Consolidate hook management in AnalyticDBSparkBaseOperator

airflow/providers/alibaba/cloud/operators/analyticdb_spark.py

@@ -21,7 +21,9 @@
 from time import sleep
 from typing import TYPE_CHECKING, Any, Sequence
 
-from airflow.exceptions import AirflowException
+from deprecated.classic import deprecated
+
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.models import BaseOperator
 from airflow.providers.alibaba.cloud.hooks.analyticdb_spark import AnalyticDBSparkHook, AppState
 
@@ -48,16 +50,15 @@ def __init__(
         self._adb_spark_conn_id = adb_spark_conn_id
         self._region = region
 
-        self._adb_spark_hook: AnalyticDBSparkHook | None = None
-
     @cached_property
+    def hook(self) -> AnalyticDBSparkHook:
+        """Get valid hook."""
+        return AnalyticDBSparkHook(adb_spark_conn_id=self._adb_spark_conn_id, region=self._region)
+
+    @deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning)
     def get_hook(self) -> AnalyticDBSparkHook:
         """Get valid hook."""
-        if self._adb_spark_hook is None or not isinstance(self._adb_spark_hook, AnalyticDBSparkHook):
-            self._adb_spark_hook = AnalyticDBSparkHook(
-                adb_spark_conn_id=self._adb_spark_conn_id, region=self._region
-            )
-        return self._adb_spark_hook
+        return self.hook
 
     def execute(self, context: Context) -> Any:
         ...
@@ -74,17 +75,18 @@ def poll_for_termination(self, app_id: str) -> None:
 
         :param app_id: id of the spark application to monitor
         """
-        hook = self.get_hook
-        state = hook.get_spark_state(app_id)
+        state = self.hook.get_spark_state(app_id)
         while AppState(state) not in AnalyticDBSparkHook.TERMINAL_STATES:
             self.log.debug("Application with id %s is in state: %s", app_id, state)
             sleep(self.polling_interval)
-            state = hook.get_spark_state(app_id)
+            state = self.hook.get_spark_state(app_id)
         self.log.info("Application with id %s terminated with state: %s", app_id, state)
         self.log.info(
-            "Web ui address is %s for application with id %s", hook.get_spark_web_ui_address(app_id), app_id
+            "Web ui address is %s for application with id %s",
+            self.hook.get_spark_web_ui_address(app_id),
+            app_id,
         )
-        self.log.info(hook.get_spark_log(app_id))
+        self.log.info(self.hook.get_spark_log(app_id))
         if AppState(state) != AppState.COMPLETED:
             raise AirflowException(f"Application {app_id} did not succeed")
 
@@ -94,7 +96,7 @@ def on_kill(self) -> None:
     def kill(self) -> None:
         """Delete the specified application."""
         if self.app_id is not None:
-            self.get_hook.kill_spark_app(self.app_id)
+            self.hook.kill_spark_app(self.app_id)
 
 
 class AnalyticDBSparkSQLOperator(AnalyticDBSparkBaseOperator):
@@ -142,7 +144,7 @@ def __init__(
         self._rg_name = rg_name
 
     def execute(self, context: Context) -> Any:
-        submit_response = self.get_hook.submit_spark_sql(
+        submit_response = self.hook.submit_spark_sql(
             cluster_id=self._cluster_id, rg_name=self._rg_name, **self.spark_params
         )
         self.app_id = submit_response.body.data.app_id
@@ -213,7 +215,7 @@ def __init__(
         self._rg_name = rg_name
 
     def execute(self, context: Context) -> Any:
-        submit_response = self.get_hook.submit_spark_app(
+        submit_response = self.hook.submit_spark_app(
             cluster_id=self._cluster_id, rg_name=self._rg_name, **self.spark_params
         )
         self.app_id = submit_response.body.data.app_id

tests/providers/alibaba/cloud/operators/test_analyticdb_spark.py

@@ -51,10 +51,10 @@ def setup_method(self):
     @mock.patch(ADB_SPARK_OPERATOR_STRING.format("AnalyticDBSparkHook"))
     def test_get_hook(self, mock_hook):
         """Test get_hook function works as expected."""
-        self.operator.get_hook()
+        self.operator.hook
         mock_hook.assert_called_once_with(adb_spark_conn_id=MOCK_ADB_SPARK_CONN_ID, region=MOCK_REGION)
 
-    @mock.patch(ADB_SPARK_OPERATOR_STRING.format("AnalyticDBSparkBaseOperator.get_hook"))
+    @mock.patch(ADB_SPARK_OPERATOR_STRING.format("AnalyticDBSparkBaseOperator.hook"))
     def test_poll_for_termination(self, mock_hook):
         """Test poll_for_termination works as expected with COMPLETED application."""
         # Given
@@ -63,7 +63,7 @@ def test_poll_for_termination(self, mock_hook):
         # When
         self.operator.poll_for_termination(MOCK_APP_ID)
 
-    @mock.patch(ADB_SPARK_OPERATOR_STRING.format("AnalyticDBSparkBaseOperator.get_hook"))
+    @mock.patch(ADB_SPARK_OPERATOR_STRING.format("AnalyticDBSparkBaseOperator.hook"))
     def test_poll_for_termination_with_exception(self, mock_hook):
         """Test poll_for_termination raises AirflowException with FATAL application."""
         # Given
@@ -107,7 +107,7 @@ def test_execute(self, mock_hook):
             name=None,
         )
 
-    @mock.patch(ADB_SPARK_OPERATOR_STRING.format("AnalyticDBSparkBaseOperator.get_hook"))
+    @mock.patch(ADB_SPARK_OPERATOR_STRING.format("AnalyticDBSparkBaseOperator.hook"))
     def test_execute_with_exception(self, mock_hook):
         """Test submit AnalyticDB Spark Batch Application raises ValueError with invalid parameter."""
         # Given
@@ -155,7 +155,7 @@ def test_execute(self, mock_hook):
             name=None,
         )
 
-    @mock.patch(ADB_SPARK_OPERATOR_STRING.format("AnalyticDBSparkBaseOperator.get_hook"))
+    @mock.patch(ADB_SPARK_OPERATOR_STRING.format("AnalyticDBSparkBaseOperator.hook"))
     def test_execute_with_exception(self, mock_hook):
         """Test submit AnalyticDB Spark SQL Application raises ValueError with invalid parameter."""
         # Given