From 5d1b136698402369ea7c5e2e77ea3c1b69d7f19f Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Fri, 13 Aug 2021 23:46:34 -0400 Subject: [PATCH 01/27] Initial commit --- .../example_dags/example_adf_run_pipeline.py | 68 +++++++ .../azure/hooks/azure_data_factory.py | 69 ++++--- .../azure/operators/azure_data_factory.py | 169 ++++++++++++++++++ .../providers/microsoft/azure/provider.yaml | 6 + .../azure/sensors/azure_data_factory.py | 89 +++++++++ .../azure/hooks/test_azure_data_factory.py | 16 +- 6 files changed, 376 insertions(+), 41 deletions(-) create mode 100644 airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py create mode 100644 airflow/providers/microsoft/azure/operators/azure_data_factory.py create mode 100644 airflow/providers/microsoft/azure/sensors/azure_data_factory.py diff --git a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py new file mode 100644 index 0000000000000..b43388d1cbc0a --- /dev/null +++ b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py @@ -0,0 +1,68 @@ +# 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. + +from datetime import datetime, timedelta + +from airflow.models import DAG +from airflow.operators.dummy import DummyOperator +from airflow.providers.microsoft.azure.operators.azure_data_factory import AzureDataFactoryRunPipelineOperator +from airflow.providers.microsoft.azure.sensors.azure_data_factory import ( + AzureDataFactoryPipelineRunStatusSensor, +) +from airflow.utils.edgemodifier import Label + +with DAG( + dag_id="example_adf_run_pipeline", + start_date=datetime(2021, 8, 13), + schedule_interval="@daily", + catchup=False, + default_args={ + "retries": 1, + "retry_delay": timedelta(minutes=3), + "conn_id": "azure_data_factory", + "factory_name": "my-data-factory", + "resource_group_name": "my-resource-group", + }, + default_view="graph", +) as dag: + begin = DummyOperator(task_id="begin") + end = DummyOperator(task_id="end") + + run_pipeline1 = AzureDataFactoryRunPipelineOperator( + task_id="run_pipeline1", + pipeline_name="pipeline1", + parameters={"myParam": "value"}, + ) + + run_pipeline2 = AzureDataFactoryRunPipelineOperator( + task_id="run_pipeline2", + pipeline_name="extractDailyExchangeRates", + do_asynchronous_wait=True, + ) + + wait_for_pipeline_run = AzureDataFactoryPipelineRunStatusSensor( + task_id="wait_for_pipeline_run", + run_id=run_pipeline2.output["run_id"], + poke_interval=10, + ) + + begin >> Label("Do async wait with sensor") >> run_pipeline2 + wait_for_pipeline_run >> end + begin >> Label("No async wait") >> run_pipeline1 >> end + + # Task dependency created via `XComArgs`: + # run_pipeline2 >> wait_for_pipeline_run diff --git a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py index 8a1cc36090ec4..daaeed0291229 100644 --- a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py @@ -23,14 +23,11 @@ from azure.mgmt.datafactory import DataFactoryManagementClient from azure.mgmt.datafactory.models import ( CreateRunResponse, - Dataset, DatasetResource, Factory, - LinkedService, LinkedServiceResource, PipelineResource, PipelineRun, - Trigger, TriggerResource, ) @@ -52,24 +49,38 @@ def wrapper(*args, **kwargs) -> Callable: bound_args = signature.bind(*args, **kwargs) def bind_argument(arg, default_key): - if arg not in bound_args.arguments: + # Check if arg was not included in the function signature or, if it is, that the value is not + # provided. + if arg not in bound_args.arguments or bound_args.arguments[arg] is None: self = args[0] conn = self.get_connection(self.conn_id) default_value = conn.extra_dejson.get(default_key) - if not default_value: raise AirflowException("Could not determine the targeted data factory.") bound_args.arguments[arg] = conn.extra_dejson[default_key] - bind_argument("resource_group_name", "resourceGroup") - bind_argument("factory_name", "factory") + bind_argument("resource_group_name", "extra__azure_data_factory__resource_group_name") + bind_argument("factory_name", "extra__azure_data_factory__factory_name") return func(*bound_args.args, **bound_args.kwargs) return wrapper +class AzureDataFactoryPipelineRunStatus: + """Azure Data Factory pipeline operation statuses.""" + + QUEUED = "Queued" + IN_PROGRESS = "InProgress" + SUCCEEDED = "Succeeded" + FAILED = "Failed" + CANCELING = "Canceling" + CANCELLED = "Cancelled" + + TERMINAL_STATUSES = {CANCELLED, FAILED, SUCCEEDED} + + class AzureDataFactoryHook(BaseHook): """ A hook to interact with Azure Data Factory. @@ -97,32 +108,22 @@ def get_connection_form_widgets() -> Dict[str, Any]: "extra__azure_data_factory__subscriptionId": StringField( lazy_gettext('Azure Subscription ID'), widget=BS3TextFieldWidget() ), + "extra__azure_data_factory__resource_group_name": StringField( + lazy_gettext('Resource Group Name'), widget=BS3TextFieldWidget() + ), + "extra__azure_data_factory__factory_name": StringField( + lazy_gettext('Factory Name'), widget=BS3TextFieldWidget() + ), } @staticmethod def get_ui_field_behaviour() -> Dict: """Returns custom field behaviour""" - import json - return { - "hidden_fields": ['schema', 'port', 'host'], + "hidden_fields": ['schema', 'port', 'host', 'extra'], "relabeling": { 'login': 'Azure Client ID', 'password': 'Azure Secret', - 'extra': 'Extra (optional)', - }, - "placeholders": { - 'extra': json.dumps( - { - "resourceGroup": "azure resource group name", - "factory": "azure data factory name", - }, - indent=1, - ), - 'login': 'client id', - 'password': 'secret', - 'extra__azure_data_factory__tenantId': 'tenant id', - 'extra__azure_data_factory__subscriptionId': 'subscription id', }, } @@ -136,12 +137,8 @@ def get_conn(self) -> DataFactoryManagementClient: return self._conn conn = self.get_connection(self.conn_id) - tenant = conn.extra_dejson.get('extra__azure_data_factory__tenantId') or conn.extra_dejson.get( - 'tenantId' - ) - subscription_id = conn.extra_dejson.get( - 'extra__azure_data_factory__subscriptionId' - ) or conn.extra_dejson.get('subscriptionId') + tenant = conn.extra_dejson.get('extra__azure_data_factory__tenantId') + subscription_id = conn.extra_dejson.get('extra__azure_data_factory__subscriptionId') self._conn = DataFactoryManagementClient( credential=ClientSecretCredential( @@ -273,7 +270,7 @@ def _linked_service_exists(self, resource_group_name, factory_name, linked_servi def update_linked_service( self, linked_service_name: str, - linked_service: LinkedService, + linked_service: LinkedServiceResource, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, **config: Any, @@ -300,7 +297,7 @@ def update_linked_service( def create_linked_service( self, linked_service_name: str, - linked_service: LinkedService, + linked_service: LinkedServiceResource, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, **config: Any, @@ -375,7 +372,7 @@ def _dataset_exists(self, resource_group_name, factory_name, dataset_name) -> bo def update_dataset( self, dataset_name: str, - dataset: Dataset, + dataset: DatasetResource, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, **config: Any, @@ -402,7 +399,7 @@ def update_dataset( def create_dataset( self, dataset_name: str, - dataset: Dataset, + dataset: DatasetResource, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, **config: Any, @@ -633,7 +630,7 @@ def _trigger_exists(self, resource_group_name, factory_name, trigger_name) -> bo def update_trigger( self, trigger_name: str, - trigger: Trigger, + trigger: TriggerResource, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, **config: Any, @@ -660,7 +657,7 @@ def update_trigger( def create_trigger( self, trigger_name: str, - trigger: Trigger, + trigger: TriggerResource, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, **config: Any, diff --git a/airflow/providers/microsoft/azure/operators/azure_data_factory.py b/airflow/providers/microsoft/azure/operators/azure_data_factory.py new file mode 100644 index 0000000000000..e13e3989dfc99 --- /dev/null +++ b/airflow/providers/microsoft/azure/operators/azure_data_factory.py @@ -0,0 +1,169 @@ +# 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. + +import time +from typing import Any, Dict, Optional + +from airflow.exceptions import AirflowException +from airflow.models import BaseOperator +from airflow.providers.microsoft.azure.hooks.azure_data_factory import ( + AzureDataFactoryHook, + AzureDataFactoryPipelineRunStatus, +) + + +class AzureDataFactoryRunPipelineOperator(BaseOperator): + """ + Executes a data factory pipeline. + + :param conn_id: The connection identifier for connecting to Azure Data Factory. + :type conn_id: str + :param pipeline_name: The name of the pipeline to execute. + :type pipeline_name: str + :param resource_group_name: The resource group name. If a value is not passed in to the operator, the + ``AzureDataFactoryHook`` will attempt to use the resource group name provided in the corresponding + connection. + :type resource_group_name: str + :param factory_name: The data factory name. If a value is not passed in to the operator, the + ``AzureDataFactoryHook`` will attempt to use the resource group name provided in the corresponding + connection. + :type factory_name: str + :param reference_pipeline_run_id: The pipeline run identifier. If this run ID is specified the parameters + of the specified run will be used to create a new run. + :type reference_pipeline_run_id: str + :param is_recovery: Recovery mode flag. If recovery mode is set to `True`, the specified referenced + pipeline run and the new run will be grouped under the same ``groupId``. + :type is_recovery: bool + :param start_activity_name: In recovery mode, the rerun will start from this activity. If not specified, + all activities will run. + :type start_activity_name: str + :param start_from_failure: In recovery mode, if set to true, the rerun will start from failed activities. + The property will be used only if ``start_activity_name`` is not specified. + :type start_from_failure: bool + :param parameters: Parameters of the pipeline run. These parameters will be used only if the + ``reference_pipeline_run_id`` is not specified. + :type start_from_failure: Dict[str, Any] + :param do_asynchronous_wait: Flag to exit after creating a pipeline run. Typically this flag would be + enabled to wait for a long-running pipeline execution using the + ``AzureDataFactoryPipelineRunStatusSensor`` rather than this operator. + :type do_asynchronous_wait: bool + :param status_check_timeout: Time in seconds to wait for a pipeline to reach a terminal status for + non-asynchronous waits. Use only if ``do_asynchronous_wait`` is False. + :type status_check_timeout: int + :param poke_interval: Time in seconds to check on a pipeline run's status for non-asynchronous waits. Use + only if ``do_asynchronous_wait`` is False. + :type poke_interval: int + """ + + template_fields = ( + "resource_group_name", + "factory_name", + "pipeline_name", + "reference_pipeline_run_id", + "parameters", + ) + template_fields_renderers = {"parameters": "json"} + + def __init__( + self, + *, + conn_id: str, + pipeline_name: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + reference_pipeline_run_id: Optional[str] = None, + is_recovery: Optional[bool] = None, + start_activity_name: Optional[str] = None, + start_from_failure: Optional[bool] = None, + parameters: Optional[Dict[str, Any]] = None, + do_asynchronous_wait: Optional[bool] = False, + status_check_timeout: Optional[int] = 60 * 60 * 5, + poke_interval: Optional[int] = 30, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.conn_id = conn_id + self.pipeline_name = pipeline_name + self.resource_group_name = resource_group_name + self.factory_name = factory_name + self.reference_pipeline_run_id = reference_pipeline_run_id + self.is_recovery = is_recovery + self.start_activity_name = start_activity_name + self.start_from_failure = start_from_failure + self.parameters = parameters + self.do_asynchronous_wait = do_asynchronous_wait + self.status_check_timeout = status_check_timeout + self.poke_interval = poke_interval + + def execute(self, context: Dict) -> None: + self.hook = AzureDataFactoryHook(conn_id=self.conn_id) + self.log.info(f"Executing the '{self.pipeline_name}' pipeline.") + response = self.hook.run_pipeline( + pipeline_name=self.pipeline_name, + resource_group_name=self.resource_group_name, + factory_name=self.factory_name, + reference_pipeline_run_id=self.reference_pipeline_run_id, + is_recovery=self.is_recovery, + start_activity_name=self.start_activity_name, + start_from_failure=self.start_from_failure, + parameters=self.parameters, + ) + self.run_id = vars(response)["run_id"] + # Push the ``run_id`` value to XCom regardless of what happens during execution. This allows users to + # retrieve the ``run_id`` as an output of this operator for downstream tasks especially if performing + # an async wait. + context["ti"].xcom_push(key="run_id", value=self.run_id) + + if not self.do_asynchronous_wait: + self.log.info(f"Waiting for run ID {self.run_id} of pipeline '{self.pipeline_name}' to complete.") + start_time = time.monotonic() + pipeline_run_status = None + while pipeline_run_status not in AzureDataFactoryPipelineRunStatus.TERMINAL_STATUSES: + # Check to see if the pipeline-run duration has exceeded the ``status_check_timeout`` + # configured. + if start_time + self.status_check_timeout < time.monotonic(): + raise AirflowException( + f"Pipeline run {self.run_id} has not reached a terminal status after " + + f"{self.status_check_timeout} seconds." + ) + + # Wait to check the status of the pipeline based on the ``poke_interval`` configured. + time.sleep(self.poke_interval) + self.log.info(f"Checking on the status of run ID {self.run_id}.") + pipeline_run = self.hook.get_pipeline_run( + run_id=self.run_id, + factory_name=self.factory_name, + resource_group_name=self.resource_group_name, + ) + pipeline_run_status = pipeline_run.status + self.log.info(f"Run ID {self.run_id} is in a status of {pipeline_run_status}.") + + if pipeline_run_status == AzureDataFactoryPipelineRunStatus.CANCELLED: + raise AirflowException(f"Pipeline run {self.run_id} has been cancelled.") + + if pipeline_run_status == AzureDataFactoryPipelineRunStatus.FAILED: + raise AirflowException(f"Pipeline run {self.run_id} has failed.") + + self.log.info(f"Pipeline run {self.run_id} has completed successfully.") + + def on_kill(self) -> None: + if self.run_id: + self.hook.cancel_pipeline_run( + run_id=self.run_id, + resource_group_name=self.resource_group_name, + factory_name=self.factory_name, + ) diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index 3622489f3cf5f..c3cc75e4b34c5 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -97,6 +97,9 @@ operators: - integration-name: Microsoft Azure Blob Storage python-modules: - airflow.providers.microsoft.azure.operators.wasb_delete_blob + - integration-name: Microsoft Azure Data Factory + python-modules: + - airflow.providers.microsoft.azure.operators.azure_data_factory sensors: - integration-name: Microsoft Azure Cosmos DB @@ -105,6 +108,9 @@ sensors: - integration-name: Microsoft Azure Blob Storage python-modules: - airflow.providers.microsoft.azure.sensors.wasb + - integration-name: Microsoft Azure Data Factory + python-modules: + - airflow.providers.microsoft.azure.sensors.azure_data_factory hooks: - integration-name: Microsoft Azure Container Instances diff --git a/airflow/providers/microsoft/azure/sensors/azure_data_factory.py b/airflow/providers/microsoft/azure/sensors/azure_data_factory.py new file mode 100644 index 0000000000000..2805f9879ed2f --- /dev/null +++ b/airflow/providers/microsoft/azure/sensors/azure_data_factory.py @@ -0,0 +1,89 @@ +# 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. + +from typing import Dict, Optional, Set, Union + +from airflow.exceptions import AirflowException +from airflow.providers.microsoft.azure.hooks.azure_data_factory import ( + AzureDataFactoryHook, + AzureDataFactoryPipelineRunStatus, +) +from airflow.sensors.base import BaseSensorOperator + + +class AzureDataFactoryPipelineRunStatusSensor(BaseSensorOperator): + """ + Checks the status of a pipeline run. + + :param conn_id: The connection identifier for connecting to Azure Data Factory. + :type conn_id: str + :param run_id: The pipeline run identifier. + :type run_id: str + :param resource_group_name: The resource group name. + :type resource_group_name: str + :param factory_name: The data factory name. + :type factory_name: str + :param expected_statuses: The status(es) which are desired for the pipeline run. + :type expected_statuses: str or Set[str] + """ + + template_fields = ("resource_group_name", "factory_name", "run_id") + + def __init__( + self, + *, + conn_id: str, + run_id: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + expected_statuses: Optional[Union[Set[str], str]] = AzureDataFactoryPipelineRunStatus.SUCCEEDED, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.conn_id = conn_id + self.run_id = run_id + self.resource_group_name = resource_group_name + self.factory_name = factory_name + self.expected_statuses = ( + {expected_statuses} if isinstance(expected_statuses, str) else expected_statuses + ) + + def poke(self, context: Dict) -> bool: + self.log.info( + f"Checking for pipeline run {self.run_id} to be in one of the following statuses: " + f"{', '.join(self.expected_statuses)}.", + ) + self.hook = AzureDataFactoryHook(conn_id=self.conn_id) + pipeline_run = self.hook.get_pipeline_run( + run_id=self.run_id, + factory_name=self.factory_name, + resource_group_name=self.resource_group_name, + ) + pipeline_run_status = pipeline_run.status + self.log.info(f"Current status for pipeline run {self.run_id}: {pipeline_run_status}.") + + if pipeline_run_status in self.expected_statuses: + return True + elif pipeline_run_status in { + AzureDataFactoryPipelineRunStatus.FAILED, + AzureDataFactoryPipelineRunStatus.CANCELLED, + }: + raise AirflowException( + f"Pipeline run {self.run_id} is in a terminal status: {pipeline_run_status}" + ) + + return False diff --git a/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py b/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py index 8c4685441bcfa..7939a798883d5 100644 --- a/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py @@ -49,10 +49,10 @@ def setup_module(): password="clientSecret", extra=json.dumps( { - "tenantId": "tenantId", - "subscriptionId": "subscriptionId", - "resourceGroup": DEFAULT_RESOURCE_GROUP, - "factory": DEFAULT_FACTORY, + "extra__azure_data_factory__tenantId": "tenantId", + "extra__azure_data_factory__subscriptionId": "subscriptionId", + "extra__azure_data_factory__resource_group_name": DEFAULT_RESOURCE_GROUP, + "extra__azure_data_factory__factory_name": DEFAULT_FACTORY, } ), ) @@ -100,8 +100,14 @@ def echo(_, resource_group_name=None, factory_name=None): conn.extra_dejson = {} assert provide_targeted_factory(echo)(hook, RESOURCE_GROUP, FACTORY) == (RESOURCE_GROUP, FACTORY) - conn.extra_dejson = {"resourceGroup": DEFAULT_RESOURCE_GROUP, "factory": DEFAULT_FACTORY} + conn.extra_dejson = { + "extra__azure_data_factory__resource_group_name": DEFAULT_RESOURCE_GROUP, + "extra__azure_data_factory__factory_name": DEFAULT_FACTORY, + } assert provide_targeted_factory(echo)(hook) == (DEFAULT_RESOURCE_GROUP, DEFAULT_FACTORY) + assert provide_targeted_factory(echo)(hook, RESOURCE_GROUP, None) == (RESOURCE_GROUP, DEFAULT_FACTORY) + assert provide_targeted_factory(echo)(hook, None, FACTORY) == (DEFAULT_RESOURCE_GROUP, FACTORY) + assert provide_targeted_factory(echo)(hook, None, None) == (DEFAULT_RESOURCE_GROUP, DEFAULT_FACTORY) with pytest.raises(AirflowException): conn.extra_dejson = {} From ef3f4938e476d493a2d8ee4933f4731298ce050c Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Sat, 14 Aug 2021 19:38:51 -0400 Subject: [PATCH 02/27] Rename operator parameter name --- .../microsoft/azure/hooks/azure_data_factory.py | 3 +-- .../azure/operators/azure_data_factory.py | 15 +++++++-------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py index daaeed0291229..c13f9751230fa 100644 --- a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py @@ -49,8 +49,7 @@ def wrapper(*args, **kwargs) -> Callable: bound_args = signature.bind(*args, **kwargs) def bind_argument(arg, default_key): - # Check if arg was not included in the function signature or, if it is, that the value is not - # provided. + # Check if arg was not included in the function signature or, if it is, the value is not provided. if arg not in bound_args.arguments or bound_args.arguments[arg] is None: self = args[0] conn = self.get_connection(self.conn_id) diff --git a/airflow/providers/microsoft/azure/operators/azure_data_factory.py b/airflow/providers/microsoft/azure/operators/azure_data_factory.py index e13e3989dfc99..999a248846b86 100644 --- a/airflow/providers/microsoft/azure/operators/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/operators/azure_data_factory.py @@ -61,9 +61,9 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): enabled to wait for a long-running pipeline execution using the ``AzureDataFactoryPipelineRunStatusSensor`` rather than this operator. :type do_asynchronous_wait: bool - :param status_check_timeout: Time in seconds to wait for a pipeline to reach a terminal status for + :param timeout: Time in seconds to wait for a pipeline to reach a terminal status for non-asynchronous waits. Use only if ``do_asynchronous_wait`` is False. - :type status_check_timeout: int + :type timeout: int :param poke_interval: Time in seconds to check on a pipeline run's status for non-asynchronous waits. Use only if ``do_asynchronous_wait`` is False. :type poke_interval: int @@ -91,7 +91,7 @@ def __init__( start_from_failure: Optional[bool] = None, parameters: Optional[Dict[str, Any]] = None, do_asynchronous_wait: Optional[bool] = False, - status_check_timeout: Optional[int] = 60 * 60 * 5, + timeout: Optional[int] = 60 * 60 * 5, poke_interval: Optional[int] = 30, **kwargs, ) -> None: @@ -106,7 +106,7 @@ def __init__( self.start_from_failure = start_from_failure self.parameters = parameters self.do_asynchronous_wait = do_asynchronous_wait - self.status_check_timeout = status_check_timeout + self.timeout = timeout self.poke_interval = poke_interval def execute(self, context: Dict) -> None: @@ -133,12 +133,11 @@ def execute(self, context: Dict) -> None: start_time = time.monotonic() pipeline_run_status = None while pipeline_run_status not in AzureDataFactoryPipelineRunStatus.TERMINAL_STATUSES: - # Check to see if the pipeline-run duration has exceeded the ``status_check_timeout`` - # configured. - if start_time + self.status_check_timeout < time.monotonic(): + # Check to see if the pipeline-run duration has exceeded the ``timeout`` configured. + if start_time + self.timeout < time.monotonic(): raise AirflowException( f"Pipeline run {self.run_id} has not reached a terminal status after " - + f"{self.status_check_timeout} seconds." + + f"{self.timeout} seconds." ) # Wait to check the status of the pipeline based on the ``poke_interval`` configured. From ac4c6bbfc051ad94c95252f9426a4ed3f90d17f0 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Sun, 15 Aug 2021 21:42:37 -0400 Subject: [PATCH 03/27] Small tweak to operator docstring --- .../providers/microsoft/azure/operators/azure_data_factory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/operators/azure_data_factory.py b/airflow/providers/microsoft/azure/operators/azure_data_factory.py index 999a248846b86..f2640f8918803 100644 --- a/airflow/providers/microsoft/azure/operators/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/operators/azure_data_factory.py @@ -39,7 +39,7 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): connection. :type resource_group_name: str :param factory_name: The data factory name. If a value is not passed in to the operator, the - ``AzureDataFactoryHook`` will attempt to use the resource group name provided in the corresponding + ``AzureDataFactoryHook`` will attempt to use the factory name name provided in the corresponding connection. :type factory_name: str :param reference_pipeline_run_id: The pipeline run identifier. If this run ID is specified the parameters From ce3df7f43598368a3e7c15b201f77d49fb598e0e Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Wed, 18 Aug 2021 23:14:38 -0400 Subject: [PATCH 04/27] Create test_azure_data_factory.py --- .../azure/sensors/test_azure_data_factory.py | 141 ++++++++++++++++++ 1 file changed, 141 insertions(+) create mode 100644 tests/providers/microsoft/azure/sensors/test_azure_data_factory.py diff --git a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py new file mode 100644 index 0000000000000..136e7607b61d4 --- /dev/null +++ b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py @@ -0,0 +1,141 @@ +# +# 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. + +import json +import unittest +from azure.mgmt.datafactory import DataFactoryManagementClient +from datetime import datetime +from parameterized import parameterized +from pytest import fixture +from unittest import mock + +from airflow.models.connection import Connection +from airflow.models.dag import DAG +from airflow.providers.microsoft.azure.hooks.azure_data_factory import ( + AzureDataFactoryHook, + AzureDataFactoryPipelineRunStatus, +) +from airflow.providers.microsoft.azure.sensors.azure_data_factory import ( + AzureDataFactoryPipelineRunStatusSensor, +) +from airflow.utils.session import create_session + +@fixture +def data_factory_client(): + client = AzureDataFactoryHook(conn_id="azure_data_factory_test") + client._conn = mock.MagicMock( + spec=[ + "factories", + "linked_services", + "datasets", + "pipelines", + "pipeline_runs", + "triggers", + "trigger_runs", + ] + ) + + return client + +class TestPipelineRunStatusSensor(unittest.TestCase): + def setUp(self): + self.dag = DAG("test", start_date=datetime(2021, 8, 16), schedule_interval=None, catchup=False) + self.config = { + "conn_id": "azure_data_factory_test", + "run_id": "run_id", + "timeout": 100, + "poke_interval": 15, + } + self.conn = Connection( + conn_id="azure_data_factory_test", + conn_type="azure_data_factory", + login="client_id", + password="client_secret", + extra=json.dumps( + { + "extra__azure_data_factory__tenantId": "tenantId", + "extra__azure_data_factory__subscriptionId": "subscriptionId", + "extra__azure_data_factory__resource_group_name": "default-resource-group-name", + "extra__azure_data_factory__factory_name": "default-factory-name", + }, + ) + ) + + with create_session() as session: + session.query(Connection).filter(Connection.conn_id == self.conn.conn_id).delete() + session.add(self.conn) + session.commit() + + @fixture(autouse=True) + def hook(self, data_factory_client): + self.hook = data_factory_client + + def test_init(self): + sensor = AzureDataFactoryPipelineRunStatusSensor( + task_id="pipeline_run_sensor", dag=self.dag, **self.config + ) + assert sensor.conn_id == self.config["conn_id"] + assert sensor.run_id == self.config["run_id"] + assert sensor.resource_group_name == None + assert sensor.factory_name == None + assert sensor.expected_statuses == {AzureDataFactoryPipelineRunStatus.SUCCEEDED} + assert sensor.timeout == self.config["timeout"] + assert sensor.poke_interval == self.config["poke_interval"] + + + @mock.patch( + "airflow.providers.microsoft.azure.hooks.azure_data_factory.AzureDataFactoryHook", + autospec=True, + ) + @mock.patch.object(AzureDataFactoryHook, 'get_pipeline_run') + def test_poke(self, mock_hook, mock_get_pipeline_run): + mock_run = mock_get_pipeline_run.return_value + + sensor = AzureDataFactoryPipelineRunStatusSensor( + task_id="pipeline_run_sensor_poke", dag=self.dag, **self.config + ) + result = sensor.poke({}) + + mock_run.get_pipeline_run.assert_called_once_with( + run_id="run_id", + factory_name="default_factory_name", + resource_group_name="default_resource_group_name", + ) + + assert result == False + + @mock.patch( + "airflow.providers.microsoft.azure.hooks.azure_data_factory.AzureDataFactoryHook", + autospec=True, + ) + def test_poke_(self, mock_hook): + mock_run = mock_hook.get_pipeline_run.return_value + mock_run.return_value = DataFactoryManagementClient + + sensor = AzureDataFactoryPipelineRunStatusSensor( + task_id="pipeline_run_sensor_poke", dag=self.dag, **self.config + ) + result = sensor.poke({}) + + mock_run.get_pipeline_run.assert_called_once_with( + run_id="run_id", + factory_name="default_factory_name", + resource_group_name="default_resource_group_name", + ) + + assert result == False From 7607d1071ce55e5b9af532fb60ce17ef02851cf3 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Thu, 26 Aug 2021 20:42:35 -0400 Subject: [PATCH 05/27] Changing operator parameter name --- .../azure/operators/azure_data_factory.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/azure_data_factory.py b/airflow/providers/microsoft/azure/operators/azure_data_factory.py index f2640f8918803..919480f9af380 100644 --- a/airflow/providers/microsoft/azure/operators/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/operators/azure_data_factory.py @@ -57,15 +57,15 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): :param parameters: Parameters of the pipeline run. These parameters will be used only if the ``reference_pipeline_run_id`` is not specified. :type start_from_failure: Dict[str, Any] - :param do_asynchronous_wait: Flag to exit after creating a pipeline run. Typically this flag would be - enabled to wait for a long-running pipeline execution using the - ``AzureDataFactoryPipelineRunStatusSensor`` rather than this operator. - :type do_asynchronous_wait: bool - :param timeout: Time in seconds to wait for a pipeline to reach a terminal status for - non-asynchronous waits. Use only if ``do_asynchronous_wait`` is False. + :param wait_for_completion: Flag to wait on a pipeline run's completion. By default, this feature is + enabled but could be disabled to wait for a long-running pipeline execution using the + ``AzureDataFactoryPipelineRunSensor`` rather than this operator. + :type wait_for_completion: bool + :param timeout: Time in seconds to wait for a pipeline to reach a terminal status for non-asynchronous + waits. Used only if ``wait_for_completion`` is True. :type timeout: int - :param poke_interval: Time in seconds to check on a pipeline run's status for non-asynchronous waits. Use - only if ``do_asynchronous_wait`` is False. + :param poke_interval: Time in seconds to check on a pipeline run's status for non-asynchronous waits. Used + only if ``wait_for_completion`` is True. :type poke_interval: int """ @@ -90,8 +90,8 @@ def __init__( start_activity_name: Optional[str] = None, start_from_failure: Optional[bool] = None, parameters: Optional[Dict[str, Any]] = None, - do_asynchronous_wait: Optional[bool] = False, - timeout: Optional[int] = 60 * 60 * 5, + wait_for_completion: Optional[bool] = True, + timeout: Optional[int] = 60 * 60 * 24 * 7, poke_interval: Optional[int] = 30, **kwargs, ) -> None: @@ -105,7 +105,7 @@ def __init__( self.start_activity_name = start_activity_name self.start_from_failure = start_from_failure self.parameters = parameters - self.do_asynchronous_wait = do_asynchronous_wait + self.wait_for_completion = wait_for_completion self.timeout = timeout self.poke_interval = poke_interval @@ -128,7 +128,7 @@ def execute(self, context: Dict) -> None: # an async wait. context["ti"].xcom_push(key="run_id", value=self.run_id) - if not self.do_asynchronous_wait: + if self.wait_for_completion: self.log.info(f"Waiting for run ID {self.run_id} of pipeline '{self.pipeline_name}' to complete.") start_time = time.monotonic() pipeline_run_status = None From 36a7840e657bfa2b53be064bc319bccebcc684d0 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Thu, 26 Aug 2021 20:57:15 -0400 Subject: [PATCH 06/27] Trivial operator formatting changes --- .../azure/operators/azure_data_factory.py | 3 +- .../azure/sensors/test_azure_data_factory.py | 42 +++++++++---------- 2 files changed, 23 insertions(+), 22 deletions(-) diff --git a/airflow/providers/microsoft/azure/operators/azure_data_factory.py b/airflow/providers/microsoft/azure/operators/azure_data_factory.py index 919480f9af380..a1f1e16ed1cd3 100644 --- a/airflow/providers/microsoft/azure/operators/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/operators/azure_data_factory.py @@ -137,11 +137,12 @@ def execute(self, context: Dict) -> None: if start_time + self.timeout < time.monotonic(): raise AirflowException( f"Pipeline run {self.run_id} has not reached a terminal status after " - + f"{self.timeout} seconds." + f"{self.timeout} seconds." ) # Wait to check the status of the pipeline based on the ``poke_interval`` configured. time.sleep(self.poke_interval) + self.log.info(f"Checking on the status of run ID {self.run_id}.") pipeline_run = self.hook.get_pipeline_run( run_id=self.run_id, diff --git a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py index 136e7607b61d4..f980a57adf37d 100644 --- a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py @@ -18,7 +18,7 @@ import json import unittest -from azure.mgmt.datafactory import DataFactoryManagementClient +# from azure.mgmt.datafactory import DataFactoryManagementClient from datetime import datetime from parameterized import parameterized from pytest import fixture @@ -119,23 +119,23 @@ def test_poke(self, mock_hook, mock_get_pipeline_run): assert result == False - @mock.patch( - "airflow.providers.microsoft.azure.hooks.azure_data_factory.AzureDataFactoryHook", - autospec=True, - ) - def test_poke_(self, mock_hook): - mock_run = mock_hook.get_pipeline_run.return_value - mock_run.return_value = DataFactoryManagementClient - - sensor = AzureDataFactoryPipelineRunStatusSensor( - task_id="pipeline_run_sensor_poke", dag=self.dag, **self.config - ) - result = sensor.poke({}) - - mock_run.get_pipeline_run.assert_called_once_with( - run_id="run_id", - factory_name="default_factory_name", - resource_group_name="default_resource_group_name", - ) - - assert result == False + # @mock.patch( + # "airflow.providers.microsoft.azure.hooks.azure_data_factory.AzureDataFactoryHook", + # autospec=True, + # ) + # def test_poke_(self, mock_hook): + # mock_run = mock_hook.get_pipeline_run.return_value + # mock_run.return_value = DataFactoryManagementClient + # + # sensor = AzureDataFactoryPipelineRunStatusSensor( + # task_id="pipeline_run_sensor_poke", dag=self.dag, **self.config + # ) + # result = sensor.poke({}) + # + # mock_run.get_pipeline_run.assert_called_once_with( + # run_id="run_id", + # factory_name="default_factory_name", + # resource_group_name="default_resource_group_name", + # ) + # + # assert result == False From 40177ab036ae4f0dc0a38f4d9a87eb3443dbf4ee Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Thu, 26 Aug 2021 21:10:33 -0400 Subject: [PATCH 07/27] Change datatype of expected_statuses for sensor --- .../microsoft/azure/sensors/azure_data_factory.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/airflow/providers/microsoft/azure/sensors/azure_data_factory.py b/airflow/providers/microsoft/azure/sensors/azure_data_factory.py index 2805f9879ed2f..e685393b848ae 100644 --- a/airflow/providers/microsoft/azure/sensors/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/sensors/azure_data_factory.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -from typing import Dict, Optional, Set, Union +from typing import Dict, List, Optional, Union from airflow.exceptions import AirflowException from airflow.providers.microsoft.azure.hooks.azure_data_factory import ( @@ -38,7 +38,7 @@ class AzureDataFactoryPipelineRunStatusSensor(BaseSensorOperator): :param factory_name: The data factory name. :type factory_name: str :param expected_statuses: The status(es) which are desired for the pipeline run. - :type expected_statuses: str or Set[str] + :type expected_statuses: str or List[str] """ template_fields = ("resource_group_name", "factory_name", "run_id") @@ -50,7 +50,7 @@ def __init__( run_id: str, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, - expected_statuses: Optional[Union[Set[str], str]] = AzureDataFactoryPipelineRunStatus.SUCCEEDED, + expected_statuses: Optional[Union[List[str], str]] = AzureDataFactoryPipelineRunStatus.SUCCEEDED, **kwargs, ) -> None: super().__init__(**kwargs) @@ -58,8 +58,9 @@ def __init__( self.run_id = run_id self.resource_group_name = resource_group_name self.factory_name = factory_name + # Normalize input of ``expected_status`` to a list. self.expected_statuses = ( - {expected_statuses} if isinstance(expected_statuses, str) else expected_statuses + [expected_statuses] if isinstance(expected_statuses, str) else expected_statuses ) def poke(self, context: Dict) -> bool: From ee453d5e03172e18150f7030c4a29e360ad35da4 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Thu, 26 Aug 2021 21:44:42 -0400 Subject: [PATCH 08/27] Removing expected_statuses from sensor --- .../azure/hooks/azure_data_factory.py | 2 -- .../azure/sensors/azure_data_factory.py | 19 +++++-------------- 2 files changed, 5 insertions(+), 16 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py index c13f9751230fa..f5fdd35534e82 100644 --- a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py @@ -77,8 +77,6 @@ class AzureDataFactoryPipelineRunStatus: CANCELING = "Canceling" CANCELLED = "Cancelled" - TERMINAL_STATUSES = {CANCELLED, FAILED, SUCCEEDED} - class AzureDataFactoryHook(BaseHook): """ diff --git a/airflow/providers/microsoft/azure/sensors/azure_data_factory.py b/airflow/providers/microsoft/azure/sensors/azure_data_factory.py index e685393b848ae..d87ca63a9befa 100644 --- a/airflow/providers/microsoft/azure/sensors/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/sensors/azure_data_factory.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -from typing import Dict, List, Optional, Union +from typing import Dict, Optional from airflow.exceptions import AirflowException from airflow.providers.microsoft.azure.hooks.azure_data_factory import ( @@ -37,8 +37,6 @@ class AzureDataFactoryPipelineRunStatusSensor(BaseSensorOperator): :type resource_group_name: str :param factory_name: The data factory name. :type factory_name: str - :param expected_statuses: The status(es) which are desired for the pipeline run. - :type expected_statuses: str or List[str] """ template_fields = ("resource_group_name", "factory_name", "run_id") @@ -50,7 +48,6 @@ def __init__( run_id: str, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, - expected_statuses: Optional[Union[List[str], str]] = AzureDataFactoryPipelineRunStatus.SUCCEEDED, **kwargs, ) -> None: super().__init__(**kwargs) @@ -58,16 +55,9 @@ def __init__( self.run_id = run_id self.resource_group_name = resource_group_name self.factory_name = factory_name - # Normalize input of ``expected_status`` to a list. - self.expected_statuses = ( - [expected_statuses] if isinstance(expected_statuses, str) else expected_statuses - ) def poke(self, context: Dict) -> bool: - self.log.info( - f"Checking for pipeline run {self.run_id} to be in one of the following statuses: " - f"{', '.join(self.expected_statuses)}.", - ) + self.log.info(f"Checking the status for pipeline run {self.run_id}.") self.hook = AzureDataFactoryHook(conn_id=self.conn_id) pipeline_run = self.hook.get_pipeline_run( run_id=self.run_id, @@ -77,14 +67,15 @@ def poke(self, context: Dict) -> bool: pipeline_run_status = pipeline_run.status self.log.info(f"Current status for pipeline run {self.run_id}: {pipeline_run_status}.") - if pipeline_run_status in self.expected_statuses: + if pipeline_run_status == AzureDataFactoryPipelineRunStatus.SUCCEEDED: + self.log.info(f"The pipeline run {self.run_id} has succeeded.") return True elif pipeline_run_status in { AzureDataFactoryPipelineRunStatus.FAILED, AzureDataFactoryPipelineRunStatus.CANCELLED, }: raise AirflowException( - f"Pipeline run {self.run_id} is in a terminal status: {pipeline_run_status}" + f"Pipeline run {self.run_id} is in a negative terminal status: {pipeline_run_status}" ) return False From a947055aa9ca8aee24b5ecd94fb962f8acbd2651 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Thu, 26 Aug 2021 21:48:20 -0400 Subject: [PATCH 09/27] Adding terminal statuses to ADF hook --- airflow/providers/microsoft/azure/hooks/azure_data_factory.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py index f5fdd35534e82..c13f9751230fa 100644 --- a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py @@ -77,6 +77,8 @@ class AzureDataFactoryPipelineRunStatus: CANCELING = "Canceling" CANCELLED = "Cancelled" + TERMINAL_STATUSES = {CANCELLED, FAILED, SUCCEEDED} + class AzureDataFactoryHook(BaseHook): """ From 34641d99dfc6b3a5fea0dd05e04a4c5dad17796e Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Fri, 27 Aug 2021 00:15:37 -0400 Subject: [PATCH 10/27] Finalizing sensor unit test --- .../azure/sensors/test_azure_data_factory.py | 119 +++++++----------- 1 file changed, 46 insertions(+), 73 deletions(-) diff --git a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py index f980a57adf37d..cf822be5cc421 100644 --- a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py @@ -18,12 +18,13 @@ import json import unittest -# from azure.mgmt.datafactory import DataFactoryManagementClient from datetime import datetime -from parameterized import parameterized -from pytest import fixture from unittest import mock +import pytest +from parameterized import parameterized + +from airflow.exceptions import AirflowException from airflow.models.connection import Connection from airflow.models.dag import DAG from airflow.providers.microsoft.azure.hooks.azure_data_factory import ( @@ -35,32 +36,10 @@ ) from airflow.utils.session import create_session -@fixture -def data_factory_client(): - client = AzureDataFactoryHook(conn_id="azure_data_factory_test") - client._conn = mock.MagicMock( - spec=[ - "factories", - "linked_services", - "datasets", - "pipelines", - "pipeline_runs", - "triggers", - "trigger_runs", - ] - ) - - return client class TestPipelineRunStatusSensor(unittest.TestCase): def setUp(self): self.dag = DAG("test", start_date=datetime(2021, 8, 16), schedule_interval=None, catchup=False) - self.config = { - "conn_id": "azure_data_factory_test", - "run_id": "run_id", - "timeout": 100, - "poke_interval": 15, - } self.conn = Connection( conn_id="azure_data_factory_test", conn_type="azure_data_factory", @@ -70,10 +49,10 @@ def setUp(self): { "extra__azure_data_factory__tenantId": "tenantId", "extra__azure_data_factory__subscriptionId": "subscriptionId", - "extra__azure_data_factory__resource_group_name": "default-resource-group-name", - "extra__azure_data_factory__factory_name": "default-factory-name", + "extra__azure_data_factory__resource_group_name": "resource-group-name", + "extra__azure_data_factory__factory_name": "factory-name", }, - ) + ), ) with create_session() as session: @@ -81,9 +60,19 @@ def setUp(self): session.add(self.conn) session.commit() - @fixture(autouse=True) - def hook(self, data_factory_client): - self.hook = data_factory_client + self.config = { + "conn_id": "azure_data_factory_test", + "run_id": "run_id", + "resource_group_name": self.conn.extra_dejson["extra__azure_data_factory__resource_group_name"], + "factory_name": self.conn.extra_dejson["extra__azure_data_factory__resource_group_name"], + "timeout": 100, + "poke_interval": 15, + } + + def create_pipeline_run(self, state: str): + run = mock.Mock() + run.status = state + return run def test_init(self): sensor = AzureDataFactoryPipelineRunStatusSensor( @@ -91,51 +80,35 @@ def test_init(self): ) assert sensor.conn_id == self.config["conn_id"] assert sensor.run_id == self.config["run_id"] - assert sensor.resource_group_name == None - assert sensor.factory_name == None - assert sensor.expected_statuses == {AzureDataFactoryPipelineRunStatus.SUCCEEDED} + assert sensor.resource_group_name == self.config["resource_group_name"] + assert sensor.factory_name == self.config["factory_name"] assert sensor.timeout == self.config["timeout"] assert sensor.poke_interval == self.config["poke_interval"] - - @mock.patch( - "airflow.providers.microsoft.azure.hooks.azure_data_factory.AzureDataFactoryHook", - autospec=True, + @parameterized.expand( + [ + (AzureDataFactoryPipelineRunStatus.SUCCEEDED, True), + (AzureDataFactoryPipelineRunStatus.FAILED, "AirflowException"), + (AzureDataFactoryPipelineRunStatus.CANCELLED, "AirflowException"), + (AzureDataFactoryPipelineRunStatus.CANCELING, False), + (AzureDataFactoryPipelineRunStatus.QUEUED, False), + (AzureDataFactoryPipelineRunStatus.IN_PROGRESS, False), + ] ) - @mock.patch.object(AzureDataFactoryHook, 'get_pipeline_run') - def test_poke(self, mock_hook, mock_get_pipeline_run): - mock_run = mock_get_pipeline_run.return_value - - sensor = AzureDataFactoryPipelineRunStatusSensor( - task_id="pipeline_run_sensor_poke", dag=self.dag, **self.config - ) - result = sensor.poke({}) + def test_poke(self, pipeline_run_status, expected_status): + mock_pipeline_run = self.create_pipeline_run(pipeline_run_status) - mock_run.get_pipeline_run.assert_called_once_with( - run_id="run_id", - factory_name="default_factory_name", - resource_group_name="default_resource_group_name", - ) - - assert result == False + with mock.patch.object(AzureDataFactoryHook, "get_pipeline_run", return_value=mock_pipeline_run): + sensor = AzureDataFactoryPipelineRunStatusSensor( + task_id="pipeline_run_sensor_poke", dag=self.dag, **self.config + ) - # @mock.patch( - # "airflow.providers.microsoft.azure.hooks.azure_data_factory.AzureDataFactoryHook", - # autospec=True, - # ) - # def test_poke_(self, mock_hook): - # mock_run = mock_hook.get_pipeline_run.return_value - # mock_run.return_value = DataFactoryManagementClient - # - # sensor = AzureDataFactoryPipelineRunStatusSensor( - # task_id="pipeline_run_sensor_poke", dag=self.dag, **self.config - # ) - # result = sensor.poke({}) - # - # mock_run.get_pipeline_run.assert_called_once_with( - # run_id="run_id", - # factory_name="default_factory_name", - # resource_group_name="default_resource_group_name", - # ) - # - # assert result == False + if expected_status == "AirflowException": + with pytest.raises( + AirflowException, + match=f"Pipeline run {self.config['run_id']} is in a negative terminal status: " + f"{pipeline_run_status}", + ): + sensor.poke({}) + else: + assert sensor.poke({}) == expected_status From 742aa9d9cf6428d9527389e4f0da34068a2a1aa7 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Fri, 27 Aug 2021 00:25:28 -0400 Subject: [PATCH 11/27] Removing Connection creation in sensor unit test --- .../azure/sensors/test_azure_data_factory.py | 24 ++----------------- 1 file changed, 2 insertions(+), 22 deletions(-) diff --git a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py index cf822be5cc421..5da03c1400f2d 100644 --- a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py @@ -40,31 +40,11 @@ class TestPipelineRunStatusSensor(unittest.TestCase): def setUp(self): self.dag = DAG("test", start_date=datetime(2021, 8, 16), schedule_interval=None, catchup=False) - self.conn = Connection( - conn_id="azure_data_factory_test", - conn_type="azure_data_factory", - login="client_id", - password="client_secret", - extra=json.dumps( - { - "extra__azure_data_factory__tenantId": "tenantId", - "extra__azure_data_factory__subscriptionId": "subscriptionId", - "extra__azure_data_factory__resource_group_name": "resource-group-name", - "extra__azure_data_factory__factory_name": "factory-name", - }, - ), - ) - - with create_session() as session: - session.query(Connection).filter(Connection.conn_id == self.conn.conn_id).delete() - session.add(self.conn) - session.commit() - self.config = { "conn_id": "azure_data_factory_test", "run_id": "run_id", - "resource_group_name": self.conn.extra_dejson["extra__azure_data_factory__resource_group_name"], - "factory_name": self.conn.extra_dejson["extra__azure_data_factory__resource_group_name"], + "resource_group_name": "resource-group-name", + "factory_name": "factory-name", "timeout": 100, "poke_interval": 15, } From 39556e9c9a399200b9c6605cb461dd51ac74e9ee Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Fri, 27 Aug 2021 12:55:21 -0400 Subject: [PATCH 12/27] Rename create_pipeline_run parameter --- .../microsoft/azure/sensors/test_azure_data_factory.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py index 5da03c1400f2d..aea5b419b8f01 100644 --- a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py @@ -49,9 +49,9 @@ def setUp(self): "poke_interval": 15, } - def create_pipeline_run(self, state: str): + def create_pipeline_run(self, status: str): run = mock.Mock() - run.status = state + run.status = status return run def test_init(self): From 2af2134eadb9bb6b96b1111759180f4cd1edd8c2 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Fri, 27 Aug 2021 19:50:11 -0400 Subject: [PATCH 13/27] Adding unit test for AzureDataFactoryRunPipelineOperator --- .../azure/operators/azure_data_factory.py | 6 +- .../operators/test_azure_data_factory.py | 202 ++++++++++++++++++ .../azure/sensors/test_azure_data_factory.py | 25 ++- 3 files changed, 222 insertions(+), 11 deletions(-) create mode 100644 tests/providers/microsoft/azure/operators/test_azure_data_factory.py diff --git a/airflow/providers/microsoft/azure/operators/azure_data_factory.py b/airflow/providers/microsoft/azure/operators/azure_data_factory.py index a1f1e16ed1cd3..a5d8bbc129b6b 100644 --- a/airflow/providers/microsoft/azure/operators/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/operators/azure_data_factory.py @@ -123,9 +123,9 @@ def execute(self, context: Dict) -> None: parameters=self.parameters, ) self.run_id = vars(response)["run_id"] - # Push the ``run_id`` value to XCom regardless of what happens during execution. This allows users to - # retrieve the ``run_id`` as an output of this operator for downstream tasks especially if performing - # an async wait. + # Push the ``run_id`` value to XCom regardless of what happens during execution. This allows users + # to retrieve the ``run_id`` as an output of this operator for downstream tasks especially if + # performing an async wait. context["ti"].xcom_push(key="run_id", value=self.run_id) if self.wait_for_completion: diff --git a/tests/providers/microsoft/azure/operators/test_azure_data_factory.py b/tests/providers/microsoft/azure/operators/test_azure_data_factory.py new file mode 100644 index 0000000000000..7b91d9a31a491 --- /dev/null +++ b/tests/providers/microsoft/azure/operators/test_azure_data_factory.py @@ -0,0 +1,202 @@ +# 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. + +import unittest +from unittest.mock import MagicMock, patch + +import pytest +from parameterized import parameterized + +from airflow.exceptions import AirflowException +from airflow.providers.microsoft.azure.hooks.azure_data_factory import ( + AzureDataFactoryHook, + AzureDataFactoryPipelineRunStatus, +) +from airflow.providers.microsoft.azure.operators.azure_data_factory import AzureDataFactoryRunPipelineOperator + +PIPELINE_RUN_RESPONSE = {"additional_properties": {}, "run_id": "run_id"} + + +class TestAzureDataFactoryRunPipelineOperator(unittest.TestCase): + def setUp(self): + self.mock_ti = MagicMock() + self.mock_context = {"ti": self.mock_ti} + self.config = { + "task_id": "run_pipeline_op", + "conn_id": "azure_data_factory_test", + "pipeline_name": "pipeline1", + "resource_group_name": "resource-group-name", + "factory_name": "factory-name", + } + + @staticmethod + def create_pipeline_run(status: str): + """Helper function to create a mock pipeline run with a given execution status.""" + + run = MagicMock() + run.status = status + return run + + @parameterized.expand( + [ + (AzureDataFactoryPipelineRunStatus.SUCCEEDED, None), + (AzureDataFactoryPipelineRunStatus.FAILED, "AirflowException"), + (AzureDataFactoryPipelineRunStatus.CANCELLED, "AirflowException"), + ] + ) + @patch.object(AzureDataFactoryHook, "run_pipeline", return_value=MagicMock(**PIPELINE_RUN_RESPONSE)) + def test_execute_wait_for_completion(self, pipeline_run_status, expected_status, mock_run_pipeline): + operator = AzureDataFactoryRunPipelineOperator(poke_interval=1, **self.config) + + assert operator.conn_id == self.config["conn_id"] + assert operator.pipeline_name == self.config["pipeline_name"] + assert operator.resource_group_name == self.config["resource_group_name"] + assert operator.factory_name == self.config["factory_name"] + assert operator.wait_for_completion + assert operator.timeout == 60 * 60 * 24 * 7 + assert operator.poke_interval == 1 + + self.assertIsNone(operator.reference_pipeline_run_id) + self.assertIsNone(operator.is_recovery) + self.assertIsNone(operator.start_activity_name) + self.assertIsNone(operator.start_from_failure) + self.assertIsNone(operator.parameters) + + with patch.object(AzureDataFactoryHook, "get_pipeline_run") as mock_get_pipeline_run: + mock_get_pipeline_run.return_value = TestAzureDataFactoryRunPipelineOperator.create_pipeline_run( + pipeline_run_status + ) + + assert mock_get_pipeline_run.return_value.status == pipeline_run_status + + if not expected_status: + self.assertIsNone(operator.execute(context=self.mock_context)) + elif expected_status == "AirflowException": + if mock_get_pipeline_run.return_value.status == AzureDataFactoryPipelineRunStatus.CANCELLED: + error_message = ( + f"Pipeline run {PIPELINE_RUN_RESPONSE['run_id']} has been " + f"{pipeline_run_status.lower()}." + ) + else: + error_message = ( + f"Pipeline run {PIPELINE_RUN_RESPONSE['run_id']} has " + f"{pipeline_run_status.lower()}." + ) + # The operator should fail if the pipeline run fails or is canceled. + with pytest.raises(AirflowException, match=error_message): + operator.execute(context=self.mock_context) + + # Check the ``run_id`` attr is assigned after executing the pipeline. + assert operator.run_id == PIPELINE_RUN_RESPONSE["run_id"] + + # Check to ensure an `XCom` is pushed regardless of pipeline run result. + self.mock_ti.xcom_push.assert_called_once_with( + key="run_id", value=PIPELINE_RUN_RESPONSE["run_id"] + ) + + mock_run_pipeline.assert_called_once_with( + pipeline_name=self.config["pipeline_name"], + resource_group_name=self.config["resource_group_name"], + factory_name=self.config["factory_name"], + reference_pipeline_run_id=None, + is_recovery=None, + start_activity_name=None, + start_from_failure=None, + parameters=None, + ) + + mock_get_pipeline_run.assert_called_once_with( + run_id=mock_run_pipeline.return_value.run_id, + factory_name=self.config["factory_name"], + resource_group_name=self.config["resource_group_name"], + ) + + @patch.object(AzureDataFactoryHook, "run_pipeline", return_value=MagicMock(**PIPELINE_RUN_RESPONSE)) + def test_execute_no_wait_for_completion(self, mock_run_pipeline): + operator = AzureDataFactoryRunPipelineOperator( + poke_interval=1, wait_for_completion=False, **self.config + ) + + assert operator.conn_id == self.config["conn_id"] + assert operator.pipeline_name == self.config["pipeline_name"] + assert operator.resource_group_name == self.config["resource_group_name"] + assert operator.factory_name == self.config["factory_name"] + assert not operator.wait_for_completion + assert operator.timeout == 60 * 60 * 24 * 7 + assert operator.poke_interval == 1 + + self.assertIsNone(operator.reference_pipeline_run_id) + self.assertIsNone(operator.is_recovery) + self.assertIsNone(operator.start_activity_name) + self.assertIsNone(operator.start_from_failure) + self.assertIsNone(operator.parameters) + + with patch.object(AzureDataFactoryHook, "get_pipeline_run") as mock_get_pipeline_run: + self.assertIsNone(operator.execute(context=self.mock_context)) + + # Check the ``run_id`` attr is assigned after executing the pipeline. + assert operator.run_id == PIPELINE_RUN_RESPONSE["run_id"] + + # Check to ensure an `XCom` is pushed regardless of pipeline run result. + self.mock_ti.xcom_push.assert_called_once_with( + key="run_id", value=PIPELINE_RUN_RESPONSE["run_id"] + ) + + mock_run_pipeline.assert_called_once_with( + pipeline_name=self.config["pipeline_name"], + resource_group_name=self.config["resource_group_name"], + factory_name=self.config["factory_name"], + reference_pipeline_run_id=None, + is_recovery=None, + start_activity_name=None, + start_from_failure=None, + parameters=None, + ) + + # Checking the pipeline run status should _not_ be called when ``wait_for_completion`` is False. + mock_get_pipeline_run.assert_not_called() + + # @mock.patch(DATAPROC_PATH.format("DataprocHook")) + # def test_on_kill(self, mock_hook): + # job = {} + # job_id = "job_id" + # mock_hook.return_value.wait_for_job.return_value = None + # mock_hook.return_value.submit_job.return_value.reference.job_id = job_id + # + # op = DataprocSubmitJobOperator( + # task_id=TASK_ID, + # region=GCP_LOCATION, + # project_id=GCP_PROJECT, + # job=job, + # gcp_conn_id=GCP_CONN_ID, + # retry=RETRY, + # timeout=TIMEOUT, + # metadata=METADATA, + # request_id=REQUEST_ID, + # impersonation_chain=IMPERSONATION_CHAIN, + # cancel_on_kill=False, + # ) + # op.execute(context=self.mock_context) + # + # op.on_kill() + # mock_hook.return_value.cancel_job.assert_not_called() + # + # op.cancel_on_kill = True + # op.on_kill() + # mock_hook.return_value.cancel_job.assert_called_once_with( + # project_id=GCP_PROJECT, region=GCP_LOCATION, job_id=job_id + # ) diff --git a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py index aea5b419b8f01..8afb10374609e 100644 --- a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py @@ -16,16 +16,14 @@ # specific language governing permissions and limitations # under the License. -import json import unittest from datetime import datetime -from unittest import mock +from unittest.mock import MagicMock, patch import pytest from parameterized import parameterized from airflow.exceptions import AirflowException -from airflow.models.connection import Connection from airflow.models.dag import DAG from airflow.providers.microsoft.azure.hooks.azure_data_factory import ( AzureDataFactoryHook, @@ -34,7 +32,6 @@ from airflow.providers.microsoft.azure.sensors.azure_data_factory import ( AzureDataFactoryPipelineRunStatusSensor, ) -from airflow.utils.session import create_session class TestPipelineRunStatusSensor(unittest.TestCase): @@ -49,8 +46,11 @@ def setUp(self): "poke_interval": 15, } - def create_pipeline_run(self, status: str): - run = mock.Mock() + @staticmethod + def create_pipeline_run(status: str): + """Helper function to create a mock pipeline run with a given execution status.""" + + run = MagicMock() run.status = status return run @@ -76,14 +76,17 @@ def test_init(self): ] ) def test_poke(self, pipeline_run_status, expected_status): - mock_pipeline_run = self.create_pipeline_run(pipeline_run_status) + mock_pipeline_run = TestPipelineRunStatusSensor.create_pipeline_run(pipeline_run_status) - with mock.patch.object(AzureDataFactoryHook, "get_pipeline_run", return_value=mock_pipeline_run): + with patch.object( + AzureDataFactoryHook, "get_pipeline_run", return_value=mock_pipeline_run + ) as mock_get_pipeline_run: sensor = AzureDataFactoryPipelineRunStatusSensor( task_id="pipeline_run_sensor_poke", dag=self.dag, **self.config ) if expected_status == "AirflowException": + # The sensor should fail if the pipeline run fails or is canceled. with pytest.raises( AirflowException, match=f"Pipeline run {self.config['run_id']} is in a negative terminal status: " @@ -92,3 +95,9 @@ def test_poke(self, pipeline_run_status, expected_status): sensor.poke({}) else: assert sensor.poke({}) == expected_status + + mock_get_pipeline_run.assert_called_once_with( + run_id=self.config["run_id"], + factory_name=self.config["factory_name"], + resource_group_name=self.config["resource_group_name"], + ) From 1a74100adc322dd3012cc58e918a4cda89bb3b9b Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Sat, 28 Aug 2021 09:19:01 -0400 Subject: [PATCH 14/27] Small update to ADF hook docstring --- airflow/providers/microsoft/azure/hooks/azure_data_factory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py index c13f9751230fa..6db7f6994b8ea 100644 --- a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py @@ -85,7 +85,7 @@ class AzureDataFactoryHook(BaseHook): A hook to interact with Azure Data Factory. :param conn_id: The :ref:`Azure Data Factory connection id`. - :type: str + :type conn_id: str """ conn_type: str = 'azure_data_factory' From b4a02fab9c6cfa54b17f05ccb06a98c9f09f0861 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Sat, 28 Aug 2021 13:26:44 -0400 Subject: [PATCH 15/27] Small refactor of operator unit test --- .../operators/test_azure_data_factory.py | 62 +++---------------- 1 file changed, 9 insertions(+), 53 deletions(-) diff --git a/tests/providers/microsoft/azure/operators/test_azure_data_factory.py b/tests/providers/microsoft/azure/operators/test_azure_data_factory.py index 7b91d9a31a491..e6aa4e8789fe6 100644 --- a/tests/providers/microsoft/azure/operators/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/operators/test_azure_data_factory.py @@ -41,6 +41,7 @@ def setUp(self): "pipeline_name": "pipeline1", "resource_group_name": "resource-group-name", "factory_name": "factory-name", + "poke_interval": 1, } @staticmethod @@ -49,6 +50,7 @@ def create_pipeline_run(status: str): run = MagicMock() run.status = status + return run @parameterized.expand( @@ -60,21 +62,14 @@ def create_pipeline_run(status: str): ) @patch.object(AzureDataFactoryHook, "run_pipeline", return_value=MagicMock(**PIPELINE_RUN_RESPONSE)) def test_execute_wait_for_completion(self, pipeline_run_status, expected_status, mock_run_pipeline): - operator = AzureDataFactoryRunPipelineOperator(poke_interval=1, **self.config) + operator = AzureDataFactoryRunPipelineOperator(**self.config) assert operator.conn_id == self.config["conn_id"] assert operator.pipeline_name == self.config["pipeline_name"] assert operator.resource_group_name == self.config["resource_group_name"] assert operator.factory_name == self.config["factory_name"] + assert operator.poke_interval == self.config["poke_interval"] assert operator.wait_for_completion - assert operator.timeout == 60 * 60 * 24 * 7 - assert operator.poke_interval == 1 - - self.assertIsNone(operator.reference_pipeline_run_id) - self.assertIsNone(operator.is_recovery) - self.assertIsNone(operator.start_activity_name) - self.assertIsNone(operator.start_from_failure) - self.assertIsNone(operator.parameters) with patch.object(AzureDataFactoryHook, "get_pipeline_run") as mock_get_pipeline_run: mock_get_pipeline_run.return_value = TestAzureDataFactoryRunPipelineOperator.create_pipeline_run( @@ -84,7 +79,8 @@ def test_execute_wait_for_completion(self, pipeline_run_status, expected_status, assert mock_get_pipeline_run.return_value.status == pipeline_run_status if not expected_status: - self.assertIsNone(operator.execute(context=self.mock_context)) + # A successful operator execution should not return any values. + assert not operator.execute(context=self.mock_context) elif expected_status == "AirflowException": if mock_get_pipeline_run.return_value.status == AzureDataFactoryPipelineRunStatus.CANCELLED: error_message = ( @@ -127,26 +123,17 @@ def test_execute_wait_for_completion(self, pipeline_run_status, expected_status, @patch.object(AzureDataFactoryHook, "run_pipeline", return_value=MagicMock(**PIPELINE_RUN_RESPONSE)) def test_execute_no_wait_for_completion(self, mock_run_pipeline): - operator = AzureDataFactoryRunPipelineOperator( - poke_interval=1, wait_for_completion=False, **self.config - ) + operator = AzureDataFactoryRunPipelineOperator(wait_for_completion=False, **self.config) assert operator.conn_id == self.config["conn_id"] assert operator.pipeline_name == self.config["pipeline_name"] assert operator.resource_group_name == self.config["resource_group_name"] assert operator.factory_name == self.config["factory_name"] + assert operator.poke_interval == self.config["poke_interval"] assert not operator.wait_for_completion - assert operator.timeout == 60 * 60 * 24 * 7 - assert operator.poke_interval == 1 - - self.assertIsNone(operator.reference_pipeline_run_id) - self.assertIsNone(operator.is_recovery) - self.assertIsNone(operator.start_activity_name) - self.assertIsNone(operator.start_from_failure) - self.assertIsNone(operator.parameters) with patch.object(AzureDataFactoryHook, "get_pipeline_run") as mock_get_pipeline_run: - self.assertIsNone(operator.execute(context=self.mock_context)) + operator.execute(context=self.mock_context) # Check the ``run_id`` attr is assigned after executing the pipeline. assert operator.run_id == PIPELINE_RUN_RESPONSE["run_id"] @@ -169,34 +156,3 @@ def test_execute_no_wait_for_completion(self, mock_run_pipeline): # Checking the pipeline run status should _not_ be called when ``wait_for_completion`` is False. mock_get_pipeline_run.assert_not_called() - - # @mock.patch(DATAPROC_PATH.format("DataprocHook")) - # def test_on_kill(self, mock_hook): - # job = {} - # job_id = "job_id" - # mock_hook.return_value.wait_for_job.return_value = None - # mock_hook.return_value.submit_job.return_value.reference.job_id = job_id - # - # op = DataprocSubmitJobOperator( - # task_id=TASK_ID, - # region=GCP_LOCATION, - # project_id=GCP_PROJECT, - # job=job, - # gcp_conn_id=GCP_CONN_ID, - # retry=RETRY, - # timeout=TIMEOUT, - # metadata=METADATA, - # request_id=REQUEST_ID, - # impersonation_chain=IMPERSONATION_CHAIN, - # cancel_on_kill=False, - # ) - # op.execute(context=self.mock_context) - # - # op.on_kill() - # mock_hook.return_value.cancel_job.assert_not_called() - # - # op.cancel_on_kill = True - # op.on_kill() - # mock_hook.return_value.cancel_job.assert_called_once_with( - # project_id=GCP_PROJECT, region=GCP_LOCATION, job_id=job_id - # ) From 5ce65da2ffe3cf0bbf0e4e9395674ed57bee955a Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Sat, 28 Aug 2021 14:07:56 -0400 Subject: [PATCH 16/27] Small updates to the operator, sensor, and example DAG --- .../azure/example_dags/example_adf_run_pipeline.py | 8 ++++---- .../azure/operators/azure_data_factory.py | 14 ++++++++------ .../microsoft/azure/sensors/azure_data_factory.py | 2 +- 3 files changed, 13 insertions(+), 11 deletions(-) diff --git a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py index b43388d1cbc0a..beab8da5a8dd7 100644 --- a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py +++ b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py @@ -34,8 +34,8 @@ "retries": 1, "retry_delay": timedelta(minutes=3), "conn_id": "azure_data_factory", - "factory_name": "my-data-factory", - "resource_group_name": "my-resource-group", + "factory_name": "my-data-factory", # This can also be specified in the ADF connection. + "resource_group_name": "my-resource-group", # This can also be specified in the ADF connection. }, default_view="graph", ) as dag: @@ -51,7 +51,7 @@ run_pipeline2 = AzureDataFactoryRunPipelineOperator( task_id="run_pipeline2", pipeline_name="extractDailyExchangeRates", - do_asynchronous_wait=True, + wait_for_completion=False, ) wait_for_pipeline_run = AzureDataFactoryPipelineRunStatusSensor( @@ -60,9 +60,9 @@ poke_interval=10, ) + begin >> Label("No async wait") >> run_pipeline1 >> end begin >> Label("Do async wait with sensor") >> run_pipeline2 wait_for_pipeline_run >> end - begin >> Label("No async wait") >> run_pipeline1 >> end # Task dependency created via `XComArgs`: # run_pipeline2 >> wait_for_pipeline_run diff --git a/airflow/providers/microsoft/azure/operators/azure_data_factory.py b/airflow/providers/microsoft/azure/operators/azure_data_factory.py index a5d8bbc129b6b..4018b504bbe77 100644 --- a/airflow/providers/microsoft/azure/operators/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/operators/azure_data_factory.py @@ -54,8 +54,9 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): :param start_from_failure: In recovery mode, if set to true, the rerun will start from failed activities. The property will be used only if ``start_activity_name`` is not specified. :type start_from_failure: bool - :param parameters: Parameters of the pipeline run. These parameters will be used only if the - ``reference_pipeline_run_id`` is not specified. + :param parameters: Parameters of the pipeline run. These parameters are referenced in a pipeline via + ``@pipeline().parameters.parameterName`` and will be used only if the ``reference_pipeline_run_id`` is + not specified. :type start_from_failure: Dict[str, Any] :param wait_for_completion: Flag to wait on a pipeline run's completion. By default, this feature is enabled but could be disabled to wait for a long-running pipeline execution using the @@ -70,6 +71,7 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): """ template_fields = ( + "conn_id", "resource_group_name", "factory_name", "pipeline_name", @@ -123,13 +125,13 @@ def execute(self, context: Dict) -> None: parameters=self.parameters, ) self.run_id = vars(response)["run_id"] - # Push the ``run_id`` value to XCom regardless of what happens during execution. This allows users - # to retrieve the ``run_id`` as an output of this operator for downstream tasks especially if - # performing an async wait. + # Push the ``run_id`` value to XCom regardless of what happens during execution. This allows for + # retrieval the executed pipeline's ``run_id`` for downstream tasks especially if performing an + # asynchronous wait. context["ti"].xcom_push(key="run_id", value=self.run_id) if self.wait_for_completion: - self.log.info(f"Waiting for run ID {self.run_id} of pipeline '{self.pipeline_name}' to complete.") + self.log.info(f"Waiting for run ID {self.run_id} of pipeline {self.pipeline_name} to complete.") start_time = time.monotonic() pipeline_run_status = None while pipeline_run_status not in AzureDataFactoryPipelineRunStatus.TERMINAL_STATUSES: diff --git a/airflow/providers/microsoft/azure/sensors/azure_data_factory.py b/airflow/providers/microsoft/azure/sensors/azure_data_factory.py index d87ca63a9befa..625bb0d71ed02 100644 --- a/airflow/providers/microsoft/azure/sensors/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/sensors/azure_data_factory.py @@ -39,7 +39,7 @@ class AzureDataFactoryPipelineRunStatusSensor(BaseSensorOperator): :type factory_name: str """ - template_fields = ("resource_group_name", "factory_name", "run_id") + template_fields = ("conn_id", "resource_group_name", "factory_name", "run_id") def __init__( self, From a282433ae0eb71ae84eaf7fd35d0a40419c0ebda Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Sat, 28 Aug 2021 14:18:33 -0400 Subject: [PATCH 17/27] Modifying custom conn field names --- .../providers/microsoft/azure/hooks/azure_data_factory.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py index 6db7f6994b8ea..c4d3648dc6ca0 100644 --- a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py @@ -102,10 +102,10 @@ def get_connection_form_widgets() -> Dict[str, Any]: return { "extra__azure_data_factory__tenantId": StringField( - lazy_gettext('Azure Tenant ID'), widget=BS3TextFieldWidget() + lazy_gettext('Tenant ID'), widget=BS3TextFieldWidget() ), "extra__azure_data_factory__subscriptionId": StringField( - lazy_gettext('Azure Subscription ID'), widget=BS3TextFieldWidget() + lazy_gettext('Subscription ID'), widget=BS3TextFieldWidget() ), "extra__azure_data_factory__resource_group_name": StringField( lazy_gettext('Resource Group Name'), widget=BS3TextFieldWidget() @@ -121,8 +121,8 @@ def get_ui_field_behaviour() -> Dict: return { "hidden_fields": ['schema', 'port', 'host', 'extra'], "relabeling": { - 'login': 'Azure Client ID', - 'password': 'Azure Secret', + 'login': 'Client ID', + 'password': 'Secret', }, } From 7816d8c2efb6d08d9a88882cba271131ec6feb50 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Sat, 28 Aug 2021 14:40:12 -0400 Subject: [PATCH 18/27] Updating ADF connection doc --- .../connections/adf.rst | 28 +++++++++++-------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/docs/apache-airflow-providers-microsoft-azure/connections/adf.rst b/docs/apache-airflow-providers-microsoft-azure/connections/adf.rst index 6c38b94e5e6cc..6750c4501b2c8 100644 --- a/docs/apache-airflow-providers-microsoft-azure/connections/adf.rst +++ b/docs/apache-airflow-providers-microsoft-azure/connections/adf.rst @@ -41,24 +41,30 @@ All hooks and operators related to Microsoft Azure Data Factory use ``azure_data Configuring the Connection -------------------------- -Login +Client ID Specify the ``client_id`` used for the initial connection. This is needed for *token credentials* authentication mechanism. -Password +Secret Specify the ``secret`` used for the initial connection. This is needed for *token credentials* authentication mechanism. -Extra (optional) - Specify the extra parameters (as json dictionary) that can be used in Azure Data Lake connection. - The following parameters are all optional: +Tenant ID (optional) + Specify the ``tenantId`` used for the initial connection. + This is needed for *token credentials* authentication mechanism. + +Subscription ID (optional) + Specify the ``subscriptionId`` used for the initial connection. + This is meeded for *token credentials* authentication mechanism. + +Factory Name (optional) + Specify the Azure Data Factory to interface with. + If not specified in the connection, this needs to be passed in directly to hooks, operators, and sensors. + +Resource Group Name (optional) + Specify the Azure Resource Group Name under which the desired data factory resides. + If not specified in the connection, this needs to be passed in directly to hooks, operators, and sensors. - * ``tenantId``: Specify the tenant to use. - This is needed for *token credentials* authentication mechanism. - * ``subscriptionId``: Specify the subscription id to use. - This is needed for *token credentials* authentication mechanism. - * ``resourceGroup``: Specify the azure resource group name. - * ``factory``: Specify the azure data factory to use When specifying the connection in environment variable you should specify it using URI syntax. From 942591af07bb8cd99d27f064aef7bf12dc6a90fb Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Sat, 28 Aug 2021 16:22:00 -0400 Subject: [PATCH 19/27] Adding operator documentation --- .../example_dags/example_adf_run_pipeline.py | 8 +-- .../azure/operators/azure_data_factory.py | 4 ++ .../connections/adf.rst | 8 +-- .../operators/adf_run_pipeline.rst | 49 +++++++++++++++++++ 4 files changed, 62 insertions(+), 7 deletions(-) create mode 100644 docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst diff --git a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py index beab8da5a8dd7..5959e595a62d2 100644 --- a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py +++ b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py @@ -34,23 +34,25 @@ "retries": 1, "retry_delay": timedelta(minutes=3), "conn_id": "azure_data_factory", - "factory_name": "my-data-factory", # This can also be specified in the ADF connection. - "resource_group_name": "my-resource-group", # This can also be specified in the ADF connection. + "factory_name": "my-data-factory", # This can also be specified in the ADF connection. + "resource_group_name": "my-resource-group", # This can also be specified in the ADF connection. }, default_view="graph", ) as dag: begin = DummyOperator(task_id="begin") end = DummyOperator(task_id="end") + # [START howto_operator_adf_run_pipeline] run_pipeline1 = AzureDataFactoryRunPipelineOperator( task_id="run_pipeline1", pipeline_name="pipeline1", parameters={"myParam": "value"}, ) + # [END howto_operator_adf_run_pipeline] run_pipeline2 = AzureDataFactoryRunPipelineOperator( task_id="run_pipeline2", - pipeline_name="extractDailyExchangeRates", + pipeline_name="run_pipeline2", wait_for_completion=False, ) diff --git a/airflow/providers/microsoft/azure/operators/azure_data_factory.py b/airflow/providers/microsoft/azure/operators/azure_data_factory.py index 4018b504bbe77..3a93c49f41978 100644 --- a/airflow/providers/microsoft/azure/operators/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/operators/azure_data_factory.py @@ -30,6 +30,10 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): """ Executes a data factory pipeline. + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:AzureDataFactoryRunPipelineOperator` + :param conn_id: The connection identifier for connecting to Azure Data Factory. :type conn_id: str :param pipeline_name: The name of the pipeline to execute. diff --git a/docs/apache-airflow-providers-microsoft-azure/connections/adf.rst b/docs/apache-airflow-providers-microsoft-azure/connections/adf.rst index 6750c4501b2c8..034f8ae806dbd 100644 --- a/docs/apache-airflow-providers-microsoft-azure/connections/adf.rst +++ b/docs/apache-airflow-providers-microsoft-azure/connections/adf.rst @@ -19,7 +19,7 @@ .. _howto/connection:adf: -Microsoft Azure Data Factory Connection +Microsoft Azure Data Factory ======================================= The Microsoft Azure Data Factory connection type enables the Azure Data Factory Integrations. @@ -49,13 +49,13 @@ Secret Specify the ``secret`` used for the initial connection. This is needed for *token credentials* authentication mechanism. -Tenant ID (optional) +Tenant ID Specify the ``tenantId`` used for the initial connection. This is needed for *token credentials* authentication mechanism. -Subscription ID (optional) +Subscription ID Specify the ``subscriptionId`` used for the initial connection. - This is meeded for *token credentials* authentication mechanism. + This is needed for *token credentials* authentication mechanism. Factory Name (optional) Specify the Azure Data Factory to interface with. diff --git a/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst b/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst new file mode 100644 index 0000000000000..9989402190ea5 --- /dev/null +++ b/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst @@ -0,0 +1,49 @@ + + .. 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. + +Azure Data Factory Operators +============================ +Azure Data Factory is Azure's cloud ETL service for scale-out serverless data integration and data transformation. +It offers a code-free UI for intuitive authoring and single-pane-of-glass monitoring and management. + +.. contents:: + :depth: 1 + :local: + +.. _howto/operator:AzureDataFactoryRunPipelineOperator: + +AzureDataFactoryRunPipelineOperator +----------------------------------- +Use the :class:`~airflow.providers.microsoft.azure.operators.azure_data_factory.AzureDataFactoryRunPipelineOperator` to execute a pipeline within a data factory. +By default, the operator will check on the status of the executed pipeline and mirror its result (i.e. succeed if the pipeline succeeds or fail if the pipeline fails or is canceled). +This functionality can be disabled for an asynchronous wait -- typically with the :class:`~airflow.providers.microsoft.azure.sensors.azure_data_factory.AzureDataFactoryPipelineRunSensor` -- by setting ``wait_for_completion`` to False. + +Below is an example of using this operator to execute an Azure Data Factory pipeline. + + .. exampleinclude:: /../../airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py + :language: python + :dedent: 0 + :start-after: [START howto_operator_adf_run_pipeline] + :end-before: [END howto_operator_adf_run_pipeline] + +Reference +--------- + +For further information, please refer to the Microsoft documentation: + + * `Azure Data Factory Documentation `__ From 5c5105cdce0d8fc1c0ed593a564bae095b3a8a99 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Sat, 28 Aug 2021 16:53:53 -0400 Subject: [PATCH 20/27] Trivial update to example DAG --- .../microsoft/azure/example_dags/example_adf_run_pipeline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py index 5959e595a62d2..45870cf05dec3 100644 --- a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py +++ b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py @@ -52,7 +52,7 @@ run_pipeline2 = AzureDataFactoryRunPipelineOperator( task_id="run_pipeline2", - pipeline_name="run_pipeline2", + pipeline_name="pipeline2", wait_for_completion=False, ) From 51cc0c1a61454cda79efcd24fe8bcb246db46c34 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Sat, 28 Aug 2021 18:19:46 -0400 Subject: [PATCH 21/27] Adding how-to doc entry in provider.yaml --- airflow/providers/microsoft/azure/provider.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index c3cc75e4b34c5..6bf887c85ab0b 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -69,6 +69,8 @@ integrations: logo: /integration-logos/azure/Microsoft-Azure-Fileshare.png tags: [azure] - integration-name: Microsoft Azure Data Factory + how-to-guide: + - /docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst external-doc-url: https://azure.microsoft.com/en-us/services/data-factory/ logo: /integration-logos/azure/Azure Data Factory.svg tags: [azure] From 2f37c55c89140b11006017125ad1618565a933e6 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Sun, 29 Aug 2021 23:20:44 -0400 Subject: [PATCH 22/27] AIP-21 + conn_id convention updates --- .../azure/example_dags/example_adf_run_pipeline.py | 8 +++----- .../{azure_data_factory.py => data_factory.py} | 12 ++++++------ .../{azure_data_factory.py => data_factory.py} | 14 +++++++------- airflow/providers/microsoft/azure/provider.yaml | 10 +++++----- .../{azure_data_factory.py => data_factory.py} | 14 +++++++------- .../operators/adf_run_pipeline.rst | 4 ++-- .../azure/hooks/test_azure_data_factory.py | 4 ++-- .../azure/operators/test_azure_data_factory.py | 10 +++++----- .../azure/sensors/test_azure_data_factory.py | 10 ++++------ 9 files changed, 41 insertions(+), 45 deletions(-) rename airflow/providers/microsoft/azure/hooks/{azure_data_factory.py => data_factory.py} (98%) rename airflow/providers/microsoft/azure/operators/{azure_data_factory.py => data_factory.py} (94%) rename airflow/providers/microsoft/azure/sensors/{azure_data_factory.py => data_factory.py} (83%) diff --git a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py index 45870cf05dec3..03212255db547 100644 --- a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py +++ b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py @@ -19,10 +19,8 @@ from airflow.models import DAG from airflow.operators.dummy import DummyOperator -from airflow.providers.microsoft.azure.operators.azure_data_factory import AzureDataFactoryRunPipelineOperator -from airflow.providers.microsoft.azure.sensors.azure_data_factory import ( - AzureDataFactoryPipelineRunStatusSensor, -) +from airflow.providers.microsoft.azure.operators.data_factory import AzureDataFactoryRunPipelineOperator +from airflow.providers.microsoft.azure.sensors.data_factory import AzureDataFactoryPipelineRunStatusSensor from airflow.utils.edgemodifier import Label with DAG( @@ -33,7 +31,7 @@ default_args={ "retries": 1, "retry_delay": timedelta(minutes=3), - "conn_id": "azure_data_factory", + "azure_data_factory_conn_id": "azure_data_factory", "factory_name": "my-data-factory", # This can also be specified in the ADF connection. "resource_group_name": "my-resource-group", # This can also be specified in the ADF connection. }, diff --git a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py b/airflow/providers/microsoft/azure/hooks/data_factory.py similarity index 98% rename from airflow/providers/microsoft/azure/hooks/azure_data_factory.py rename to airflow/providers/microsoft/azure/hooks/data_factory.py index c4d3648dc6ca0..182ca80f8ca29 100644 --- a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/hooks/data_factory.py @@ -52,7 +52,7 @@ def bind_argument(arg, default_key): # Check if arg was not included in the function signature or, if it is, the value is not provided. if arg not in bound_args.arguments or bound_args.arguments[arg] is None: self = args[0] - conn = self.get_connection(self.conn_id) + conn = self.get_connection(self.azure_data_factory_conn_id) default_value = conn.extra_dejson.get(default_key) if not default_value: raise AirflowException("Could not determine the targeted data factory.") @@ -84,8 +84,8 @@ class AzureDataFactoryHook(BaseHook): """ A hook to interact with Azure Data Factory. - :param conn_id: The :ref:`Azure Data Factory connection id`. - :type conn_id: str + :param azure_data_factory_conn_id: The :ref:`Azure Data Factory connection id`. + :type azure_data_factory_conn_id: str """ conn_type: str = 'azure_data_factory' @@ -126,16 +126,16 @@ def get_ui_field_behaviour() -> Dict: }, } - def __init__(self, conn_id: Optional[str] = default_conn_name): + def __init__(self, azure_data_factory_conn_id: Optional[str] = default_conn_name): self._conn: DataFactoryManagementClient = None - self.conn_id = conn_id + self.azure_data_factory_conn_id = azure_data_factory_conn_id super().__init__() def get_conn(self) -> DataFactoryManagementClient: if self._conn is not None: return self._conn - conn = self.get_connection(self.conn_id) + conn = self.get_connection(self.azure_data_factory_conn_id) tenant = conn.extra_dejson.get('extra__azure_data_factory__tenantId') subscription_id = conn.extra_dejson.get('extra__azure_data_factory__subscriptionId') diff --git a/airflow/providers/microsoft/azure/operators/azure_data_factory.py b/airflow/providers/microsoft/azure/operators/data_factory.py similarity index 94% rename from airflow/providers/microsoft/azure/operators/azure_data_factory.py rename to airflow/providers/microsoft/azure/operators/data_factory.py index 3a93c49f41978..72fb29171a89f 100644 --- a/airflow/providers/microsoft/azure/operators/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/operators/data_factory.py @@ -20,7 +20,7 @@ from airflow.exceptions import AirflowException from airflow.models import BaseOperator -from airflow.providers.microsoft.azure.hooks.azure_data_factory import ( +from airflow.providers.microsoft.azure.hooks.data_factory import ( AzureDataFactoryHook, AzureDataFactoryPipelineRunStatus, ) @@ -34,8 +34,8 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:AzureDataFactoryRunPipelineOperator` - :param conn_id: The connection identifier for connecting to Azure Data Factory. - :type conn_id: str + :param azure_data_factory_conn_id: The connection identifier for connecting to Azure Data Factory. + :type azure_data_factory_conn_id: str :param pipeline_name: The name of the pipeline to execute. :type pipeline_name: str :param resource_group_name: The resource group name. If a value is not passed in to the operator, the @@ -75,7 +75,7 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): """ template_fields = ( - "conn_id", + "azure_data_factory_conn_id", "resource_group_name", "factory_name", "pipeline_name", @@ -87,7 +87,7 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): def __init__( self, *, - conn_id: str, + azure_data_factory_conn_id: str, pipeline_name: str, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, @@ -102,7 +102,7 @@ def __init__( **kwargs, ) -> None: super().__init__(**kwargs) - self.conn_id = conn_id + self.azure_data_factory_conn_id = azure_data_factory_conn_id self.pipeline_name = pipeline_name self.resource_group_name = resource_group_name self.factory_name = factory_name @@ -116,7 +116,7 @@ def __init__( self.poke_interval = poke_interval def execute(self, context: Dict) -> None: - self.hook = AzureDataFactoryHook(conn_id=self.conn_id) + self.hook = AzureDataFactoryHook(azure_data_factory_conn_id=self.azure_data_factory_conn_id) self.log.info(f"Executing the '{self.pipeline_name}' pipeline.") response = self.hook.run_pipeline( pipeline_name=self.pipeline_name, diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index 6bf887c85ab0b..63ede42a4f8ea 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -101,7 +101,7 @@ operators: - airflow.providers.microsoft.azure.operators.wasb_delete_blob - integration-name: Microsoft Azure Data Factory python-modules: - - airflow.providers.microsoft.azure.operators.azure_data_factory + - airflow.providers.microsoft.azure.operators.data_factory sensors: - integration-name: Microsoft Azure Cosmos DB @@ -112,7 +112,7 @@ sensors: - airflow.providers.microsoft.azure.sensors.wasb - integration-name: Microsoft Azure Data Factory python-modules: - - airflow.providers.microsoft.azure.sensors.azure_data_factory + - airflow.providers.microsoft.azure.sensors.data_factory hooks: - integration-name: Microsoft Azure Container Instances @@ -143,7 +143,7 @@ hooks: - airflow.providers.microsoft.azure.hooks.wasb - integration-name: Microsoft Azure Data Factory python-modules: - - airflow.providers.microsoft.azure.hooks.azure_data_factory + - airflow.providers.microsoft.azure.hooks.data_factory transfers: - source-integration-name: Local @@ -171,7 +171,7 @@ hook-class-names: # deprecated - to be removed after providers add dependency o - airflow.providers.microsoft.azure.hooks.azure_container_volume.AzureContainerVolumeHook - airflow.providers.microsoft.azure.hooks.azure_container_instance.AzureContainerInstanceHook - airflow.providers.microsoft.azure.hooks.wasb.WasbHook - - airflow.providers.microsoft.azure.hooks.azure_data_factory.AzureDataFactoryHook + - airflow.providers.microsoft.azure.hooks.data_factory.AzureDataFactoryHook - airflow.providers.microsoft.azure.hooks.azure_container_registry.AzureContainerRegistryHook connection-types: @@ -194,7 +194,7 @@ connection-types: connection-type: azure_container_instance - hook-class-name: airflow.providers.microsoft.azure.hooks.wasb.WasbHook connection-type: wasb - - hook-class-name: airflow.providers.microsoft.azure.hooks.azure_data_factory.AzureDataFactoryHook + - hook-class-name: airflow.providers.microsoft.azure.hooks.data_factory.AzureDataFactoryHook connection-type: azure_data_factory - hook-class-name: >- airflow.providers.microsoft.azure.hooks.azure_container_registry.AzureContainerRegistryHook diff --git a/airflow/providers/microsoft/azure/sensors/azure_data_factory.py b/airflow/providers/microsoft/azure/sensors/data_factory.py similarity index 83% rename from airflow/providers/microsoft/azure/sensors/azure_data_factory.py rename to airflow/providers/microsoft/azure/sensors/data_factory.py index 625bb0d71ed02..201cc6f462ddb 100644 --- a/airflow/providers/microsoft/azure/sensors/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/sensors/data_factory.py @@ -18,7 +18,7 @@ from typing import Dict, Optional from airflow.exceptions import AirflowException -from airflow.providers.microsoft.azure.hooks.azure_data_factory import ( +from airflow.providers.microsoft.azure.hooks.data_factory import ( AzureDataFactoryHook, AzureDataFactoryPipelineRunStatus, ) @@ -29,8 +29,8 @@ class AzureDataFactoryPipelineRunStatusSensor(BaseSensorOperator): """ Checks the status of a pipeline run. - :param conn_id: The connection identifier for connecting to Azure Data Factory. - :type conn_id: str + :param azure_data_factory_conn_id: The connection identifier for connecting to Azure Data Factory. + :type azure_data_factory_conn_id: str :param run_id: The pipeline run identifier. :type run_id: str :param resource_group_name: The resource group name. @@ -39,26 +39,26 @@ class AzureDataFactoryPipelineRunStatusSensor(BaseSensorOperator): :type factory_name: str """ - template_fields = ("conn_id", "resource_group_name", "factory_name", "run_id") + template_fields = ("azure_data_factory_conn_id", "resource_group_name", "factory_name", "run_id") def __init__( self, *, - conn_id: str, + azure_data_factory_conn_id: str, run_id: str, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, **kwargs, ) -> None: super().__init__(**kwargs) - self.conn_id = conn_id + self.azure_data_factory_conn_id = azure_data_factory_conn_id self.run_id = run_id self.resource_group_name = resource_group_name self.factory_name = factory_name def poke(self, context: Dict) -> bool: self.log.info(f"Checking the status for pipeline run {self.run_id}.") - self.hook = AzureDataFactoryHook(conn_id=self.conn_id) + self.hook = AzureDataFactoryHook(azure_data_factory_conn_id=self.azure_data_factory_conn_id) pipeline_run = self.hook.get_pipeline_run( run_id=self.run_id, factory_name=self.factory_name, diff --git a/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst b/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst index 9989402190ea5..eea04822a5a79 100644 --- a/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst +++ b/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst @@ -29,9 +29,9 @@ It offers a code-free UI for intuitive authoring and single-pane-of-glass monito AzureDataFactoryRunPipelineOperator ----------------------------------- -Use the :class:`~airflow.providers.microsoft.azure.operators.azure_data_factory.AzureDataFactoryRunPipelineOperator` to execute a pipeline within a data factory. +Use the :class:`~airflow.providers.microsoft.azure.operators.data_factory.AzureDataFactoryRunPipelineOperator` to execute a pipeline within a data factory. By default, the operator will check on the status of the executed pipeline and mirror its result (i.e. succeed if the pipeline succeeds or fail if the pipeline fails or is canceled). -This functionality can be disabled for an asynchronous wait -- typically with the :class:`~airflow.providers.microsoft.azure.sensors.azure_data_factory.AzureDataFactoryPipelineRunSensor` -- by setting ``wait_for_completion`` to False. +This functionality can be disabled for an asynchronous wait -- typically with the :class:`~airflow.providers.microsoft.azure.sensors.data_factory.AzureDataFactoryPipelineRunSensor` -- by setting ``wait_for_completion`` to False. Below is an example of using this operator to execute an Azure Data Factory pipeline. diff --git a/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py b/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py index 7939a798883d5..6cafa5b9cb209 100644 --- a/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py @@ -24,7 +24,7 @@ from airflow.exceptions import AirflowException from airflow.models.connection import Connection -from airflow.providers.microsoft.azure.hooks.azure_data_factory import ( +from airflow.providers.microsoft.azure.hooks.data_factory import ( AzureDataFactoryHook, provide_targeted_factory, ) @@ -62,7 +62,7 @@ def setup_module(): @fixture def hook(): - client = AzureDataFactoryHook(conn_id="azure_data_factory_test") + client = AzureDataFactoryHook(azure_data_factory_conn_id="azure_data_factory_test") client._conn = MagicMock( spec=[ "factories", diff --git a/tests/providers/microsoft/azure/operators/test_azure_data_factory.py b/tests/providers/microsoft/azure/operators/test_azure_data_factory.py index e6aa4e8789fe6..4f72447c6ad31 100644 --- a/tests/providers/microsoft/azure/operators/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/operators/test_azure_data_factory.py @@ -22,11 +22,11 @@ from parameterized import parameterized from airflow.exceptions import AirflowException -from airflow.providers.microsoft.azure.hooks.azure_data_factory import ( +from airflow.providers.microsoft.azure.hooks.data_factory import ( AzureDataFactoryHook, AzureDataFactoryPipelineRunStatus, ) -from airflow.providers.microsoft.azure.operators.azure_data_factory import AzureDataFactoryRunPipelineOperator +from airflow.providers.microsoft.azure.operators.data_factory import AzureDataFactoryRunPipelineOperator PIPELINE_RUN_RESPONSE = {"additional_properties": {}, "run_id": "run_id"} @@ -37,7 +37,7 @@ def setUp(self): self.mock_context = {"ti": self.mock_ti} self.config = { "task_id": "run_pipeline_op", - "conn_id": "azure_data_factory_test", + "azure_data_factory_conn_id": "azure_data_factory_test", "pipeline_name": "pipeline1", "resource_group_name": "resource-group-name", "factory_name": "factory-name", @@ -64,7 +64,7 @@ def create_pipeline_run(status: str): def test_execute_wait_for_completion(self, pipeline_run_status, expected_status, mock_run_pipeline): operator = AzureDataFactoryRunPipelineOperator(**self.config) - assert operator.conn_id == self.config["conn_id"] + assert operator.azure_data_factory_conn_id == self.config["azure_data_factory_conn_id"] assert operator.pipeline_name == self.config["pipeline_name"] assert operator.resource_group_name == self.config["resource_group_name"] assert operator.factory_name == self.config["factory_name"] @@ -125,7 +125,7 @@ def test_execute_wait_for_completion(self, pipeline_run_status, expected_status, def test_execute_no_wait_for_completion(self, mock_run_pipeline): operator = AzureDataFactoryRunPipelineOperator(wait_for_completion=False, **self.config) - assert operator.conn_id == self.config["conn_id"] + assert operator.azure_data_factory_conn_id == self.config["azure_data_factory_conn_id"] assert operator.pipeline_name == self.config["pipeline_name"] assert operator.resource_group_name == self.config["resource_group_name"] assert operator.factory_name == self.config["factory_name"] diff --git a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py index 8afb10374609e..67de2952ddb5c 100644 --- a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py @@ -25,20 +25,18 @@ from airflow.exceptions import AirflowException from airflow.models.dag import DAG -from airflow.providers.microsoft.azure.hooks.azure_data_factory import ( +from airflow.providers.microsoft.azure.hooks.data_factory import ( AzureDataFactoryHook, AzureDataFactoryPipelineRunStatus, ) -from airflow.providers.microsoft.azure.sensors.azure_data_factory import ( - AzureDataFactoryPipelineRunStatusSensor, -) +from airflow.providers.microsoft.azure.sensors.data_factory import AzureDataFactoryPipelineRunStatusSensor class TestPipelineRunStatusSensor(unittest.TestCase): def setUp(self): self.dag = DAG("test", start_date=datetime(2021, 8, 16), schedule_interval=None, catchup=False) self.config = { - "conn_id": "azure_data_factory_test", + "azure_data_factory_conn_id": "azure_data_factory_test", "run_id": "run_id", "resource_group_name": "resource-group-name", "factory_name": "factory-name", @@ -58,7 +56,7 @@ def test_init(self): sensor = AzureDataFactoryPipelineRunStatusSensor( task_id="pipeline_run_sensor", dag=self.dag, **self.config ) - assert sensor.conn_id == self.config["conn_id"] + assert sensor.azure_data_factory_conn_id == self.config["azure_data_factory_conn_id"] assert sensor.run_id == self.config["run_id"] assert sensor.resource_group_name == self.config["resource_group_name"] assert sensor.factory_name == self.config["factory_name"] From deccf389a40a781f49321065fcd91f7d0737970e Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Mon, 30 Aug 2021 16:19:51 -0400 Subject: [PATCH 23/27] Refactoring the check of a pipeline run status --- .../example_dags/example_adf_run_pipeline.py | 5 +- .../microsoft/azure/hooks/data_factory.py | 89 ++++++++++++++++++- .../microsoft/azure/operators/data_factory.py | 71 +++++---------- .../microsoft/azure/sensors/data_factory.py | 27 ++---- .../operators/adf_run_pipeline.rst | 10 ++- 5 files changed, 126 insertions(+), 76 deletions(-) diff --git a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py index 03212255db547..614decde169ad 100644 --- a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py +++ b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py @@ -48,17 +48,18 @@ ) # [END howto_operator_adf_run_pipeline] + # [START howto_operator_adf_run_pipeline_async] run_pipeline2 = AzureDataFactoryRunPipelineOperator( task_id="run_pipeline2", pipeline_name="pipeline2", - wait_for_completion=False, + wait_for_pipeline_run=False, ) wait_for_pipeline_run = AzureDataFactoryPipelineRunStatusSensor( task_id="wait_for_pipeline_run", run_id=run_pipeline2.output["run_id"], - poke_interval=10, ) + # [END howto_operator_adf_run_pipeline_async] begin >> Label("No async wait") >> run_pipeline1 >> end begin >> Label("Do async wait with sensor") >> run_pipeline2 diff --git a/airflow/providers/microsoft/azure/hooks/data_factory.py b/airflow/providers/microsoft/azure/hooks/data_factory.py index 182ca80f8ca29..65d5161ef8771 100644 --- a/airflow/providers/microsoft/azure/hooks/data_factory.py +++ b/airflow/providers/microsoft/azure/hooks/data_factory.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. import inspect +import time from functools import wraps from typing import Any, Callable, Dict, Optional @@ -33,6 +34,7 @@ from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook +from airflow.typing_compat import Literal def provide_targeted_factory(func: Callable) -> Callable: @@ -52,7 +54,7 @@ def bind_argument(arg, default_key): # Check if arg was not included in the function signature or, if it is, the value is not provided. if arg not in bound_args.arguments or bound_args.arguments[arg] is None: self = args[0] - conn = self.get_connection(self.azure_data_factory_conn_id) + conn = self.get_connection(self.conn_id) default_value = conn.extra_dejson.get(default_key) if not default_value: raise AirflowException("Could not determine the targeted data factory.") @@ -80,6 +82,10 @@ class AzureDataFactoryPipelineRunStatus: TERMINAL_STATUSES = {CANCELLED, FAILED, SUCCEEDED} +class AzureDataFactoryPipelineRunException(AirflowException): + """An exception that indicates a pipeline run failed to complete.""" + + class AzureDataFactoryHook(BaseHook): """ A hook to interact with Azure Data Factory. @@ -128,14 +134,14 @@ def get_ui_field_behaviour() -> Dict: def __init__(self, azure_data_factory_conn_id: Optional[str] = default_conn_name): self._conn: DataFactoryManagementClient = None - self.azure_data_factory_conn_id = azure_data_factory_conn_id + self.conn_id = azure_data_factory_conn_id super().__init__() def get_conn(self) -> DataFactoryManagementClient: if self._conn is not None: return self._conn - conn = self.get_connection(self.azure_data_factory_conn_id) + conn = self.get_connection(self.conn_id) tenant = conn.extra_dejson.get('extra__azure_data_factory__tenantId') subscription_id = conn.extra_dejson.get('extra__azure_data_factory__subscriptionId') @@ -579,6 +585,83 @@ def get_pipeline_run( """ return self.get_conn().pipeline_runs.get(resource_group_name, factory_name, run_id, **config) + @provide_targeted_factory + def check_pipeline_run_status( + self, + run_id: str, + mode: Literal["poke", "wait"] = "wait", + expected_status: str = AzureDataFactoryPipelineRunStatus.SUCCEEDED, + check_interval: Optional[int] = 60, + timeout: Optional[int] = 60 * 60 * 24 * 7, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> bool: + """ + Checks a pipeline run's status with a configurable option to do so until the run has reached a + terminal status or the desired status. + + :param run_id: The pipeline run identifier. + :param mode: If set to "wait", periodically check a pipeline run's status to see if it has reached a + terminal or desired status. If set to "poke", check a pipeline run's status only once and return + whether or not the status matches the ``expected_status``. + :param expected_status: The desired status of a pipeline run. + :param check_interval: Time in seconds to check a pipeline run's status when ``mode`` is set to + "wait". + :param timeout: Time in seconds to wait for a pipeline to reach a terminal status when ``mode`` is set + to "wait". + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :return: Boolean value indicating if the current pipeline run's status matches the + ``expected_status``. + """ + if mode not in ("wait", "poke"): + raise ValueError( + f"The configured mode, '{mode}', is not supported. Please use 'wait' or 'poke'." + ) + + pipeline_run = self.get_pipeline_run( + run_id=run_id, + factory_name=factory_name, + resource_group_name=resource_group_name, + **config, + ) + pipeline_run_status = pipeline_run.status + + if mode == "wait": + start_time = time.monotonic() + + while ( + pipeline_run_status not in AzureDataFactoryPipelineRunStatus.TERMINAL_STATUSES + and pipeline_run_status != expected_status + ): + # Check if the pipeline-run duration has exceeded the ``timeout`` configured. + if start_time + timeout < time.monotonic(): + raise AzureDataFactoryPipelineRunException( + f"Pipeline run {run_id} has not reached a terminal status after {timeout} seconds." + ) + + # Wait to check the status of the pipeline based on the ``check_interval`` configured. + time.sleep(check_interval) + + self.log.info(f"Checking on the status of run ID {run_id}.") + pipeline_run = self.get_pipeline_run( + run_id=run_id, + factory_name=factory_name, + resource_group_name=resource_group_name, + **config, + ) + pipeline_run_status = pipeline_run.status + self.log.info(f"Current status of run ID {run_id}: {pipeline_run_status}") + + if pipeline_run_status == AzureDataFactoryPipelineRunStatus.FAILED: + raise AzureDataFactoryPipelineRunException(f"Pipeline run {run_id} has failed.") + + if pipeline_run_status == AzureDataFactoryPipelineRunStatus.CANCELLED: + raise AzureDataFactoryPipelineRunException(f"Pipeline run {run_id} has been cancelled.") + + return pipeline_run_status == expected_status + @provide_targeted_factory def cancel_pipeline_run( self, diff --git a/airflow/providers/microsoft/azure/operators/data_factory.py b/airflow/providers/microsoft/azure/operators/data_factory.py index 72fb29171a89f..3ca1a14d4f848 100644 --- a/airflow/providers/microsoft/azure/operators/data_factory.py +++ b/airflow/providers/microsoft/azure/operators/data_factory.py @@ -15,15 +15,10 @@ # specific language governing permissions and limitations # under the License. -import time from typing import Any, Dict, Optional -from airflow.exceptions import AirflowException from airflow.models import BaseOperator -from airflow.providers.microsoft.azure.hooks.data_factory import ( - AzureDataFactoryHook, - AzureDataFactoryPipelineRunStatus, -) +from airflow.providers.microsoft.azure.hooks.data_factory import AzureDataFactoryHook class AzureDataFactoryRunPipelineOperator(BaseOperator): @@ -38,6 +33,10 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): :type azure_data_factory_conn_id: str :param pipeline_name: The name of the pipeline to execute. :type pipeline_name: str + :param wait_for_pipeline_run: Flag to wait on a pipeline run's completion. By default, this feature is + enabled but could be disabled to wait for a long-running pipeline execution using the + ``AzureDataFactoryPipelineRunSensor`` rather than this operator. + :type wait_for_pipeline_run: bool :param resource_group_name: The resource group name. If a value is not passed in to the operator, the ``AzureDataFactoryHook`` will attempt to use the resource group name provided in the corresponding connection. @@ -62,16 +61,12 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): ``@pipeline().parameters.parameterName`` and will be used only if the ``reference_pipeline_run_id`` is not specified. :type start_from_failure: Dict[str, Any] - :param wait_for_completion: Flag to wait on a pipeline run's completion. By default, this feature is - enabled but could be disabled to wait for a long-running pipeline execution using the - ``AzureDataFactoryPipelineRunSensor`` rather than this operator. - :type wait_for_completion: bool :param timeout: Time in seconds to wait for a pipeline to reach a terminal status for non-asynchronous - waits. Used only if ``wait_for_completion`` is True. + waits. Used only if ``wait_for_pipeline_run`` is True. :type timeout: int - :param poke_interval: Time in seconds to check on a pipeline run's status for non-asynchronous waits. Used - only if ``wait_for_completion`` is True. - :type poke_interval: int + :param check_interval: Time in seconds to check on a pipeline run's status for non-asynchronous waits. + Used only if ``wait_for_pipeline_run`` is True. + :type check_interval: int """ template_fields = ( @@ -89,6 +84,7 @@ def __init__( *, azure_data_factory_conn_id: str, pipeline_name: str, + wait_for_pipeline_run: bool = True, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, reference_pipeline_run_id: Optional[str] = None, @@ -96,14 +92,14 @@ def __init__( start_activity_name: Optional[str] = None, start_from_failure: Optional[bool] = None, parameters: Optional[Dict[str, Any]] = None, - wait_for_completion: Optional[bool] = True, timeout: Optional[int] = 60 * 60 * 24 * 7, - poke_interval: Optional[int] = 30, + check_interval: Optional[int] = 60, **kwargs, ) -> None: super().__init__(**kwargs) self.azure_data_factory_conn_id = azure_data_factory_conn_id self.pipeline_name = pipeline_name + self.wait_for_pipeline_run = wait_for_pipeline_run self.resource_group_name = resource_group_name self.factory_name = factory_name self.reference_pipeline_run_id = reference_pipeline_run_id @@ -111,9 +107,8 @@ def __init__( self.start_activity_name = start_activity_name self.start_from_failure = start_from_failure self.parameters = parameters - self.wait_for_completion = wait_for_completion self.timeout = timeout - self.poke_interval = poke_interval + self.check_interval = check_interval def execute(self, context: Dict) -> None: self.hook = AzureDataFactoryHook(azure_data_factory_conn_id=self.azure_data_factory_conn_id) @@ -134,36 +129,16 @@ def execute(self, context: Dict) -> None: # asynchronous wait. context["ti"].xcom_push(key="run_id", value=self.run_id) - if self.wait_for_completion: - self.log.info(f"Waiting for run ID {self.run_id} of pipeline {self.pipeline_name} to complete.") - start_time = time.monotonic() - pipeline_run_status = None - while pipeline_run_status not in AzureDataFactoryPipelineRunStatus.TERMINAL_STATUSES: - # Check to see if the pipeline-run duration has exceeded the ``timeout`` configured. - if start_time + self.timeout < time.monotonic(): - raise AirflowException( - f"Pipeline run {self.run_id} has not reached a terminal status after " - f"{self.timeout} seconds." - ) - - # Wait to check the status of the pipeline based on the ``poke_interval`` configured. - time.sleep(self.poke_interval) - - self.log.info(f"Checking on the status of run ID {self.run_id}.") - pipeline_run = self.hook.get_pipeline_run( - run_id=self.run_id, - factory_name=self.factory_name, - resource_group_name=self.resource_group_name, - ) - pipeline_run_status = pipeline_run.status - self.log.info(f"Run ID {self.run_id} is in a status of {pipeline_run_status}.") - - if pipeline_run_status == AzureDataFactoryPipelineRunStatus.CANCELLED: - raise AirflowException(f"Pipeline run {self.run_id} has been cancelled.") - - if pipeline_run_status == AzureDataFactoryPipelineRunStatus.FAILED: - raise AirflowException(f"Pipeline run {self.run_id} has failed.") - + if self.wait_for_pipeline_run: + self.log.info(f"Waiting for pipeline run {self.run_id} to complete.") + self.hook.check_pipeline_run_status( + run_id=self.run_id, + mode="wait", + check_interval=self.check_interval, + timeout=self.timeout, + resource_group_name=self.resource_group_name, + factory_name=self.factory_name, + ) self.log.info(f"Pipeline run {self.run_id} has completed successfully.") def on_kill(self) -> None: diff --git a/airflow/providers/microsoft/azure/sensors/data_factory.py b/airflow/providers/microsoft/azure/sensors/data_factory.py index 201cc6f462ddb..b52412aac6d15 100644 --- a/airflow/providers/microsoft/azure/sensors/data_factory.py +++ b/airflow/providers/microsoft/azure/sensors/data_factory.py @@ -17,11 +17,7 @@ from typing import Dict, Optional -from airflow.exceptions import AirflowException -from airflow.providers.microsoft.azure.hooks.data_factory import ( - AzureDataFactoryHook, - AzureDataFactoryPipelineRunStatus, -) +from airflow.providers.microsoft.azure.hooks.data_factory import AzureDataFactoryHook from airflow.sensors.base import BaseSensorOperator @@ -59,23 +55,10 @@ def __init__( def poke(self, context: Dict) -> bool: self.log.info(f"Checking the status for pipeline run {self.run_id}.") self.hook = AzureDataFactoryHook(azure_data_factory_conn_id=self.azure_data_factory_conn_id) - pipeline_run = self.hook.get_pipeline_run( + + return self.hook.check_pipeline_run_status( run_id=self.run_id, - factory_name=self.factory_name, + mode="poke", resource_group_name=self.resource_group_name, + factory_name=self.factory_name, ) - pipeline_run_status = pipeline_run.status - self.log.info(f"Current status for pipeline run {self.run_id}: {pipeline_run_status}.") - - if pipeline_run_status == AzureDataFactoryPipelineRunStatus.SUCCEEDED: - self.log.info(f"The pipeline run {self.run_id} has succeeded.") - return True - elif pipeline_run_status in { - AzureDataFactoryPipelineRunStatus.FAILED, - AzureDataFactoryPipelineRunStatus.CANCELLED, - }: - raise AirflowException( - f"Pipeline run {self.run_id} is in a negative terminal status: {pipeline_run_status}" - ) - - return False diff --git a/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst b/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst index eea04822a5a79..7cf74589e98b6 100644 --- a/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst +++ b/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst @@ -31,7 +31,7 @@ AzureDataFactoryRunPipelineOperator ----------------------------------- Use the :class:`~airflow.providers.microsoft.azure.operators.data_factory.AzureDataFactoryRunPipelineOperator` to execute a pipeline within a data factory. By default, the operator will check on the status of the executed pipeline and mirror its result (i.e. succeed if the pipeline succeeds or fail if the pipeline fails or is canceled). -This functionality can be disabled for an asynchronous wait -- typically with the :class:`~airflow.providers.microsoft.azure.sensors.data_factory.AzureDataFactoryPipelineRunSensor` -- by setting ``wait_for_completion`` to False. +This functionality can be disabled for an asynchronous wait -- typically with the :class:`~airflow.providers.microsoft.azure.sensors.data_factory.AzureDataFactoryPipelineRunSensor` -- by setting ``wait_for_pipeline_run`` to False. Below is an example of using this operator to execute an Azure Data Factory pipeline. @@ -41,6 +41,14 @@ Below is an example of using this operator to execute an Azure Data Factory pipe :start-after: [START howto_operator_adf_run_pipeline] :end-before: [END howto_operator_adf_run_pipeline] +Here is a different example of using this operator to execute a pipeline but coupled with the :class:`~airflow.providers.microsoft.azure.sensors.data_factory.AzureDataFactoryPipelineRunSensor` to perform an asynchronous wait. + + .. exampleinclude:: /../../airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py + :language: python + :dedent: 0 + :start-after: [START howto_operator_adf_run_pipeline_async] + :end-before: [END howto_operator_adf_run_pipeline_async] + Reference --------- From 052ed37cf98dab6642940c0e76473976b2fcc285 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Wed, 1 Sep 2021 21:33:06 -0400 Subject: [PATCH 24/27] Refactoring unit test + parameter rename --- .../example_dags/example_adf_run_pipeline.py | 2 +- .../microsoft/azure/operators/data_factory.py | 17 +++-- .../microsoft/azure/sensors/data_factory.py | 2 +- .../operators/adf_run_pipeline.rst | 2 +- .../azure/hooks/test_azure_data_factory.py | 42 ++++++++++- .../operators/test_azure_data_factory.py | 27 ++++---- .../azure/sensors/test_azure_data_factory.py | 69 ++++++------------- 7 files changed, 87 insertions(+), 74 deletions(-) diff --git a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py index 614decde169ad..740c92b5bdf4c 100644 --- a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py +++ b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py @@ -52,7 +52,7 @@ run_pipeline2 = AzureDataFactoryRunPipelineOperator( task_id="run_pipeline2", pipeline_name="pipeline2", - wait_for_pipeline_run=False, + wait_for_completion=False, ) wait_for_pipeline_run = AzureDataFactoryPipelineRunStatusSensor( diff --git a/airflow/providers/microsoft/azure/operators/data_factory.py b/airflow/providers/microsoft/azure/operators/data_factory.py index 3ca1a14d4f848..3c7f93c7c5c1a 100644 --- a/airflow/providers/microsoft/azure/operators/data_factory.py +++ b/airflow/providers/microsoft/azure/operators/data_factory.py @@ -33,9 +33,9 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): :type azure_data_factory_conn_id: str :param pipeline_name: The name of the pipeline to execute. :type pipeline_name: str - :param wait_for_pipeline_run: Flag to wait on a pipeline run's completion. By default, this feature is - enabled but could be disabled to wait for a long-running pipeline execution using the - ``AzureDataFactoryPipelineRunSensor`` rather than this operator. + :param wait_for_completion: Flag to wait on a pipeline run's completion. By default, this feature is + enabled but could be disabled to perform an asynchronous wait for a long-running pipeline execution + using the ``AzureDataFactoryPipelineRunSensor``. :type wait_for_pipeline_run: bool :param resource_group_name: The resource group name. If a value is not passed in to the operator, the ``AzureDataFactoryHook`` will attempt to use the resource group name provided in the corresponding @@ -62,10 +62,10 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): not specified. :type start_from_failure: Dict[str, Any] :param timeout: Time in seconds to wait for a pipeline to reach a terminal status for non-asynchronous - waits. Used only if ``wait_for_pipeline_run`` is True. + waits. Used only if ``wait_for_completion`` is True. :type timeout: int :param check_interval: Time in seconds to check on a pipeline run's status for non-asynchronous waits. - Used only if ``wait_for_pipeline_run`` is True. + Used only if ``wait_for_completion`` is True. :type check_interval: int """ @@ -84,7 +84,7 @@ def __init__( *, azure_data_factory_conn_id: str, pipeline_name: str, - wait_for_pipeline_run: bool = True, + wait_for_completion: bool = True, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, reference_pipeline_run_id: Optional[str] = None, @@ -99,7 +99,7 @@ def __init__( super().__init__(**kwargs) self.azure_data_factory_conn_id = azure_data_factory_conn_id self.pipeline_name = pipeline_name - self.wait_for_pipeline_run = wait_for_pipeline_run + self.wait_for_completion = wait_for_completion self.resource_group_name = resource_group_name self.factory_name = factory_name self.reference_pipeline_run_id = reference_pipeline_run_id @@ -129,11 +129,10 @@ def execute(self, context: Dict) -> None: # asynchronous wait. context["ti"].xcom_push(key="run_id", value=self.run_id) - if self.wait_for_pipeline_run: + if self.wait_for_completion: self.log.info(f"Waiting for pipeline run {self.run_id} to complete.") self.hook.check_pipeline_run_status( run_id=self.run_id, - mode="wait", check_interval=self.check_interval, timeout=self.timeout, resource_group_name=self.resource_group_name, diff --git a/airflow/providers/microsoft/azure/sensors/data_factory.py b/airflow/providers/microsoft/azure/sensors/data_factory.py index b52412aac6d15..2ef66597caedb 100644 --- a/airflow/providers/microsoft/azure/sensors/data_factory.py +++ b/airflow/providers/microsoft/azure/sensors/data_factory.py @@ -58,7 +58,7 @@ def poke(self, context: Dict) -> bool: return self.hook.check_pipeline_run_status( run_id=self.run_id, - mode="poke", + wait_for_completion=False, resource_group_name=self.resource_group_name, factory_name=self.factory_name, ) diff --git a/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst b/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst index 7cf74589e98b6..6978dbbb32610 100644 --- a/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst +++ b/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst @@ -31,7 +31,7 @@ AzureDataFactoryRunPipelineOperator ----------------------------------- Use the :class:`~airflow.providers.microsoft.azure.operators.data_factory.AzureDataFactoryRunPipelineOperator` to execute a pipeline within a data factory. By default, the operator will check on the status of the executed pipeline and mirror its result (i.e. succeed if the pipeline succeeds or fail if the pipeline fails or is canceled). -This functionality can be disabled for an asynchronous wait -- typically with the :class:`~airflow.providers.microsoft.azure.sensors.data_factory.AzureDataFactoryPipelineRunSensor` -- by setting ``wait_for_pipeline_run`` to False. +This functionality can be disabled for an asynchronous wait -- typically with the :class:`~airflow.providers.microsoft.azure.sensors.data_factory.AzureDataFactoryPipelineRunSensor` -- by setting ``wait_for_completion`` to False. Below is an example of using this operator to execute an Azure Data Factory pipeline. diff --git a/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py b/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py index 6cafa5b9cb209..2be2cf55b0092 100644 --- a/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py @@ -17,7 +17,7 @@ import json -from unittest.mock import MagicMock, Mock +from unittest.mock import MagicMock, Mock, patch import pytest from pytest import fixture @@ -26,6 +26,8 @@ from airflow.models.connection import Connection from airflow.providers.microsoft.azure.hooks.data_factory import ( AzureDataFactoryHook, + AzureDataFactoryPipelineRunException, + AzureDataFactoryPipelineRunStatus, provide_targeted_factory, ) from airflow.utils import db @@ -342,6 +344,44 @@ def test_get_pipeline_run(hook: AzureDataFactoryHook, user_args, sdk_args): hook._conn.pipeline_runs.get.assert_called_with(*sdk_args) +def test_check_pipeline_run_status(hook): + config = {"run_id": ID, "timeout": 5, "check_interval": 1} + + with patch.object(AzureDataFactoryHook, "get_pipeline_run") as mock_pipeline_run: + # Begin tests when ``wait_for_completion`` is enabled -- the default behavior. + mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.SUCCEEDED + assert hook.check_pipeline_run_status(**config) + + mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.FAILED + with pytest.raises(AzureDataFactoryPipelineRunException): + hook.check_pipeline_run_status(**config) + + mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.CANCELLED + with pytest.raises(AzureDataFactoryPipelineRunException): + hook.check_pipeline_run_status(**config) + + # Begin tests when ``wait_for_completion`` is disabled (aka poking for status). + mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.SUCCEEDED + assert hook.check_pipeline_run_status(wait_for_completion=False, **config) + + mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.IN_PROGRESS + assert not hook.check_pipeline_run_status(wait_for_completion=False, **config) + + mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.QUEUED + assert not hook.check_pipeline_run_status(wait_for_completion=False, **config) + + mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.CANCELING + assert not hook.check_pipeline_run_status(wait_for_completion=False, **config) + + mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.FAILED + with pytest.raises(AzureDataFactoryPipelineRunException): + hook.check_pipeline_run_status(wait_for_completion=False, **config) + + mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.CANCELLED + with pytest.raises(AzureDataFactoryPipelineRunException): + hook.check_pipeline_run_status(wait_for_completion=False, **config) + + @parametrize( explicit_factory=((ID, RESOURCE_GROUP, FACTORY), (RESOURCE_GROUP, FACTORY, ID)), implicit_factory=((ID,), (DEFAULT_RESOURCE_GROUP, DEFAULT_FACTORY, ID)), diff --git a/tests/providers/microsoft/azure/operators/test_azure_data_factory.py b/tests/providers/microsoft/azure/operators/test_azure_data_factory.py index 4f72447c6ad31..ab718f657ce75 100644 --- a/tests/providers/microsoft/azure/operators/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/operators/test_azure_data_factory.py @@ -21,9 +21,9 @@ import pytest from parameterized import parameterized -from airflow.exceptions import AirflowException from airflow.providers.microsoft.azure.hooks.data_factory import ( AzureDataFactoryHook, + AzureDataFactoryPipelineRunException, AzureDataFactoryPipelineRunStatus, ) from airflow.providers.microsoft.azure.operators.data_factory import AzureDataFactoryRunPipelineOperator @@ -41,7 +41,7 @@ def setUp(self): "pipeline_name": "pipeline1", "resource_group_name": "resource-group-name", "factory_name": "factory-name", - "poke_interval": 1, + "check_interval": 1, } @staticmethod @@ -56,19 +56,20 @@ def create_pipeline_run(status: str): @parameterized.expand( [ (AzureDataFactoryPipelineRunStatus.SUCCEEDED, None), - (AzureDataFactoryPipelineRunStatus.FAILED, "AirflowException"), - (AzureDataFactoryPipelineRunStatus.CANCELLED, "AirflowException"), + (AzureDataFactoryPipelineRunStatus.FAILED, "exception"), + (AzureDataFactoryPipelineRunStatus.CANCELLED, "exception"), ] ) @patch.object(AzureDataFactoryHook, "run_pipeline", return_value=MagicMock(**PIPELINE_RUN_RESPONSE)) - def test_execute_wait_for_completion(self, pipeline_run_status, expected_status, mock_run_pipeline): + def test_execute_wait_for_pipeline_run(self, pipeline_run_status, expected_output, mock_run_pipeline): operator = AzureDataFactoryRunPipelineOperator(**self.config) assert operator.azure_data_factory_conn_id == self.config["azure_data_factory_conn_id"] assert operator.pipeline_name == self.config["pipeline_name"] assert operator.resource_group_name == self.config["resource_group_name"] assert operator.factory_name == self.config["factory_name"] - assert operator.poke_interval == self.config["poke_interval"] + assert operator.check_interval == self.config["check_interval"] + assert operator.timeout == 604800 assert operator.wait_for_completion with patch.object(AzureDataFactoryHook, "get_pipeline_run") as mock_get_pipeline_run: @@ -76,12 +77,10 @@ def test_execute_wait_for_completion(self, pipeline_run_status, expected_status, pipeline_run_status ) - assert mock_get_pipeline_run.return_value.status == pipeline_run_status - - if not expected_status: + if not expected_output: # A successful operator execution should not return any values. assert not operator.execute(context=self.mock_context) - elif expected_status == "AirflowException": + elif expected_output == "exception": if mock_get_pipeline_run.return_value.status == AzureDataFactoryPipelineRunStatus.CANCELLED: error_message = ( f"Pipeline run {PIPELINE_RUN_RESPONSE['run_id']} has been " @@ -93,7 +92,7 @@ def test_execute_wait_for_completion(self, pipeline_run_status, expected_status, f"{pipeline_run_status.lower()}." ) # The operator should fail if the pipeline run fails or is canceled. - with pytest.raises(AirflowException, match=error_message): + with pytest.raises(AzureDataFactoryPipelineRunException, match=error_message): operator.execute(context=self.mock_context) # Check the ``run_id`` attr is assigned after executing the pipeline. @@ -122,14 +121,14 @@ def test_execute_wait_for_completion(self, pipeline_run_status, expected_status, ) @patch.object(AzureDataFactoryHook, "run_pipeline", return_value=MagicMock(**PIPELINE_RUN_RESPONSE)) - def test_execute_no_wait_for_completion(self, mock_run_pipeline): + def test_execute_no_wait_for_pipeline_run(self, mock_run_pipeline): operator = AzureDataFactoryRunPipelineOperator(wait_for_completion=False, **self.config) assert operator.azure_data_factory_conn_id == self.config["azure_data_factory_conn_id"] assert operator.pipeline_name == self.config["pipeline_name"] assert operator.resource_group_name == self.config["resource_group_name"] assert operator.factory_name == self.config["factory_name"] - assert operator.poke_interval == self.config["poke_interval"] + assert operator.check_interval == self.config["check_interval"] assert not operator.wait_for_completion with patch.object(AzureDataFactoryHook, "get_pipeline_run") as mock_get_pipeline_run: @@ -154,5 +153,5 @@ def test_execute_no_wait_for_completion(self, mock_run_pipeline): parameters=None, ) - # Checking the pipeline run status should _not_ be called when ``wait_for_completion`` is False. + # Checking the pipeline run status should _not_ be called when ``wait_for_pipeline_run`` is False. mock_get_pipeline_run.assert_not_called() diff --git a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py index 67de2952ddb5c..76d3094ab6d6e 100644 --- a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py @@ -17,16 +17,14 @@ # under the License. import unittest -from datetime import datetime -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest from parameterized import parameterized -from airflow.exceptions import AirflowException -from airflow.models.dag import DAG from airflow.providers.microsoft.azure.hooks.data_factory import ( AzureDataFactoryHook, + AzureDataFactoryPipelineRunException, AzureDataFactoryPipelineRunStatus, ) from airflow.providers.microsoft.azure.sensors.data_factory import AzureDataFactoryPipelineRunStatusSensor @@ -34,7 +32,6 @@ class TestPipelineRunStatusSensor(unittest.TestCase): def setUp(self): - self.dag = DAG("test", start_date=datetime(2021, 8, 16), schedule_interval=None, catchup=False) self.config = { "azure_data_factory_conn_id": "azure_data_factory_test", "run_id": "run_id", @@ -43,59 +40,37 @@ def setUp(self): "timeout": 100, "poke_interval": 15, } - - @staticmethod - def create_pipeline_run(status: str): - """Helper function to create a mock pipeline run with a given execution status.""" - - run = MagicMock() - run.status = status - return run + self.sensor = AzureDataFactoryPipelineRunStatusSensor(task_id="pipeline_run_sensor", **self.config) def test_init(self): - sensor = AzureDataFactoryPipelineRunStatusSensor( - task_id="pipeline_run_sensor", dag=self.dag, **self.config - ) - assert sensor.azure_data_factory_conn_id == self.config["azure_data_factory_conn_id"] - assert sensor.run_id == self.config["run_id"] - assert sensor.resource_group_name == self.config["resource_group_name"] - assert sensor.factory_name == self.config["factory_name"] - assert sensor.timeout == self.config["timeout"] - assert sensor.poke_interval == self.config["poke_interval"] + assert self.sensor.azure_data_factory_conn_id == self.config["azure_data_factory_conn_id"] + assert self.sensor.run_id == self.config["run_id"] + assert self.sensor.resource_group_name == self.config["resource_group_name"] + assert self.sensor.factory_name == self.config["factory_name"] + assert self.sensor.timeout == self.config["timeout"] + assert self.sensor.poke_interval == self.config["poke_interval"] @parameterized.expand( [ (AzureDataFactoryPipelineRunStatus.SUCCEEDED, True), - (AzureDataFactoryPipelineRunStatus.FAILED, "AirflowException"), - (AzureDataFactoryPipelineRunStatus.CANCELLED, "AirflowException"), + (AzureDataFactoryPipelineRunStatus.FAILED, "exception"), + (AzureDataFactoryPipelineRunStatus.CANCELLED, "exception"), (AzureDataFactoryPipelineRunStatus.CANCELING, False), (AzureDataFactoryPipelineRunStatus.QUEUED, False), (AzureDataFactoryPipelineRunStatus.IN_PROGRESS, False), ] ) - def test_poke(self, pipeline_run_status, expected_status): - mock_pipeline_run = TestPipelineRunStatusSensor.create_pipeline_run(pipeline_run_status) - - with patch.object( - AzureDataFactoryHook, "get_pipeline_run", return_value=mock_pipeline_run - ) as mock_get_pipeline_run: - sensor = AzureDataFactoryPipelineRunStatusSensor( - task_id="pipeline_run_sensor_poke", dag=self.dag, **self.config - ) + @patch.object(AzureDataFactoryHook, "get_pipeline_run") + def test_poke(self, pipeline_run_status, expected_status, mock_pipeline_run): + mock_pipeline_run.return_value.status = pipeline_run_status - if expected_status == "AirflowException": - # The sensor should fail if the pipeline run fails or is canceled. - with pytest.raises( - AirflowException, - match=f"Pipeline run {self.config['run_id']} is in a negative terminal status: " - f"{pipeline_run_status}", - ): - sensor.poke({}) + if expected_status == "exception": + if pipeline_run_status == AzureDataFactoryPipelineRunStatus.FAILED: + error_message = f"Pipeline run {self.config['run_id']} has failed." else: - assert sensor.poke({}) == expected_status + error_message = f"Pipeline run {self.config['run_id']} has been cancelled." - mock_get_pipeline_run.assert_called_once_with( - run_id=self.config["run_id"], - factory_name=self.config["factory_name"], - resource_group_name=self.config["resource_group_name"], - ) + with pytest.raises(AzureDataFactoryPipelineRunException, match=error_message): + self.sensor.poke({}) + else: + assert self.sensor.poke({}) == expected_status From 15491d96906fdcb5a5bd10e7fc359eeff20b62c6 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Wed, 1 Sep 2021 21:33:36 -0400 Subject: [PATCH 25/27] Parameter rename in hook --- .../microsoft/azure/hooks/data_factory.py | 21 +++++++------------ 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/data_factory.py b/airflow/providers/microsoft/azure/hooks/data_factory.py index 65d5161ef8771..e23f03d593b95 100644 --- a/airflow/providers/microsoft/azure/hooks/data_factory.py +++ b/airflow/providers/microsoft/azure/hooks/data_factory.py @@ -34,7 +34,6 @@ from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook -from airflow.typing_compat import Literal def provide_targeted_factory(func: Callable) -> Callable: @@ -585,11 +584,10 @@ def get_pipeline_run( """ return self.get_conn().pipeline_runs.get(resource_group_name, factory_name, run_id, **config) - @provide_targeted_factory def check_pipeline_run_status( self, run_id: str, - mode: Literal["poke", "wait"] = "wait", + wait_for_completion: bool = True, expected_status: str = AzureDataFactoryPipelineRunStatus.SUCCEEDED, check_interval: Optional[int] = 60, timeout: Optional[int] = 60 * 60 * 24 * 7, @@ -602,12 +600,12 @@ def check_pipeline_run_status( terminal status or the desired status. :param run_id: The pipeline run identifier. - :param mode: If set to "wait", periodically check a pipeline run's status to see if it has reached a - terminal or desired status. If set to "poke", check a pipeline run's status only once and return - whether or not the status matches the ``expected_status``. + :param wait_for_completion: If enabled, periodically check a pipeline run's status to see if it has + reached a terminal or the desired status. Otherwise, check a pipeline run's status only once and + return whether or not the status matches the ``expected_status``. :param expected_status: The desired status of a pipeline run. - :param check_interval: Time in seconds to check a pipeline run's status when ``mode`` is set to - "wait". + :param check_interval: Time in seconds to check a pipeline run's status when ``wait_for_completion`` + is enabled. :param timeout: Time in seconds to wait for a pipeline to reach a terminal status when ``mode`` is set to "wait". :param resource_group_name: The resource group name. @@ -615,11 +613,6 @@ def check_pipeline_run_status( :return: Boolean value indicating if the current pipeline run's status matches the ``expected_status``. """ - if mode not in ("wait", "poke"): - raise ValueError( - f"The configured mode, '{mode}', is not supported. Please use 'wait' or 'poke'." - ) - pipeline_run = self.get_pipeline_run( run_id=run_id, factory_name=factory_name, @@ -628,7 +621,7 @@ def check_pipeline_run_status( ) pipeline_run_status = pipeline_run.status - if mode == "wait": + if wait_for_completion: start_time = time.monotonic() while ( From e0e417384f5f0b351cc9e656cf4bbfdd258f51ed Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Wed, 8 Sep 2021 12:38:23 -0400 Subject: [PATCH 26/27] Refactoring pipeline run status and waiting checks --- .../example_dags/example_adf_run_pipeline.py | 12 +- .../azure/hooks/azure_data_factory.py | 27 +++++ .../microsoft/azure/hooks/data_factory.py | 108 +++++++++--------- .../microsoft/azure/operators/data_factory.py | 55 ++++++--- .../providers/microsoft/azure/provider.yaml | 1 + .../microsoft/azure/sensors/data_factory.py | 23 +++- .../operators/adf_run_pipeline.rst | 4 +- .../azure/hooks/test_azure_data_factory.py | 64 +++++------ .../operators/test_azure_data_factory.py | 68 +++++++---- .../azure/sensors/test_azure_data_factory.py | 8 +- 10 files changed, 226 insertions(+), 144 deletions(-) create mode 100644 airflow/providers/microsoft/azure/hooks/azure_data_factory.py diff --git a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py index 740c92b5bdf4c..21269f112009a 100644 --- a/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py +++ b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py @@ -52,18 +52,18 @@ run_pipeline2 = AzureDataFactoryRunPipelineOperator( task_id="run_pipeline2", pipeline_name="pipeline2", - wait_for_completion=False, + wait_for_termination=False, ) - wait_for_pipeline_run = AzureDataFactoryPipelineRunStatusSensor( - task_id="wait_for_pipeline_run", + pipeline_run_sensor = AzureDataFactoryPipelineRunStatusSensor( + task_id="pipeline_run_sensor", run_id=run_pipeline2.output["run_id"], ) # [END howto_operator_adf_run_pipeline_async] - begin >> Label("No async wait") >> run_pipeline1 >> end + begin >> Label("No async wait") >> run_pipeline1 begin >> Label("Do async wait with sensor") >> run_pipeline2 - wait_for_pipeline_run >> end + [run_pipeline1, pipeline_run_sensor] >> end # Task dependency created via `XComArgs`: - # run_pipeline2 >> wait_for_pipeline_run + # run_pipeline2 >> pipeline_run_sensor diff --git a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py new file mode 100644 index 0000000000000..52faa0b91182e --- /dev/null +++ b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py @@ -0,0 +1,27 @@ +# 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. +"""This module is deprecated. Please use :mod:`airflow.providers.microsoft.azure.hooks.data_factory`.""" + +import warnings + +from airflow.providers.microsoft.azure.hooks.data_factory import AzureDataFactoryHook # noqa + +warnings.warn( + "This module is deprecated. Please use `airflow.providers.microsoft.azure.hooks.data_factory`.", + DeprecationWarning, + stacklevel=2, +) diff --git a/airflow/providers/microsoft/azure/hooks/data_factory.py b/airflow/providers/microsoft/azure/hooks/data_factory.py index e23f03d593b95..588b687d30103 100644 --- a/airflow/providers/microsoft/azure/hooks/data_factory.py +++ b/airflow/providers/microsoft/azure/hooks/data_factory.py @@ -17,7 +17,7 @@ import inspect import time from functools import wraps -from typing import Any, Callable, Dict, Optional +from typing import Any, Callable, Dict, Optional, Set, Union from azure.core.polling import LROPoller from azure.identity import ClientSecretCredential @@ -584,76 +584,76 @@ def get_pipeline_run( """ return self.get_conn().pipeline_runs.get(resource_group_name, factory_name, run_id, **config) - def check_pipeline_run_status( + def get_pipeline_run_status( self, run_id: str, - wait_for_completion: bool = True, - expected_status: str = AzureDataFactoryPipelineRunStatus.SUCCEEDED, - check_interval: Optional[int] = 60, - timeout: Optional[int] = 60 * 60 * 24 * 7, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, - **config: Any, - ) -> bool: + ) -> str: """ - Checks a pipeline run's status with a configurable option to do so until the run has reached a - terminal status or the desired status. + Get a pipeline run's current status. :param run_id: The pipeline run identifier. - :param wait_for_completion: If enabled, periodically check a pipeline run's status to see if it has - reached a terminal or the desired status. Otherwise, check a pipeline run's status only once and - return whether or not the status matches the ``expected_status``. - :param expected_status: The desired status of a pipeline run. - :param check_interval: Time in seconds to check a pipeline run's status when ``wait_for_completion`` - is enabled. - :param timeout: Time in seconds to wait for a pipeline to reach a terminal status when ``mode`` is set - to "wait". :param resource_group_name: The resource group name. :param factory_name: The factory name. - :return: Boolean value indicating if the current pipeline run's status matches the - ``expected_status``. + :return: The status of the pipeline run. """ - pipeline_run = self.get_pipeline_run( + self.log.info(f"Getting the status of run ID {run_id}.") + pipeline_run_status = self.get_pipeline_run( run_id=run_id, factory_name=factory_name, resource_group_name=resource_group_name, - **config, - ) - pipeline_run_status = pipeline_run.status - - if wait_for_completion: - start_time = time.monotonic() - - while ( - pipeline_run_status not in AzureDataFactoryPipelineRunStatus.TERMINAL_STATUSES - and pipeline_run_status != expected_status - ): - # Check if the pipeline-run duration has exceeded the ``timeout`` configured. - if start_time + timeout < time.monotonic(): - raise AzureDataFactoryPipelineRunException( - f"Pipeline run {run_id} has not reached a terminal status after {timeout} seconds." - ) - - # Wait to check the status of the pipeline based on the ``check_interval`` configured. - time.sleep(check_interval) - - self.log.info(f"Checking on the status of run ID {run_id}.") - pipeline_run = self.get_pipeline_run( - run_id=run_id, - factory_name=factory_name, - resource_group_name=resource_group_name, - **config, + ).status + self.log.info(f"Current status of pipeline run {run_id}: {pipeline_run_status}") + + return pipeline_run_status + + def wait_for_pipeline_run_status( + self, + run_id: str, + expected_statuses: Union[str, Set[str]], + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + check_interval: Optional[int] = 60, + timeout: Optional[int] = 60 * 60 * 24 * 7, + ) -> bool: + """ + Waits for a pipeline run to match an expected status. + + :param run_id: The pipeline run identifier. + :param expected_statuses: The desired status(es) to check against a pipeline run's current status. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param check_interval: Time in seconds to check on a pipeline run's status. + :param timeout: Time in seconds to wait for a pipeline to reach a terminal status or the expected + status. + :return: Boolean indicating if the pipeline run has reached the ``expected_status``. + """ + pipeline_run_info = { + "run_id": run_id, + "factory_name": factory_name, + "resource_group_name": resource_group_name, + } + pipeline_run_status = self.get_pipeline_run_status(**pipeline_run_info) + + start_time = time.monotonic() + + while ( + pipeline_run_status not in AzureDataFactoryPipelineRunStatus.TERMINAL_STATUSES + and pipeline_run_status not in expected_statuses + ): + # Check if the pipeline-run duration has exceeded the ``timeout`` configured. + if start_time + timeout < time.monotonic(): + raise AzureDataFactoryPipelineRunException( + f"Pipeline run {run_id} has not reached a terminal status after {timeout} seconds." ) - pipeline_run_status = pipeline_run.status - self.log.info(f"Current status of run ID {run_id}: {pipeline_run_status}") - if pipeline_run_status == AzureDataFactoryPipelineRunStatus.FAILED: - raise AzureDataFactoryPipelineRunException(f"Pipeline run {run_id} has failed.") + # Wait to check the status of the pipeline run based on the ``check_interval`` configured. + time.sleep(check_interval) - if pipeline_run_status == AzureDataFactoryPipelineRunStatus.CANCELLED: - raise AzureDataFactoryPipelineRunException(f"Pipeline run {run_id} has been cancelled.") + pipeline_run_status = self.get_pipeline_run_status(**pipeline_run_info) - return pipeline_run_status == expected_status + return pipeline_run_status in expected_statuses @provide_targeted_factory def cancel_pipeline_run( diff --git a/airflow/providers/microsoft/azure/operators/data_factory.py b/airflow/providers/microsoft/azure/operators/data_factory.py index 3c7f93c7c5c1a..ddcf489b93438 100644 --- a/airflow/providers/microsoft/azure/operators/data_factory.py +++ b/airflow/providers/microsoft/azure/operators/data_factory.py @@ -18,7 +18,11 @@ from typing import Any, Dict, Optional from airflow.models import BaseOperator -from airflow.providers.microsoft.azure.hooks.data_factory import AzureDataFactoryHook +from airflow.providers.microsoft.azure.hooks.data_factory import ( + AzureDataFactoryHook, + AzureDataFactoryPipelineRunException, + AzureDataFactoryPipelineRunStatus, +) class AzureDataFactoryRunPipelineOperator(BaseOperator): @@ -33,10 +37,10 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): :type azure_data_factory_conn_id: str :param pipeline_name: The name of the pipeline to execute. :type pipeline_name: str - :param wait_for_completion: Flag to wait on a pipeline run's completion. By default, this feature is + :param wait_for_termination: Flag to wait on a pipeline run's termination. By default, this feature is enabled but could be disabled to perform an asynchronous wait for a long-running pipeline execution using the ``AzureDataFactoryPipelineRunSensor``. - :type wait_for_pipeline_run: bool + :type wait_for_termination: bool :param resource_group_name: The resource group name. If a value is not passed in to the operator, the ``AzureDataFactoryHook`` will attempt to use the resource group name provided in the corresponding connection. @@ -60,12 +64,12 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): :param parameters: Parameters of the pipeline run. These parameters are referenced in a pipeline via ``@pipeline().parameters.parameterName`` and will be used only if the ``reference_pipeline_run_id`` is not specified. - :type start_from_failure: Dict[str, Any] + :type parameters: Dict[str, Any] :param timeout: Time in seconds to wait for a pipeline to reach a terminal status for non-asynchronous - waits. Used only if ``wait_for_completion`` is True. + waits. Used only if ``wait_for_termination`` is True. :type timeout: int :param check_interval: Time in seconds to check on a pipeline run's status for non-asynchronous waits. - Used only if ``wait_for_completion`` is True. + Used only if ``wait_for_termination`` is True. :type check_interval: int """ @@ -79,12 +83,14 @@ class AzureDataFactoryRunPipelineOperator(BaseOperator): ) template_fields_renderers = {"parameters": "json"} + ui_color = "#0678d4" + def __init__( self, *, - azure_data_factory_conn_id: str, pipeline_name: str, - wait_for_completion: bool = True, + azure_data_factory_conn_id: str = AzureDataFactoryHook.default_conn_name, + wait_for_termination: bool = True, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, reference_pipeline_run_id: Optional[str] = None, @@ -99,7 +105,7 @@ def __init__( super().__init__(**kwargs) self.azure_data_factory_conn_id = azure_data_factory_conn_id self.pipeline_name = pipeline_name - self.wait_for_completion = wait_for_completion + self.wait_for_termination = wait_for_termination self.resource_group_name = resource_group_name self.factory_name = factory_name self.reference_pipeline_run_id = reference_pipeline_run_id @@ -112,7 +118,7 @@ def __init__( def execute(self, context: Dict) -> None: self.hook = AzureDataFactoryHook(azure_data_factory_conn_id=self.azure_data_factory_conn_id) - self.log.info(f"Executing the '{self.pipeline_name}' pipeline.") + self.log.info(f"Executing the {self.pipeline_name} pipeline.") response = self.hook.run_pipeline( pipeline_name=self.pipeline_name, resource_group_name=self.resource_group_name, @@ -129,16 +135,22 @@ def execute(self, context: Dict) -> None: # asynchronous wait. context["ti"].xcom_push(key="run_id", value=self.run_id) - if self.wait_for_completion: - self.log.info(f"Waiting for pipeline run {self.run_id} to complete.") - self.hook.check_pipeline_run_status( + if self.wait_for_termination: + self.log.info(f"Waiting for pipeline run {self.run_id} to terminate.") + + if self.hook.wait_for_pipeline_run_status( run_id=self.run_id, + expected_statuses=AzureDataFactoryPipelineRunStatus.SUCCEEDED, check_interval=self.check_interval, timeout=self.timeout, resource_group_name=self.resource_group_name, factory_name=self.factory_name, - ) - self.log.info(f"Pipeline run {self.run_id} has completed successfully.") + ): + self.log.info(f"Pipeline run {self.run_id} has completed successfully.") + else: + raise AzureDataFactoryPipelineRunException( + f"Pipeline run {self.run_id} has failed or has been cancelled." + ) def on_kill(self) -> None: if self.run_id: @@ -147,3 +159,16 @@ def on_kill(self) -> None: resource_group_name=self.resource_group_name, factory_name=self.factory_name, ) + + # Check to ensure the pipeline run was cancelled as expected. + if self.hook.wait_for_pipeline_run_status( + run_id=self.run_id, + expected_statuses=AzureDataFactoryPipelineRunStatus.CANCELLED, + check_interval=self.check_interval, + timeout=self.timeout, + resource_group_name=self.resource_group_name, + factory_name=self.factory_name, + ): + self.log.info(f"Pipeline run {self.run_id} has been cancelled successfully.") + else: + raise AzureDataFactoryPipelineRunException(f"Pipeline run {self.run_id} was not cancelled.") diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index 63ede42a4f8ea..9025d8c662978 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -144,6 +144,7 @@ hooks: - integration-name: Microsoft Azure Data Factory python-modules: - airflow.providers.microsoft.azure.hooks.data_factory + - airflow.providers.microsoft.azure.hooks.azure_data_factory transfers: - source-integration-name: Local diff --git a/airflow/providers/microsoft/azure/sensors/data_factory.py b/airflow/providers/microsoft/azure/sensors/data_factory.py index 2ef66597caedb..de2f3938c57af 100644 --- a/airflow/providers/microsoft/azure/sensors/data_factory.py +++ b/airflow/providers/microsoft/azure/sensors/data_factory.py @@ -17,7 +17,11 @@ from typing import Dict, Optional -from airflow.providers.microsoft.azure.hooks.data_factory import AzureDataFactoryHook +from airflow.providers.microsoft.azure.hooks.data_factory import ( + AzureDataFactoryHook, + AzureDataFactoryPipelineRunException, + AzureDataFactoryPipelineRunStatus, +) from airflow.sensors.base import BaseSensorOperator @@ -37,11 +41,13 @@ class AzureDataFactoryPipelineRunStatusSensor(BaseSensorOperator): template_fields = ("azure_data_factory_conn_id", "resource_group_name", "factory_name", "run_id") + ui_color = "#50e6ff" + def __init__( self, *, - azure_data_factory_conn_id: str, run_id: str, + azure_data_factory_conn_id: str = AzureDataFactoryHook.default_conn_name, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, **kwargs, @@ -53,12 +59,17 @@ def __init__( self.factory_name = factory_name def poke(self, context: Dict) -> bool: - self.log.info(f"Checking the status for pipeline run {self.run_id}.") self.hook = AzureDataFactoryHook(azure_data_factory_conn_id=self.azure_data_factory_conn_id) - - return self.hook.check_pipeline_run_status( + pipeline_run_status = self.hook.get_pipeline_run_status( run_id=self.run_id, - wait_for_completion=False, resource_group_name=self.resource_group_name, factory_name=self.factory_name, ) + + if pipeline_run_status == AzureDataFactoryPipelineRunStatus.FAILED: + raise AzureDataFactoryPipelineRunException(f"Pipeline run {self.run_id} has failed.") + + if pipeline_run_status == AzureDataFactoryPipelineRunStatus.CANCELLED: + raise AzureDataFactoryPipelineRunException(f"Pipeline run {self.run_id} has been cancelled.") + + return pipeline_run_status == AzureDataFactoryPipelineRunStatus.SUCCEEDED diff --git a/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst b/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst index 6978dbbb32610..1462648b8b691 100644 --- a/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst +++ b/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst @@ -30,8 +30,8 @@ It offers a code-free UI for intuitive authoring and single-pane-of-glass monito AzureDataFactoryRunPipelineOperator ----------------------------------- Use the :class:`~airflow.providers.microsoft.azure.operators.data_factory.AzureDataFactoryRunPipelineOperator` to execute a pipeline within a data factory. -By default, the operator will check on the status of the executed pipeline and mirror its result (i.e. succeed if the pipeline succeeds or fail if the pipeline fails or is canceled). -This functionality can be disabled for an asynchronous wait -- typically with the :class:`~airflow.providers.microsoft.azure.sensors.data_factory.AzureDataFactoryPipelineRunSensor` -- by setting ``wait_for_completion`` to False. +By default, the operator will periodically check on the status of the executed pipeline to terminate with a "Succeeded" status. +This functionality can be disabled for an asynchronous wait -- typically with the :class:`~airflow.providers.microsoft.azure.sensors.data_factory.AzureDataFactoryPipelineRunSensor` -- by setting ``wait_for_termination`` to False. Below is an example of using this operator to execute an Azure Data Factory pipeline. diff --git a/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py b/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py index 2be2cf55b0092..670d3254aea7c 100644 --- a/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py @@ -344,42 +344,40 @@ def test_get_pipeline_run(hook: AzureDataFactoryHook, user_args, sdk_args): hook._conn.pipeline_runs.get.assert_called_with(*sdk_args) -def test_check_pipeline_run_status(hook): - config = {"run_id": ID, "timeout": 5, "check_interval": 1} +_wait_for_pipeline_run_status_test_args = [ + (AzureDataFactoryPipelineRunStatus.SUCCEEDED, AzureDataFactoryPipelineRunStatus.SUCCEEDED, True), + (AzureDataFactoryPipelineRunStatus.FAILED, AzureDataFactoryPipelineRunStatus.SUCCEEDED, False), + (AzureDataFactoryPipelineRunStatus.CANCELLED, AzureDataFactoryPipelineRunStatus.SUCCEEDED, False), + (AzureDataFactoryPipelineRunStatus.IN_PROGRESS, AzureDataFactoryPipelineRunStatus.SUCCEEDED, "timeout"), + (AzureDataFactoryPipelineRunStatus.QUEUED, AzureDataFactoryPipelineRunStatus.SUCCEEDED, "timeout"), + (AzureDataFactoryPipelineRunStatus.CANCELING, AzureDataFactoryPipelineRunStatus.SUCCEEDED, "timeout"), + (AzureDataFactoryPipelineRunStatus.SUCCEEDED, AzureDataFactoryPipelineRunStatus.TERMINAL_STATUSES, True), + (AzureDataFactoryPipelineRunStatus.FAILED, AzureDataFactoryPipelineRunStatus.TERMINAL_STATUSES, True), + (AzureDataFactoryPipelineRunStatus.CANCELLED, AzureDataFactoryPipelineRunStatus.TERMINAL_STATUSES, True), +] + + +@pytest.mark.parametrize( + argnames=("pipeline_run_status", "expected_status", "expected_output"), + argvalues=_wait_for_pipeline_run_status_test_args, + ids=[ + f"run_status_{argval[0]}_expected_{argval[1]}" + if isinstance(argval[1], str) + else f"run_status_{argval[0]}_expected_AnyTerminalStatus" + for argval in _wait_for_pipeline_run_status_test_args + ], +) +def test_wait_for_pipeline_run_status(hook, pipeline_run_status, expected_status, expected_output): + config = {"run_id": ID, "timeout": 3, "check_interval": 1, "expected_statuses": expected_status} with patch.object(AzureDataFactoryHook, "get_pipeline_run") as mock_pipeline_run: - # Begin tests when ``wait_for_completion`` is enabled -- the default behavior. - mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.SUCCEEDED - assert hook.check_pipeline_run_status(**config) + mock_pipeline_run.return_value.status = pipeline_run_status - mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.FAILED - with pytest.raises(AzureDataFactoryPipelineRunException): - hook.check_pipeline_run_status(**config) - - mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.CANCELLED - with pytest.raises(AzureDataFactoryPipelineRunException): - hook.check_pipeline_run_status(**config) - - # Begin tests when ``wait_for_completion`` is disabled (aka poking for status). - mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.SUCCEEDED - assert hook.check_pipeline_run_status(wait_for_completion=False, **config) - - mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.IN_PROGRESS - assert not hook.check_pipeline_run_status(wait_for_completion=False, **config) - - mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.QUEUED - assert not hook.check_pipeline_run_status(wait_for_completion=False, **config) - - mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.CANCELING - assert not hook.check_pipeline_run_status(wait_for_completion=False, **config) - - mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.FAILED - with pytest.raises(AzureDataFactoryPipelineRunException): - hook.check_pipeline_run_status(wait_for_completion=False, **config) - - mock_pipeline_run.return_value.status = AzureDataFactoryPipelineRunStatus.CANCELLED - with pytest.raises(AzureDataFactoryPipelineRunException): - hook.check_pipeline_run_status(wait_for_completion=False, **config) + if expected_output != "timeout": + assert hook.wait_for_pipeline_run_status(**config) == expected_output + else: + with pytest.raises(AzureDataFactoryPipelineRunException): + hook.wait_for_pipeline_run_status(**config) @parametrize( diff --git a/tests/providers/microsoft/azure/operators/test_azure_data_factory.py b/tests/providers/microsoft/azure/operators/test_azure_data_factory.py index ab718f657ce75..25080b5a1a9b8 100644 --- a/tests/providers/microsoft/azure/operators/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/operators/test_azure_data_factory.py @@ -42,6 +42,7 @@ def setUp(self): "resource_group_name": "resource-group-name", "factory_name": "factory-name", "check_interval": 1, + "timeout": 3, } @staticmethod @@ -58,10 +59,13 @@ def create_pipeline_run(status: str): (AzureDataFactoryPipelineRunStatus.SUCCEEDED, None), (AzureDataFactoryPipelineRunStatus.FAILED, "exception"), (AzureDataFactoryPipelineRunStatus.CANCELLED, "exception"), + (AzureDataFactoryPipelineRunStatus.IN_PROGRESS, "timeout"), + (AzureDataFactoryPipelineRunStatus.QUEUED, "timeout"), + (AzureDataFactoryPipelineRunStatus.CANCELING, "timeout"), ] ) @patch.object(AzureDataFactoryHook, "run_pipeline", return_value=MagicMock(**PIPELINE_RUN_RESPONSE)) - def test_execute_wait_for_pipeline_run(self, pipeline_run_status, expected_output, mock_run_pipeline): + def test_execute_wait_for_termination(self, pipeline_run_status, expected_output, mock_run_pipeline): operator = AzureDataFactoryRunPipelineOperator(**self.config) assert operator.azure_data_factory_conn_id == self.config["azure_data_factory_conn_id"] @@ -69,8 +73,8 @@ def test_execute_wait_for_pipeline_run(self, pipeline_run_status, expected_outpu assert operator.resource_group_name == self.config["resource_group_name"] assert operator.factory_name == self.config["factory_name"] assert operator.check_interval == self.config["check_interval"] - assert operator.timeout == 604800 - assert operator.wait_for_completion + assert operator.timeout == self.config["timeout"] + assert operator.wait_for_termination with patch.object(AzureDataFactoryHook, "get_pipeline_run") as mock_get_pipeline_run: mock_get_pipeline_run.return_value = TestAzureDataFactoryRunPipelineOperator.create_pipeline_run( @@ -81,18 +85,21 @@ def test_execute_wait_for_pipeline_run(self, pipeline_run_status, expected_outpu # A successful operator execution should not return any values. assert not operator.execute(context=self.mock_context) elif expected_output == "exception": - if mock_get_pipeline_run.return_value.status == AzureDataFactoryPipelineRunStatus.CANCELLED: - error_message = ( - f"Pipeline run {PIPELINE_RUN_RESPONSE['run_id']} has been " - f"{pipeline_run_status.lower()}." - ) - else: - error_message = ( - f"Pipeline run {PIPELINE_RUN_RESPONSE['run_id']} has " - f"{pipeline_run_status.lower()}." - ) # The operator should fail if the pipeline run fails or is canceled. - with pytest.raises(AzureDataFactoryPipelineRunException, match=error_message): + with pytest.raises( + AzureDataFactoryPipelineRunException, + match=f"Pipeline run {PIPELINE_RUN_RESPONSE['run_id']} has failed or has been cancelled.", + ): + operator.execute(context=self.mock_context) + else: + # Demonstrating the operator timing out after surpassing the configured timeout value. + with pytest.raises( + AzureDataFactoryPipelineRunException, + match=( + f"Pipeline run {PIPELINE_RUN_RESPONSE['run_id']} has not reached a terminal status " + f"after {self.config['timeout']} seconds." + ), + ): operator.execute(context=self.mock_context) # Check the ``run_id`` attr is assigned after executing the pipeline. @@ -114,24 +121,37 @@ def test_execute_wait_for_pipeline_run(self, pipeline_run_status, expected_outpu parameters=None, ) - mock_get_pipeline_run.assert_called_once_with( - run_id=mock_run_pipeline.return_value.run_id, - factory_name=self.config["factory_name"], - resource_group_name=self.config["resource_group_name"], - ) + if pipeline_run_status in AzureDataFactoryPipelineRunStatus.TERMINAL_STATUSES: + mock_get_pipeline_run.assert_called_once_with( + run_id=mock_run_pipeline.return_value.run_id, + factory_name=self.config["factory_name"], + resource_group_name=self.config["resource_group_name"], + ) + else: + # When the pipeline run status is not in a terminal status or "Succeeded", the operator will + # continue to call ``get_pipeline_run()`` until a ``timeout`` number of seconds has passed + # (3 seconds for this test). Therefore, there should be 4 calls of this function: one + # initially and 3 for each check done at a 1 second interval. + assert mock_get_pipeline_run.call_count == 4 + + mock_get_pipeline_run.assert_called_with( + run_id=mock_run_pipeline.return_value.run_id, + factory_name=self.config["factory_name"], + resource_group_name=self.config["resource_group_name"], + ) @patch.object(AzureDataFactoryHook, "run_pipeline", return_value=MagicMock(**PIPELINE_RUN_RESPONSE)) - def test_execute_no_wait_for_pipeline_run(self, mock_run_pipeline): - operator = AzureDataFactoryRunPipelineOperator(wait_for_completion=False, **self.config) + def test_execute_no_wait_for_termination(self, mock_run_pipeline): + operator = AzureDataFactoryRunPipelineOperator(wait_for_termination=False, **self.config) assert operator.azure_data_factory_conn_id == self.config["azure_data_factory_conn_id"] assert operator.pipeline_name == self.config["pipeline_name"] assert operator.resource_group_name == self.config["resource_group_name"] assert operator.factory_name == self.config["factory_name"] assert operator.check_interval == self.config["check_interval"] - assert not operator.wait_for_completion + assert not operator.wait_for_termination - with patch.object(AzureDataFactoryHook, "get_pipeline_run") as mock_get_pipeline_run: + with patch.object(AzureDataFactoryHook, "get_pipeline_run", autospec=True) as mock_get_pipeline_run: operator.execute(context=self.mock_context) # Check the ``run_id`` attr is assigned after executing the pipeline. @@ -153,5 +173,5 @@ def test_execute_no_wait_for_pipeline_run(self, mock_run_pipeline): parameters=None, ) - # Checking the pipeline run status should _not_ be called when ``wait_for_pipeline_run`` is False. + # Checking the pipeline run status should _not_ be called when ``wait_for_termination`` is False. mock_get_pipeline_run.assert_not_called() diff --git a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py index 76d3094ab6d6e..4f626af3dd4f0 100644 --- a/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py @@ -1,4 +1,3 @@ -# # 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 @@ -64,7 +63,10 @@ def test_init(self): def test_poke(self, pipeline_run_status, expected_status, mock_pipeline_run): mock_pipeline_run.return_value.status = pipeline_run_status - if expected_status == "exception": + if expected_status != "exception": + assert self.sensor.poke({}) == expected_status + else: + # The sensor should fail if the pipeline run status is "Failed" or "Cancelled". if pipeline_run_status == AzureDataFactoryPipelineRunStatus.FAILED: error_message = f"Pipeline run {self.config['run_id']} has failed." else: @@ -72,5 +74,3 @@ def test_poke(self, pipeline_run_status, expected_status, mock_pipeline_run): with pytest.raises(AzureDataFactoryPipelineRunException, match=error_message): self.sensor.poke({}) - else: - assert self.sensor.poke({}) == expected_status From e5cdc2e10dce61ef4f8e5d47533556956069da29 Mon Sep 17 00:00:00 2001 From: Josh Fell Date: Wed, 8 Sep 2021 13:22:05 -0400 Subject: [PATCH 27/27] Adding hook deprecation messge to KNOWN_DEPRECATED_DIRECT_IMPORTS --- dev/provider_packages/prepare_provider_packages.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dev/provider_packages/prepare_provider_packages.py b/dev/provider_packages/prepare_provider_packages.py index b45ee3c5f42cb..006266a4c4263 100755 --- a/dev/provider_packages/prepare_provider_packages.py +++ b/dev/provider_packages/prepare_provider_packages.py @@ -2116,6 +2116,7 @@ def summarise_total_vs_bad_and_warnings(total: int, bad: int, warns: List[warnin # ignore those messages when the warnings are generated directly by importlib - which means that # we imported it directly during module walk by the importlib library KNOWN_DEPRECATED_DIRECT_IMPORTS: Set[str] = { + "This module is deprecated. Please use `airflow.providers.microsoft.azure.hooks.data_factory`.", "This module is deprecated. Please use `airflow.providers.amazon.aws.hooks.dynamodb`.", "This module is deprecated. Please use `airflow.providers.tableau.operators.tableau_refresh_workbook`.", "This module is deprecated. Please use `airflow.providers.tableau.sensors.tableau_job_status`.",