kaxil commented on code in PR #76270:
URL: https://github.com/apache/airflow/pull/76270#discussion_r4446275605
##########
providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py:
##########
@@ -80,14 +89,12 @@ def get_hook(self) -> DbApiHook:
return hook
async def _get_records(self) -> Any:
- from asgiref.sync import sync_to_async
-
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(hook.get_records)(self.sql,
parameters=self.parameters)
Review Comment:
Good catch, changed in 02a760b. With no `parameters` the trigger calls
`get_records(sql)` exactly as before, so `GenericTransfer`'s call is unchanged
(I reverted the mock change in `test_generic_transfer.py`, and it passes as it
was). With parameters it passes them positionally, the same call `poke()` makes.
##########
providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py:
##########
@@ -58,6 +66,7 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
"sql": self.sql,
"conn_id": self.conn_id,
"hook_params": self.hook_params,
+ "parameters": self.parameters,
Review Comment:
Checked both serializers. On Airflow 3 (`airflow.sdk.serde`) a tuple,
`date`, `datetime` and `Decimal` all round-trip unchanged, and a live run with
`parameters=(3, date(2026, 11, 2), Decimal("1.5"))` showed the triggerer
running with the same tuple. On Airflow 2 they don't: `date` and `Decimal` come
back as strings and a naive `datetime` picks up UTC. So on Airflow 2 the sensor
now stays on the worker, with a warning, when `parameters` holds anything other
than strings, numbers, booleans or `None`. Tests cover both.
##########
providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py:
##########
@@ -103,3 +110,63 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
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.
+
+ :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.get_hook().get_records(self.sql,
parameters=self.parameters)
Review Comment:
Kept this on purpose. Poke and reschedule mode also build the hook on every
poke, and rebuilding it means a long wait picks up rotated credentials (IAM
tokens, Vault leases) instead of failing hours in. At `poke_interval=60` it's
one lookup a minute per sensor.
##########
providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py:
##########
@@ -103,3 +110,63 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
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.
+
+ :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.get_hook().get_records(self.sql,
parameters=self.parameters)
+
+ 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:
Review Comment:
That's intended, and it matches poke mode: an error fails the attempt and
the task's `retries` apply (`SqlSensorTriggerError` is a plain
`AirflowException`, so it retries). Added that to the trigger docstring.
##########
providers/common/sql/src/airflow/providers/common/sql/sensors/sql.py:
##########
@@ -121,3 +136,41 @@ def poke(self, context: Context) -> bool:
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)
+ 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),
Review Comment:
This is already handled by `BaseSensorOperator.resume_execution`. On 3.2.2+,
the only Airflow 3 versions where this defers, a trigger timeout becomes
`AirflowSensorTimeout`, which the task runner treats as terminal (no retries),
and `soft_fail` turns it into a skip. Added
`test_trigger_timeout_fails_like_a_sensor_timeout` for both. Airflow 2.11 is
the exception: its core raises a retryable `TaskDeferralError` on a trigger
timeout, which is in the PR's gotchas.
##########
providers/common/sql/src/airflow/providers/common/sql/sensors/sql.py:
##########
@@ -121,3 +136,41 @@ def poke(self, context: Context) -> bool:
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)
+ 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
self.selector is not _first_cell:
+ return "success, failure and selector only run on a worker"
+ if type(self).poke is not SqlSensor.poke:
+ return f"{type(self).__name__} overrides poke"
+ 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"
+ return None
+
+ def execute_complete(self, context: Context, event: dict[str, Any]) ->
None:
+ if event["status"] == "failure":
+ raise SqlSensorTriggerError(event["message"])
Review Comment:
This one already skips: `resume_execution` catches any `AirflowException`
(`SqlSensorTriggerError` subclasses it) and raises `AirflowSkipException` under
`soft_fail`, see `test_failure_event_skips_with_soft_fail`. It's poke mode that
doesn't skip here, since its `soft_fail` only covers timeouts. On `.get()`: the
event only ever comes from `SqlSensorTrigger`, so I'd rather not guard against
a shape it never sends.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
