Vamsi-klu commented on code in PR #70088:
URL: https://github.com/apache/airflow/pull/70088#discussion_r3746565847
##########
providers/databricks/src/airflow/providers/databricks/operators/warehouse.py:
##########
Review Comment:
Renamed in `727bfc3`. It is `operators/warehouse.py` and
`tests/unit/databricks/operators/test_warehouse.py` now, and I updated
everything that pointed at the old name: `provider.yaml`, the regenerated
`get_provider_info`, the two `:class:` references in the docs, the `mock.patch`
targets in the unit tests and the system example import. `git grep
databricks_warehouse` across the provider comes back empty.
##########
providers/databricks/tests/unit/databricks/operators/test_warehouse.py:
##########
Review Comment:
Renamed alongside the operator module in `727bfc3`, see the reply on the
other file-name thread for the full list of what got updated.
##########
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:
+ """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__)
+
+ def to_json(self) -> str:
+ return json.dumps(self.__dict__)
+
+ @classmethod
+ def from_json(cls, data: str) -> WarehouseState:
+ return WarehouseState(**json.loads(data))
Review Comment:
Nothing uses them, you are right. I added them for symmetry with `RunState`
and `SQLStatementState`, whose `to_json` and `from_json` are used by the
deferrable triggers, but `WarehouseState` has no trigger yet so they were dead
code from day one. Removed both along with the round-trip test in `727bfc3`.
They will come back with `DatabricksWarehouseStateTrigger` when the deferrable
work lands. `is_deleted` stays, since `_wait_for_state` uses it now.
##########
providers/databricks/docs/operators/sql_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.databricks_warehouse.DatabricksStartWarehouseOperator`
+and
:class:`~airflow.providers.databricks.operators.databricks_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 system-test example marks the stop task with the ``all_done`` trigger rule
so it runs as cleanup
Review Comment:
Reworded in `727bfc3`. It now just says the stop task uses the `all_done`
trigger rule so the warehouse is stopped even when upstream tasks fail, with no
mention of system tests.
--
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]
