Skip to content

Commit 759ce2a

Browse files
turbaszekANiteckiPmik-laj
authored
[AIRFLOW-6978] Add PubSubPullOperator (#7766)
* Fix typo. * Type checking: pass empty dict instead of None to operators in test scenarios. * Reformat some whitespace. * Fix PubSubHook.pull return value type annotation. * PubSubHook.acknowledge: Implement acknowledging by passing list of ReceivedMessage objects. * Refactor PubSubPullSensor. * Implement PubSubPullOperator. * Remove unused constant from PubSub tests. * Fix Sensor contract: some do have a return value. * Reformat whitespace in PubSub tests. * PubSub: Add tests for PubSubPullOperator. * Fix PubSubPullSensor tests after refactoring. * Sort imports. * Fix pylint complaining about callback interface. * Update airflow/providers/google/cloud/operators/pubsub.py Co-Authored-By: Kamil Breguła <mik-laj@users.noreply.github.com> * Update airflow/providers/google/cloud/hooks/pubsub.py Co-Authored-By: Kamil Breguła <mik-laj@users.noreply.github.com> * Update airflow/providers/google/cloud/hooks/pubsub.py Co-Authored-By: Kamil Breguła <mik-laj@users.noreply.github.com> * Update airflow/providers/google/cloud/operators/pubsub.py Co-Authored-By: Kamil Breguła <mik-laj@users.noreply.github.com> * Reorder PubSub messages_callback argument. * Update docstring. * Reformat mutually exclusive argument handling logic in PubSubHook. * PubSub: Deprecate return_immediately argument. * Implement example DAG and system test for PubSubPullOperator. * Fix docstring formatting. * PubSubPullOperator: Fix docs. * Apply suggestions from code review Co-Authored-By: Tomek Urbaszek <turbaszek@gmail.com> Co-authored-by: Aleksander Nitecki <aleksander.nitecki@polidea.com> Co-authored-by: ANiteckiP <60920935+ANiteckiP@users.noreply.github.com> Co-authored-by: Kamil Breguła <mik-laj@users.noreply.github.com>
1 parent 6b9b214 commit 759ce2a

10 files changed

Lines changed: 534 additions & 70 deletions

File tree

‎airflow/providers/google/cloud/example_dags/example_pubsub.py‎

Lines changed: 74 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,14 @@
2525
from airflow.operators.bash import BashOperator
2626
from airflow.providers.google.cloud.operators.pubsub import (
2727
PubSubCreateSubscriptionOperator, PubSubCreateTopicOperator, PubSubDeleteSubscriptionOperator,
28-
PubSubDeleteTopicOperator, PubSubPublishMessageOperator,
28+
PubSubDeleteTopicOperator, PubSubPublishMessageOperator, PubSubPullOperator,
2929
)
3030
from airflow.providers.google.cloud.sensors.pubsub import PubSubPullSensor
3131
from airflow.utils.dates import days_ago
3232

3333
GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "your-project-id")
34-
TOPIC = "PubSubTestTopic"
34+
TOPIC_FOR_SENSOR_DAG = "PubSubSensorTestTopic"
35+
TOPIC_FOR_OPERATOR_DAG = "PubSubOperatorTestTopic"
3536
MESSAGE = {"data": b"Tool", "attributes": {"name": "wrench", "mass": "1.3kg", "count": "3"}}
3637

3738
default_args = {"start_date": days_ago(1)}
@@ -45,23 +46,23 @@
4546
# [END howto_operator_gcp_pubsub_pull_messages_result_cmd]
4647

4748
with models.DAG(
48-
"example_gcp_pubsub",
49+
"example_gcp_pubsub_sensor",
4950
default_args=default_args,
5051
schedule_interval=None, # Override to match your needs
51-
) as example_dag:
52+
) as example_sensor_dag:
5253
# [START howto_operator_gcp_pubsub_create_topic]
5354
create_topic = PubSubCreateTopicOperator(
54-
task_id="create_topic", topic=TOPIC, project_id=GCP_PROJECT_ID
55+
task_id="create_topic", topic=TOPIC_FOR_SENSOR_DAG, project_id=GCP_PROJECT_ID
5556
)
5657
# [END howto_operator_gcp_pubsub_create_topic]
5758

5859
# [START howto_operator_gcp_pubsub_create_subscription]
5960
subscribe_task = PubSubCreateSubscriptionOperator(
60-
task_id="subscribe_task", project_id=GCP_PROJECT_ID, topic=TOPIC
61+
task_id="subscribe_task", project_id=GCP_PROJECT_ID, topic=TOPIC_FOR_SENSOR_DAG
6162
)
6263
# [END howto_operator_gcp_pubsub_create_subscription]
6364

