Vamsi-klu commented on code in PR #71752:
URL: https://github.com/apache/airflow/pull/71752#discussion_r3813930229
##########
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:
Followed DatabricksSQLStatementExecutionTrigger. The operator now records an
absolute end_time from timeout when it defers, serialize carries that
timestamp, and run() compares it with time.time(). A triggerer restart no
longer resets the wait. The how-to and the trigger docstring now describe that.
--
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]
