Vamsi-klu commented on code in PR #71752:
URL: https://github.com/apache/airflow/pull/71752#discussion_r3813938306
##########
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:
Kept the serialized state key so the payload matches the RunState and
SQLStatementState triggers, which is also why WarehouseState.to_json and
from_json came back after the Phase 1 removal. last_state stays as the plain
string used in the timeout message. The PR description now calls out that
restore from 70088.
--
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]