64-
# [START howto_operator_gcp_pubsub_pull_message]
65+
# [START howto_operator_gcp_pubsub_pull_message_with_sensor]
6566
subscription = "{{ task_instance.xcom_pull('subscribe_task') }}"
6667

6768
pull_messages = PubSubPullSensor(
@@ -70,7 +71,7 @@
7071
project_id=GCP_PROJECT_ID,
7172
subscription=subscription,
7273
)
73-
# [END howto_operator_gcp_pubsub_pull_message]
74+
# [END howto_operator_gcp_pubsub_pull_message_with_sensor]
7475

7576
# [START howto_operator_gcp_pubsub_pull_messages_result]
7677
pull_messages_result = BashOperator(
@@ -82,7 +83,7 @@
8283
publish_task = PubSubPublishMessageOperator(
8384
task_id="publish_task",
8485
project_id=GCP_PROJECT_ID,
85-
topic=TOPIC,
86+
topic=TOPIC_FOR_SENSOR_DAG,
8687
messages=[MESSAGE, MESSAGE, MESSAGE],
8788
)
8889
# [END howto_operator_gcp_pubsub_publish]
@@ -97,9 +98,72 @@
9798

9899
# [START howto_operator_gcp_pubsub_delete_topic]
99100
delete_topic = PubSubDeleteTopicOperator(
100-
task_id="delete_topic", topic=TOPIC, project_id=GCP_PROJECT_ID
101+
task_id="delete_topic", topic=TOPIC_FOR_SENSOR_DAG, project_id=GCP_PROJECT_ID
101102
)
102103
# [END howto_operator_gcp_pubsub_delete_topic]
103104

104105
create_topic >> subscribe_task >> publish_task
105106
subscribe_task >> pull_messages >> pull_messages_result >> unsubscribe_task >> delete_topic
107+
108+
109+
with models.DAG(
110+
"example_gcp_pubsub_operator",
111+
default_args=default_args,
112+
schedule_interval=None, # Override to match your needs
113+
) as example_operator_dag:
114+
# [START howto_operator_gcp_pubsub_create_topic]
115+
create_topic = PubSubCreateTopicOperator(
116+
task_id="create_topic", topic=TOPIC_FOR_OPERATOR_DAG, project_id=GCP_PROJECT_ID
117+
)
118+
# [END howto_operator_gcp_pubsub_create_topic]
119+
120+
# [START howto_operator_gcp_pubsub_create_subscription]
121+
subscribe_task = PubSubCreateSubscriptionOperator(
122+
task_id="subscribe_task", project_id=GCP_PROJECT_ID, topic=TOPIC_FOR_OPERATOR_DAG
123+
)
124+
# [END howto_operator_gcp_pubsub_create_subscription]
125+
126+
# [START howto_operator_gcp_pubsub_pull_message_with_operator]
127+
subscription = "{{ task_instance.xcom_pull('subscribe_task') }}"
128+
129+
pull_messages = PubSubPullOperator(
130+
task_id="pull_messages",
131+
ack_messages=True,
132+
project_id=GCP_PROJECT_ID,
133+
subscription=subscription,
134+
)
135+
# [END howto_operator_gcp_pubsub_pull_message_with_operator]
136+
137+
# [START howto_operator_gcp_pubsub_pull_messages_result]
138+
pull_messages_result = BashOperator(
139+
task_id="pull_messages_result", bash_command=echo_cmd
140+
)
141+
# [END howto_operator_gcp_pubsub_pull_messages_result]
142+
143+
# [START howto_operator_gcp_pubsub_publish]
144+
publish_task = PubSubPublishMessageOperator(
145+
task_id="publish_task",
146+
project_id=GCP_PROJECT_ID,
147+
topic=TOPIC_FOR_OPERATOR_DAG,
148+
messages=[MESSAGE, MESSAGE, MESSAGE],
149+
)
150+
# [END howto_operator_gcp_pubsub_publish]
151+
152+
# [START howto_operator_gcp_pubsub_unsubscribe]
153+
unsubscribe_task = PubSubDeleteSubscriptionOperator(
154+
task_id="unsubscribe_task",
155+
project_id=GCP_PROJECT_ID,
156+
subscription="{{ task_instance.xcom_pull('subscribe_task') }}",
157+
)
158+
# [END howto_operator_gcp_pubsub_unsubscribe]
159+
160+
# [START howto_operator_gcp_pubsub_delete_topic]
161+
delete_topic = PubSubDeleteTopicOperator(
162+
task_id="delete_topic", topic=TOPIC_FOR_OPERATOR_DAG, project_id=GCP_PROJECT_ID
163+
)
164+
# [END howto_operator_gcp_pubsub_delete_topic]
165+
166+
(
167+
create_topic >> subscribe_task >> publish_task
168+
>> pull_messages >> pull_messages_result >> unsubscribe_task >> delete_topic
169+
)

‎airflow/providers/google/cloud/hooks/pubsub.py‎

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
from google.api_core.retry import Retry
2929
from google.cloud.exceptions import NotFound
3030
from google.cloud.pubsub_v1 import PublisherClient, SubscriberClient
31-
from google.cloud.pubsub_v1.types import Duration, MessageStoragePolicy, PushConfig
31+
from google.cloud.pubsub_v1.types import Duration, MessageStoragePolicy, PushConfig, ReceivedMessage
3232
from googleapiclient.errors import HttpError
3333

3434
from airflow.providers.google.cloud.hooks.base import CloudBaseHook
@@ -460,7 +460,7 @@ def pull(
460460
retry: Optional[Retry] = None,
461461
timeout: Optional[float] = None,
462462
metadata: Optional[Sequence[Tuple[str, str]]] = None,
463-
) -> List[Dict]:
463+
) -> List[ReceivedMessage]:
464464
"""
465465
Pulls up to ``max_messages`` messages from Pub/Sub subscription.
466466
@@ -496,7 +496,7 @@ def pull(
496496
subscriber = self.subscriber_client
497497
subscription_path = SubscriberClient.subscription_path(project_id, subscription) # noqa E501 # pylint: disable=no-member,line-too-long
498498

499-
self.log.info("Pulling mex %d messages from subscription (path) %s", max_messages, subscription_path)
499+
self.log.info("Pulling max %d messages from subscription (path) %s", max_messages, subscription_path)
500500
try:
501501
# pylint: disable=no-member
502502
response = subscriber.pull(
@@ -517,7 +517,8 @@ def pull(
517517
def acknowledge(
518518
self,
519519
subscription: str,
520-
ack_ids: List[str],
520+
ack_ids: Optional[List[str]] = None,
521+
messages: Optional[List[ReceivedMessage]] = None,
521522
project_id: Optional[str] = None,
522523
retry: Optional[Retry] = None,
523524
timeout: Optional[float] = None,
@@ -529,9 +530,12 @@ def acknowledge(
529530
:param subscription: the Pub/Sub subscription name to delete; do not
530531
include the 'projects/{project}/topics/' prefix.
531532
:type subscription: str
532-
:param ack_ids: List of ReceivedMessage ackIds from a previous pull
533-
response
533+
:param ack_ids: List of ReceivedMessage ackIds from a previous pull response.
534+
Mutually exclusive with ``messages`` argument.
534535
:type ack_ids: list
536+
:param messages: List of ReceivedMessage objects to acknowledge.
537+
Mutually exclusive with ``ack_ids`` argument.
538+
:type messages: list
535539
:param project_id: Optional, the GCP project name or ID in which to create the topic
536540
If set to None or missing, the default project_id from the GCP connection is used.
537541
:type project_id: str
@@ -545,8 +549,20 @@ def acknowledge(
545549
:param metadata: (Optional) Additional metadata that is provided to the method.
546550
:type metadata: Sequence[Tuple[str, str]]]
547551
"""
552+
548553
if not project_id:
549554
raise ValueError("Project ID should be set.")
555+
556+
if ack_ids is not None and messages is None:
557+
pass
558+
elif ack_ids is None and messages is not None:
559+
ack_ids = [
560+
message.ack_id
561+
for message in messages
562+
]
563+
else:
564+
raise ValueError("One and only one of 'ack_ids' and 'messages' arguments have to be provided")
565+
550566
subscriber = self.subscriber_client
551567
subscription_path = SubscriberClient.subscription_path(project_id, subscription) # noqa E501 # pylint: disable=no-member,line-too-long
552568

