SameerMesiah97 commented on code in PR #76504: URL: https://github.com/apache/airflow/pull/76504#discussion_r4503833229
########## 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``. + +The connection ID, warehouse ID, and nested ``json`` values support templating. +For numeric or boolean Jinja expressions, +set ``render_template_as_native_obj=True`` on the Dag. + +The task returns ``None`` without a ``return_value`` XCom, so pass the warehouse ID directly to +downstream tasks. It completes when the edit request succeeds, without waiting for compute changes +to finish. If the workflow needs a running warehouse, +chain a :class:`~airflow.providers.databricks.operators.warehouse.DatabricksStartWarehouseOperator` +after the edit task. Reaching ``RUNNING`` checks availability; it does not verify that all requested +configuration changes have finished applying. Review Comment: I think this part could be a bit more readable. Please see the below: ``` The operator returns ``None`` and completes when the edit request succeeds, without waiting for configuration changes to finish applying. Pass the warehouse ID explicitly to downstream tasks. If needed, follow it with :class:`~airflow.providers.databricks.operators.warehouse.DatabricksStartWarehouseOperator` to wait for warehouse availability. A ``RUNNING`` state does not guarantee that all configuration changes have been applied. ``` ########## 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: Agreed. ########## 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: Agreed. And to add to that. I think there is too much unnecessary context in this entire docstring. A simple wrapper operator should not require an entire essay in the docstring. I would try to trim it to 2/3 of its current length and retain only information the user would defnitely need to know. You dont need to dive deep into the rationale/implementation details bejind each flag. ########## 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: Agreed. Also, why not restore the warehouse to its existing auto stop setting in the DAG run? It appears stopping a warehouse wont delete it so you dont want to inadvertently change an existing warehouse's settings indefinitely just to run the example DAG. ########## 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: Agreed. -- 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]
