Adding Google Pub/Sub as a viable queue to be used by the MessageQueueTrigger

providers/google/src/airflow/providers/google/cloud/queues/pubsub.py

@@ -0,0 +1,52 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING
+
+from airflow.exceptions import AirflowOptionalProviderFeatureException
+from airflow.providers.google.cloud.triggers.pubsub import PubsubPullTrigger
+
+try:
+    from airflow.providers.common.messaging.providers.base_provider import BaseMessageQueueProvider
+except ImportError:
+    raise AirflowOptionalProviderFeatureException(
+        "This feature requires the 'common.messaging' provider to be installed in version >= 1.0.1."
+    )
+
+if TYPE_CHECKING:
+    from airflow.triggers.base import BaseEventTrigger
+
+# [START queue_regexp]
+QUEUE_REGEXP = r".+"
+# [END queue_regexp]
+
+
+class PubSubMessageQueueProvider(BaseMessageQueueProvider):
+    """Configuration for Pub/Sub integration with common-messaging."""
+
+    def queue_matches(self, queue: str) -> bool:
+        return bool(re.match(QUEUE_REGEXP, queue))
+
+    def trigger_class(self) -> type[BaseEventTrigger]:
+        return PubsubPullTrigger
+
+    def trigger_kwargs(self, queue: str, **kwargs) -> dict:
+        return {
+            "subscription": queue
+        }

providers/google/src/airflow/providers/google/cloud/triggers/pubsub.py

@@ -33,55 +33,60 @@ class PubsubPullTrigger(BaseTrigger):
     """
     Initialize the Pubsub Pull Trigger with needed parameters.
 
+    :param subscription: the Pub/Sub subscription name. Do not include the full subscription path
     :param project_id: the Google Cloud project ID for the subscription (templated)
-    :param subscription: the Pub/Sub subscription name. Do not include the full subscription path.
-    :param max_messages: The maximum number of messages to retrieve per
-        PubSub pull request
-    :param ack_messages: If True, each message will be acknowledged
-        immediately rather than by any downstream tasks
     :param gcp_conn_id: Reference to google cloud connection id
-    :param poke_interval: polling period in seconds to check for the status
     :param impersonation_chain: Optional service account to impersonate using short-term
         credentials, or chained list of accounts required to get the access_token
         of the last account in the list, which will be impersonated in the request.
         If set as a string, the account must grant the originating account
         the Service Account Token Creator IAM role.
         If set as a sequence, the identities from the list must grant
         Service Account Token Creator IAM role to the directly preceding identity, with first
-        account from the list granting this role to the originating account (templated).
+        account from the list granting this role to the originating account (templated)
+
+    :param max_messages: The maximum number of messages to retrieve per
+        PubSub pull request
+    :param ack_messages: If True, each message will be acknowledged
+        immediately rather than by any downstream tasks
+    :param poke_interval: polling period in seconds to check for the status
+    :param waiter_delay: The time in seconds to wait between calls to the SQS API to receive messages
     """
 
     def __init__(
         self,
         project_id: str,
         subscription: str,
-        max_messages: int,
-        ack_messages: bool,
-        gcp_conn_id: str,
-        poke_interval: float = 10.0,
+        gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
+        max_messages: int = 100,
+        ack_messages: bool = True,
+        poke_interval: float = 10.0,
+        waiter_delay: int = 60,
     ):
         super().__init__()
         self.project_id = project_id
         self.subscription = subscription
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
         self.max_messages = max_messages
         self.ack_messages = ack_messages
         self.poke_interval = poke_interval
-        self.gcp_conn_id = gcp_conn_id
-        self.impersonation_chain = impersonation_chain
+        self.waiter_delay = waiter_delay
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         """Serialize PubsubPullTrigger arguments and classpath."""
         return (
             "airflow.providers.google.cloud.triggers.pubsub.PubsubPullTrigger",
             {
-                "project_id": self.project_id,
                 "subscription": self.subscription,
+                "project_id": self.project_id,
+                "gcp_conn_id": self.gcp_conn_id,
+                "impersonation_chain": self.impersonation_chain,
                 "max_messages": self.max_messages,
                 "ack_messages": self.ack_messages,
                 "poke_interval": self.poke_interval,
-                "gcp_conn_id": self.gcp_conn_id,
-                "impersonation_chain": self.impersonation_chain,
+                "waiter_delay": self.waiter_delay,
             },
         )
 
@@ -101,8 +106,10 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
 
                     yield TriggerEvent({"status": "success", "message": messages_json})
                     return
+
                 self.log.info("Sleeping for %s seconds.", self.poke_interval)
                 await asyncio.sleep(self.poke_interval)
+
         except Exception as e:
             yield TriggerEvent({"status": "error", "message": str(e)})
             return

providers/google/src/airflow/providers/google/get_provider_info.py

@@ -1517,4 +1517,5 @@ def get_provider_info():
             "airflow.providers.google.cloud.log.gcs_task_handler.GCSTaskHandler",
             "airflow.providers.google.cloud.log.stackdriver_task_handler.StackdriverTaskHandler",
         ],
+        "queues": ["airflow.providers.google.cloud.queues.pubsub.PubSubMessageQueueProvider"]
     }