‎airflow/providers/google/cloud/operators/pubsub.py‎

Lines changed: 123 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,11 @@
1919
This module contains Google PubSub operators.
2020
"""
2121
import warnings
22-
from typing import Dict, List, Optional, Sequence, Tuple, Union
22+
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
2323

2424
from google.api_core.retry import Retry
25-
from google.cloud.pubsub_v1.types import Duration, MessageStoragePolicy, PushConfig
25+
from google.cloud.pubsub_v1.types import Duration, MessageStoragePolicy, PushConfig, ReceivedMessage
26+
from google.protobuf.json_format import MessageToDict
2627

2728
from airflow.models import BaseOperator
2829
from airflow.providers.google.cloud.hooks.pubsub import PubSubHook
@@ -666,3 +667,123 @@ def execute(self, context):
666667
self.log.info("Publishing to topic %s", self.topic)
667668
hook.publish(project_id=self.project_id, topic=self.topic, messages=self.messages)
668669
self.log.info("Published to topic %s", self.topic)
670+
671+
672+
class PubSubPullOperator(BaseOperator):
673+
"""Pulls messages from a PubSub subscription and passes them through XCom.
674+
If the queue is empty, returns empty list - never waits for messages.
675+
If you do need to wait, please use :class:`airflow.providers.google.cloud.sensors.PubSubPullSensor`
676+
instead.
677+
678+
.. seealso::
679+
For more information on how to use this operator, take a look at the guide:
680+
:ref:`howto/operator:PubSubPullSensor`
681+
682+
This sensor operator will pull up to ``max_messages`` messages from the
683+
specified PubSub subscription. When the subscription returns messages,
684+
the poke method's criteria will be fulfilled and the messages will be
685+
returned from the operator and passed through XCom for downstream tasks.
686+
687+
If ``ack_messages`` is set to True, messages will be immediately
688+
acknowledged before being returned, otherwise, downstream tasks will be
689+
responsible for acknowledging them.
690+
691+
``project`` and ``subscription`` are templated so you can use
692+
variables in them.
693+
694+
:param project: the GCP project ID for the subscription (templated)
695+
:type project: str
696+
:param subscription: the Pub/Sub subscription name. Do not include the
697+
full subscription path.
698+
:type subscription: str
699+
:param max_messages: The maximum number of messages to retrieve per
700+
PubSub pull request
701+
:type max_messages: int
702+
:param ack_messages: If True, each message will be acknowledged
703+
immediately rather than by any downstream tasks
704+
:type ack_messages: bool
705+
:param gcp_conn_id: The connection ID to use connecting to
706+
Google Cloud Platform.
707+
:type gcp_conn_id: str
708+
:param delegate_to: The account to impersonate, if any.
709+
For this to work, the service account making the request
710+
must have domain-wide delegation enabled.
711+
:type delegate_to: str
712+
:param messages_callback: (Optional) Callback to process received messages.
713+
It's return value will be saved to XCom.
714+
If you are pulling large messages, you probably want to provide a custom callback.
715+
If not provided, the default implementation will convert `ReceivedMessage` objects
716+
into JSON-serializable dicts using `google.protobuf.json_format.MessageToDict` function.
717+
:type messages_callback: Optional[Callable[[List[ReceivedMessage], Dict[str, Any]], Any]]
718+
"""
719+
template_fields = ['project_id', 'subscription']
720+
721+
@apply_defaults
722+
def __init__(
723+
self,
724+
project_id: str,
725+
subscription: str,
726+
max_messages: int = 5,
727+
ack_messages: bool = False,
728+
messages_callback: Optional[Callable[[List[ReceivedMessage], Dict[str, Any]], Any]] = None,
729+
gcp_conn_id: str = 'google_cloud_default',
730+
delegate_to: Optional[str] = None,
731+
*args,
732+
**kwargs
733+
) -> None:
734+
super().__init__(*args, **kwargs)
735+
self.gcp_conn_id = gcp_conn_id
736+
self.delegate_to = delegate_to
737+
self.project_id = project_id
738+
self.subscription = subscription
739+
self.max_messages = max_messages
740+
self.ack_messages = ack_messages
741+
self.messages_callback = messages_callback
742+
743+
def execute(self, context):
744+
hook = PubSubHook(
745+
gcp_conn_id=self.gcp_conn_id,
746+
delegate_to=self.delegate_to,
747+
)
748+
749+
pulled_messages = hook.pull(
750+
project_id=self.project_id,
751+
subscription=self.subscription,
752+
max_messages=self.max_messages,
753+
return_immediately=True,
754+
)
755+
756+
handle_messages = self.messages_callback or self._default_message_callback
757+
758+
ret = handle_messages(pulled_messages, context)
759+
760+
if pulled_messages and self.ack_messages:
761+
hook.acknowledge(
762+
project_id=self.project_id,
763+
subscription=self.subscription,
764+
messages=pulled_messages,
765+
)
766+
767+
return ret
768+
769+
def _default_message_callback(
770+
self,
771+
pulled_messages: List[ReceivedMessage],
772+
context: Dict[str, Any], # pylint: disable=unused-argument
773+
):
774+
"""
775+
This method can be overridden by subclasses or by `messages_callback` constructor argument.
776+
This default implementation converts `ReceivedMessage` objects into JSON-serializable dicts.
777+
778+
:param pulled_messages: messages received from the topic.
779+
:type pulled_messages: List[ReceivedMessage]
780+
:param context: same as in `execute`
781+
:return: value to be saved to XCom.
782+
"""
783+
784+
messages_json = [
785+
MessageToDict(m)
786+
for m in pulled_messages
787+
]
788+
789+
return messages_json

0 commit comments

Comments
 (0)