moomindani commented on code in PR #71752:
URL: https://github.com/apache/airflow/pull/71752#discussion_r3813498696
##########
providers/databricks/src/airflow/providers/databricks/triggers/databricks.py:
##########
@@ -336,3 +336,128 @@ async def run(self):
}
)
return
+
+
+class DatabricksWarehouseStateTrigger(BaseTrigger):
+ """
+ Poll a Databricks SQL warehouse until it reaches a target lifecycle state.
+
+ :param warehouse_id: ID of the Databricks SQL warehouse.
+ :param target_state: Lifecycle state to wait for (``RUNNING`` or
``STOPPED``).
+ :param databricks_conn_id: Reference to the :ref:`Databricks connection
<howto/connection:databricks>`.
+ :param timeout: Maximum number of seconds to wait after the trigger starts
polling.
+ The deadline uses ``time.monotonic()`` locally so it survives trigger
serialization
+ without depending on wall-clock time.
+ :param polling_period_seconds: Controls the rate of the poll for the
warehouse state.
+ By default, the trigger will poll every 30 seconds.
+ :param retry_limit: The number of times to retry the connection in case of
service outages.
+ :param retry_delay: Minimum wait in seconds between retryable attempts
when using the
+ default retry strategy. The wait uses exponential backoff (doubling
after each
+ failure, capped at ``2 ** retry_limit`` seconds). May be a floating
point number.
+ :param retry_args: An optional dictionary with arguments passed to
``tenacity.Retrying`` class.
+ :param caller: The name of the operator that is calling the hook.
+ """
+
+ def __init__(
+ self,
+ warehouse_id: str,
+ target_state: str,
+ databricks_conn_id: str,
+ timeout: float,
+ polling_period_seconds: int = 30,
+ retry_limit: int = 3,
+ retry_delay: int = 10,
+ retry_args: dict[Any, Any] | None = None,
+ caller: str = "DatabricksWarehouseStateTrigger",
+ ) -> None:
+ super().__init__()
+ # Trigger kwargs cross Airflow's serialization boundary, so fail
before storing invalid
+ # trigger state or surfacing a generic serializer error without
Databricks-specific guidance.
+ validate_deferrable_databricks_retry_args(retry_args, owner=caller)
+ self.warehouse_id = warehouse_id
+ self.target_state = target_state
+ self.databricks_conn_id = databricks_conn_id
+ self.timeout = timeout
+ self.polling_period_seconds = polling_period_seconds
+ self.retry_limit = retry_limit
+ self.retry_delay = retry_delay
+ self.retry_args = retry_args
+ self.caller = caller
+ self.hook = DatabricksHook(
+ databricks_conn_id,
+ retry_limit=self.retry_limit,
+ retry_delay=self.retry_delay,
+ retry_args=retry_args,
+ caller=caller,
+ )
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ return (
+
"airflow.providers.databricks.triggers.databricks.DatabricksWarehouseStateTrigger",
+ {
+ "warehouse_id": self.warehouse_id,
+ "target_state": self.target_state,
+ "databricks_conn_id": self.databricks_conn_id,
+ "timeout": self.timeout,
+ "polling_period_seconds": self.polling_period_seconds,
+ "retry_limit": self.retry_limit,
+ "retry_delay": self.retry_delay,
+ "retry_args": self.retry_args,
+ "caller": self.caller,
+ },
+ )
+
+ async def on_kill(self) -> None:
+ # Warehouses have no cancel-start/stop API. Clearing a deferred start
must not stop the
+ # warehouse — the Dag author may still want it running after the wait
is abandoned.
+ self.log.info(
+ "Databricks SQL warehouse %s wait cancelled; leaving warehouse
state unchanged.",
+ self.warehouse_id,
+ )
+
+ def _build_trigger_event(
+ self, *, status: str, last_state: str, state: WarehouseState | None =
None
+ ) -> TriggerEvent:
+ payload: dict[str, Any] = {
+ "status": status,
+ "warehouse_id": self.warehouse_id,
+ "target_state": self.target_state,
+ "last_state": last_state,
+ }
+ if state is not None:
+ payload["state"] = state.to_json()
+ return TriggerEvent(payload)
+
+ async def run(self):
+ async with self.hook:
+ deadline = time.monotonic() + self.timeout
Review Comment:
The deadline is computed here, inside `run()`, and `serialize()` carries
`timeout` as a duration — so after a triggerer restart, deploy, or HA rebalance
the wait starts a fresh full `timeout` window, and a warehouse that never
reaches its target can stay deferred well past the configured value. The
synchronous path fails deterministically instead.
`DatabricksSQLStatementExecutionTrigger` in this same file takes an absolute
`end_time` (serialized at line 281, consumed at line 299) that the operator
computes, precisely so the deadline survives serialization. Following that
pattern here would make the two consistent; if the per-run semantics are
deliberate instead, they are worth stating in the docstring and the how-to.
Related: the docstring above says the monotonic deadline "survives trigger
serialization", and the PR description repeats it. It is the other way round —
a monotonic value cannot be serialized, which is why this line recomputes it.
##########
providers/databricks/src/airflow/providers/databricks/triggers/databricks.py:
##########
@@ -336,3 +336,128 @@ async def run(self):
}
)
return
+
+
+class DatabricksWarehouseStateTrigger(BaseTrigger):
+ """
+ Poll a Databricks SQL warehouse until it reaches a target lifecycle state.
+
+ :param warehouse_id: ID of the Databricks SQL warehouse.
+ :param target_state: Lifecycle state to wait for (``RUNNING`` or
``STOPPED``).
+ :param databricks_conn_id: Reference to the :ref:`Databricks connection
<howto/connection:databricks>`.
+ :param timeout: Maximum number of seconds to wait after the trigger starts
polling.
+ The deadline uses ``time.monotonic()`` locally so it survives trigger
serialization
+ without depending on wall-clock time.
+ :param polling_period_seconds: Controls the rate of the poll for the
warehouse state.
+ By default, the trigger will poll every 30 seconds.
+ :param retry_limit: The number of times to retry the connection in case of
service outages.
+ :param retry_delay: Minimum wait in seconds between retryable attempts
when using the
+ default retry strategy. The wait uses exponential backoff (doubling
after each
+ failure, capped at ``2 ** retry_limit`` seconds). May be a floating
point number.
+ :param retry_args: An optional dictionary with arguments passed to
``tenacity.Retrying`` class.
+ :param caller: The name of the operator that is calling the hook.
+ """
+
+ def __init__(
+ self,
+ warehouse_id: str,
+ target_state: str,
+ databricks_conn_id: str,
+ timeout: float,
+ polling_period_seconds: int = 30,
+ retry_limit: int = 3,
+ retry_delay: int = 10,
+ retry_args: dict[Any, Any] | None = None,
+ caller: str = "DatabricksWarehouseStateTrigger",
+ ) -> None:
+ super().__init__()
+ # Trigger kwargs cross Airflow's serialization boundary, so fail
before storing invalid
+ # trigger state or surfacing a generic serializer error without
Databricks-specific guidance.
+ validate_deferrable_databricks_retry_args(retry_args, owner=caller)
+ self.warehouse_id = warehouse_id
+ self.target_state = target_state
+ self.databricks_conn_id = databricks_conn_id
+ self.timeout = timeout
+ self.polling_period_seconds = polling_period_seconds
+ self.retry_limit = retry_limit
+ self.retry_delay = retry_delay
+ self.retry_args = retry_args
+ self.caller = caller
+ self.hook = DatabricksHook(
+ databricks_conn_id,
+ retry_limit=self.retry_limit,
+ retry_delay=self.retry_delay,
+ retry_args=retry_args,
+ caller=caller,
+ )
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ return (
+
"airflow.providers.databricks.triggers.databricks.DatabricksWarehouseStateTrigger",
+ {
+ "warehouse_id": self.warehouse_id,
+ "target_state": self.target_state,
+ "databricks_conn_id": self.databricks_conn_id,
+ "timeout": self.timeout,
+ "polling_period_seconds": self.polling_period_seconds,
+ "retry_limit": self.retry_limit,
+ "retry_delay": self.retry_delay,
+ "retry_args": self.retry_args,
+ "caller": self.caller,
+ },
+ )
+
+ async def on_kill(self) -> None:
Review Comment:
`BaseTrigger.on_kill` is already a documented no-op, so this override only
adds a log line. I deleted it entirely and
`test_on_kill_leaves_warehouse_unchanged` still passed, which means the test
does not pin the behaviour the description credits it with.
Either drop both the override and the test and keep the reasoning as a
comment, or keep the override and stop describing it as tested. Worth noting
the base-class method only exists on newer cores while the provider supports
`apache-airflow>=2.11.0`.
##########
providers/databricks/docs/changelog.rst:
##########
@@ -29,6 +29,11 @@ Changelog
7.18.1
......
+Features
+~~~~~~~~
+
+* ``Add deferrable mode to Databricks SQL warehouse start and stop operators``
Review Comment:
This block sits under the `7.18.1` header, which the release manager created
for the 2026-08-06 wave — a released version, so its changelog would advertise
a feature that release does not contain, and it will collide with the entry
regenerated from `git log` for the next version.
`providers/AGENTS.md` and the `NOTE TO CONTRIBUTORS` block at the top of
this file both say routine feature entries are collected automatically and
direct edits are reserved for breaking/important behaviour notes. Simplest fix
is to drop all five lines; the commit message already carries the note.
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -799,6 +806,15 @@ def get_warehouse_state(self, warehouse_id: str) ->
WarehouseState:
"""
return WarehouseState(self.get_warehouse(warehouse_id)["state"])
+ async def a_get_warehouse_state(self, warehouse_id: str) -> WarehouseState:
Review Comment:
This inlines the endpoint tuple, while the synchronous `get_warehouse_state`
goes through `get_warehouse()`. They are identical today — same `2.0` constant,
same `WarehouseState(response["state"])` parsing — but that is now two copies,
so a later change to `get_warehouse` (query params, error mapping, an API
version bump) would silently miss the async path. An `a_get_warehouse` mirror,
or a small shared endpoint builder, keeps them in step.
##########
providers/databricks/src/airflow/providers/databricks/triggers/databricks.py:
##########
@@ -336,3 +336,128 @@ async def run(self):
}
)
return
+
+
+class DatabricksWarehouseStateTrigger(BaseTrigger):
+ """
+ Poll a Databricks SQL warehouse until it reaches a target lifecycle state.
+
+ :param warehouse_id: ID of the Databricks SQL warehouse.
+ :param target_state: Lifecycle state to wait for (``RUNNING`` or
``STOPPED``).
+ :param databricks_conn_id: Reference to the :ref:`Databricks connection
<howto/connection:databricks>`.
+ :param timeout: Maximum number of seconds to wait after the trigger starts
polling.
+ The deadline uses ``time.monotonic()`` locally so it survives trigger
serialization
+ without depending on wall-clock time.
+ :param polling_period_seconds: Controls the rate of the poll for the
warehouse state.
+ By default, the trigger will poll every 30 seconds.
+ :param retry_limit: The number of times to retry the connection in case of
service outages.
+ :param retry_delay: Minimum wait in seconds between retryable attempts
when using the
+ default retry strategy. The wait uses exponential backoff (doubling
after each
+ failure, capped at ``2 ** retry_limit`` seconds). May be a floating
point number.
+ :param retry_args: An optional dictionary with arguments passed to
``tenacity.Retrying`` class.
+ :param caller: The name of the operator that is calling the hook.
+ """
+
+ def __init__(
+ self,
+ warehouse_id: str,
+ target_state: str,
+ databricks_conn_id: str,
+ timeout: float,
+ polling_period_seconds: int = 30,
+ retry_limit: int = 3,
+ retry_delay: int = 10,
+ retry_args: dict[Any, Any] | None = None,
+ caller: str = "DatabricksWarehouseStateTrigger",
+ ) -> None:
+ super().__init__()
+ # Trigger kwargs cross Airflow's serialization boundary, so fail
before storing invalid
+ # trigger state or surfacing a generic serializer error without
Databricks-specific guidance.
+ validate_deferrable_databricks_retry_args(retry_args, owner=caller)
+ self.warehouse_id = warehouse_id
+ self.target_state = target_state
+ self.databricks_conn_id = databricks_conn_id
+ self.timeout = timeout
+ self.polling_period_seconds = polling_period_seconds
+ self.retry_limit = retry_limit
+ self.retry_delay = retry_delay
+ self.retry_args = retry_args
+ self.caller = caller
+ self.hook = DatabricksHook(
+ databricks_conn_id,
+ retry_limit=self.retry_limit,
+ retry_delay=self.retry_delay,
+ retry_args=retry_args,
+ caller=caller,
+ )
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ return (
+
"airflow.providers.databricks.triggers.databricks.DatabricksWarehouseStateTrigger",
+ {
+ "warehouse_id": self.warehouse_id,
+ "target_state": self.target_state,
+ "databricks_conn_id": self.databricks_conn_id,
+ "timeout": self.timeout,
+ "polling_period_seconds": self.polling_period_seconds,
+ "retry_limit": self.retry_limit,
+ "retry_delay": self.retry_delay,
+ "retry_args": self.retry_args,
+ "caller": self.caller,
+ },
+ )
+
+ async def on_kill(self) -> None:
+ # Warehouses have no cancel-start/stop API. Clearing a deferred start
must not stop the
+ # warehouse — the Dag author may still want it running after the wait
is abandoned.
+ self.log.info(
+ "Databricks SQL warehouse %s wait cancelled; leaving warehouse
state unchanged.",
+ self.warehouse_id,
+ )
+
+ def _build_trigger_event(
+ self, *, status: str, last_state: str, state: WarehouseState | None =
None
+ ) -> TriggerEvent:
+ payload: dict[str, Any] = {
+ "status": status,
+ "warehouse_id": self.warehouse_id,
+ "target_state": self.target_state,
+ "last_state": last_state,
+ }
+ if state is not None:
+ payload["state"] = state.to_json()
Review Comment:
`payload["state"]` is always `json.dumps({"state": payload["last_state"]})`,
and `execute_complete` then does
`WarehouseState.from_json(event["state"]).state` where `event["last_state"]`
would do — while hard-indexing `event["state"]` and `.get()`-ing everything
else. Worth either dropping the redundant key or keeping it deliberately for
symmetry with the sibling triggers and saying so.
On the `to_json`/`from_json` restore: this closes the loop from #70088,
where their removal was requested because nothing used them and you noted they
would return with this trigger. That is the strongest justification for the
change and the description does not mention it.
--
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]
