diff --git a/airflow/providers/alibaba/cloud/operators/analyticdb_spark.py b/airflow/providers/alibaba/cloud/operators/analyticdb_spark.py index 0fd65e1ac6c80..79834dbe65765 100644 --- a/airflow/providers/alibaba/cloud/operators/analyticdb_spark.py +++ b/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 diff --git a/tests/providers/alibaba/cloud/operators/test_analyticdb_spark.py b/tests/providers/alibaba/cloud/operators/test_analyticdb_spark.py index eb2db3ff39567..c6d960b3ef337 100644 --- a/tests/providers/alibaba/cloud/operators/test_analyticdb_spark.py +++ b/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