From 6cc987e27ec0d7578ab586d7a23d6f2360b3d228 Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Sun, 17 Sep 2023 02:03:37 +0200 Subject: [PATCH 1/2] Consolidate hook management in AzureBatchOperator --- airflow/providers/microsoft/azure/operators/batch.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/batch.py b/airflow/providers/microsoft/azure/operators/batch.py index 06122ed1f9f00..e54177676b8a6 100644 --- a/airflow/providers/microsoft/azure/operators/batch.py +++ b/airflow/providers/microsoft/azure/operators/batch.py @@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, Any, Sequence from azure.batch import models as batch_models +from deprecated.classic import deprecated from airflow.exceptions import AirflowException from airflow.models import BaseOperator @@ -181,7 +182,12 @@ def __init__( @cached_property def hook(self) -> AzureBatchHook: """Create and return an AzureBatchHook (cached).""" - return self.get_hook() + return AzureBatchHook(self.azure_batch_conn_id) + + @deprecated(reason="use `hook` property instead.") + def get_hook(self) -> AzureBatchHook: + """Create and return an AzureBatchHook.""" + return self.hook def _check_inputs(self) -> Any: if not self.os_family and not self.vm_publisher: @@ -315,10 +321,6 @@ def on_kill(self) -> None: ) self.log.info("Azure Batch job (%s) terminated: %s", self.batch_job_id, response) - def get_hook(self) -> AzureBatchHook: - """Create and return an AzureBatchHook.""" - return AzureBatchHook(azure_batch_conn_id=self.azure_batch_conn_id) - def clean_up(self, pool_id: str | None = None, job_id: str | None = None) -> None: """ Delete the given pool and job in the batch account. From 8b15cbf0543d71c7fcb7c96203b1aab05fb237c5 Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Sun, 17 Sep 2023 12:19:38 +0200 Subject: [PATCH 2/2] use AirflowProviderDeprecationWarning --- airflow/providers/microsoft/azure/operators/batch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/batch.py b/airflow/providers/microsoft/azure/operators/batch.py index e54177676b8a6..0c1a0849b5b05 100644 --- a/airflow/providers/microsoft/azure/operators/batch.py +++ b/airflow/providers/microsoft/azure/operators/batch.py @@ -23,7 +23,7 @@ from azure.batch import models as batch_models from deprecated.classic import deprecated -from airflow.exceptions import AirflowException +from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.models import BaseOperator from airflow.providers.microsoft.azure.hooks.batch import AzureBatchHook @@ -184,7 +184,7 @@ def hook(self) -> AzureBatchHook: """Create and return an AzureBatchHook (cached).""" return AzureBatchHook(self.azure_batch_conn_id) - @deprecated(reason="use `hook` property instead.") + @deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning) def get_hook(self) -> AzureBatchHook: """Create and return an AzureBatchHook.""" return self.hook