kaxil commented on code in PR #76270:
URL: https://github.com/apache/airflow/pull/76270#discussion_r4446277441
##########
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:
Yes, intended. `silent_fail` and `never_fail` send the sensor back to poke
mode before this, and a DB error raises from the first poke of
`super().execute()` as well, so retries behave the same. Added a comment. The
one difference: if `execution_timeout` fires during that first query with
`soft_fail` set, the task fails instead of skipping. I don't think that's worth
more code.
##########
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:
Good catch, confirmed: a `copy.deepcopy` of the sensor and an explicit
`itemgetter(0)` both made it fall back to poking. It now compares by value (an
`itemgetter` with the same arguments), with tests for both. I kept the default
as it is rather than a `None` sentinel, so `selector` keeps its type. Mapped
tasks build the operator through the constructor (`MappedOperator.unmap`), so
they get the default.
##########
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:
Agreed, it now also falls back when a subclass overrides `_get_hook`.
`hook_params` is passed to the trigger, which builds the hook the same way
`_get_hook` does.
##########
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 no-rows message is shared (`NO_ROWS_MESSAGE`), and the default condition
is a single truthiness check on the first cell in both places. A shared helper
would have to live in the trigger module for the sensor to import it, and I'd
rather not couple the two over one expression. Happy to if you feel strongly.
--
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]
