moomindani commented on code in PR #70088: URL: https://github.com/apache/airflow/pull/70088#discussion_r3654302440
########## providers/databricks/src/airflow/providers/databricks/operators/databricks_warehouse.py: ########## @@ -0,0 +1,159 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Operators for managing Databricks SQL warehouse lifecycle state.""" + +from __future__ import annotations + +import time +from collections.abc import Sequence +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.sdk import BaseOperator +from airflow.providers.databricks.exceptions import DatabricksWarehouseError +from airflow.providers.databricks.hooks.databricks import DatabricksHook + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context + + +class _DatabricksWarehouseBaseOperator(BaseOperator): + """Share Databricks SQL warehouse connection and polling behavior.""" + + template_fields: Sequence[str] = ("databricks_conn_id", "warehouse_id") + ui_color = "#1CB1C2" + ui_fgcolor = "#fff" + + def __init__( + self, + warehouse_id: str, + *, + databricks_conn_id: str = "databricks_default", + wait_for_termination: bool = True, + polling_period_seconds: int = 30, + timeout: float = 3600, + databricks_retry_limit: int = 3, + databricks_retry_delay: int = 1, + databricks_retry_args: dict[Any, Any] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.warehouse_id = warehouse_id + self.databricks_conn_id = databricks_conn_id + self.wait_for_termination = wait_for_termination + self.polling_period_seconds = polling_period_seconds + self.timeout = timeout + self.databricks_retry_limit = databricks_retry_limit + self.databricks_retry_delay = databricks_retry_delay + self.databricks_retry_args = databricks_retry_args + + def _validate_warehouse_id(self) -> None: + if not self.warehouse_id: + raise ValueError("warehouse_id must be provided.") + + @cached_property + def _hook(self) -> DatabricksHook: + return self._get_hook(caller=self.__class__.__name__) + + def _get_hook(self, caller: str) -> DatabricksHook: + return DatabricksHook( + self.databricks_conn_id, + retry_limit=self.databricks_retry_limit, + retry_delay=self.databricks_retry_delay, + retry_args=self.databricks_retry_args, + caller=caller, + ) + + def _wait_for_state(self, target: str, failure_states: set[str]) -> None: + deadline = time.monotonic() + self.timeout + last_state = "unknown" + while time.monotonic() < deadline: + state = self._hook.get_warehouse_state(self.warehouse_id) + last_state = state.state + now = time.monotonic() + if state.state == target: + return + if state.state in failure_states: + raise DatabricksWarehouseError( + f"Databricks SQL warehouse {self.warehouse_id} entered {state.state} " + f"while waiting for {target}." + ) + if now >= deadline: + break + self.log.info( + "Databricks SQL warehouse %s is %s; waiting for %s.", + self.warehouse_id, + state.state, + target, + ) + time.sleep(min(self.polling_period_seconds, deadline - now)) + raise DatabricksWarehouseError( + f"Databricks SQL warehouse {self.warehouse_id} did not reach {target} " + f"within {self.timeout}s; last state: {last_state}." + ) + + +class DatabricksStartWarehouseOperator(_DatabricksWarehouseBaseOperator): + """ + Start a Databricks SQL warehouse and optionally wait for it to run. + + :param warehouse_id: ID of the Databricks SQL warehouse. (templated) + :param databricks_conn_id: Reference to the Databricks connection. (templated) + :param wait_for_termination: Wait until the warehouse reaches ``RUNNING``. + :param polling_period_seconds: Number of seconds between state checks. + :param timeout: Maximum number of seconds to wait for the target state. + :param databricks_retry_limit: Number of times to retry unavailable Databricks requests. + :param databricks_retry_delay: Number of seconds between Databricks request retries. + :param databricks_retry_args: Additional arguments for ``tenacity.Retrying``. + """ + + def execute(self, context: Context) -> None: + self._validate_warehouse_id() + state = self._hook.get_warehouse_state(self.warehouse_id) + if state.is_running: + self.log.info("Databricks SQL warehouse %s is already running.", self.warehouse_id) + return + if state.state != "STARTING": + self._hook.start_warehouse(self.warehouse_id) + if self.wait_for_termination: + self._wait_for_state("RUNNING", {"STOPPED", "DELETING", "DELETED"}) Review Comment: `STOPPED` should not be a failure state for the start path. After the pre-check issues `start_warehouse()`, the very first poll happens with no `time.sleep()` in between — so if that `GET` still reports the pre-transition `STOPPED`, the task fails immediately even though the start succeeded. Reproduction: ```python op = DatabricksStartWarehouseOperator(task_id="t", warehouse_id="w") # default polling_period_seconds=30 hook.get_warehouse_state.side_effect = [ WarehouseState("STOPPED"), # pre-check -> start_warehouse() issued WarehouseState("STOPPED"), # first poll, still lagging WarehouseState("RUNNING"), ] op.execute(None) # DatabricksWarehouseError: warehouse w entered STOPPED while waiting for RUNNING # (time.sleep call count at failure: 0) ``` I measured the real window: `POST /start` -> `GET` reporting `STARTING` took 0.45-0.46s across three trials, so in practice this will usually pass. But `polling_period_seconds` does not protect against it, because the first poll precedes the first sleep — raising the interval does not shrink the window. Dropping `STOPPED` from the start path's failure set is enough; `DELETING`/`DELETED` remain genuinely terminal, and the timeout still catches a start that never takes effect. A warehouse that legitimately re-stops mid-wait (auto-stop, or a concurrent stop) would then surface as a timeout rather than an immediate error, which seems like the better trade against false failures — but if you'd rather keep the fast failure, tolerating `STOPPED` only until the first observed non-`STOPPED` state would also work. Worth a test pinning whichever semantics you choose, since the current suite's `STOPPED`-as-failure case (`test_execute_raises_on_failure_state`) asserts the behaviour I'm questioning. ########## providers/databricks/src/airflow/providers/databricks/hooks/databricks.py: ########## @@ -267,6 +268,53 @@ def from_json(cls, data: str) -> SQLStatementState: return SQLStatementState(**json.loads(data)) +class WarehouseState: + """Utility class for the state of a Databricks SQL warehouse.""" + + WAREHOUSE_STATES = ["STARTING", "RUNNING", "STOPPING", "STOPPED", "DELETING", "DELETED"] + + def __init__(self, state: str = "", *args, **kwargs) -> None: + if state not in self.WAREHOUSE_STATES: + raise ValueError( + f"Unexpected warehouse state: {state}: If the state has been introduced recently, " + "please check the Databricks user guide for troubleshooting information" + ) + self.state = state + + @property + def is_running(self) -> bool: + """Return whether the warehouse is running.""" + return self.state == "RUNNING" + + @property + def is_stopped(self) -> bool: + """Return whether the warehouse is stopped.""" + return self.state == "STOPPED" + + @property + def is_deleted(self) -> bool: Review Comment: `WarehouseState.is_deleted` is not used by any production code path — only `test_warehouse_state_valid_and_properties` references it. The operators compare against the `"DELETING"` / `"DELETED"` literals in their `failure_states` sets instead. Either drop the property, or use it in `_wait_for_state` so the terminal-state check goes through it rather than duplicating the literals. The latter is probably nicer, since it puts the definition of "deleted" in one place. ########## providers/databricks/src/airflow/providers/databricks/operators/databricks_warehouse.py: ########## @@ -0,0 +1,159 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Operators for managing Databricks SQL warehouse lifecycle state.""" + +from __future__ import annotations + +import time +from collections.abc import Sequence +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.sdk import BaseOperator +from airflow.providers.databricks.exceptions import DatabricksWarehouseError +from airflow.providers.databricks.hooks.databricks import DatabricksHook + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context + + +class _DatabricksWarehouseBaseOperator(BaseOperator): + """Share Databricks SQL warehouse connection and polling behavior.""" + + template_fields: Sequence[str] = ("databricks_conn_id", "warehouse_id") + ui_color = "#1CB1C2" + ui_fgcolor = "#fff" + + def __init__( + self, + warehouse_id: str, + *, + databricks_conn_id: str = "databricks_default", + wait_for_termination: bool = True, + polling_period_seconds: int = 30, + timeout: float = 3600, + databricks_retry_limit: int = 3, + databricks_retry_delay: int = 1, + databricks_retry_args: dict[Any, Any] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.warehouse_id = warehouse_id + self.databricks_conn_id = databricks_conn_id + self.wait_for_termination = wait_for_termination + self.polling_period_seconds = polling_period_seconds + self.timeout = timeout + self.databricks_retry_limit = databricks_retry_limit + self.databricks_retry_delay = databricks_retry_delay + self.databricks_retry_args = databricks_retry_args + + def _validate_warehouse_id(self) -> None: + if not self.warehouse_id: + raise ValueError("warehouse_id must be provided.") + + @cached_property + def _hook(self) -> DatabricksHook: + return self._get_hook(caller=self.__class__.__name__) + + def _get_hook(self, caller: str) -> DatabricksHook: Review Comment: The `_get_hook(caller=...)` indirection has a single call site (`_hook`, one line above) and no override in either subclass, so it can be inlined into the `cached_property` and the method dropped. Non-blocking — flagging it because a private one-caller wrapper tends to read as a leftover from an earlier iteration. ########## providers/databricks/tests/system/databricks/example_databricks_sql_warehouse.py: ########## @@ -0,0 +1,68 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Example Dag for starting and stopping an existing Databricks SQL warehouse.""" + +from __future__ import annotations + +import os +from datetime import datetime + +from airflow.providers.common.compat.sdk import DAG +from airflow.providers.databricks.operators.databricks_warehouse import ( + DatabricksStartWarehouseOperator, + DatabricksStopWarehouseOperator, +) + +ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID") Review Comment: `ENV_ID` is assigned but never used. The other Databricks system examples that define it use it, e.g. to namespace resources per environment. Either drop the line or fold it into `DAG_ID` the way the sibling examples do. -- 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]
