This is an automated email from the ASF dual-hosted git repository.
eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 263fafcec05 Add Databricks SQL warehouse lifecycle operators (#70088)
263fafcec05 is described below
commit 263fafcec054c5463a76225937dd8e1b528a691c
Author: deepinsight coder <[email protected]>
AuthorDate: Mon Aug 17 10:38:01 2026 -0700
Add Databricks SQL warehouse lifecycle operators (#70088)
* Add Databricks SQL warehouse lifecycle operators
* Honor observed Databricks warehouse states at timeout
A slow final status request can finish after the deadline even when it
confirms the requested state. Treating that observation as a timeout can fail
an otherwise successful Dag.
* Avoid false failures when starting Databricks warehouses
A warehouse can still report STOPPED immediately after the start request
because the API response is eventually consistent. Keep polling until RUNNING
or deletion/timeout so valid starts do not fail spuriously.
* Address review feedback on Databricks SQL warehouse operators
Rename the operator module to warehouse.py to follow provider naming
conventions, drop the unused WarehouseState.to_json/from_json helpers
that have no production caller until the deferrable trigger lands, and
describe the all_done trigger rule behavior in the docs without
referring to system tests.
Co-Authored-By: Claude Fable 5 <[email protected]>
* Rename Databricks SQL warehouse how-to page to warehouse.rst
Match the operator module name so the how-to page follows the
same provider naming convention as warehouse.py.
Co-authored-by: nrvamsi13 <[email protected]>
---------
Co-authored-by: Claude Fable 5 <[email protected]>
Co-authored-by: Cursor Agent <[email protected]>
---
providers/databricks/docs/operators/warehouse.rst | 55 +++
providers/databricks/provider.yaml | 2 +
.../src/airflow/providers/databricks/exceptions.py | 4 +
.../providers/databricks/get_provider_info.py | 6 +-
.../providers/databricks/hooks/databricks.py | 73 ++++
.../providers/databricks/operators/warehouse.py | 156 ++++++++
.../databricks/example_databricks_sql_warehouse.py | 67 ++++
.../tests/unit/databricks/hooks/test_databricks.py | 63 ++++
.../unit/databricks/operators/test_warehouse.py | 406 +++++++++++++++++++++
9 files changed, 831 insertions(+), 1 deletion(-)
diff --git a/providers/databricks/docs/operators/warehouse.rst
b/providers/databricks/docs/operators/warehouse.rst
new file mode 100644
index 00000000000..1467bbb7f45
--- /dev/null
+++ b/providers/databricks/docs/operators/warehouse.rst
@@ -0,0 +1,55 @@
+ .. 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.
+
+.. _howto/operator:DatabricksStartWarehouseOperator:
+.. _howto/operator:DatabricksStopWarehouseOperator:
+
+Databricks SQL warehouse lifecycle operators
+============================================
+
+Use
:class:`~airflow.providers.databricks.operators.warehouse.DatabricksStartWarehouseOperator`
+and
:class:`~airflow.providers.databricks.operators.warehouse.DatabricksStopWarehouseOperator`
+to start and stop an existing Databricks SQL warehouse through the
+`Databricks SQL Warehouses API
<https://docs.databricks.com/api/workspace/warehouses>`_.
+Both operators require the warehouse ID and use the :ref:`Databricks connection
+<howto/connection:databricks>` for authentication.
+
+By default, each operator waits for the requested state: ``RUNNING`` when
starting and ``STOPPED``
+when stopping. Use ``polling_period_seconds`` to control the polling interval,
``timeout`` to limit
+the wait, or ``wait_for_termination=False`` to return after requesting the
transition. Repeated task
+attempts are safe: an already running warehouse is not started again, and an
already stopped warehouse
+is not stopped again. If a start is requested while a warehouse is stopping,
any transition rejection
+from Databricks is propagated to the task.
+
+Start a SQL warehouse
+---------------------
+
+.. exampleinclude::
/../../databricks/tests/system/databricks/example_databricks_sql_warehouse.py
+ :language: python
+ :start-after: [START howto_operator_databricks_start_sql_warehouse]
+ :end-before: [END howto_operator_databricks_start_sql_warehouse]
+
+Stop a SQL warehouse
+--------------------
+
+The stop task uses the ``all_done`` trigger rule so the warehouse is stopped
even when an upstream
+task fails.
+
+.. exampleinclude::
/../../databricks/tests/system/databricks/example_databricks_sql_warehouse.py
+ :language: python
+ :start-after: [START howto_operator_databricks_stop_sql_warehouse]
+ :end-before: [END howto_operator_databricks_stop_sql_warehouse]
diff --git a/providers/databricks/provider.yaml
b/providers/databricks/provider.yaml
index f2dc2d7a2de..14e46e4fe1f 100644
--- a/providers/databricks/provider.yaml
+++ b/providers/databricks/provider.yaml
@@ -125,6 +125,7 @@ integrations:
how-to-guide:
- /docs/apache-airflow-providers-databricks/operators/sql.rst
- /docs/apache-airflow-providers-databricks/operators/sql_statements.rst
+ - /docs/apache-airflow-providers-databricks/operators/warehouse.rst
- /docs/apache-airflow-providers-databricks/operators/copy_into.rst
tags: [service]
- integration-name: Databricks Repos
@@ -147,6 +148,7 @@ operators:
- integration-name: Databricks SQL
python-modules:
- airflow.providers.databricks.operators.databricks_sql
+ - airflow.providers.databricks.operators.warehouse
- integration-name: Databricks Repos
python-modules:
- airflow.providers.databricks.operators.databricks_repos
diff --git
a/providers/databricks/src/airflow/providers/databricks/exceptions.py
b/providers/databricks/src/airflow/providers/databricks/exceptions.py
index 0831810f70f..a0f52a6077c 100644
--- a/providers/databricks/src/airflow/providers/databricks/exceptions.py
+++ b/providers/databricks/src/airflow/providers/databricks/exceptions.py
@@ -36,6 +36,10 @@ class DatabricksOperatorPayloadError(AirflowException):
"""Raised when a Databricks operator payload is invalid."""
+class DatabricksWarehouseError(AirflowException):
+ """Raised when a SQL warehouse fails to reach or times out waiting for a
target state."""
+
+
class DatabricksApiError(AirflowException):
"""Raised when a Databricks REST API call returns an error response."""
diff --git
a/providers/databricks/src/airflow/providers/databricks/get_provider_info.py
b/providers/databricks/src/airflow/providers/databricks/get_provider_info.py
index f80d54a130d..a218ce35464 100644
--- a/providers/databricks/src/airflow/providers/databricks/get_provider_info.py
+++ b/providers/databricks/src/airflow/providers/databricks/get_provider_info.py
@@ -46,6 +46,7 @@ def get_provider_info():
"how-to-guide": [
"/docs/apache-airflow-providers-databricks/operators/sql.rst",
"/docs/apache-airflow-providers-databricks/operators/sql_statements.rst",
+
"/docs/apache-airflow-providers-databricks/operators/warehouse.rst",
"/docs/apache-airflow-providers-databricks/operators/copy_into.rst",
],
"tags": ["service"],
@@ -74,7 +75,10 @@ def get_provider_info():
},
{
"integration-name": "Databricks SQL",
- "python-modules":
["airflow.providers.databricks.operators.databricks_sql"],
+ "python-modules": [
+ "airflow.providers.databricks.operators.databricks_sql",
+ "airflow.providers.databricks.operators.warehouse",
+ ],
},
{
"integration-name": "Databricks Repos",
diff --git
a/providers/databricks/src/airflow/providers/databricks/hooks/databricks.py
b/providers/databricks/src/airflow/providers/databricks/hooks/databricks.py
index d780816f9fa..4a0f6454983 100644
--- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks.py
+++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks.py
@@ -69,6 +69,7 @@ WORKSPACE_GET_STATUS_ENDPOINT = ("GET",
"2.0/workspace/get-status")
SPARK_VERSIONS_ENDPOINT = ("GET", "2.1/clusters/spark-versions")
SQL_STATEMENTS_ENDPOINT = "2.0/sql/statements"
+SQL_WAREHOUSES_ENDPOINT = "2.0/sql/warehouses"
class RunLifeCycleState(Enum):
@@ -267,6 +268,46 @@ class 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:
+ """Return whether the warehouse is deleting or deleted."""
+ return self.state in ("DELETING", "DELETED")
+
+ def __eq__(self, other: object) -> bool:
+ if not isinstance(other, WarehouseState):
+ return NotImplemented
+ return self.state == other.state
+
+ def __hash__(self):
+ return hash(self.state)
+
+ def __repr__(self) -> str:
+ return str(self.__dict__)
+
+
class DatabricksHook(BaseDatabricksHook):
"""
Interact with Databricks.
@@ -742,6 +783,38 @@ class DatabricksHook(BaseDatabricksHook):
"""
self._do_api_call(TERMINATE_CLUSTER_ENDPOINT, json)
+ def get_warehouse(self, warehouse_id: str) -> dict[str, Any]:
+ """
+ Retrieve a Databricks SQL warehouse.
+
+ :param warehouse_id: ID of the SQL warehouse.
+ """
+ return self._do_api_call(("GET",
f"{SQL_WAREHOUSES_ENDPOINT}/{warehouse_id}"))
+
+ def get_warehouse_state(self, warehouse_id: str) -> WarehouseState:
+ """
+ Retrieve the state of a Databricks SQL warehouse.
+
+ :param warehouse_id: ID of the SQL warehouse.
+ """
+ return WarehouseState(self.get_warehouse(warehouse_id)["state"])
+
+ def start_warehouse(self, warehouse_id: str) -> None:
+ """
+ Start a Databricks SQL warehouse.
+
+ :param warehouse_id: ID of the SQL warehouse.
+ """
+ self._do_api_call(("POST",
f"{SQL_WAREHOUSES_ENDPOINT}/{warehouse_id}/start"))
+
+ def stop_warehouse(self, warehouse_id: str) -> None:
+ """
+ Stop a Databricks SQL warehouse.
+
+ :param warehouse_id: ID of the SQL warehouse.
+ """
+ self._do_api_call(("POST",
f"{SQL_WAREHOUSES_ENDPOINT}/{warehouse_id}/stop"))
+
def install(self, json: dict) -> None:
"""
Install libraries on the cluster.
diff --git
a/providers/databricks/src/airflow/providers/databricks/operators/warehouse.py
b/providers/databricks/src/airflow/providers/databricks/operators/warehouse.py
new file mode 100644
index 00000000000..e2d0ce9a69b
--- /dev/null
+++
b/providers/databricks/src/airflow/providers/databricks/operators/warehouse.py
@@ -0,0 +1,156 @@
+# 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 DatabricksHook(
+ self.databricks_conn_id,
+ retry_limit=self.databricks_retry_limit,
+ retry_delay=self.databricks_retry_delay,
+ retry_args=self.databricks_retry_args,
+ caller=self.__class__.__name__,
+ )
+
+ def _wait_for_state(self, target: 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.is_deleted:
+ 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")
+
+
+class DatabricksStopWarehouseOperator(_DatabricksWarehouseBaseOperator):
+ """
+ Stop a Databricks SQL warehouse and optionally wait for it to stop.
+
+ :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 ``STOPPED``.
+ :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_stopped:
+ self.log.info("Databricks SQL warehouse %s is already stopped.",
self.warehouse_id)
+ return
+ if state.state != "STOPPING":
+ self._hook.stop_warehouse(self.warehouse_id)
+ if self.wait_for_termination:
+ self._wait_for_state("STOPPED")
diff --git
a/providers/databricks/tests/system/databricks/example_databricks_sql_warehouse.py
b/providers/databricks/tests/system/databricks/example_databricks_sql_warehouse.py
new file mode 100644
index 00000000000..81f2f5591b5
--- /dev/null
+++
b/providers/databricks/tests/system/databricks/example_databricks_sql_warehouse.py
@@ -0,0 +1,67 @@
+# 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.warehouse import (
+ DatabricksStartWarehouseOperator,
+ DatabricksStopWarehouseOperator,
+)
+
+DAG_ID = "example_databricks_sql_warehouse"
+WAREHOUSE_ID = os.environ.get("WAREHOUSE_ID", "your-warehouse-id")
+
+with DAG(
+ dag_id=DAG_ID,
+ schedule="@daily",
+ start_date=datetime(2021, 1, 1),
+ tags=["example"],
+ catchup=False,
+) as dag:
+ # [START howto_operator_databricks_start_sql_warehouse]
+ start_warehouse = DatabricksStartWarehouseOperator(
+ task_id="start_warehouse",
+ databricks_conn_id="databricks_default",
+ warehouse_id=WAREHOUSE_ID,
+ wait_for_termination=True,
+ )
+ # [END howto_operator_databricks_start_sql_warehouse]
+
+ # [START howto_operator_databricks_stop_sql_warehouse]
+ stop_warehouse = DatabricksStopWarehouseOperator(
+ task_id="stop_warehouse",
+ databricks_conn_id="databricks_default",
+ warehouse_id=WAREHOUSE_ID,
+ wait_for_termination=True,
+ trigger_rule="all_done",
+ )
+ # [END howto_operator_databricks_stop_sql_warehouse]
+
+ start_warehouse >> stop_warehouse
+
+ from tests_common.test_utils.watcher import watcher
+
+ list(dag.tasks) >> watcher()
+
+from tests_common.test_utils.system_tests import get_test_run # noqa: E402
+
+test_run = get_test_run(dag)
diff --git
a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
index 9086e291986..be87a82483c 100644
--- a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
@@ -43,6 +43,7 @@ from airflow.providers.databricks.hooks.databricks import (
DatabricksHook,
RunState,
SQLStatementState,
+ WarehouseState,
)
from airflow.providers.databricks.hooks.databricks_base import (
AZURE_MANAGEMENT_ENDPOINT,
@@ -1552,6 +1553,68 @@ class
TestDatabricksHookConnSettings(TestDatabricksHookToken):
assert mock_get.call_args.args ==
(f"http://{HOST}:7908/api/2.1/foo/bar";,)
+class TestWarehouseLifecycle:
+ @mock.patch.object(DatabricksHook, "_do_api_call", autospec=True)
+ def test_get_warehouse_calls_correct_endpoint(self, mock_do_api_call):
+ mock_do_api_call.return_value = {"id": "wh-1", "state": "RUNNING"}
+ hook = DatabricksHook()
+
+ result = hook.get_warehouse("wh-1")
+
+ assert result == {"id": "wh-1", "state": "RUNNING"}
+ mock_do_api_call.assert_called_once_with(hook, ("GET",
"2.0/sql/warehouses/wh-1"))
+
+ @mock.patch.object(DatabricksHook, "_do_api_call", autospec=True)
+ def test_get_warehouse_state_wraps_state(self, mock_do_api_call):
+ mock_do_api_call.return_value = {"state": "RUNNING"}
+ hook = DatabricksHook()
+
+ state = hook.get_warehouse_state("wh-1")
+
+ assert state == WarehouseState("RUNNING")
+ assert state.is_running
+ mock_do_api_call.assert_called_once_with(hook, ("GET",
"2.0/sql/warehouses/wh-1"))
+
+ @mock.patch.object(DatabricksHook, "_do_api_call", autospec=True)
+ def test_start_warehouse_endpoint(self, mock_do_api_call):
+ hook = DatabricksHook()
+
+ hook.start_warehouse("wh-1")
+
+ mock_do_api_call.assert_called_once_with(hook, ("POST",
"2.0/sql/warehouses/wh-1/start"))
+
+ @mock.patch.object(DatabricksHook, "_do_api_call", autospec=True)
+ def test_stop_warehouse_endpoint(self, mock_do_api_call):
+ hook = DatabricksHook()
+
+ hook.stop_warehouse("wh-1")
+
+ mock_do_api_call.assert_called_once_with(hook, ("POST",
"2.0/sql/warehouses/wh-1/stop"))
+
+ @pytest.mark.parametrize(
+ ("state", "is_running", "is_stopped", "is_deleted"),
+ [
+ ("STARTING", False, False, False),
+ ("RUNNING", True, False, False),
+ ("STOPPING", False, False, False),
+ ("STOPPED", False, True, False),
+ ("DELETING", False, False, True),
+ ("DELETED", False, False, True),
+ ],
+ )
+ def test_warehouse_state_valid_and_properties(self, state, is_running,
is_stopped, is_deleted):
+ warehouse_state = WarehouseState(state)
+
+ assert warehouse_state.state == state
+ assert warehouse_state.is_running is is_running
+ assert warehouse_state.is_stopped is is_stopped
+ assert warehouse_state.is_deleted is is_deleted
+
+ def test_warehouse_state_unexpected_raises_value_error(self):
+ with pytest.raises(ValueError, match="Unexpected warehouse state:
FOO"):
+ WarehouseState("FOO")
+
+
class TestRunState:
def test_is_terminal_true(self):
terminal_states = ["TERMINATED", "SKIPPED", "INTERNAL_ERROR"]
diff --git
a/providers/databricks/tests/unit/databricks/operators/test_warehouse.py
b/providers/databricks/tests/unit/databricks/operators/test_warehouse.py
new file mode 100644
index 00000000000..aa99e4d54cc
--- /dev/null
+++ b/providers/databricks/tests/unit/databricks/operators/test_warehouse.py
@@ -0,0 +1,406 @@
+# 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.
+from __future__ import annotations
+
+from unittest import mock
+
+import pytest
+
+from airflow.providers.databricks.exceptions import DatabricksWarehouseError
+from airflow.providers.databricks.hooks.databricks import DatabricksHook,
WarehouseState
+from airflow.providers.databricks.operators.warehouse import (
+ DatabricksStartWarehouseOperator,
+ DatabricksStopWarehouseOperator,
+)
+
+TASK_ID = "warehouse-lifecycle"
+WAREHOUSE_ID = "wh-1"
+
+
+class TestDatabricksStartWarehouseOperator:
+ @pytest.mark.parametrize(
+ "first_polled_state",
+ ["STARTING", "STOPPED"],
+ ids=["transitioning", "stale-stopped"],
+ )
+ @mock.patch("airflow.providers.databricks.operators.warehouse.time.sleep")
+ @mock.patch.object(DatabricksStartWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ def test_starts_then_waits_until_running(self, mock_hook_property,
mock_sleep, first_polled_state):
+ hook = mock.MagicMock(spec=DatabricksHook)
+ mock_hook_property.return_value = hook
+ hook.get_warehouse_state.side_effect = [
+ WarehouseState("STOPPED"),
+ WarehouseState(first_polled_state),
+ WarehouseState("RUNNING"),
+ ]
+ operator = DatabricksStartWarehouseOperator(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ polling_period_seconds=0,
+ )
+
+ operator.execute(None)
+
+ hook.start_warehouse.assert_called_once_with(WAREHOUSE_ID)
+ assert hook.get_warehouse_state.call_count == 3
+ mock_sleep.assert_called_once_with(0)
+
+ @mock.patch.object(DatabricksStartWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ def test_skips_start_when_already_running(self, mock_hook_property):
+ hook = mock.MagicMock(spec=DatabricksHook)
+ mock_hook_property.return_value = hook
+ hook.get_warehouse_state.return_value = WarehouseState("RUNNING")
+ operator = DatabricksStartWarehouseOperator(task_id=TASK_ID,
warehouse_id=WAREHOUSE_ID)
+
+ operator.execute(None)
+
+ hook.start_warehouse.assert_not_called()
+ hook.get_warehouse_state.assert_called_once_with(WAREHOUSE_ID)
+
+ @mock.patch.object(DatabricksStartWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ def test_skips_start_when_starting_without_waiting(self,
mock_hook_property):
+ hook = mock.MagicMock(spec=DatabricksHook)
+ mock_hook_property.return_value = hook
+ hook.get_warehouse_state.return_value = WarehouseState("STARTING")
+ operator = DatabricksStartWarehouseOperator(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ wait_for_termination=False,
+ )
+
+ operator.execute(None)
+
+ hook.start_warehouse.assert_not_called()
+ hook.get_warehouse_state.assert_called_once_with(WAREHOUSE_ID)
+
+ @mock.patch("airflow.providers.databricks.operators.warehouse.time.sleep")
+ @mock.patch.object(DatabricksStartWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ def test_starts_while_stopping_and_waits_until_running(self,
mock_hook_property, mock_sleep):
+ hook = mock.MagicMock(spec=DatabricksHook)
+ mock_hook_property.return_value = hook
+ hook.get_warehouse_state.side_effect = [
+ WarehouseState("STOPPING"),
+ WarehouseState("STOPPING"),
+ WarehouseState("RUNNING"),
+ ]
+ operator = DatabricksStartWarehouseOperator(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ polling_period_seconds=0,
+ )
+
+ operator.execute(None)
+
+ hook.start_warehouse.assert_called_once_with(WAREHOUSE_ID)
+ assert hook.get_warehouse_state.call_count == 3
+ mock_sleep.assert_called_once_with(0)
+
+
+class TestDatabricksStopWarehouseOperator:
+ @mock.patch("airflow.providers.databricks.operators.warehouse.time.sleep")
+ @mock.patch.object(DatabricksStopWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ def test_stops_then_waits_until_stopped(self, mock_hook_property,
mock_sleep):
+ hook = mock.MagicMock(spec=DatabricksHook)
+ mock_hook_property.return_value = hook
+ hook.get_warehouse_state.side_effect = [
+ WarehouseState("RUNNING"),
+ WarehouseState("STOPPING"),
+ WarehouseState("STOPPED"),
+ ]
+ operator = DatabricksStopWarehouseOperator(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ polling_period_seconds=0,
+ )
+
+ operator.execute(None)
+
+ hook.stop_warehouse.assert_called_once_with(WAREHOUSE_ID)
+ assert hook.get_warehouse_state.call_count == 3
+ mock_sleep.assert_called_once_with(0)
+
+ @mock.patch.object(DatabricksStopWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ def test_skips_stop_when_already_stopped(self, mock_hook_property):
+ hook = mock.MagicMock(spec=DatabricksHook)
+ mock_hook_property.return_value = hook
+ hook.get_warehouse_state.return_value = WarehouseState("STOPPED")
+ operator = DatabricksStopWarehouseOperator(task_id=TASK_ID,
warehouse_id=WAREHOUSE_ID)
+
+ operator.execute(None)
+
+ hook.stop_warehouse.assert_not_called()
+ hook.get_warehouse_state.assert_called_once_with(WAREHOUSE_ID)
+
+ @mock.patch.object(DatabricksStopWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ def test_skips_stop_when_stopping_without_waiting(self,
mock_hook_property):
+ hook = mock.MagicMock(spec=DatabricksHook)
+ mock_hook_property.return_value = hook
+ hook.get_warehouse_state.return_value = WarehouseState("STOPPING")
+ operator = DatabricksStopWarehouseOperator(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ wait_for_termination=False,
+ )
+
+ operator.execute(None)
+
+ hook.stop_warehouse.assert_not_called()
+ hook.get_warehouse_state.assert_called_once_with(WAREHOUSE_ID)
+
+
+class TestDatabricksWarehouseOperatorBase:
+ @pytest.mark.parametrize(
+ ("operator_class", "initial_state", "failure_state",
"transition_method", "target_state"),
+ [
+ (DatabricksStartWarehouseOperator, "STOPPED", "DELETING",
"start_warehouse", "RUNNING"),
+ (DatabricksStopWarehouseOperator, "RUNNING", "DELETED",
"stop_warehouse", "STOPPED"),
+ ],
+ )
+ @mock.patch.object(DatabricksStartWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ @mock.patch.object(DatabricksStopWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ def test_execute_raises_on_failure_state(
+ self,
+ mock_stop_hook_property,
+ mock_start_hook_property,
+ operator_class,
+ initial_state,
+ failure_state,
+ transition_method,
+ target_state,
+ ):
+ hook = mock.MagicMock(spec=DatabricksHook)
+ mock_start_hook_property.return_value = hook
+ mock_stop_hook_property.return_value = hook
+ hook.get_warehouse_state.side_effect = [
+ WarehouseState(initial_state),
+ WarehouseState(failure_state),
+ ]
+ operator = operator_class(task_id=TASK_ID, warehouse_id=WAREHOUSE_ID)
+
+ with pytest.raises(
+ DatabricksWarehouseError,
+ match=f"entered {failure_state} while waiting for {target_state}",
+ ):
+ operator.execute(None)
+
+ getattr(hook, transition_method).assert_called_once_with(WAREHOUSE_ID)
+
+ @pytest.mark.parametrize(
+ ("operator_class", "initial_state", "target_state",
"transition_method"),
+ [
+ (DatabricksStartWarehouseOperator, "STARTING", "RUNNING",
"start_warehouse"),
+ (DatabricksStopWarehouseOperator, "STOPPING", "STOPPED",
"stop_warehouse"),
+ ],
+ )
+ @mock.patch("airflow.providers.databricks.operators.warehouse.time.sleep")
+ @mock.patch.object(DatabricksStartWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ @mock.patch.object(DatabricksStopWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ def test_transition_in_progress_skips_request_and_continues_waiting(
+ self,
+ mock_stop_hook_property,
+ mock_start_hook_property,
+ mock_sleep,
+ operator_class,
+ initial_state,
+ target_state,
+ transition_method,
+ ):
+ hook = mock.MagicMock(spec=DatabricksHook)
+ mock_start_hook_property.return_value = hook
+ mock_stop_hook_property.return_value = hook
+ hook.get_warehouse_state.side_effect = [
+ WarehouseState(initial_state),
+ WarehouseState(initial_state),
+ WarehouseState(target_state),
+ ]
+ operator = operator_class(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ polling_period_seconds=0,
+ )
+
+ operator.execute(None)
+
+ getattr(hook, transition_method).assert_not_called()
+ assert hook.get_warehouse_state.call_count == 3
+ mock_sleep.assert_called_once_with(0)
+
+ @pytest.mark.parametrize(
+ ("operator_class", "initial_state", "transition_method"),
+ [
+ (DatabricksStartWarehouseOperator, "STOPPED", "start_warehouse"),
+ (DatabricksStopWarehouseOperator, "RUNNING", "stop_warehouse"),
+ ],
+ )
+ @mock.patch.object(DatabricksStartWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ @mock.patch.object(DatabricksStopWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ def test_transition_without_waiting(
+ self,
+ mock_stop_hook_property,
+ mock_start_hook_property,
+ operator_class,
+ initial_state,
+ transition_method,
+ ):
+ hook = mock.MagicMock(spec=DatabricksHook)
+ mock_start_hook_property.return_value = hook
+ mock_stop_hook_property.return_value = hook
+ hook.get_warehouse_state.return_value = WarehouseState(initial_state)
+ operator = operator_class(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ wait_for_termination=False,
+ )
+
+ operator.execute(None)
+
+ getattr(hook, transition_method).assert_called_once_with(WAREHOUSE_ID)
+ hook.get_warehouse_state.assert_called_once_with(WAREHOUSE_ID)
+
+ @mock.patch("airflow.providers.databricks.operators.warehouse.time.sleep")
+ @mock.patch(
+ "airflow.providers.databricks.operators.warehouse.time.monotonic",
+ return_value=0,
+ )
+ @mock.patch.object(DatabricksStartWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ def test_wait_caps_sleep_and_does_not_poll_after_deadline(
+ self, mock_hook_property, mock_monotonic, mock_sleep
+ ):
+ hook = mock.MagicMock(spec=DatabricksHook)
+ mock_hook_property.return_value = hook
+ hook.get_warehouse_state.side_effect = [WarehouseState("STARTING"),
WarehouseState("RUNNING")]
+ operator = DatabricksStartWarehouseOperator(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ polling_period_seconds=30,
+ timeout=10,
+ )
+
+ def advance_to_deadline(seconds):
+ mock_monotonic.return_value = seconds
+
+ mock_sleep.side_effect = advance_to_deadline
+
+ with pytest.raises(
+ DatabricksWarehouseError,
+ match="did not reach RUNNING within 10s; last state: STARTING",
+ ):
+ operator._wait_for_state("RUNNING")
+
+ hook.get_warehouse_state.assert_called_once_with(WAREHOUSE_ID)
+ mock_sleep.assert_called_once_with(10)
+
+ @mock.patch("airflow.providers.databricks.operators.warehouse.time.sleep")
+ @mock.patch(
+ "airflow.providers.databricks.operators.warehouse.time.monotonic",
+ return_value=0,
+ )
+ @mock.patch.object(DatabricksStartWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ def test_wait_accepts_target_received_after_deadline(
+ self, mock_hook_property, mock_monotonic, mock_sleep
+ ):
+ hook = mock.MagicMock(spec=DatabricksHook)
+ mock_hook_property.return_value = hook
+
+ def get_state_after_deadline(_):
+ mock_monotonic.return_value = 11
+ return WarehouseState("RUNNING")
+
+ hook.get_warehouse_state.side_effect = get_state_after_deadline
+ operator = DatabricksStartWarehouseOperator(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ timeout=10,
+ )
+
+ operator._wait_for_state("RUNNING")
+
+ hook.get_warehouse_state.assert_called_once_with(WAREHOUSE_ID)
+ mock_sleep.assert_not_called()
+
+ @mock.patch("airflow.providers.databricks.operators.warehouse.time.sleep")
+ @mock.patch(
+ "airflow.providers.databricks.operators.warehouse.time.monotonic",
+ return_value=0,
+ )
+ @mock.patch.object(DatabricksStartWarehouseOperator, "_hook",
new_callable=mock.PropertyMock)
+ def test_wait_rejects_deleted_state_received_after_deadline(
+ self, mock_hook_property, mock_monotonic, mock_sleep
+ ):
+ hook = mock.MagicMock(spec=DatabricksHook)
+ mock_hook_property.return_value = hook
+
+ def get_state_after_deadline(_):
+ mock_monotonic.return_value = 11
+ return WarehouseState("DELETED")
+
+ hook.get_warehouse_state.side_effect = get_state_after_deadline
+ operator = DatabricksStartWarehouseOperator(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ timeout=10,
+ )
+
+ with pytest.raises(
+ DatabricksWarehouseError,
+ match="entered DELETED while waiting for RUNNING",
+ ):
+ operator._wait_for_state("RUNNING")
+
+ hook.get_warehouse_state.assert_called_once_with(WAREHOUSE_ID)
+ mock_sleep.assert_not_called()
+
+ def test_operator_template_fields(self):
+ expected = ("databricks_conn_id", "warehouse_id")
+ assert DatabricksStartWarehouseOperator.template_fields == expected
+ assert DatabricksStopWarehouseOperator.template_fields == expected
+
+ @pytest.mark.parametrize(
+ ("operator_class", "warehouse_id"),
+ [
+ (DatabricksStartWarehouseOperator, ""),
+ (DatabricksStartWarehouseOperator, None),
+ (DatabricksStopWarehouseOperator, ""),
+ (DatabricksStopWarehouseOperator, None),
+ ],
+ )
+ def test_invalid_warehouse_id(self, operator_class, warehouse_id):
+ operator = operator_class(task_id=TASK_ID, warehouse_id=warehouse_id)
+
+ with pytest.raises(ValueError, match="warehouse_id must be provided"):
+ operator.execute(None)
+
+
@mock.patch("airflow.providers.databricks.operators.warehouse.DatabricksHook",
autospec=True)
+ def test_operator_builds_hook(self, mock_hook_class):
+ retry_args = {"reraise": True}
+ operator = DatabricksStartWarehouseOperator(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ databricks_conn_id="custom_conn",
+ databricks_retry_limit=7,
+ databricks_retry_delay=4,
+ databricks_retry_args=retry_args,
+ )
+
+ assert operator._hook is mock_hook_class.return_value
+ mock_hook_class.assert_called_once_with(
+ "custom_conn",
+ retry_limit=7,
+ retry_delay=4,
+ retry_args=retry_args,
+ caller="DatabricksStartWarehouseOperator",
+ )
