Vamsi-klu commented on code in PR #71752:
URL: https://github.com/apache/airflow/pull/71752#discussion_r3813932727
##########
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:
Removed the on_kill override and the test that only asserted the base no-op.
The class docstring now says we leave warehouse state unchanged because
Databricks has no cancel API for start or stop.
--
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]
