pankajastro commented on code in PR #76270:
URL: https://github.com/apache/airflow/pull/76270#discussion_r4444280355
##########
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:
`_get_records` now passes `parameters=` as a keyword on every call.
`SQLExecuteQueryOperator` and `GenericTransfer` deferred paths previously
called `get_records(sql)`, and the `test_generic_transfer.py` mock had to
change to accept `parameters`. Could a third-party `DbApiHook` whose
`get_records` override doesn't accept `parameters` now raise `TypeError` on
those paths? Could we pass `parameters` only when it is not None so existing
callers keep the same call shape?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
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:
`parameters` is now serialized into trigger kwargs. A tuple such as `(3,)`
would come back as a list, and a `date` or `Decimal` may fail to serialize at
`defer()`. Is there a test with tuple and date parameters that checks the
trigger runs the same query as the worker's first `poke`? If not, could we add
one, or validate `parameters` before deferring and raise a clear error?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
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:
In poke and reschedule mode a timeout raises `AirflowSensorTimeout`, which
skips retries and honors `soft_fail`. Here the timeout is a trigger timeout.
With `retries=3`, does a deferred timeout retry and restart the full timeout
each time? With `soft_fail=True`, does it fail instead of skip? Can
`execute_complete` map the timeout to the worker behavior? Could we add tests
for `soft_fail` and `retries` with `deferrable=True`?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
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):
Review Comment:
`execute` calls `self.poke(context)` once before deferring. If that call
hits a transient DB error, does it bypass the base sensor's error handling and
retry accounting? Is that intended, and if so could we note it in a comment?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
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:
`get_hook()` is called on every poll, so with `poke_interval=60` and the
default 7-day `timeout` that is about 10k connection lookups and
connect/disconnect cycles. Can the hook be built once outside the loop?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
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:
A single transient DB or network error here ends the whole wait with a
failure event, and a bug in `_poll` is reported the same way. On a worker, the
task retries. Should transient errors get a limited retry in the loop, or is
failing the task and relying on task `retries` the intended behavior? If so,
could we say so in the docstring?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
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)
Review Comment:
The default condition is `bool(condition)` of the selected cell here,
`records[0][0]` in `SqlSensorTrigger.run`, and the empty-row handling and
`NO_ROWS_MESSAGE` are also in both places. Could the default truthiness check
and empty-row handling live in one helper so the worker's first check and the
triggerer's later checks cannot diverge?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
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:
Review Comment:
`_reason_not_to_defer` guards against a `poke` override, but not an
overridden `_get_hook` or custom `hook_params` handling. The trigger builds its
hook through `BaseHook.get_connection` and `get_hook`, so it bypasses those
overrides. Should a `_get_hook` override also fall back to the worker, or is it
out of scope?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
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:
Review Comment:
`self.selector is not _first_cell` is an identity check, and the default in
the signature changed from `itemgetter(0)` to `_first_cell`. If a user passes
`itemgetter(0)` explicitly, or a mapped or copied sensor ends up with a
different object, would it refuse to defer? Is there a test for mapped sensors
(`expand`) and an explicit `itemgetter(0)`? Would a `None` sentinel default be
safer?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
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:
Every failure event raises `SqlSensorTriggerError`. With `soft_fail=True`, a
`fail_on_empty` or DB error from the trigger would fail the task, where the
worker path would skip it. Can this go through the same `soft_fail` handling?
Also, `event["status"]` is read directly. Should it use `.get()` so a malformed
event gives a clear error instead of `KeyError`?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
--
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]
