diff --git a/astronomer/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py b/astronomer/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py index d25007bba..72c17d5a6 100644 --- a/astronomer/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py +++ b/astronomer/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py @@ -7,6 +7,7 @@ from airflow import DAG from airflow.models.xcom_arg import XComArg from airflow.operators.python import PythonOperator + from astronomer.providers.microsoft.azure.operators.data_factory import ( AzureDataFactoryRunPipelineOperatorAsync, ) diff --git a/astronomer/providers/microsoft/azure/hooks/data_factory.py b/astronomer/providers/microsoft/azure/hooks/data_factory.py index 2504dbdc4..65fa7cec9 100644 --- a/astronomer/providers/microsoft/azure/hooks/data_factory.py +++ b/astronomer/providers/microsoft/azure/hooks/data_factory.py @@ -82,7 +82,7 @@ def __init__(self, azure_data_factory_conn_id: str): async def get_async_conn(self) -> DataFactoryManagementClient: """Get async connection and connect to azure data factory.""" if self._conn is not None: - return self._conn + return cast(DataFactoryManagementClient, self._conn) conn = await sync_to_async(self.get_connection)(self.conn_id) extras = conn.extra_dejson @@ -147,7 +147,7 @@ async def get_adf_pipeline_run_status( factory_name=factory_name, resource_group_name=resource_group_name, ) - status: str = pipeline_run.status + status: str = cast(str, pipeline_run.status) return status except Exception as e: raise AirflowException(e) diff --git a/astronomer/providers/microsoft/azure/triggers/data_factory.py b/astronomer/providers/microsoft/azure/triggers/data_factory.py index cc64d4512..1628ddd6b 100644 --- a/astronomer/providers/microsoft/azure/triggers/data_factory.py +++ b/astronomer/providers/microsoft/azure/triggers/data_factory.py @@ -6,6 +6,7 @@ AzureDataFactoryPipelineRunStatus, ) from airflow.triggers.base import BaseTrigger, TriggerEvent + from astronomer.providers.microsoft.azure.hooks.data_factory import ( AzureDataFactoryHookAsync, )