diff --git a/providers/google/src/airflow/providers/google/cloud/operators/pubsub.py b/providers/google/src/airflow/providers/google/cloud/operators/pubsub.py index 33a1b302786cc..207ca834004b0 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/pubsub.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/pubsub.py @@ -25,6 +25,7 @@ from __future__ import annotations +import warnings from collections.abc import Callable, Sequence from functools import cached_property from typing import TYPE_CHECKING, Any @@ -42,6 +43,7 @@ SchemaSettings, ) +from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.common.compat.sdk import AirflowException, conf from airflow.providers.google.cloud.hooks.pubsub import PubSubHook from airflow.providers.google.cloud.links.pubsub import PubSubSubscriptionLink, PubSubTopicLink @@ -799,6 +801,13 @@ class PubSubPullOperator(GoogleCloudBaseOperator): :param deferrable: If True, run the task in the deferrable mode. :param poll_interval: Time (seconds) to wait between two consecutive calls to check the job. The default is 300 seconds. + :param return_immediately: If this field set to true, the system will + respond immediately even if it there are no messages available to + return in the ``Pull`` response. Otherwise, the system may wait + (for a bounded amount of time) until at least one message is available, + rather than returning no messages. Warning: setting this field to + ``true`` is discouraged because it adversely impacts the performance + of ``Pull`` operations. We recommend that users do not set this field. """ template_fields: Sequence[str] = ( @@ -819,6 +828,7 @@ def __init__( impersonation_chain: str | Sequence[str] | None = None, deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), poll_interval: int = 300, + return_immediately: bool | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -831,6 +841,15 @@ def __init__( self.impersonation_chain = impersonation_chain self.deferrable = deferrable self.poll_interval = poll_interval + if return_immediately is not None: + warnings.warn( + "The default value of `return_immediately` will be changed to `False` in a future major release.", + AirflowProviderDeprecationWarning, + stacklevel=2, + ) + self.return_immediately = return_immediately + else: + self.return_immediately = True def execute(self, context: Context) -> list: if self.deferrable: @@ -843,6 +862,7 @@ def execute(self, context: Context) -> list: gcp_conn_id=self.gcp_conn_id, poke_interval=self.poll_interval, impersonation_chain=self.impersonation_chain, + return_immediately=self.return_immediately, ), method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME, ) @@ -855,7 +875,7 @@ def execute(self, context: Context) -> list: project_id=self.project_id, subscription=self.subscription, max_messages=self.max_messages, - return_immediately=True, + return_immediately=self.return_immediately, ) handle_messages = self.messages_callback or self._default_message_callback diff --git a/providers/google/src/airflow/providers/google/cloud/sensors/pubsub.py b/providers/google/src/airflow/providers/google/cloud/sensors/pubsub.py index f138271b66e4c..7c9ceccd2774a 100644 --- a/providers/google/src/airflow/providers/google/cloud/sensors/pubsub.py +++ b/providers/google/src/airflow/providers/google/cloud/sensors/pubsub.py @@ -19,6 +19,7 @@ from __future__ import annotations +import warnings from collections.abc import Callable, Sequence from datetime import timedelta from typing import TYPE_CHECKING, Any @@ -26,6 +27,7 @@ from google.cloud import pubsub_v1 from google.cloud.pubsub_v1.types import ReceivedMessage +from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator, conf from airflow.providers.google.cloud.hooks.pubsub import PubSubHook from airflow.providers.google.cloud.triggers.pubsub import PubsubPullTrigger @@ -113,7 +115,7 @@ def __init__( project_id: str, subscription: str, max_messages: int = 5, - return_immediately: bool = True, + return_immediately: bool | None = None, ack_messages: bool = False, gcp_conn_id: str = "google_cloud_default", messages_callback: Callable[[list[ReceivedMessage], Context], Any] | None = None, @@ -127,13 +129,21 @@ def __init__( self.project_id = project_id self.subscription = subscription self.max_messages = max_messages - self.return_immediately = return_immediately self.ack_messages = ack_messages self.messages_callback = messages_callback self.impersonation_chain = impersonation_chain self.deferrable = deferrable self.poke_interval = poke_interval self._return_value = None + if return_immediately is not None: + warnings.warn( + "The default value of `return_immediately` will be changed to `False` in a future major release.", + AirflowProviderDeprecationWarning, + stacklevel=2, + ) + self.return_immediately = return_immediately + else: + self.return_immediately = True def poke(self, context: Context) -> bool: hook = PubSubHook( @@ -176,6 +186,7 @@ def execute(self, context: Context) -> None: poke_interval=self.poke_interval, gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, + return_immediately=self.return_immediately, ), method_name="execute_complete", ) diff --git a/providers/google/src/airflow/providers/google/cloud/triggers/pubsub.py b/providers/google/src/airflow/providers/google/cloud/triggers/pubsub.py index 22314838c4649..4c8aff5d27ac4 100644 --- a/providers/google/src/airflow/providers/google/cloud/triggers/pubsub.py +++ b/providers/google/src/airflow/providers/google/cloud/triggers/pubsub.py @@ -19,12 +19,14 @@ from __future__ import annotations import asyncio +import warnings from collections.abc import AsyncIterator, Sequence from functools import cached_property from typing import Any from google.cloud.pubsub_v1.types import ReceivedMessage +from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.google.cloud.hooks.pubsub import PubSubAsyncHook from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS from airflow.triggers.base import TriggerEvent @@ -55,6 +57,13 @@ class PubsubPullTrigger(BaseEventTrigger): 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). + :param return_immediately: If this field set to true, the system will + respond immediately even if it there are no messages available to + return in the ``Pull`` response. Otherwise, the system may wait + (for a bounded amount of time) until at least one message is available, + rather than returning no messages. Warning: setting this field to + ``true`` is discouraged because it adversely impacts the performance + of ``Pull`` operations. We recommend that users do not set this field. """ def __init__( @@ -66,6 +75,7 @@ def __init__( gcp_conn_id: str, poke_interval: float = 10.0, impersonation_chain: str | Sequence[str] | None = None, + return_immediately: bool | None = None, ): super().__init__() self.project_id = project_id @@ -75,6 +85,15 @@ def __init__( self.poke_interval = poke_interval self.gcp_conn_id = gcp_conn_id self.impersonation_chain = impersonation_chain + if return_immediately is not None: + warnings.warn( + "The default value of `return_immediately` will be changed to `False` in a future major release.", + AirflowProviderDeprecationWarning, + stacklevel=2, + ) + self.return_immediately = return_immediately + else: + self.return_immediately = True def serialize(self) -> tuple[str, dict[str, Any]]: """Serialize PubsubPullTrigger arguments and classpath.""" @@ -88,6 +107,7 @@ def serialize(self) -> tuple[str, dict[str, Any]]: "poke_interval": self.poke_interval, "gcp_conn_id": self.gcp_conn_id, "impersonation_chain": self.impersonation_chain, + "return_immediately": self.return_immediately, }, ) @@ -97,7 +117,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: project_id=self.project_id, subscription=self.subscription, max_messages=self.max_messages, - return_immediately=True, + return_immediately=self.return_immediately, ): if self.ack_messages: await self.message_acknowledgement(pulled_messages) diff --git a/providers/google/tests/unit/google/cloud/operators/test_pubsub.py b/providers/google/tests/unit/google/cloud/operators/test_pubsub.py index 3537c5266db2e..f8570d14dd253 100644 --- a/providers/google/tests/unit/google/cloud/operators/test_pubsub.py +++ b/providers/google/tests/unit/google/cloud/operators/test_pubsub.py @@ -25,6 +25,7 @@ from google.cloud import pubsub_v1 from google.cloud.pubsub_v1.types import ReceivedMessage +from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.common.compat.sdk import TaskDeferred from airflow.providers.google.cloud.operators.pubsub import ( PubSubCreateSubscriptionOperator, @@ -34,6 +35,9 @@ PubSubPublishMessageOperator, PubSubPullOperator, ) +from airflow.providers.google.cloud.triggers.pubsub import PubsubPullTrigger + +pytestmark = pytest.mark.filterwarnings("ignore::airflow.exceptions.AirflowProviderDeprecationWarning") TASK_ID = "test-task-id" TEST_PROJECT = "test-project" @@ -513,6 +517,21 @@ def messages_callback( assert response == messages_callback_return_value + @mock.patch("airflow.providers.google.cloud.operators.pubsub.PubSubHook") + def test_execute_with_return_immediately_false(self, mock_hook): + operator = PubSubPullOperator( + task_id=TASK_ID, + project_id=TEST_PROJECT, + subscription=TEST_SUBSCRIPTION, + return_immediately=False, + ) + + mock_hook.return_value.pull.return_value = [] + operator.execute({}) + mock_hook.return_value.pull.assert_called_once_with( + project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, max_messages=5, return_immediately=False + ) + @pytest.mark.db_test @mock.patch("airflow.providers.google.cloud.operators.pubsub.PubSubHook") def test_execute_deferred(self, mock_hook): @@ -526,9 +545,32 @@ def test_execute_deferred(self, mock_hook): subscription=TEST_SUBSCRIPTION, deferrable=True, ) - with pytest.raises(TaskDeferred) as _: + with pytest.raises(TaskDeferred) as exc: + task.execute(mock.MagicMock()) + + assert isinstance(exc.value.trigger, PubsubPullTrigger) + assert exc.value.trigger.return_immediately is True + + @pytest.mark.db_test + @mock.patch("airflow.providers.google.cloud.operators.pubsub.PubSubHook") + def test_execute_deferred_with_return_immediately_false(self, mock_hook): + """ + Asserts that a task is deferred and a PubSubPullOperator will be fired + when the PubSubPullOperator is executed with deferrable=True. + """ + task = PubSubPullOperator( + task_id=TASK_ID, + project_id=TEST_PROJECT, + subscription=TEST_SUBSCRIPTION, + return_immediately=False, + deferrable=True, + ) + with pytest.raises(TaskDeferred) as exc: task.execute(mock.MagicMock()) + assert isinstance(exc.value.trigger, PubsubPullTrigger) + assert exc.value.trigger.return_immediately is False + @mock.patch("airflow.providers.google.cloud.operators.pubsub.PubSubHook") def test_get_openlineage_facets(self, mock_hook): operator = PubSubPullOperator( @@ -631,3 +673,12 @@ def test_execute_complete_use_default_message_callback(self, mock_hook): resp = operator.execute_complete(context={}, event={"status": "success", "message": test_message}) mock_log_info.assert_called_with("Sensor pulls messages: %s", test_message) assert resp == [ReceivedMessage.to_dict(m) for m in received_messages] + + def test_pubsub_pull_operator_deprecation_warning(self): + with pytest.warns(AirflowProviderDeprecationWarning, match="return_immediately"): + PubSubPullOperator( + task_id=TASK_ID, + project_id=TEST_PROJECT, + subscription=TEST_SUBSCRIPTION, + return_immediately=False, + ) diff --git a/providers/google/tests/unit/google/cloud/sensors/test_pubsub.py b/providers/google/tests/unit/google/cloud/sensors/test_pubsub.py index 4cd1b48fbfb60..64a6d72dd511a 100644 --- a/providers/google/tests/unit/google/cloud/sensors/test_pubsub.py +++ b/providers/google/tests/unit/google/cloud/sensors/test_pubsub.py @@ -24,10 +24,13 @@ from google.cloud import pubsub_v1 from google.cloud.pubsub_v1.types import ReceivedMessage +from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred from airflow.providers.google.cloud.sensors.pubsub import PubSubPullSensor from airflow.providers.google.cloud.triggers.pubsub import PubsubPullTrigger +pytestmark = pytest.mark.filterwarnings("ignore::airflow.exceptions.AirflowProviderDeprecationWarning") + TASK_ID = "test-task-id" TEST_PROJECT = "test-project" TEST_SUBSCRIPTION = "test-subscription" @@ -99,6 +102,26 @@ def test_execute(self, mock_hook): ) assert generated_dicts == response + @mock.patch("airflow.providers.google.cloud.sensors.pubsub.PubSubHook") + def test_execute_with_return_immediately_false(self, mock_hook): + operator = PubSubPullSensor( + task_id=TASK_ID, + project_id=TEST_PROJECT, + subscription=TEST_SUBSCRIPTION, + poke_interval=0, + return_immediately=False, + ) + + generated_messages = self._generate_messages(5) + generated_dicts = self._generate_dicts(5) + mock_hook.return_value.pull.return_value = generated_messages + + response = operator.execute({}) + mock_hook.return_value.pull.assert_called_once_with( + project_id=TEST_PROJECT, subscription=TEST_SUBSCRIPTION, max_messages=5, return_immediately=False + ) + assert generated_dicts == response + @mock.patch("airflow.providers.google.cloud.sensors.pubsub.PubSubHook") def test_execute_timeout(self, mock_hook): operator = PubSubPullSensor( @@ -167,6 +190,25 @@ def test_pubsub_pull_sensor_async(self): with pytest.raises(TaskDeferred) as exc: task.execute(context={}) assert isinstance(exc.value.trigger, PubsubPullTrigger), "Trigger is not a PubsubPullTrigger" + assert exc.value.trigger.return_immediately is True + + def test_pubsub_pull_sensor_async_with_return_immediately_false(self): + """ + Asserts that a task is deferred and a PubsubPullTrigger will be fired + with custom return_immediately value. + """ + task = PubSubPullSensor( + task_id="test_task_id", + ack_messages=True, + project_id=TEST_PROJECT, + subscription=TEST_SUBSCRIPTION, + deferrable=True, + return_immediately=False, + ) + with pytest.raises(TaskDeferred) as exc: + task.execute(context={}) + assert isinstance(exc.value.trigger, PubsubPullTrigger), "Trigger is not a PubsubPullTrigger" + assert exc.value.trigger.return_immediately is False def test_pubsub_pull_sensor_async_execute_should_throw_exception(self): """Tests that an AirflowException is raised in case of error event""" @@ -245,3 +287,12 @@ def messages_callback( resp = operator.execute_complete(context={}, event={"status": "success", "message": test_message}) mock_log_info.assert_called_with("Sensor pulls messages: %s", test_message) assert resp == messages_callback_return_value + + def test_pubsub_pull_sensor_deprecation_warning(self): + with pytest.warns(AirflowProviderDeprecationWarning, match="return_immediately"): + PubSubPullSensor( + task_id=TASK_ID, + project_id=TEST_PROJECT, + subscription=TEST_SUBSCRIPTION, + return_immediately=False, + ) diff --git a/providers/google/tests/unit/google/cloud/triggers/test_pubsub.py b/providers/google/tests/unit/google/cloud/triggers/test_pubsub.py index 7fb95260b17fa..9ff897be74bb9 100644 --- a/providers/google/tests/unit/google/cloud/triggers/test_pubsub.py +++ b/providers/google/tests/unit/google/cloud/triggers/test_pubsub.py @@ -22,9 +22,12 @@ from google.api_core.exceptions import GoogleAPICallError from google.cloud.pubsub_v1.types import ReceivedMessage +from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.google.cloud.triggers.pubsub import PubsubPullTrigger from airflow.triggers.base import TriggerEvent +pytestmark = pytest.mark.filterwarnings("ignore::airflow.exceptions.AirflowProviderDeprecationWarning") + TEST_POLL_INTERVAL = 10 TEST_GCP_CONN_ID = "google_cloud_default" PROJECT_ID = "test_project_id" @@ -74,8 +77,34 @@ def test_async_pubsub_pull_trigger_serialization_should_execute_successfully(sel "poke_interval": TEST_POLL_INTERVAL, "gcp_conn_id": TEST_GCP_CONN_ID, "impersonation_chain": None, + "return_immediately": True, } + @pytest.mark.asyncio + @mock.patch("airflow.providers.google.cloud.hooks.pubsub.PubSubAsyncHook.pull") + async def test_async_pubsub_pull_trigger_passes_return_immediately_false(self, mock_pull): + """Test that return_immediately is passed to the hook.""" + mock_pull.return_value = generate_messages(1) + trigger = PubsubPullTrigger( + project_id=PROJECT_ID, + subscription="subscription", + max_messages=MAX_MESSAGES, + ack_messages=False, + poke_interval=TEST_POLL_INTERVAL, + gcp_conn_id=TEST_GCP_CONN_ID, + impersonation_chain=None, + return_immediately=False, + ) + + await trigger.run().asend(None) + + mock_pull.assert_called_once_with( + project_id=PROJECT_ID, + subscription="subscription", + max_messages=MAX_MESSAGES, + return_immediately=False, + ) + @pytest.mark.asyncio @mock.patch("airflow.providers.google.cloud.hooks.pubsub.PubSubAsyncHook.pull") async def test_async_pubsub_pull_trigger_return_event(self, mock_pull): @@ -172,3 +201,16 @@ async def test_async_pubsub_pull_trigger_exception_during_ack(self, mock_pull, m with pytest.raises(GoogleAPICallError, match="Acknowledgement failed"): await trigger.run().asend(None) + + def test_pubsub_pull_trigger_deprecation_warning(self): + with pytest.warns(AirflowProviderDeprecationWarning, match="return_immediately"): + PubsubPullTrigger( + project_id=PROJECT_ID, + subscription="subscription", + max_messages=MAX_MESSAGES, + ack_messages=ACK_MESSAGES, + poke_interval=TEST_POLL_INTERVAL, + gcp_conn_id=TEST_GCP_CONN_ID, + impersonation_chain=None, + return_immediately=False, + )