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 <noreply@anthropic.com>

Author: Vamsi-klu

providers/databricks/docs/operators/sql_warehouse.rst

@@ -21,8 +21,8 @@
 Databricks SQL warehouse lifecycle operators
 ============================================
 
-Use :class:`~airflow.providers.databricks.operators.databricks_warehouse.DatabricksStartWarehouseOperator`
-and :class:`~airflow.providers.databricks.operators.databricks_warehouse.DatabricksStopWarehouseOperator`
+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
@@ -46,8 +46,8 @@ Start a SQL warehouse
 Stop a SQL warehouse
 --------------------
 
-The system-test example marks the stop task with the ``all_done`` trigger rule so it runs as cleanup
-even when an upstream task fails.
+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

providers/databricks/provider.yaml

@@ -148,7 +148,7 @@ operators:
   - integration-name: Databricks SQL
     python-modules:
       - airflow.providers.databricks.operators.databricks_sql
-      - airflow.providers.databricks.operators.databricks_warehouse
+      - airflow.providers.databricks.operators.warehouse
   - integration-name: Databricks Repos
     python-modules:
       - airflow.providers.databricks.operators.databricks_repos

providers/databricks/src/airflow/providers/databricks/get_provider_info.py

@@ -77,7 +77,7 @@ def get_provider_info():
                 "integration-name": "Databricks SQL",
                 "python-modules": [
                     "airflow.providers.databricks.operators.databricks_sql",
-                    "airflow.providers.databricks.operators.databricks_warehouse",
+                    "airflow.providers.databricks.operators.warehouse",
                 ],
             },
             {

providers/databricks/src/airflow/providers/databricks/hooks/databricks.py

@@ -307,13 +307,6 @@ def __hash__(self):
     def __repr__(self) -> str:
         return str(self.__dict__)
 
-    def to_json(self) -> str:
-        return json.dumps(self.__dict__)
-
-    @classmethod
-    def from_json(cls, data: str) -> WarehouseState:
-        return WarehouseState(**json.loads(data))
-
 
 class DatabricksHook(BaseDatabricksHook):
     """

…bricks/operators/databricks_warehouse.py …viders/databricks/operators/warehouse.pyproviders/databricks/src/airflow/providers/databricks/operators/databricks_warehouse.py renamed to providers/databricks/src/airflow/providers/databricks/operators/warehouse.py providers/databricks/src/airflow/providers/databricks/operators/databricks_warehouse.py renamed to providers/databricks/src/airflow/providers/databricks/operators/warehouse.py

@@ -22,7 +22,7 @@
 from datetime import datetime
 
 from airflow.providers.common.compat.sdk import DAG
-from airflow.providers.databricks.operators.databricks_warehouse import (
+from airflow.providers.databricks.operators.warehouse import (
     DatabricksStartWarehouseOperator,
     DatabricksStopWarehouseOperator,
 )

providers/databricks/tests/system/databricks/example_databricks_sql_warehouse.py

@@ -1614,15 +1614,6 @@ def test_warehouse_state_unexpected_raises_value_error(self):
         with pytest.raises(ValueError, match="Unexpected warehouse state: FOO"):
             WarehouseState("FOO")
 
-    def test_warehouse_state_json_roundtrip(self):
-        state = WarehouseState("STARTING")
-        restored = WarehouseState.from_json(state.to_json())
-
-        assert restored == state
-        assert restored != object()
-        assert hash(restored) == hash(state)
-        assert repr(restored) == "{'state': 'STARTING'}"
-
 
 class TestRunState:
     def test_is_terminal_true(self):

providers/databricks/tests/unit/databricks/hooks/test_databricks.py

@@ -22,7 +22,7 @@
 
 from airflow.providers.databricks.exceptions import DatabricksWarehouseError
 from airflow.providers.databricks.hooks.databricks import DatabricksHook, WarehouseState
-from airflow.providers.databricks.operators.databricks_warehouse import (
+from airflow.providers.databricks.operators.warehouse import (
     DatabricksStartWarehouseOperator,
     DatabricksStopWarehouseOperator,
 )
@@ -37,7 +37,7 @@ class TestDatabricksStartWarehouseOperator:
         ["STARTING", "STOPPED"],
         ids=["transitioning", "stale-stopped"],
     )
-    @mock.patch("airflow.providers.databricks.operators.databricks_warehouse.time.sleep")
+    @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)
@@ -87,7 +87,7 @@ def test_skips_start_when_starting_without_waiting(self, mock_hook_property):
         hook.start_warehouse.assert_not_called()
         hook.get_warehouse_state.assert_called_once_with(WAREHOUSE_ID)
 
-    @mock.patch("airflow.providers.databricks.operators.databricks_warehouse.time.sleep")
+    @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)
@@ -111,7 +111,7 @@ def test_starts_while_stopping_and_waits_until_running(self, mock_hook_property,
 
 
 class TestDatabricksStopWarehouseOperator:
-    @mock.patch("airflow.providers.databricks.operators.databricks_warehouse.time.sleep")
+    @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)
@@ -206,7 +206,7 @@ def test_execute_raises_on_failure_state(
             (DatabricksStopWarehouseOperator, "STOPPING", "STOPPED", "stop_warehouse"),
         ],
     )
-    @mock.patch("airflow.providers.databricks.operators.databricks_warehouse.time.sleep")
+    @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(
@@ -271,9 +271,9 @@ def test_transition_without_waiting(
         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.databricks_warehouse.time.sleep")
+    @mock.patch("airflow.providers.databricks.operators.warehouse.time.sleep")
     @mock.patch(
-        "airflow.providers.databricks.operators.databricks_warehouse.time.monotonic",
+        "airflow.providers.databricks.operators.warehouse.time.monotonic",
         return_value=0,
     )
     @mock.patch.object(DatabricksStartWarehouseOperator, "_hook", new_callable=mock.PropertyMock)
@@ -304,9 +304,9 @@ def advance_to_deadline(seconds):
         hook.get_warehouse_state.assert_called_once_with(WAREHOUSE_ID)
         mock_sleep.assert_called_once_with(10)
 
-    @mock.patch("airflow.providers.databricks.operators.databricks_warehouse.time.sleep")
+    @mock.patch("airflow.providers.databricks.operators.warehouse.time.sleep")
     @mock.patch(
-        "airflow.providers.databricks.operators.databricks_warehouse.time.monotonic",
+        "airflow.providers.databricks.operators.warehouse.time.monotonic",
         return_value=0,
     )
     @mock.patch.object(DatabricksStartWarehouseOperator, "_hook", new_callable=mock.PropertyMock)
@@ -332,9 +332,9 @@ def get_state_after_deadline(_):
         hook.get_warehouse_state.assert_called_once_with(WAREHOUSE_ID)
         mock_sleep.assert_not_called()
 
-    @mock.patch("airflow.providers.databricks.operators.databricks_warehouse.time.sleep")
+    @mock.patch("airflow.providers.databricks.operators.warehouse.time.sleep")
     @mock.patch(
-        "airflow.providers.databricks.operators.databricks_warehouse.time.monotonic",
+        "airflow.providers.databricks.operators.warehouse.time.monotonic",
         return_value=0,
     )
     @mock.patch.object(DatabricksStartWarehouseOperator, "_hook", new_callable=mock.PropertyMock)
@@ -384,7 +384,7 @@ def test_invalid_warehouse_id(self, operator_class, warehouse_id):
         with pytest.raises(ValueError, match="warehouse_id must be provided"):
             operator.execute(None)
 
-    @mock.patch("airflow.providers.databricks.operators.databricks_warehouse.DatabricksHook", autospec=True)
+    @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(