From 60498102544005ff534fa978f47cd271c5778a97 Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Sun, 17 Sep 2023 01:53:35 +0200 Subject: [PATCH 1/2] Consolidate hook management in AnalyticDBSparkSensor --- .../alibaba/cloud/sensors/analyticdb_spark.py | 16 +++++++++------- .../cloud/sensors/test_analyticdb_spark.py | 6 +++--- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/airflow/providers/alibaba/cloud/sensors/analyticdb_spark.py b/airflow/providers/alibaba/cloud/sensors/analyticdb_spark.py index fb6a962d43798..fc78747b9f844 100644 --- a/airflow/providers/alibaba/cloud/sensors/analyticdb_spark.py +++ b/airflow/providers/alibaba/cloud/sensors/analyticdb_spark.py @@ -20,6 +20,8 @@ from functools import cached_property from typing import TYPE_CHECKING, Any, Sequence +from deprecated.classic import deprecated + from airflow.providers.alibaba.cloud.hooks.analyticdb_spark import AnalyticDBSparkHook, AppState from airflow.sensors.base import BaseSensorOperator @@ -50,19 +52,19 @@ def __init__( self.app_id = app_id self._region = region self._adb_spark_conn_id = adb_spark_conn_id - 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.") 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 poke(self, context: Context) -> bool: app_id = self.app_id - state = self.get_hook.get_spark_state(app_id) + state = self.hook.get_spark_state(app_id) return AppState(state) in AnalyticDBSparkHook.TERMINAL_STATES diff --git a/tests/providers/alibaba/cloud/sensors/test_analyticdb_spark.py b/tests/providers/alibaba/cloud/sensors/test_analyticdb_spark.py index 8cef5175005f3..86c67190635ac 100644 --- a/tests/providers/alibaba/cloud/sensors/test_analyticdb_spark.py +++ b/tests/providers/alibaba/cloud/sensors/test_analyticdb_spark.py @@ -42,10 +42,10 @@ def setup_method(self): @mock.patch(ADB_SPARK_SENSOR_STRING.format("AnalyticDBSparkHook")) def test_get_hook(self, mock_service): """Test get_hook function works as expected.""" - self.sensor.get_hook() + self.sensor.hook mock_service.assert_called_once_with(adb_spark_conn_id=MOCK_ADB_SPARK_CONN_ID, region=MOCK_REGION) - @mock.patch(ADB_SPARK_SENSOR_STRING.format("AnalyticDBSparkSensor.get_hook")) + @mock.patch(ADB_SPARK_SENSOR_STRING.format("AnalyticDBSparkSensor.hook")) def test_poke_terminal_state(self, mock_service): """Test poke_terminal_state works as expected with COMPLETED application.""" # Given @@ -58,7 +58,7 @@ def test_poke_terminal_state(self, mock_service): assert res is True mock_service.get_spark_state.assert_called_once_with(MOCK_ADB_SPARK_ID) - @mock.patch(ADB_SPARK_SENSOR_STRING.format("AnalyticDBSparkSensor.get_hook")) + @mock.patch(ADB_SPARK_SENSOR_STRING.format("AnalyticDBSparkSensor.hook")) def test_poke_non_terminal_state(self, mock_service): """Test poke_terminal_state works as expected with RUNNING application.""" # Given From d4a47a3e75dd68d9097c2633e4e3da2d62f9aee0 Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Sun, 17 Sep 2023 12:43:17 +0200 Subject: [PATCH 2/2] use AirflowProviderDeprecationWarning --- airflow/providers/alibaba/cloud/sensors/analyticdb_spark.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow/providers/alibaba/cloud/sensors/analyticdb_spark.py b/airflow/providers/alibaba/cloud/sensors/analyticdb_spark.py index fc78747b9f844..02bab417b6d7a 100644 --- a/airflow/providers/alibaba/cloud/sensors/analyticdb_spark.py +++ b/airflow/providers/alibaba/cloud/sensors/analyticdb_spark.py @@ -22,6 +22,7 @@ from deprecated.classic import deprecated +from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.alibaba.cloud.hooks.analyticdb_spark import AnalyticDBSparkHook, AppState from airflow.sensors.base import BaseSensorOperator @@ -58,7 +59,7 @@ 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.") + @deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning) def get_hook(self) -> AnalyticDBSparkHook: """Get valid hook.""" return self.hook