Skip to content

Make return_immediately configurable for all Pub/Sub modules and mark the setting as deprecated #67621

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
michaelpri10 wants to merge 10 commits into from
Closed

Make return_immediately configurable for all Pub/Sub modules and mark the setting as deprecated #67621

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
3ec7a55
Make return_immediately default to False and be configurable for all …
michaelpri10 May 27, 2026
3b5e6dd
Update test expectations for having return_immediately default to fal…
michaelpri10 Jun 2, 2026
c488202
Merge branch 'main' into main
michaelpri10 Jun 19, 2026
ea7f6c3
Change return_immediately default value back to True and add deprecat…
michaelpri10 Jul 6, 2026
e9d18e6
test: Suppress test failures caused by use of the deprecated return_i…
michaelpri10 Jul 13, 2026
9ce49fd
Only show the deprecation warning if the return_immediately option i…
michaelpri10 Jul 21, 2026
478e400
Update deprecation warning to only include plan to change default val…
michaelpri10 Jul 23, 2026
f2d9b11
Update wording for return_immediately default value deprecation warning
michaelpri10 Jul 23, 2026
33fdf92
Merge branch 'main' into main
michaelpri10 Jul 28, 2026
6e9eb67
Merge branch 'main' into main
michaelpri10 Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 21 additions & 1 deletion providers/google/src/airflow/providers/google/cloud/operators/pubsub.py
View file
Open in desktop

@MaksYermak MaksYermak Aug 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@michaelpri10 did you run the system tests for these changes with return_immediately=False?
I am not sure that we should make this parameter configurable for PubSubPullOperator as I understand the idea of this operator was return any value immediately include 0 if user does not have any messages. I think it was a reason why we have hardcoded return_immediately=True here. Otherwise with return_immediately=False it blocks whole task until message will not appear in my opinion this behavior contradicts with operator logic. If users need blocking waiting for this reason we have PubSubPullSensor.

It is the message from docstring to this operator:

Pulls messages from a PubSub subscription and passes them through XCom.

    If the queue is empty, returns empty list - never waits for messages.
    If you do need to wait, please use :class:`airflow.providers.google.cloud.sensors.PubSubPullSensor`
    instead.

Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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] = (
Expand All @@ -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)
Expand All @@ -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:
Expand All @@ -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,
)
Expand All @@ -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
Expand Down
15 changes: 13 additions & 2 deletions providers/google/src/airflow/providers/google/cloud/sensors/pubsub.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,15 @@

from __future__ import annotations

import warnings
from collections.abc import Callable, Sequence
from datetime import timedelta
from typing import TYPE_CHECKING, Any

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
Expand Down Expand Up @@ -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,
Expand All @@ -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,
)
Comment on lines +139 to +143

@MaksYermak MaksYermak Aug 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@michaelpri10 in google provider we have deprecation policy and by this policy you need to specify planned_removal_date and by default all deprecations should allow a 6 months time period until they will be removed and not available. Could you please add the removal date to this warning message?

self.return_immediately = return_immediately
else:
self.return_immediately = True

def poke(self, context: Context) -> bool:
hook = PubSubHook(
Expand Down Expand Up @@ -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",
)
Expand Down
22 changes: 21 additions & 1 deletion providers/google/src/airflow/providers/google/cloud/triggers/pubsub.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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__(
Expand All @@ -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
Expand All @@ -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."""
Expand All @@ -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,
},
)

Expand All @@ -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)
Expand Down
53 changes: 52 additions & 1 deletion providers/google/tests/unit/google/cloud/operators/test_pubsub.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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"
Expand Down Expand Up @@ -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):
Expand All @@ -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
Comment thread
michaelpri10 marked this conversation as resolved.

@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(
Expand Down Expand Up @@ -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,
)
51 changes: 51 additions & 0 deletions providers/google/tests/unit/google/cloud/sensors/test_pubsub.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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"""
Expand Down Expand Up @@ -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,
)
Loading
Loading