This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new fc172d09a5b Add deferrable mode to `SqlSensor` (#76270)
fc172d09a5b is described below
commit fc172d09a5bc0544df3d86d581d20de4f06f564d
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Nov 4 12:11:24 2026 +0000
Add deferrable mode to `SqlSensor` (#76270)
With deferrable=True the triggerer polls the query every poke_interval
until the first cell is truthy, so the task holds no worker slot while it
waits. Works with any DbApiHook connection.
When success, failure or selector callables, silent_fail, never_fail or
exponential_backoff are set, a subclass overrides poke, or Airflow 3 is
older than 3.2.2, the sensor logs a warning and runs in its usual mode.
* Run deferrable SqlSensor tests against a deferral-capable Airflow version
The provider's compat jobs run these tests on Airflow 3.0 and 3.1, where
the sensor falls back to poking, so the deferral and fallback-reason tests
failed or hung there.
* Keep SQLExecuteQueryTrigger's call shape and widen SqlSensor's deferral
checks
SQLExecuteQueryTrigger calls get_records(sql) unchanged when no parameters
are given, so GenericTransfer's call is the same as before.
SqlSensor now treats any itemgetter(0) selector, including one copied by
deepcopy, as the default condition, and stays on the worker when a subclass
overrides _get_hook or, on Airflow 2, when parameters hold values that its
JSON trigger storage would turn into strings.
---
providers/common/sql/docs/operators.rst | 47 +++++
.../airflow/providers/common/sql/sensors/sql.py | 86 ++++++++-
.../airflow/providers/common/sql/sensors/sql.pyi | 10 +-
.../airflow/providers/common/sql/triggers/sql.py | 85 ++++++++-
.../airflow/providers/common/sql/triggers/sql.pyi | 26 ++-
.../airflow/providers/common/sql/version_compat.py | 2 +
.../tests/system/common/sql/example_sql_sensor.py | 50 ++++++
.../sql/tests/unit/common/sql/sensors/test_sql.py | 200 ++++++++++++++++++++-
.../sql/tests/unit/common/sql/triggers/test_sql.py | 94 +++++++++-
9 files changed, 585 insertions(+), 15 deletions(-)
diff --git a/providers/common/sql/docs/operators.rst
b/providers/common/sql/docs/operators.rst
index 5b59b805d71..7b7f3f49eb0 100644
--- a/providers/common/sql/docs/operators.rst
+++ b/providers/common/sql/docs/operators.rst
@@ -282,6 +282,53 @@ between two connections.
:start-after: [START howto_operator_generic_transfer]
:end-before: [END howto_operator_generic_transfer]
+.. _howto/sensor:SqlSensor:
+
+Wait for a SQL condition
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+Use the :class:`~airflow.providers.common.sql.sensors.sql.SqlSensor` to wait
until a query returns a row
+whose first cell is truthy, for example until today's orders have landed. It
works with any connection
+whose hook is a :class:`~airflow.providers.common.sql.hooks.sql.DbApiHook`.
The examples below use
+PostgreSQL; write ``sql`` and ``parameters`` in your database's dialect and
placeholder style.
+
+.. exampleinclude:: /../tests/system/common/sql/example_sql_sensor.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_sensor_sql]
+ :end-before: [END howto_sensor_sql]
+
+For a condition you cannot write in SQL, pass ``success`` or ``failure``
callables. They receive the first
+cell of the first row, or whatever ``selector`` returns for that row.
``fail_on_empty=True`` fails the task
+when the query returns no rows instead of checking again.
+
+Deferrable mode
+---------------
+
+With ``deferrable=True``, the sensor checks once on the worker. If the
condition is not met, it defers and
+the triggerer runs the query every ``poke_interval`` seconds, so the task
holds no worker slot while it
+waits. ``timeout`` caps the time spent deferred.
+
+.. exampleinclude:: /../tests/system/common/sql/example_sql_sensor.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_sensor_sql_deferrable]
+ :end-before: [END howto_sensor_sql_deferrable]
+
+The triggerer only checks the default condition, that the first cell of the
first row is truthy, so write
+the condition in SQL as above. The sensor logs a warning and runs in its usual
``mode`` instead of
+deferring when:
+
+- ``success``, ``failure`` or ``selector`` is set, because these Python
callables only run on a worker;
+- ``silent_fail``, ``never_fail`` or ``exponential_backoff`` is set, because
they only apply to pokes on a
+ worker;
+- a subclass overrides ``poke`` or ``_get_hook``, because the triggerer would
not run that code;
+- it runs on Airflow 3 before 3.2.2, whose triggerer cannot safely look up
connections from the thread
+ that runs the query. Deferring works on Airflow 2.11 and on 3.2.2 and later;
+- on Airflow 2, ``parameters`` holds values other than strings, numbers,
booleans and ``None``.
+ Airflow 2 stores trigger arguments as JSON, so dates, datetimes and
``Decimal`` values would reach
+ the triggerer as strings.
+
Analytics Operator
~~~~~~~~~~~~~~~~~~
diff --git
a/providers/common/sql/src/airflow/providers/common/sql/sensors/sql.py
b/providers/common/sql/src/airflow/providers/common/sql/sensors/sql.py
index c9c4ab3b76e..510eb60317a 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/sensors/sql.py
+++ b/providers/common/sql/src/airflow/providers/common/sql/sensors/sql.py
@@ -17,15 +17,41 @@
from __future__ import annotations
from collections.abc import Callable, Mapping, Sequence
+from datetime import timedelta
from operator import itemgetter
from typing import TYPE_CHECKING, Any
-from airflow.providers.common.compat.sdk import AirflowException, BaseHook,
BaseSensorOperator
+from airflow.providers.common.compat.sdk import AirflowException, BaseHook,
BaseSensorOperator, conf
from airflow.providers.common.sql.hooks.sql import DbApiHook
+from airflow.providers.common.sql.triggers.sql import NO_ROWS_MESSAGE,
SqlSensorTrigger
+from airflow.providers.common.sql.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_2_2_PLUS
if TYPE_CHECKING:
from airflow.providers.common.compat.sdk import Context
+_first_cell = itemgetter(0)
+
+# Values that survive Airflow 2's JSON storage of trigger arguments unchanged.
+_JSON_SCALARS = (str, int, float, bool, type(None))
+
+
+def _is_first_cell(selector: Callable[[tuple[Any]], Any]) -> bool:
+ # Compare by value rather than identity: a deep copy of the sensor, or an
explicit
+ # ``selector=itemgetter(0)``, must still count as the default condition.
+ return isinstance(selector, itemgetter) and selector.__reduce__() ==
_first_cell.__reduce__()
+
+
+def _survives_json(value: Any) -> bool:
+ if isinstance(value, Mapping):
+ return all(isinstance(k, str) and _survives_json(v) for k, v in
value.items())
+ if isinstance(value, (list, tuple)):
+ return all(_survives_json(v) for v in value)
+ return isinstance(value, _JSON_SCALARS)
+
+
+class SqlSensorTriggerError(AirflowException):
+ """Raised when the deferred ``SqlSensor`` trigger reports a failure."""
+
class SqlSensor(BaseSensorOperator):
"""
@@ -55,6 +81,12 @@ class SqlSensor(BaseSensorOperator):
:param fail_on_empty: Explicitly fail on no rows returned.
:param hook_params: Extra config params to be passed to the underlying
hook.
Should match the desired hook constructor params.
+ :param deferrable: Wait in the triggerer instead of on a worker. The
triggerer only
+ checks the default condition (the first cell is truthy), so the sensor
logs a
+ warning and runs in its usual ``mode`` instead when ``success``,
``failure``,
+ ``selector``, ``silent_fail``, ``never_fail`` or
``exponential_backoff`` is set,
+ when a subclass overrides ``poke`` or ``_get_hook``, on Airflow 3
before 3.2.2, or
+ on Airflow 2 when ``parameters`` holds values JSON cannot keep, such
as dates.
"""
template_fields: Sequence[str] = ("sql", "hook_params", "parameters",
"conn_id")
@@ -69,9 +101,10 @@ class SqlSensor(BaseSensorOperator):
parameters: Mapping[str, Any] | None = None,
success: Callable[[Any], bool] | None = None,
failure: Callable[[Any], bool] | None = None,
- selector: Callable[[tuple[Any]], Any] = itemgetter(0),
+ selector: Callable[[tuple[Any]], Any] = _first_cell,
fail_on_empty: bool = False,
hook_params: Mapping[str, Any] | None = None,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
**kwargs,
) -> None:
self.conn_id = conn_id
@@ -82,6 +115,7 @@ class SqlSensor(BaseSensorOperator):
self.selector = selector
self.fail_on_empty = fail_on_empty
self.hook_params = hook_params
+ self.deferrable = deferrable
super().__init__(**kwargs)
def _get_hook(self) -> DbApiHook:
@@ -101,8 +135,7 @@ class SqlSensor(BaseSensorOperator):
records = hook.get_records(self.sql, self.parameters)
if not records:
if self.fail_on_empty:
- message = "No rows returned, raising as per fail_on_empty flag"
- raise AirflowException(message)
+ raise AirflowException(NO_ROWS_MESSAGE)
return False
condition = self.selector(records[0])
@@ -121,3 +154,48 @@ class SqlSensor(BaseSensorOperator):
message = f"self.success is present, but not callable ->
{self.success}"
raise AirflowException(message)
return bool(condition)
+
+ def execute(self, context: Context) -> Any:
+ if not self.deferrable:
+ return super().execute(context)
+ if reason := self._reason_not_to_defer():
+ self.log.warning("Cannot defer because %s; running in %s mode
instead.", reason, self.mode)
+ return super().execute(context)
+ # silent_fail and never_fail fall back above, so an error here
propagates exactly as
+ # it would from the first poke of super().execute().
+ if self.poke(context):
+ self.log.info("Success criteria met. Exiting.")
+ return None
+ self.defer(
+ trigger=SqlSensorTrigger(
+ sql=self.sql,
+ conn_id=self.conn_id,
+ poke_interval=self.poke_interval,
+ parameters=self.parameters,
+ hook_params=dict(self.hook_params) if self.hook_params is not
None else None,
+ fail_on_empty=self.fail_on_empty,
+ ),
+ method_name=self.execute_complete.__name__,
+ timeout=timedelta(seconds=self.timeout),
+ )
+
+ def _reason_not_to_defer(self) -> str | None:
+ if AIRFLOW_V_3_0_PLUS and not AIRFLOW_V_3_2_2_PLUS:
+ return "the triggerer before Airflow 3.2.2 cannot look up
connections from a thread safely"
+ if self.success is not None or self.failure is not None or not
_is_first_cell(self.selector):
+ return "success, failure and selector only run on a worker"
+ if type(self).poke is not SqlSensor.poke or type(self)._get_hook is
not SqlSensor._get_hook:
+ return f"{type(self).__name__} overrides poke or _get_hook, which
the triggerer does not run"
+ if self.silent_fail or self.never_fail or self.exponential_backoff:
+ return "silent_fail, never_fail and exponential_backoff only apply
when poking on a worker"
+ if not AIRFLOW_V_3_0_PLUS and not _survives_json(self.parameters):
+ return (
+ "Airflow 2 stores trigger arguments as JSON, which would hand
the triggerer "
+ "date, datetime or Decimal parameters as strings"
+ )
+ return None
+
+ def execute_complete(self, context: Context, event: dict[str, Any]) ->
None:
+ if event["status"] == "failure":
+ raise SqlSensorTriggerError(event["message"])
+ self.log.info("Success criteria met. Exiting.")
diff --git
a/providers/common/sql/src/airflow/providers/common/sql/sensors/sql.pyi
b/providers/common/sql/src/airflow/providers/common/sql/sensors/sql.pyi
index dc5f916ffb3..b9a4d5799aa 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/sensors/sql.pyi
+++ b/providers/common/sql/src/airflow/providers/common/sql/sensors/sql.pyi
@@ -37,9 +37,11 @@ from typing import Any
from _typeshed import Incomplete as Incomplete
-from airflow.providers.common.compat.sdk import BaseSensorOperator
+from airflow.providers.common.compat.sdk import AirflowException,
BaseSensorOperator
from airflow.sdk import Context
+class SqlSensorTriggerError(AirflowException): ...
+
class SqlSensor(BaseSensorOperator):
template_fields: Sequence[str]
template_ext: Sequence[str]
@@ -49,8 +51,10 @@ class SqlSensor(BaseSensorOperator):
parameters: Incomplete
success: Incomplete
failure: Incomplete
+ selector: Incomplete
fail_on_empty: Incomplete
hook_params: Incomplete
+ deferrable: Incomplete
def __init__(
self,
*,
@@ -59,8 +63,12 @@ class SqlSensor(BaseSensorOperator):
parameters: Mapping[str, Any] | None = None,
success: Callable[[Any], bool] | None = None,
failure: Callable[[Any], bool] | None = None,
+ selector: Callable[[tuple[Any]], Any] = ...,
fail_on_empty: bool = False,
hook_params: Mapping[str, Any] | None = None,
+ deferrable: bool = ...,
**kwargs,
) -> None: ...
def poke(self, context: Context) -> bool: ...
+ def execute(self, context: Context) -> Any: ...
+ def execute_complete(self, context: Context, event: dict[str, Any]) ->
None: ...
diff --git
a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py
b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py
index f9d8980f332..6167821279f 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py
+++ b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py
@@ -17,17 +17,22 @@
# under the License.
from __future__ import annotations
+import asyncio
from typing import TYPE_CHECKING
+from asgiref.sync import sync_to_async
+
from airflow.providers.common.compat.sdk import AirflowException, BaseHook
from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS
from airflow.providers.common.sql.hooks.sql import DbApiHook
from airflow.triggers.base import BaseTrigger, TriggerEvent
if TYPE_CHECKING:
- from collections.abc import AsyncIterator
+ from collections.abc import AsyncIterator, Iterable, Mapping
from typing import Any
+NO_ROWS_MESSAGE = "No rows returned, raising as per fail_on_empty flag"
+
class SQLExecuteQueryTrigger(BaseTrigger):
"""
@@ -36,6 +41,7 @@ class SQLExecuteQueryTrigger(BaseTrigger):
:param sql: the sql statement to be executed (str) or a list of sql
statements to execute
:param conn_id: the connection ID used to connect to the database
:param hook_params: hook parameters
+ :param parameters: The parameters to render the SQL query with (optional).
"""
def __init__(
@@ -43,12 +49,14 @@ class SQLExecuteQueryTrigger(BaseTrigger):
sql: str | list[str],
conn_id: str,
hook_params: dict | None = None,
+ parameters: Iterable | Mapping[str, Any] | None = None,
**kwargs,
):
super().__init__(**kwargs)
self.sql = sql
self.conn_id = conn_id
self.hook_params = hook_params
+ self.parameters = parameters
def serialize(self) -> tuple[str, dict[str, Any]]:
"""Serialize the SQLExecuteQueryTrigger arguments and classpath."""
@@ -58,6 +66,7 @@ class SQLExecuteQueryTrigger(BaseTrigger):
"sql": self.sql,
"conn_id": self.conn_id,
"hook_params": self.hook_params,
+ "parameters": self.parameters,
},
)
@@ -79,15 +88,18 @@ class SQLExecuteQueryTrigger(BaseTrigger):
)
return hook
- async def _get_records(self) -> Any:
- from asgiref.sync import sync_to_async
+ def _fetch_records(self, hook: DbApiHook) -> Any:
+ if self.parameters is None:
+ return hook.get_records(self.sql)
+ return hook.get_records(self.sql, self.parameters)
+ async def _get_records(self) -> Any:
hook = self.get_hook()
if AIRFLOW_V_3_2_PLUS:
# This is only supported from Airflow 3.2 or higher due to added
async support in CommsDecoder
- return await sync_to_async(hook.get_records)(self.sql)
- return hook.get_records(self.sql)
+ return await sync_to_async(self._fetch_records)(hook)
+ return self._fetch_records(hook)
async def run(self) -> AsyncIterator[TriggerEvent]:
try:
@@ -103,3 +115,66 @@ class SQLExecuteQueryTrigger(BaseTrigger):
except Exception as e:
self.log.exception("An error occurred: %s", e)
yield TriggerEvent({"status": "failure", "message": str(e)})
+
+
+class SqlSensorTrigger(SQLExecuteQueryTrigger):
+ """
+ Poll a SQL query until the first cell of the first row is truthy.
+
+ This is the default ``SqlSensor`` condition. The trigger waits
``poke_interval``
+ seconds before each query, because the sensor checks once on the worker
before
+ deferring. Each poll, connection lookup included, runs in a thread so a
slow query
+ does not block other triggers. The triggerer only serves connection
lookups from
+ threads safely on Airflow 2 and on 3.2.2 or later, so ``SqlSensor`` only
defers there.
+
+ An error during a poll ends the wait with a failure event, so the task
fails and
+ retries as it would if the same error happened in poke mode.
+
+ :param sql: The SQL query to poll.
+ :param conn_id: The connection ID used to connect to the database.
+ :param poke_interval: Seconds to wait before each query.
+ :param parameters: The parameters to render the SQL query with (optional).
+ :param hook_params: Extra config params to be passed to the underlying
hook.
+ :param fail_on_empty: Yield a failure event when the query returns no rows.
+ """
+
+ def __init__(
+ self,
+ *,
+ sql: str,
+ conn_id: str,
+ poke_interval: float,
+ parameters: Iterable | Mapping[str, Any] | None = None,
+ hook_params: dict | None = None,
+ fail_on_empty: bool = False,
+ **kwargs,
+ ):
+ super().__init__(sql=sql, conn_id=conn_id, hook_params=hook_params,
parameters=parameters, **kwargs)
+ self.poke_interval = poke_interval
+ self.fail_on_empty = fail_on_empty
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize the SqlSensorTrigger arguments and classpath."""
+ classpath, kwargs = super().serialize()
+ return classpath, {**kwargs, "poke_interval": self.poke_interval,
"fail_on_empty": self.fail_on_empty}
+
+ def _poll(self) -> Any:
+ return self._fetch_records(self.get_hook())
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ try:
+ while True:
+ await asyncio.sleep(self.poke_interval)
+ self.log.info("Poking: %s (with parameters %s)", self.sql,
self.parameters)
+ records = await sync_to_async(self._poll,
thread_sensitive=False)()
+ if not records:
+ if self.fail_on_empty:
+ yield TriggerEvent({"status": "failure", "message":
NO_ROWS_MESSAGE})
+ return
+ elif records[0][0]:
+ yield TriggerEvent({"status": "success"})
+ return
+ self.log.info("Criteria not met; checking again in %s
seconds.", self.poke_interval)
+ except Exception as e:
+ self.log.exception("An error occurred: %s", e)
+ yield TriggerEvent({"status": "failure", "message": str(e)})
diff --git
a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.pyi
b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.pyi
index 3a14c65998d..a1d53fabb13 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.pyi
+++ b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.pyi
@@ -29,16 +29,38 @@
#
"""Definition of the public interface for
airflow.providers.common.sql.triggers.sql."""
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Iterable, Mapping
from typing import Any
from airflow.providers.common.sql.hooks.sql import DbApiHook as DbApiHook
from airflow.triggers.base import BaseTrigger as BaseTrigger, TriggerEvent as
TriggerEvent
+NO_ROWS_MESSAGE: str
+
class SQLExecuteQueryTrigger(BaseTrigger):
def __init__(
- self, sql: str | list[str], conn_id: str, hook_params: dict | None =
None, **kwargs
+ self,
+ sql: str | list[str],
+ conn_id: str,
+ hook_params: dict | None = None,
+ parameters: Iterable | Mapping[str, Any] | None = None,
+ **kwargs,
) -> None: ...
def serialize(self) -> tuple[str, dict[str, Any]]: ...
def get_hook(self) -> DbApiHook: ...
async def run(self) -> AsyncIterator[TriggerEvent]: ... # type: ignore
+
+class SqlSensorTrigger(SQLExecuteQueryTrigger):
+ def __init__(
+ self,
+ *,
+ sql: str,
+ conn_id: str,
+ poke_interval: float,
+ parameters: Iterable | Mapping[str, Any] | None = None,
+ hook_params: dict | None = None,
+ fail_on_empty: bool = False,
+ **kwargs,
+ ) -> None: ...
+ def serialize(self) -> tuple[str, dict[str, Any]]: ...
+ async def run(self) -> AsyncIterator[TriggerEvent]: ... # type: ignore
diff --git
a/providers/common/sql/src/airflow/providers/common/sql/version_compat.py
b/providers/common/sql/src/airflow/providers/common/sql/version_compat.py
index a4fae50317e..81bdd0f6029 100644
--- a/providers/common/sql/src/airflow/providers/common/sql/version_compat.py
+++ b/providers/common/sql/src/airflow/providers/common/sql/version_compat.py
@@ -34,8 +34,10 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)
AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0)
+AIRFLOW_V_3_2_2_PLUS: bool = get_base_airflow_version_tuple() >= (3, 2, 2)
__all__ = [
"AIRFLOW_V_3_0_PLUS",
"AIRFLOW_V_3_1_PLUS",
+ "AIRFLOW_V_3_2_2_PLUS",
]
diff --git a/providers/common/sql/tests/system/common/sql/example_sql_sensor.py
b/providers/common/sql/tests/system/common/sql/example_sql_sensor.py
new file mode 100644
index 00000000000..c997299f32e
--- /dev/null
+++ b/providers/common/sql/tests/system/common/sql/example_sql_sensor.py
@@ -0,0 +1,50 @@
+#
+# 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
+
+from airflow import DAG
+from airflow.providers.common.sql.sensors.sql import SqlSensor
+
+with DAG(dag_id="example_sql_sensor", schedule=None, tags=["example"]) as dag:
+ # [START howto_sensor_sql]
+ wait_for_orders = SqlSensor(
+ task_id="wait_for_orders",
+ conn_id="sales_db",
+ sql="SELECT COUNT(*) FROM orders WHERE order_date = CURRENT_DATE",
+ poke_interval=300,
+ mode="reschedule",
+ )
+ # [END howto_sensor_sql]
+
+ # [START howto_sensor_sql_deferrable]
+ wait_for_orders_deferrable = SqlSensor(
+ task_id="wait_for_orders_deferrable",
+ conn_id="sales_db",
+ sql="SELECT COUNT(*) >= %(min_orders)s FROM orders WHERE order_date =
CURRENT_DATE",
+ parameters={"min_orders": 100},
+ poke_interval=300,
+ timeout=6 * 60 * 60,
+ deferrable=True,
+ )
+ # [END howto_sensor_sql_deferrable]
+
+
+from tests_common.test_utils.system_tests import get_test_run
+
+# Needed to run the example DAG with pytest (see:
contributing-docs/testing/system_tests.rst)
+test_run = get_test_run(dag)
diff --git a/providers/common/sql/tests/unit/common/sql/sensors/test_sql.py
b/providers/common/sql/tests/unit/common/sql/sensors/test_sql.py
index b5b23f71906..5253780a4f3 100644
--- a/providers/common/sql/tests/unit/common/sql/sensors/test_sql.py
+++ b/providers/common/sql/tests/unit/common/sql/sensors/test_sql.py
@@ -17,14 +17,30 @@
# under the License.
from __future__ import annotations
+import copy
+from datetime import date, datetime, timedelta
+from decimal import Decimal
+from operator import itemgetter
from unittest import mock
import pytest
from airflow.models.dag import DAG
-from airflow.providers.common.compat.sdk import AirflowException, timezone
+from airflow.providers.common.compat.sdk import (
+ AirflowException,
+ AirflowSensorTimeout,
+ AirflowSkipException,
+ TaskDeferred,
+ timezone,
+)
from airflow.providers.common.sql.hooks.sql import DbApiHook
-from airflow.providers.common.sql.sensors.sql import SqlSensor
+from airflow.providers.common.sql.sensors.sql import SqlSensor,
SqlSensorTriggerError
+from airflow.providers.common.sql.triggers.sql import SqlSensorTrigger
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_2_2_PLUS
+
+if AIRFLOW_V_3_2_2_PLUS:
+ from airflow.sdk.bases.operator import TriggerFailureReason
DEFAULT_DATE = timezone.datetime(2015, 1, 1)
TEST_DAG_ID = "unit_test_sql_dag"
@@ -308,3 +324,183 @@ class TestSqlSensor:
op.execute(context=mock.MagicMock())
mock_get_records.assert_called_once_with("SELECT %(something)s",
{"something": "1970-01-01"})
+
+
[email protected]
+def mock_get_records():
+ with mock.patch("airflow.providers.common.sql.sensors.sql.BaseHook") as
mock_base_hook:
+ hook = mock.MagicMock(spec=DbApiHook)
+ mock_base_hook.get_connection.return_value.get_hook.return_value = hook
+ yield hook.get_records
+
+
+def make_sensor(deferrable: bool = True, sensor_class: type[SqlSensor] =
SqlSensor, **kwargs) -> SqlSensor:
+ defaults = {
+ "task_id": "sql_sensor_deferrable",
+ "conn_id": "postgres_default",
+ "sql": "SELECT count(*) FROM orders WHERE ds = %(ds)s",
+ "parameters": {"ds": "2026-01-01"},
+ "hook_params": {"log_sql": False},
+ "poke_interval": 0,
+ "timeout": 600,
+ }
+ return sensor_class(**{**defaults, **kwargs}, deferrable=deferrable)
+
+
+class PokeOverridingSensor(SqlSensor):
+ def poke(self, context):
+ return super().poke(context)
+
+
+class HookOverridingSensor(SqlSensor):
+ def _get_hook(self):
+ return super()._get_hook()
+
+
+class TestSqlSensorDeferrable:
+ @pytest.fixture(autouse=True)
+ def _airflow_version_that_can_defer(self):
+ # The provider's compat jobs run these tests on Airflow 3 releases
before 3.2.2, where the
+ # sensor never defers; pin a version that can, and let the version
tests override it.
+ with (
+
mock.patch("airflow.providers.common.sql.sensors.sql.AIRFLOW_V_3_0_PLUS", True),
+
mock.patch("airflow.providers.common.sql.sensors.sql.AIRFLOW_V_3_2_2_PLUS",
True),
+ ):
+ yield
+
+ def test_returns_without_deferring_when_first_poke_succeeds(self,
mock_get_records):
+ mock_get_records.return_value = [(5,)]
+
+ assert make_sensor().execute(context={}) is None
+ assert mock_get_records.call_count == 1
+
+ def test_defers_to_trigger_when_first_poke_fails(self, mock_get_records):
+ mock_get_records.side_effect = [[(0,)], [(1,)]]
+
+ with pytest.raises(TaskDeferred) as deferred:
+ make_sensor(fail_on_empty=True).execute(context={})
+
+ assert isinstance(deferred.value.trigger, SqlSensorTrigger)
+ assert deferred.value.trigger.serialize()[1] == {
+ "sql": "SELECT count(*) FROM orders WHERE ds = %(ds)s",
+ "conn_id": "postgres_default",
+ "hook_params": {"log_sql": False},
+ "parameters": {"ds": "2026-01-01"},
+ "poke_interval": 0,
+ "fail_on_empty": True,
+ }
+ assert deferred.value.method_name == "execute_complete"
+ assert deferred.value.timeout == timedelta(seconds=600)
+
+ def test_does_not_warn_when_not_deferrable(self, mock_get_records, caplog):
+ mock_get_records.side_effect = [[(0,)], [(1,)]]
+
+ assert make_sensor(deferrable=False).execute(context={}) is None
+ assert mock_get_records.call_count == 2
+ assert "Cannot defer" not in caplog.text
+
+ @pytest.mark.parametrize(
+ ("sensor_kwargs", "reason"),
+ [
+ pytest.param(
+ {"success": lambda value: value > 10}, "success, failure and
selector", id="success"
+ ),
+ pytest.param({"failure": lambda value: value < 0}, "success,
failure and selector", id="failure"),
+ pytest.param({"selector": itemgetter(1)}, "success, failure and
selector", id="selector"),
+ pytest.param({"silent_fail": True}, "silent_fail, never_fail",
id="silent_fail"),
+ pytest.param({"never_fail": True}, "silent_fail, never_fail",
id="never_fail"),
+ pytest.param({"exponential_backoff": True}, "silent_fail,
never_fail", id="exponential_backoff"),
+ pytest.param(
+ {"sensor_class": PokeOverridingSensor},
+ "PokeOverridingSensor overrides poke or _get_hook",
+ id="poke",
+ ),
+ pytest.param(
+ {"sensor_class": HookOverridingSensor},
+ "HookOverridingSensor overrides poke or _get_hook",
+ id="get_hook",
+ ),
+ ],
+ )
+ def test_pokes_on_the_worker_when_the_triggerer_cannot_check(
+ self, mock_get_records, caplog, sensor_kwargs, reason
+ ):
+ mock_get_records.side_effect = [[(0, 0)], [(20, 20)]]
+
+ assert make_sensor(**sensor_kwargs).execute(context={}) is None
+ assert mock_get_records.call_count == 2
+ assert f"Cannot defer because {reason}" in caplog.text
+
+ @pytest.mark.parametrize(
+ "build_sensor",
+ [
+ pytest.param(lambda: make_sensor(selector=itemgetter(0)),
id="explicit-first-cell-selector"),
+ pytest.param(lambda: copy.deepcopy(make_sensor()),
id="deep-copied-sensor"),
+ ],
+ )
+ def test_default_condition_defers_however_the_selector_was_built(self,
mock_get_records, build_sensor):
+ # The second result meets the condition, so a wrong fallback to poking
returns instead of hanging.
+ mock_get_records.side_effect = [[(0,)], [(1,)]]
+
+ with pytest.raises(TaskDeferred):
+ build_sensor().execute(context={})
+
+
@mock.patch("airflow.providers.common.sql.sensors.sql.AIRFLOW_V_3_2_2_PLUS",
False)
+ @mock.patch("airflow.providers.common.sql.sensors.sql.AIRFLOW_V_3_0_PLUS",
True)
+ def test_pokes_on_the_worker_on_airflow_3_before_3_2_2(self,
mock_get_records, caplog):
+ mock_get_records.side_effect = [[(0,)], [(1,)]]
+
+ assert make_sensor().execute(context={}) is None
+ assert mock_get_records.call_count == 2
+ assert "Cannot defer because the triggerer before Airflow 3.2.2" in
caplog.text
+
+
@mock.patch("airflow.providers.common.sql.sensors.sql.AIRFLOW_V_3_2_2_PLUS",
False)
+ @mock.patch("airflow.providers.common.sql.sensors.sql.AIRFLOW_V_3_0_PLUS",
False)
+ def test_defers_on_airflow_2(self, mock_get_records):
+ mock_get_records.side_effect = [[(0,)], [(1,)]]
+
+ with pytest.raises(TaskDeferred):
+ make_sensor(parameters=("2026-01-01", 3, 1.5, True,
None)).execute(context={})
+
+
@mock.patch("airflow.providers.common.sql.sensors.sql.AIRFLOW_V_3_2_2_PLUS",
False)
+ @mock.patch("airflow.providers.common.sql.sensors.sql.AIRFLOW_V_3_0_PLUS",
False)
+ @pytest.mark.parametrize(
+ "value",
+ [date(2026, 1, 1), datetime(2026, 1, 1, 12), Decimal("1.5")],
+ ids=["date", "datetime", "decimal"],
+ )
+ def
test_pokes_on_the_worker_on_airflow_2_when_parameters_would_change_in_json(
+ self, mock_get_records, caplog, value
+ ):
+ mock_get_records.side_effect = [[(0,)], [(1,)]]
+
+ assert make_sensor(parameters={"ds": value}).execute(context={}) is
None
+ assert mock_get_records.call_count == 2
+ assert "Cannot defer because Airflow 2 stores trigger arguments as
JSON" in caplog.text
+
+ def test_success_event_completes(self):
+ assert make_sensor().execute_complete(context={}, event={"status":
"success"}) is None
+
+ def test_failure_event_raises(self):
+ with pytest.raises(SqlSensorTriggerError, match="connection refused"):
+ make_sensor().resume_execution(
+ "execute_complete", {"event": {"status": "failure", "message":
"connection refused"}}, {}
+ )
+
+ def test_failure_event_skips_with_soft_fail(self):
+ with pytest.raises(AirflowSkipException):
+ make_sensor(soft_fail=True).resume_execution(
+ "execute_complete", {"event": {"status": "failure", "message":
"connection refused"}}, {}
+ )
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_2_PLUS, reason="SqlSensor only
defers on Airflow 3 from 3.2.2")
+ @pytest.mark.parametrize(
+ ("soft_fail", "expected"), [(False, AirflowSensorTimeout), (True,
AirflowSkipException)]
+ )
+ def test_trigger_timeout_fails_like_a_sensor_timeout(self, soft_fail,
expected):
+ with pytest.raises(expected) as excinfo:
+ make_sensor(soft_fail=soft_fail).resume_execution(
+ "__fail__", {"error": TriggerFailureReason.TRIGGER_TIMEOUT}, {}
+ )
+
+ assert excinfo.type is expected
diff --git a/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py
b/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py
index e09fba590be..c73ff32db01 100644
--- a/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py
+++ b/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py
@@ -18,9 +18,12 @@ from __future__ import annotations
from unittest import mock
+import pytest
+from asgiref.sync import sync_to_async
+
from airflow.models.connection import Connection
from airflow.providers.common.sql.hooks.sql import DbApiHook
-from airflow.providers.common.sql.triggers.sql import SQLExecuteQueryTrigger
+from airflow.providers.common.sql.triggers.sql import SQLExecuteQueryTrigger,
SqlSensorTrigger
from airflow.triggers.base import TriggerEvent
try:
@@ -35,6 +38,19 @@ except ImportError:
from tests_common.test_utils.operators.run_deferrable import run_trigger
[email protected]
+def mock_get_connection():
+ with mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection") as get_connection:
+ yield get_connection
+
+
[email protected]
+def db_hook(mock_get_connection):
+ hook = mock.MagicMock(spec=DbApiHook)
+ mock_get_connection.return_value.get_hook.return_value = hook
+ return hook
+
+
class TestSQLExecuteQueryTrigger:
@mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection")
def test_run(self, mock_get_connection):
@@ -52,3 +68,79 @@ class TestSQLExecuteQueryTrigger:
assert isinstance(actual[0], TriggerEvent)
assert actual[0].payload["status"] == "success"
assert actual[0].payload["results"] == data
+
+ def test_run_passes_parameters(self, db_hook):
+ db_hook.get_records.return_value = [(1,)]
+
+ run_trigger(SQLExecuteQueryTrigger(sql="SELECT %(id)s",
conn_id="test_conn_id", parameters={"id": 1}))
+
+ db_hook.get_records.assert_called_once_with("SELECT %(id)s", {"id": 1})
+
+
+def make_sensor_trigger(**kwargs) -> SqlSensorTrigger:
+ return SqlSensorTrigger(
+ sql="SELECT count(*) FROM orders WHERE ds = %(ds)s",
+ conn_id="test_conn_id",
+ poke_interval=30,
+ parameters={"ds": "2026-01-01"},
+ **kwargs,
+ )
+
+
+class TestSqlSensorTrigger:
+ def test_serialize_round_trip(self):
+ trigger = make_sensor_trigger(hook_params={"schema": "s"},
fail_on_empty=True)
+
+ classpath, kwargs = trigger.serialize()
+
+ assert classpath ==
"airflow.providers.common.sql.triggers.sql.SqlSensorTrigger"
+ assert kwargs == {
+ "sql": "SELECT count(*) FROM orders WHERE ds = %(ds)s",
+ "conn_id": "test_conn_id",
+ "hook_params": {"schema": "s"},
+ "parameters": {"ds": "2026-01-01"},
+ "poke_interval": 30,
+ "fail_on_empty": True,
+ }
+ assert SqlSensorTrigger(**kwargs).serialize() == (classpath, kwargs)
+
+ def test_polls_in_a_thread_until_first_cell_is_truthy(self,
mock_get_connection, db_hook):
+ db_hook.get_records.side_effect = [[], [(0,)], [(None,)], [(7,), (8,)]]
+ module = "airflow.providers.common.sql.triggers.sql"
+ trigger = make_sensor_trigger()
+
+ with (
+ mock.patch(f"{module}.sync_to_async", wraps=sync_to_async) as
mock_sync_to_async,
+ mock.patch(f"{module}.asyncio.sleep") as sleep,
+ ):
+ events = run_trigger(trigger)
+
+ assert [event.payload for event in events] == [{"status": "success"}]
+ assert (
+ db_hook.get_records.call_args_list
+ == [mock.call("SELECT count(*) FROM orders WHERE ds = %(ds)s",
{"ds": "2026-01-01"})] * 4
+ )
+ assert sleep.await_args_list == [mock.call(30)] * 4
+ assert mock_get_connection.call_count == 4
+ assert mock_sync_to_async.call_args_list == [mock.call(trigger._poll,
thread_sensitive=False)] * 4
+
+ @pytest.mark.parametrize(
+ ("get_records", "message"),
+ [
+ pytest.param(
+ {"side_effect": [[], [(1,)]]},
+ "No rows returned, raising as per fail_on_empty flag",
+ id="empty",
+ ),
+ pytest.param(
+ {"side_effect": RuntimeError("connection refused")},
"connection refused", id="error"
+ ),
+ ],
+ )
+ def test_yields_failure(self, db_hook, get_records, message):
+ db_hook.get_records.configure_mock(**get_records)
+
+ with
mock.patch("airflow.providers.common.sql.triggers.sql.asyncio.sleep"):
+ events = run_trigger(make_sensor_trigger(fail_on_empty=True))
+
+ assert [event.payload for event in events] == [{"status": "failure",
"message": message}]
