shahar1 commented on code in PR #76504:
URL: https://github.com/apache/airflow/pull/76504#discussion_r4489941513
##########
providers/databricks/docs/operators/warehouse.rst:
##########
@@ -17,6 +17,7 @@
.. _howto/operator:DatabricksStartWarehouseOperator:
.. _howto/operator:DatabricksStopWarehouseOperator:
+.. _howto/operator:DatabricksEditWarehouseOperator:
Databricks SQL warehouse lifecycle operators
Review Comment:
nit: The page now covers editing too, but the title still says "lifecycle"
and the intro below only introduces the start and stop operators ("Both
operators...", "each operator waits for the requested state"). Those paragraphs
then sit above the Edit section, where they don't apply. Could the title become
"Databricks SQL warehouse operators", with the start/stop waiting and retry
paragraphs moved under their own heading? `providers/AGENTS.md` asks to "Keep
`provider.yaml` metadata, docs, and tests in sync."
##########
providers/databricks/src/airflow/providers/databricks/operators/warehouse.py:
##########
@@ -207,3 +230,65 @@ def execute(self, context: Context) -> None:
if state.state != "STOPPING":
self._hook.stop_warehouse(self.warehouse_id)
self._wait_or_defer("STOPPED")
+
+
+class DatabricksEditWarehouseOperator(_DatabricksWarehouseBaseOperator):
+ """
+ Update the configuration of an existing Databricks SQL warehouse.
+
+ Calls the SQL Warehouses edit API and returns once the request succeeds,
without
+ waiting for the warehouse to reach a particular state. Use
+ :class:`DatabricksStartWarehouseOperator` to start the warehouse and wait
for ``RUNNING``.
+
+ :param warehouse_id: Required ID of the Databricks SQL warehouse. Has no
default. (templated)
+ :param json: Required, non-empty dictionary containing the SQL Warehouses
edit API request body,
+ with no default. Requiring at least one setting is an operator
validation. For example,
Review Comment:
nit: "Requiring at least one setting is an operator validation" and the "Has
no default" phrases (also in the hook docstring) restate what the signature and
`execute` already show. The start and stop docstrings in this file are much
shorter, and trimming this one to match would make it easier to read. (Reviewer
observation, not a project rule.)
##########
providers/databricks/tests/system/databricks/example_databricks_sql_warehouse.py:
##########
@@ -37,6 +38,15 @@
tags=["example"],
catchup=False,
) as dag:
+ # [START howto_operator_databricks_edit_sql_warehouse]
+ edit_warehouse = DatabricksEditWarehouseOperator(
+ task_id="edit_warehouse",
+ databricks_conn_id="databricks_default",
+ warehouse_id=WAREHOUSE_ID,
+ json={"auto_stop_mins": 10},
Review Comment:
nit: Running this system test now leaves `auto_stop_mins=10` on the
warehouse named by the environment variable, and nothing restores it. That
warehouse is pre-existing rather than created by the test, so it may be worth a
short note in the module docstring. (Reviewer observation, not a project rule.)
##########
providers/databricks/docs/operators/warehouse.rst:
##########
@@ -42,6 +43,31 @@ Repeated task attempts are safe: an already running
warehouse is not started aga
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.
+Edit a SQL warehouse
+--------------------
+
+Use
:class:`~airflow.providers.databricks.operators.warehouse.DatabricksEditWarehouseOperator`
+to change the configuration of an existing warehouse. Pass its
``warehouse_id`` and a non-empty
+``json`` dictionary containing the request body for the
+`SQL Warehouses edit API
<https://docs.databricks.com/api/warehouses/v1/edit-warehouse>`_.
+For example, you can set ``cluster_size``, ``auto_stop_mins``, or ``tags``.
Review Comment:
nit: It would help to state that settings omitted from `json` keep their
current values. The REST reference lists defaults for the same fields (for
example `auto_stop_mins` "Defaults to 120 mins") without saying how an edit
treats omitted ones, so a reader may worry a partial body resets them. The
[Databricks SDK
docs](https://databricks.github.io/sdk-js/api/interfaces/_databricks_sdk-warehouses.v1.EditWarehouseRequest.html)
do spell it out: "If a field is unset, the existing configuration value in the
SQL warehouse is retained."
##########
providers/databricks/src/airflow/providers/databricks/operators/warehouse.py:
##########
@@ -83,6 +71,41 @@ def _hook(self) -> DatabricksHook:
caller=self.__class__.__name__,
)
+
+class _DatabricksWarehouseStateOperator(_DatabricksWarehouseBaseOperator):
+ """Add state polling and deferred execution to warehouse operations that
support waiting."""
+
+ def __init__(
Review Comment:
nit: This class defines no `template_fields` of its own, so it could declare
only `wait_for_termination`, `polling_period_seconds`, `timeout` and
`deferrable`, and forward the rest with `super().__init__(**kwargs)`. That is
the `AwsBaseOperator` subclass pattern the PR description cites, and it avoids
repeating the base signature and its defaults. (Reviewer observation, not a
project rule.)
##########
providers/databricks/tests/unit/databricks/operators/test_warehouse.py:
##########
@@ -36,6 +37,111 @@
WAREHOUSE_ID = "wh-1"
+class TestDatabricksEditWarehouseOperator:
+
@mock.patch("airflow.providers.databricks.operators.warehouse.DatabricksHook",
autospec=True)
+ def test_execute(self, mock_hook):
+ payload = {"cluster_size": "Small", "auto_stop_mins": 0,
"enable_photon": False}
+ retry_args = {"stop": stop_after_attempt(5)}
+ operator = DatabricksEditWarehouseOperator(
+ task_id=TASK_ID,
+ warehouse_id=WAREHOUSE_ID,
+ json=payload,
+ databricks_conn_id="custom_databricks",
+ databricks_retry_limit=5,
+ databricks_retry_delay=2,
+ databricks_retry_args=retry_args,
+ )
+
+ assert operator.execute(None) is None
+
+ mock_hook.assert_called_once_with(
+ "custom_databricks",
+ retry_limit=5,
+ retry_delay=2,
+ retry_args=retry_args,
+ caller="DatabricksEditWarehouseOperator",
+ )
+ assert mock_hook.return_value.mock_calls ==
[mock.call.edit_warehouse(WAREHOUSE_ID, payload)]
+
+
@mock.patch("airflow.providers.databricks.operators.warehouse.DatabricksHook",
autospec=True)
+ def test_execute_renders_templates(self, mock_hook):
+ with DAG(dag_id="edit_warehouse", render_template_as_native_obj=True):
+ operator = DatabricksEditWarehouseOperator(
+ task_id=TASK_ID,
+ warehouse_id="{{ params.warehouse_id }}",
+ databricks_conn_id="{{ params.conn_id }}",
+ json={
+ "auto_stop_mins": "{{ params.auto_stop_mins }}",
+ "enable_photon": "{{ params.enable_photon }}",
+ "tags": {"custom_tags": [{"key": "run_date", "value": "{{
ds }}"}]},
+ },
+ )
+ operator.render_template_fields(
+ context={
+ "ds": "2026-11-04",
+ "params": {
+ "warehouse_id": WAREHOUSE_ID,
+ "conn_id": "rendered_databricks",
+ "auto_stop_mins": 0,
+ "enable_photon": False,
+ },
+ }
+ )
+
+ operator.execute(None)
+
+ assert mock_hook.call_args.args == ("rendered_databricks",)
+ mock_hook.return_value.edit_warehouse.assert_called_once_with(
+ WAREHOUSE_ID,
+ {
+ "auto_stop_mins": 0,
+ "enable_photon": False,
+ "tags": {"custom_tags": [{"key": "run_date", "value":
"2026-11-04"}]},
+ },
+ )
+
+
@mock.patch("airflow.providers.databricks.operators.warehouse.DatabricksHook",
autospec=True)
+ def test_rejects_empty_rendered_warehouse_id(self, mock_hook):
+ operator = DatabricksEditWarehouseOperator(
+ task_id=TASK_ID, warehouse_id="{{ params.warehouse_id }}",
json={"auto_stop_mins": 10}
+ )
+ operator.render_template_fields(context={"params": {"warehouse_id":
""}})
+
+ with pytest.raises(ValueError, match="warehouse_id must be provided"):
+ operator.execute(None)
+
+ mock_hook.assert_not_called()
+
+ @pytest.mark.parametrize("payload", [None, [], '{"auto_stop_mins": 10}'])
+
@mock.patch("airflow.providers.databricks.operators.warehouse.DatabricksHook",
autospec=True)
+ def test_rejects_non_dictionary_payload(self, mock_hook, payload):
+ operator = DatabricksEditWarehouseOperator(task_id=TASK_ID,
warehouse_id=WAREHOUSE_ID, json=payload)
+
+ with pytest.raises(TypeError, match="json must be a dictionary"):
+ operator.execute(None)
+
+ mock_hook.assert_not_called()
+
+
@mock.patch("airflow.providers.databricks.operators.warehouse.DatabricksHook",
autospec=True)
+ def test_rejects_empty_payload(self, mock_hook):
+ operator = DatabricksEditWarehouseOperator(task_id=TASK_ID,
warehouse_id=WAREHOUSE_ID, json={})
+
+ with pytest.raises(ValueError, match="json must contain at least one
warehouse setting"):
+ operator.execute(None)
+
+ mock_hook.assert_not_called()
+
+
@mock.patch("airflow.providers.databricks.operators.warehouse.DatabricksHook",
autospec=True)
+ def test_propagates_edit_failure(self, mock_hook):
Review Comment:
nit: The operator has no exception handling, so this test only exercises
Python's own propagation and could be dropped. `AGENTS.md`: "Do not add tests
for pre-existing logic that was already present before the PR, and do not test
standard-library or third-party functions."
--
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]
