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..21269f112009a --- /dev/null +++ b/airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py @@ -0,0 +1,69 @@ +# 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.data_factory import AzureDataFactoryRunPipelineOperator +from airflow.providers.microsoft.azure.sensors.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), + "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. + }, + 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] + + # [START howto_operator_adf_run_pipeline_async] + run_pipeline2 = AzureDataFactoryRunPipelineOperator( + task_id="run_pipeline2", + pipeline_name="pipeline2", + wait_for_termination=False, + ) + + 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 + begin >> Label("Do async wait with sensor") >> run_pipeline2 + [run_pipeline1, pipeline_run_sensor] >> end + + # Task dependency created via `XComArgs`: + # 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 index 8a1cc36090ec4..52faa0b91182e 100644 --- a/airflow/providers/microsoft/azure/hooks/azure_data_factory.py +++ b/airflow/providers/microsoft/azure/hooks/azure_data_factory.py @@ -14,769 +14,14 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -import inspect -from functools import wraps -from typing import Any, Callable, Dict, Optional +"""This module is deprecated. Please use :mod:`airflow.providers.microsoft.azure.hooks.data_factory`.""" -from azure.core.polling import LROPoller -from azure.identity import ClientSecretCredential -from azure.mgmt.datafactory import DataFactoryManagementClient -from azure.mgmt.datafactory.models import ( - CreateRunResponse, - Dataset, - DatasetResource, - Factory, - LinkedService, - LinkedServiceResource, - PipelineResource, - PipelineRun, - Trigger, - TriggerResource, -) - -from airflow.exceptions import AirflowException -from airflow.hooks.base import BaseHook - - -def provide_targeted_factory(func: Callable) -> Callable: - """ - Provide the targeted factory to the decorated function in case it isn't specified. - - If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in - the connection extras. - """ - signature = inspect.signature(func) - - @wraps(func) - def wrapper(*args, **kwargs) -> Callable: - bound_args = signature.bind(*args, **kwargs) - - def bind_argument(arg, default_key): - if arg not in bound_args.arguments: - 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") - - return func(*bound_args.args, **bound_args.kwargs) - - return wrapper - - -class AzureDataFactoryHook(BaseHook): - """ - A hook to interact with Azure Data Factory. - - :param conn_id: The :ref:`Azure Data Factory connection id`. - :type: str - """ - - conn_type: str = 'azure_data_factory' - conn_name_attr: str = 'azure_data_factory_conn_id' - default_conn_name: str = 'azure_data_factory_default' - hook_name: str = 'Azure Data Factory' - - @staticmethod - def get_connection_form_widgets() -> Dict[str, Any]: - """Returns connection widgets to add to connection form""" - from flask_appbuilder.fieldwidgets import BS3TextFieldWidget - from flask_babel import lazy_gettext - from wtforms import StringField - - return { - "extra__azure_data_factory__tenantId": StringField( - lazy_gettext('Azure Tenant ID'), widget=BS3TextFieldWidget() - ), - "extra__azure_data_factory__subscriptionId": StringField( - lazy_gettext('Azure Subscription ID'), widget=BS3TextFieldWidget() - ), - } - - @staticmethod - def get_ui_field_behaviour() -> Dict: - """Returns custom field behaviour""" - import json - - return { - "hidden_fields": ['schema', 'port', 'host'], - "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', - }, - } - - def __init__(self, conn_id: Optional[str] = default_conn_name): - self._conn: DataFactoryManagementClient = None - self.conn_id = 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) - 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') - - self._conn = DataFactoryManagementClient( - credential=ClientSecretCredential( - client_id=conn.login, client_secret=conn.password, tenant_id=tenant - ), - subscription_id=subscription_id, - ) - - return self._conn - - @provide_targeted_factory - def get_factory( - self, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, **config: Any - ) -> Factory: - """ - Get the factory. - - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :return: The factory. - """ - return self.get_conn().factories.get(resource_group_name, factory_name, **config) - - def _factory_exists(self, resource_group_name, factory_name) -> bool: - """Return whether or not the factory already exists.""" - factories = { - factory.name for factory in self.get_conn().factories.list_by_resource_group(resource_group_name) - } - - return factory_name in factories - - @provide_targeted_factory - def update_factory( - self, - factory: Factory, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> Factory: - """ - Update the factory. - - :param factory: The factory resource definition. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :raise AirflowException: If the factory does not exist. - :return: The factory. - """ - if not self._factory_exists(resource_group_name, factory_name): - raise AirflowException(f"Factory {factory!r} does not exist.") - - return self.get_conn().factories.create_or_update( - resource_group_name, factory_name, factory, **config - ) - - @provide_targeted_factory - def create_factory( - self, - factory: Factory, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> Factory: - """ - Create the factory. - - :param factory: The factory resource definition. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :raise AirflowException: If the factory already exists. - :return: The factory. - """ - if self._factory_exists(resource_group_name, factory_name): - raise AirflowException(f"Factory {factory!r} already exists.") - - return self.get_conn().factories.create_or_update( - resource_group_name, factory_name, factory, **config - ) - - @provide_targeted_factory - def delete_factory( - self, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, **config: Any - ) -> None: - """ - Delete the factory. - - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - """ - self.get_conn().factories.delete(resource_group_name, factory_name, **config) - - @provide_targeted_factory - def get_linked_service( - self, - linked_service_name: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> LinkedServiceResource: - """ - Get the linked service. - - :param linked_service_name: The linked service name. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :return: The linked service. - """ - return self.get_conn().linked_services.get( - resource_group_name, factory_name, linked_service_name, **config - ) - - def _linked_service_exists(self, resource_group_name, factory_name, linked_service_name) -> bool: - """Return whether or not the linked service already exists.""" - linked_services = { - linked_service.name - for linked_service in self.get_conn().linked_services.list_by_factory( - resource_group_name, factory_name - ) - } - - return linked_service_name in linked_services - - @provide_targeted_factory - def update_linked_service( - self, - linked_service_name: str, - linked_service: LinkedService, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> LinkedServiceResource: - """ - Update the linked service. - - :param linked_service_name: The linked service name. - :param linked_service: The linked service resource definition. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :raise AirflowException: If the linked service does not exist. - :return: The linked service. - """ - if not self._linked_service_exists(resource_group_name, factory_name, linked_service_name): - raise AirflowException(f"Linked service {linked_service_name!r} does not exist.") - - return self.get_conn().linked_services.create_or_update( - resource_group_name, factory_name, linked_service_name, linked_service, **config - ) - - @provide_targeted_factory - def create_linked_service( - self, - linked_service_name: str, - linked_service: LinkedService, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> LinkedServiceResource: - """ - Create the linked service. - - :param linked_service_name: The linked service name. - :param linked_service: The linked service resource definition. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :raise AirflowException: If the linked service already exists. - :return: The linked service. - """ - if self._linked_service_exists(resource_group_name, factory_name, linked_service_name): - raise AirflowException(f"Linked service {linked_service_name!r} already exists.") - - return self.get_conn().linked_services.create_or_update( - resource_group_name, factory_name, linked_service_name, linked_service, **config - ) - - @provide_targeted_factory - def delete_linked_service( - self, - linked_service_name: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> None: - """ - Delete the linked service. - - :param linked_service_name: The linked service name. - :param resource_group_name: The linked service name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - """ - self.get_conn().linked_services.delete( - resource_group_name, factory_name, linked_service_name, **config - ) +import warnings - @provide_targeted_factory - def get_dataset( - self, - dataset_name: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> DatasetResource: - """ - Get the dataset. +from airflow.providers.microsoft.azure.hooks.data_factory import AzureDataFactoryHook # noqa - :param dataset_name: The dataset name. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :return: The dataset. - """ - return self.get_conn().datasets.get(resource_group_name, factory_name, dataset_name, **config) - - def _dataset_exists(self, resource_group_name, factory_name, dataset_name) -> bool: - """Return whether or not the dataset already exists.""" - datasets = { - dataset.name - for dataset in self.get_conn().datasets.list_by_factory(resource_group_name, factory_name) - } - - return dataset_name in datasets - - @provide_targeted_factory - def update_dataset( - self, - dataset_name: str, - dataset: Dataset, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> DatasetResource: - """ - Update the dataset. - - :param dataset_name: The dataset name. - :param dataset: The dataset resource definition. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :raise AirflowException: If the dataset does not exist. - :return: The dataset. - """ - if not self._dataset_exists(resource_group_name, factory_name, dataset_name): - raise AirflowException(f"Dataset {dataset_name!r} does not exist.") - - return self.get_conn().datasets.create_or_update( - resource_group_name, factory_name, dataset_name, dataset, **config - ) - - @provide_targeted_factory - def create_dataset( - self, - dataset_name: str, - dataset: Dataset, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> DatasetResource: - """ - Create the dataset. - - :param dataset_name: The dataset name. - :param dataset: The dataset resource definition. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :raise AirflowException: If the dataset already exists. - :return: The dataset. - """ - if self._dataset_exists(resource_group_name, factory_name, dataset_name): - raise AirflowException(f"Dataset {dataset_name!r} already exists.") - - return self.get_conn().datasets.create_or_update( - resource_group_name, factory_name, dataset_name, dataset, **config - ) - - @provide_targeted_factory - def delete_dataset( - self, - dataset_name: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> None: - """ - Delete the dataset. - - :param dataset_name: The dataset name. - :param resource_group_name: The dataset name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - """ - self.get_conn().datasets.delete(resource_group_name, factory_name, dataset_name, **config) - - @provide_targeted_factory - def get_pipeline( - self, - pipeline_name: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> PipelineResource: - """ - Get the pipeline. - - :param pipeline_name: The pipeline name. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :return: The pipeline. - """ - return self.get_conn().pipelines.get(resource_group_name, factory_name, pipeline_name, **config) - - def _pipeline_exists(self, resource_group_name, factory_name, pipeline_name) -> bool: - """Return whether or not the pipeline already exists.""" - pipelines = { - pipeline.name - for pipeline in self.get_conn().pipelines.list_by_factory(resource_group_name, factory_name) - } - - return pipeline_name in pipelines - - @provide_targeted_factory - def update_pipeline( - self, - pipeline_name: str, - pipeline: PipelineResource, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> PipelineResource: - """ - Update the pipeline. - - :param pipeline_name: The pipeline name. - :param pipeline: The pipeline resource definition. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :raise AirflowException: If the pipeline does not exist. - :return: The pipeline. - """ - if not self._pipeline_exists(resource_group_name, factory_name, pipeline_name): - raise AirflowException(f"Pipeline {pipeline_name!r} does not exist.") - - return self.get_conn().pipelines.create_or_update( - resource_group_name, factory_name, pipeline_name, pipeline, **config - ) - - @provide_targeted_factory - def create_pipeline( - self, - pipeline_name: str, - pipeline: PipelineResource, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> PipelineResource: - """ - Create the pipeline. - - :param pipeline_name: The pipeline name. - :param pipeline: The pipeline resource definition. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :raise AirflowException: If the pipeline already exists. - :return: The pipeline. - """ - if self._pipeline_exists(resource_group_name, factory_name, pipeline_name): - raise AirflowException(f"Pipeline {pipeline_name!r} already exists.") - - return self.get_conn().pipelines.create_or_update( - resource_group_name, factory_name, pipeline_name, pipeline, **config - ) - - @provide_targeted_factory - def delete_pipeline( - self, - pipeline_name: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> None: - """ - Delete the pipeline. - - :param pipeline_name: The pipeline name. - :param resource_group_name: The pipeline name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - """ - self.get_conn().pipelines.delete(resource_group_name, factory_name, pipeline_name, **config) - - @provide_targeted_factory - def run_pipeline( - self, - pipeline_name: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> CreateRunResponse: - """ - Run a pipeline. - - :param pipeline_name: The pipeline name. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :return: The pipeline run. - """ - return self.get_conn().pipelines.create_run( - resource_group_name, factory_name, pipeline_name, **config - ) - - @provide_targeted_factory - def get_pipeline_run( - self, - run_id: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> PipelineRun: - """ - Get the pipeline run. - - :param run_id: The pipeline run identifier. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :return: The pipeline run. - """ - return self.get_conn().pipeline_runs.get(resource_group_name, factory_name, run_id, **config) - - @provide_targeted_factory - def cancel_pipeline_run( - self, - run_id: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> None: - """ - Cancel the pipeline run. - - :param run_id: The pipeline run identifier. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - """ - self.get_conn().pipeline_runs.cancel(resource_group_name, factory_name, run_id, **config) - - @provide_targeted_factory - def get_trigger( - self, - trigger_name: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> TriggerResource: - """ - Get the trigger. - - :param trigger_name: The trigger name. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :return: The trigger. - """ - return self.get_conn().triggers.get(resource_group_name, factory_name, trigger_name, **config) - - def _trigger_exists(self, resource_group_name, factory_name, trigger_name) -> bool: - """Return whether or not the trigger already exists.""" - triggers = { - trigger.name - for trigger in self.get_conn().triggers.list_by_factory(resource_group_name, factory_name) - } - - return trigger_name in triggers - - @provide_targeted_factory - def update_trigger( - self, - trigger_name: str, - trigger: Trigger, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> TriggerResource: - """ - Update the trigger. - - :param trigger_name: The trigger name. - :param trigger: The trigger resource definition. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :raise AirflowException: If the trigger does not exist. - :return: The trigger. - """ - if not self._trigger_exists(resource_group_name, factory_name, trigger_name): - raise AirflowException(f"Trigger {trigger_name!r} does not exist.") - - return self.get_conn().triggers.create_or_update( - resource_group_name, factory_name, trigger_name, trigger, **config - ) - - @provide_targeted_factory - def create_trigger( - self, - trigger_name: str, - trigger: Trigger, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> TriggerResource: - """ - Create the trigger. - - :param trigger_name: The trigger name. - :param trigger: The trigger resource definition. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :raise AirflowException: If the trigger already exists. - :return: The trigger. - """ - if self._trigger_exists(resource_group_name, factory_name, trigger_name): - raise AirflowException(f"Trigger {trigger_name!r} already exists.") - - return self.get_conn().triggers.create_or_update( - resource_group_name, factory_name, trigger_name, trigger, **config - ) - - @provide_targeted_factory - def delete_trigger( - self, - trigger_name: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> None: - """ - Delete the trigger. - - :param trigger_name: The trigger name. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - """ - self.get_conn().triggers.delete(resource_group_name, factory_name, trigger_name, **config) - - @provide_targeted_factory - def start_trigger( - self, - trigger_name: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> LROPoller: - """ - Start the trigger. - - :param trigger_name: The trigger name. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :return: An Azure operation poller. - """ - return self.get_conn().triggers.begin_start(resource_group_name, factory_name, trigger_name, **config) - - @provide_targeted_factory - def stop_trigger( - self, - trigger_name: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> LROPoller: - """ - Stop the trigger. - - :param trigger_name: The trigger name. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - :return: An Azure operation poller. - """ - return self.get_conn().triggers.begin_stop(resource_group_name, factory_name, trigger_name, **config) - - @provide_targeted_factory - def rerun_trigger( - self, - trigger_name: str, - run_id: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> None: - """ - Rerun the trigger. - - :param trigger_name: The trigger name. - :param run_id: The trigger run identifier. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - """ - return self.get_conn().trigger_runs.rerun( - resource_group_name, factory_name, trigger_name, run_id, **config - ) - - @provide_targeted_factory - def cancel_trigger( - self, - trigger_name: str, - run_id: str, - resource_group_name: Optional[str] = None, - factory_name: Optional[str] = None, - **config: Any, - ) -> None: - """ - Cancel the trigger. - - :param trigger_name: The trigger name. - :param run_id: The trigger run identifier. - :param resource_group_name: The resource group name. - :param factory_name: The factory name. - :param config: Extra parameters for the ADF client. - """ - self.get_conn().trigger_runs.cancel(resource_group_name, factory_name, trigger_name, run_id, **config) +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 new file mode 100644 index 0000000000000..588b687d30103 --- /dev/null +++ b/airflow/providers/microsoft/azure/hooks/data_factory.py @@ -0,0 +1,854 @@ +# 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 inspect +import time +from functools import wraps +from typing import Any, Callable, Dict, Optional, Set, Union + +from azure.core.polling import LROPoller +from azure.identity import ClientSecretCredential +from azure.mgmt.datafactory import DataFactoryManagementClient +from azure.mgmt.datafactory.models import ( + CreateRunResponse, + DatasetResource, + Factory, + LinkedServiceResource, + PipelineResource, + PipelineRun, + TriggerResource, +) + +from airflow.exceptions import AirflowException +from airflow.hooks.base import BaseHook + + +def provide_targeted_factory(func: Callable) -> Callable: + """ + Provide the targeted factory to the decorated function in case it isn't specified. + + If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in + the connection extras. + """ + signature = inspect.signature(func) + + @wraps(func) + 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, 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", "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 AzureDataFactoryPipelineRunException(AirflowException): + """An exception that indicates a pipeline run failed to complete.""" + + +class AzureDataFactoryHook(BaseHook): + """ + A hook to interact with Azure Data Factory. + + :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' + conn_name_attr: str = 'azure_data_factory_conn_id' + default_conn_name: str = 'azure_data_factory_default' + hook_name: str = 'Azure Data Factory' + + @staticmethod + def get_connection_form_widgets() -> Dict[str, Any]: + """Returns connection widgets to add to connection form""" + from flask_appbuilder.fieldwidgets import BS3TextFieldWidget + from flask_babel import lazy_gettext + from wtforms import StringField + + return { + "extra__azure_data_factory__tenantId": StringField( + lazy_gettext('Tenant ID'), widget=BS3TextFieldWidget() + ), + "extra__azure_data_factory__subscriptionId": StringField( + lazy_gettext('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""" + return { + "hidden_fields": ['schema', 'port', 'host', 'extra'], + "relabeling": { + 'login': 'Client ID', + 'password': 'Secret', + }, + } + + def __init__(self, azure_data_factory_conn_id: Optional[str] = default_conn_name): + self._conn: DataFactoryManagementClient = None + 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.conn_id) + 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( + client_id=conn.login, client_secret=conn.password, tenant_id=tenant + ), + subscription_id=subscription_id, + ) + + return self._conn + + @provide_targeted_factory + def get_factory( + self, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, **config: Any + ) -> Factory: + """ + Get the factory. + + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :return: The factory. + """ + return self.get_conn().factories.get(resource_group_name, factory_name, **config) + + def _factory_exists(self, resource_group_name, factory_name) -> bool: + """Return whether or not the factory already exists.""" + factories = { + factory.name for factory in self.get_conn().factories.list_by_resource_group(resource_group_name) + } + + return factory_name in factories + + @provide_targeted_factory + def update_factory( + self, + factory: Factory, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> Factory: + """ + Update the factory. + + :param factory: The factory resource definition. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :raise AirflowException: If the factory does not exist. + :return: The factory. + """ + if not self._factory_exists(resource_group_name, factory_name): + raise AirflowException(f"Factory {factory!r} does not exist.") + + return self.get_conn().factories.create_or_update( + resource_group_name, factory_name, factory, **config + ) + + @provide_targeted_factory + def create_factory( + self, + factory: Factory, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> Factory: + """ + Create the factory. + + :param factory: The factory resource definition. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :raise AirflowException: If the factory already exists. + :return: The factory. + """ + if self._factory_exists(resource_group_name, factory_name): + raise AirflowException(f"Factory {factory!r} already exists.") + + return self.get_conn().factories.create_or_update( + resource_group_name, factory_name, factory, **config + ) + + @provide_targeted_factory + def delete_factory( + self, resource_group_name: Optional[str] = None, factory_name: Optional[str] = None, **config: Any + ) -> None: + """ + Delete the factory. + + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + """ + self.get_conn().factories.delete(resource_group_name, factory_name, **config) + + @provide_targeted_factory + def get_linked_service( + self, + linked_service_name: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> LinkedServiceResource: + """ + Get the linked service. + + :param linked_service_name: The linked service name. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :return: The linked service. + """ + return self.get_conn().linked_services.get( + resource_group_name, factory_name, linked_service_name, **config + ) + + def _linked_service_exists(self, resource_group_name, factory_name, linked_service_name) -> bool: + """Return whether or not the linked service already exists.""" + linked_services = { + linked_service.name + for linked_service in self.get_conn().linked_services.list_by_factory( + resource_group_name, factory_name + ) + } + + return linked_service_name in linked_services + + @provide_targeted_factory + def update_linked_service( + self, + linked_service_name: str, + linked_service: LinkedServiceResource, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> LinkedServiceResource: + """ + Update the linked service. + + :param linked_service_name: The linked service name. + :param linked_service: The linked service resource definition. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :raise AirflowException: If the linked service does not exist. + :return: The linked service. + """ + if not self._linked_service_exists(resource_group_name, factory_name, linked_service_name): + raise AirflowException(f"Linked service {linked_service_name!r} does not exist.") + + return self.get_conn().linked_services.create_or_update( + resource_group_name, factory_name, linked_service_name, linked_service, **config + ) + + @provide_targeted_factory + def create_linked_service( + self, + linked_service_name: str, + linked_service: LinkedServiceResource, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> LinkedServiceResource: + """ + Create the linked service. + + :param linked_service_name: The linked service name. + :param linked_service: The linked service resource definition. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :raise AirflowException: If the linked service already exists. + :return: The linked service. + """ + if self._linked_service_exists(resource_group_name, factory_name, linked_service_name): + raise AirflowException(f"Linked service {linked_service_name!r} already exists.") + + return self.get_conn().linked_services.create_or_update( + resource_group_name, factory_name, linked_service_name, linked_service, **config + ) + + @provide_targeted_factory + def delete_linked_service( + self, + linked_service_name: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> None: + """ + Delete the linked service. + + :param linked_service_name: The linked service name. + :param resource_group_name: The linked service name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + """ + self.get_conn().linked_services.delete( + resource_group_name, factory_name, linked_service_name, **config + ) + + @provide_targeted_factory + def get_dataset( + self, + dataset_name: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> DatasetResource: + """ + Get the dataset. + + :param dataset_name: The dataset name. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :return: The dataset. + """ + return self.get_conn().datasets.get(resource_group_name, factory_name, dataset_name, **config) + + def _dataset_exists(self, resource_group_name, factory_name, dataset_name) -> bool: + """Return whether or not the dataset already exists.""" + datasets = { + dataset.name + for dataset in self.get_conn().datasets.list_by_factory(resource_group_name, factory_name) + } + + return dataset_name in datasets + + @provide_targeted_factory + def update_dataset( + self, + dataset_name: str, + dataset: DatasetResource, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> DatasetResource: + """ + Update the dataset. + + :param dataset_name: The dataset name. + :param dataset: The dataset resource definition. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :raise AirflowException: If the dataset does not exist. + :return: The dataset. + """ + if not self._dataset_exists(resource_group_name, factory_name, dataset_name): + raise AirflowException(f"Dataset {dataset_name!r} does not exist.") + + return self.get_conn().datasets.create_or_update( + resource_group_name, factory_name, dataset_name, dataset, **config + ) + + @provide_targeted_factory + def create_dataset( + self, + dataset_name: str, + dataset: DatasetResource, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> DatasetResource: + """ + Create the dataset. + + :param dataset_name: The dataset name. + :param dataset: The dataset resource definition. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :raise AirflowException: If the dataset already exists. + :return: The dataset. + """ + if self._dataset_exists(resource_group_name, factory_name, dataset_name): + raise AirflowException(f"Dataset {dataset_name!r} already exists.") + + return self.get_conn().datasets.create_or_update( + resource_group_name, factory_name, dataset_name, dataset, **config + ) + + @provide_targeted_factory + def delete_dataset( + self, + dataset_name: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> None: + """ + Delete the dataset. + + :param dataset_name: The dataset name. + :param resource_group_name: The dataset name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + """ + self.get_conn().datasets.delete(resource_group_name, factory_name, dataset_name, **config) + + @provide_targeted_factory + def get_pipeline( + self, + pipeline_name: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> PipelineResource: + """ + Get the pipeline. + + :param pipeline_name: The pipeline name. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :return: The pipeline. + """ + return self.get_conn().pipelines.get(resource_group_name, factory_name, pipeline_name, **config) + + def _pipeline_exists(self, resource_group_name, factory_name, pipeline_name) -> bool: + """Return whether or not the pipeline already exists.""" + pipelines = { + pipeline.name + for pipeline in self.get_conn().pipelines.list_by_factory(resource_group_name, factory_name) + } + + return pipeline_name in pipelines + + @provide_targeted_factory + def update_pipeline( + self, + pipeline_name: str, + pipeline: PipelineResource, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> PipelineResource: + """ + Update the pipeline. + + :param pipeline_name: The pipeline name. + :param pipeline: The pipeline resource definition. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :raise AirflowException: If the pipeline does not exist. + :return: The pipeline. + """ + if not self._pipeline_exists(resource_group_name, factory_name, pipeline_name): + raise AirflowException(f"Pipeline {pipeline_name!r} does not exist.") + + return self.get_conn().pipelines.create_or_update( + resource_group_name, factory_name, pipeline_name, pipeline, **config + ) + + @provide_targeted_factory + def create_pipeline( + self, + pipeline_name: str, + pipeline: PipelineResource, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> PipelineResource: + """ + Create the pipeline. + + :param pipeline_name: The pipeline name. + :param pipeline: The pipeline resource definition. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :raise AirflowException: If the pipeline already exists. + :return: The pipeline. + """ + if self._pipeline_exists(resource_group_name, factory_name, pipeline_name): + raise AirflowException(f"Pipeline {pipeline_name!r} already exists.") + + return self.get_conn().pipelines.create_or_update( + resource_group_name, factory_name, pipeline_name, pipeline, **config + ) + + @provide_targeted_factory + def delete_pipeline( + self, + pipeline_name: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> None: + """ + Delete the pipeline. + + :param pipeline_name: The pipeline name. + :param resource_group_name: The pipeline name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + """ + self.get_conn().pipelines.delete(resource_group_name, factory_name, pipeline_name, **config) + + @provide_targeted_factory + def run_pipeline( + self, + pipeline_name: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> CreateRunResponse: + """ + Run a pipeline. + + :param pipeline_name: The pipeline name. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :return: The pipeline run. + """ + return self.get_conn().pipelines.create_run( + resource_group_name, factory_name, pipeline_name, **config + ) + + @provide_targeted_factory + def get_pipeline_run( + self, + run_id: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> PipelineRun: + """ + Get the pipeline run. + + :param run_id: The pipeline run identifier. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :return: The pipeline run. + """ + return self.get_conn().pipeline_runs.get(resource_group_name, factory_name, run_id, **config) + + def get_pipeline_run_status( + self, + run_id: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + ) -> str: + """ + Get a pipeline run's current status. + + :param run_id: The pipeline run identifier. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :return: The status of the 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, + ).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." + ) + + # Wait to check the status of the pipeline run based on the ``check_interval`` configured. + time.sleep(check_interval) + + pipeline_run_status = self.get_pipeline_run_status(**pipeline_run_info) + + return pipeline_run_status in expected_statuses + + @provide_targeted_factory + def cancel_pipeline_run( + self, + run_id: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> None: + """ + Cancel the pipeline run. + + :param run_id: The pipeline run identifier. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + """ + self.get_conn().pipeline_runs.cancel(resource_group_name, factory_name, run_id, **config) + + @provide_targeted_factory + def get_trigger( + self, + trigger_name: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> TriggerResource: + """ + Get the trigger. + + :param trigger_name: The trigger name. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :return: The trigger. + """ + return self.get_conn().triggers.get(resource_group_name, factory_name, trigger_name, **config) + + def _trigger_exists(self, resource_group_name, factory_name, trigger_name) -> bool: + """Return whether or not the trigger already exists.""" + triggers = { + trigger.name + for trigger in self.get_conn().triggers.list_by_factory(resource_group_name, factory_name) + } + + return trigger_name in triggers + + @provide_targeted_factory + def update_trigger( + self, + trigger_name: str, + trigger: TriggerResource, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> TriggerResource: + """ + Update the trigger. + + :param trigger_name: The trigger name. + :param trigger: The trigger resource definition. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :raise AirflowException: If the trigger does not exist. + :return: The trigger. + """ + if not self._trigger_exists(resource_group_name, factory_name, trigger_name): + raise AirflowException(f"Trigger {trigger_name!r} does not exist.") + + return self.get_conn().triggers.create_or_update( + resource_group_name, factory_name, trigger_name, trigger, **config + ) + + @provide_targeted_factory + def create_trigger( + self, + trigger_name: str, + trigger: TriggerResource, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> TriggerResource: + """ + Create the trigger. + + :param trigger_name: The trigger name. + :param trigger: The trigger resource definition. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :raise AirflowException: If the trigger already exists. + :return: The trigger. + """ + if self._trigger_exists(resource_group_name, factory_name, trigger_name): + raise AirflowException(f"Trigger {trigger_name!r} already exists.") + + return self.get_conn().triggers.create_or_update( + resource_group_name, factory_name, trigger_name, trigger, **config + ) + + @provide_targeted_factory + def delete_trigger( + self, + trigger_name: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> None: + """ + Delete the trigger. + + :param trigger_name: The trigger name. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + """ + self.get_conn().triggers.delete(resource_group_name, factory_name, trigger_name, **config) + + @provide_targeted_factory + def start_trigger( + self, + trigger_name: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> LROPoller: + """ + Start the trigger. + + :param trigger_name: The trigger name. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :return: An Azure operation poller. + """ + return self.get_conn().triggers.begin_start(resource_group_name, factory_name, trigger_name, **config) + + @provide_targeted_factory + def stop_trigger( + self, + trigger_name: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> LROPoller: + """ + Stop the trigger. + + :param trigger_name: The trigger name. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + :return: An Azure operation poller. + """ + return self.get_conn().triggers.begin_stop(resource_group_name, factory_name, trigger_name, **config) + + @provide_targeted_factory + def rerun_trigger( + self, + trigger_name: str, + run_id: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> None: + """ + Rerun the trigger. + + :param trigger_name: The trigger name. + :param run_id: The trigger run identifier. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + """ + return self.get_conn().trigger_runs.rerun( + resource_group_name, factory_name, trigger_name, run_id, **config + ) + + @provide_targeted_factory + def cancel_trigger( + self, + trigger_name: str, + run_id: str, + resource_group_name: Optional[str] = None, + factory_name: Optional[str] = None, + **config: Any, + ) -> None: + """ + Cancel the trigger. + + :param trigger_name: The trigger name. + :param run_id: The trigger run identifier. + :param resource_group_name: The resource group name. + :param factory_name: The factory name. + :param config: Extra parameters for the ADF client. + """ + self.get_conn().trigger_runs.cancel(resource_group_name, factory_name, trigger_name, run_id, **config) diff --git a/airflow/providers/microsoft/azure/operators/data_factory.py b/airflow/providers/microsoft/azure/operators/data_factory.py new file mode 100644 index 0000000000000..ddcf489b93438 --- /dev/null +++ b/airflow/providers/microsoft/azure/operators/data_factory.py @@ -0,0 +1,174 @@ +# 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 Any, Dict, Optional + +from airflow.models import BaseOperator +from airflow.providers.microsoft.azure.hooks.data_factory import ( + AzureDataFactoryHook, + AzureDataFactoryPipelineRunException, + AzureDataFactoryPipelineRunStatus, +) + + +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 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 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_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. + :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 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 + 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 are referenced in a pipeline via + ``@pipeline().parameters.parameterName`` and will be used only if the ``reference_pipeline_run_id`` is + not specified. + :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_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_termination`` is True. + :type check_interval: int + """ + + template_fields = ( + "azure_data_factory_conn_id", + "resource_group_name", + "factory_name", + "pipeline_name", + "reference_pipeline_run_id", + "parameters", + ) + template_fields_renderers = {"parameters": "json"} + + ui_color = "#0678d4" + + def __init__( + self, + *, + pipeline_name: str, + 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, + is_recovery: Optional[bool] = None, + start_activity_name: Optional[str] = None, + start_from_failure: Optional[bool] = None, + parameters: Optional[Dict[str, Any]] = None, + timeout: Optional[int] = 60 * 60 * 24 * 7, + 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_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 + self.is_recovery = is_recovery + self.start_activity_name = start_activity_name + self.start_from_failure = start_from_failure + self.parameters = parameters + self.timeout = timeout + 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) + 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 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_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.") + else: + raise AzureDataFactoryPipelineRunException( + f"Pipeline run {self.run_id} has failed or has been cancelled." + ) + + 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, + ) + + # 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 3622489f3cf5f..9025d8c662978 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] @@ -97,6 +99,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.data_factory sensors: - integration-name: Microsoft Azure Cosmos DB @@ -105,6 +110,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.data_factory hooks: - integration-name: Microsoft Azure Container Instances @@ -135,6 +143,7 @@ hooks: - airflow.providers.microsoft.azure.hooks.wasb - integration-name: Microsoft Azure Data Factory python-modules: + - airflow.providers.microsoft.azure.hooks.data_factory - airflow.providers.microsoft.azure.hooks.azure_data_factory transfers: @@ -163,7 +172,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: @@ -186,7 +195,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/data_factory.py b/airflow/providers/microsoft/azure/sensors/data_factory.py new file mode 100644 index 0000000000000..de2f3938c57af --- /dev/null +++ b/airflow/providers/microsoft/azure/sensors/data_factory.py @@ -0,0 +1,75 @@ +# 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 + +from airflow.providers.microsoft.azure.hooks.data_factory import ( + AzureDataFactoryHook, + AzureDataFactoryPipelineRunException, + AzureDataFactoryPipelineRunStatus, +) +from airflow.sensors.base import BaseSensorOperator + + +class AzureDataFactoryPipelineRunStatusSensor(BaseSensorOperator): + """ + Checks the status of a pipeline run. + + :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. + :type resource_group_name: str + :param factory_name: The data factory name. + :type factory_name: str + """ + + template_fields = ("azure_data_factory_conn_id", "resource_group_name", "factory_name", "run_id") + + ui_color = "#50e6ff" + + def __init__( + self, + *, + 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, + ) -> None: + super().__init__(**kwargs) + 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.hook = AzureDataFactoryHook(azure_data_factory_conn_id=self.azure_data_factory_conn_id) + pipeline_run_status = self.hook.get_pipeline_run_status( + run_id=self.run_id, + 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/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`.", diff --git a/docs/apache-airflow-providers-microsoft-azure/connections/adf.rst b/docs/apache-airflow-providers-microsoft-azure/connections/adf.rst index 6c38b94e5e6cc..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. @@ -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 + Specify the ``tenantId`` used for the initial connection. + This is needed for *token credentials* authentication mechanism. + +Subscription ID + Specify the ``subscriptionId`` used for the initial connection. + This is needed 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. 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..1462648b8b691 --- /dev/null +++ b/docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst @@ -0,0 +1,57 @@ + + .. 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.data_factory.AzureDataFactoryRunPipelineOperator` to execute a pipeline within a data factory. +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. + + .. 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] + +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 +--------- + +For further information, please refer to the Microsoft documentation: + + * `Azure Data Factory Documentation `__ 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..670d3254aea7c 100644 --- a/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/hooks/test_azure_data_factory.py @@ -17,15 +17,17 @@ import json -from unittest.mock import MagicMock, Mock +from unittest.mock import MagicMock, Mock, patch import pytest from pytest import fixture 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, + AzureDataFactoryPipelineRunException, + AzureDataFactoryPipelineRunStatus, provide_targeted_factory, ) from airflow.utils import db @@ -49,10 +51,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, } ), ) @@ -62,7 +64,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", @@ -100,8 +102,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 = {} @@ -336,6 +344,42 @@ def test_get_pipeline_run(hook: AzureDataFactoryHook, user_args, sdk_args): hook._conn.pipeline_runs.get.assert_called_with(*sdk_args) +_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: + mock_pipeline_run.return_value.status = pipeline_run_status + + 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( 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 new file mode 100644 index 0000000000000..25080b5a1a9b8 --- /dev/null +++ b/tests/providers/microsoft/azure/operators/test_azure_data_factory.py @@ -0,0 +1,177 @@ +# 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.providers.microsoft.azure.hooks.data_factory import ( + AzureDataFactoryHook, + AzureDataFactoryPipelineRunException, + AzureDataFactoryPipelineRunStatus, +) +from airflow.providers.microsoft.azure.operators.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", + "azure_data_factory_conn_id": "azure_data_factory_test", + "pipeline_name": "pipeline1", + "resource_group_name": "resource-group-name", + "factory_name": "factory-name", + "check_interval": 1, + "timeout": 3, + } + + @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, "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_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"] + 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 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( + pipeline_run_status + ) + + if not expected_output: + # A successful operator execution should not return any values. + assert not operator.execute(context=self.mock_context) + elif expected_output == "exception": + # The operator should fail if the pipeline run fails or is canceled. + 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. + 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, + ) + + 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_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_termination + + 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. + 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_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 new file mode 100644 index 0000000000000..4f626af3dd4f0 --- /dev/null +++ b/tests/providers/microsoft/azure/sensors/test_azure_data_factory.py @@ -0,0 +1,76 @@ +# 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 patch + +import pytest +from parameterized import parameterized + +from airflow.providers.microsoft.azure.hooks.data_factory import ( + AzureDataFactoryHook, + AzureDataFactoryPipelineRunException, + AzureDataFactoryPipelineRunStatus, +) +from airflow.providers.microsoft.azure.sensors.data_factory import AzureDataFactoryPipelineRunStatusSensor + + +class TestPipelineRunStatusSensor(unittest.TestCase): + def setUp(self): + self.config = { + "azure_data_factory_conn_id": "azure_data_factory_test", + "run_id": "run_id", + "resource_group_name": "resource-group-name", + "factory_name": "factory-name", + "timeout": 100, + "poke_interval": 15, + } + self.sensor = AzureDataFactoryPipelineRunStatusSensor(task_id="pipeline_run_sensor", **self.config) + + def test_init(self): + 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, "exception"), + (AzureDataFactoryPipelineRunStatus.CANCELLED, "exception"), + (AzureDataFactoryPipelineRunStatus.CANCELING, False), + (AzureDataFactoryPipelineRunStatus.QUEUED, False), + (AzureDataFactoryPipelineRunStatus.IN_PROGRESS, False), + ] + ) + @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 != "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: + error_message = f"Pipeline run {self.config['run_id']} has been cancelled." + + with pytest.raises(AzureDataFactoryPipelineRunException, match=error_message): + self.sensor.poke({})