diff --git a/airflow/providers/microsoft/azure/hooks/data_factory.py b/airflow/providers/microsoft/azure/hooks/data_factory.py index da8455daa29a7..b9e7feea0a13a 100644 --- a/airflow/providers/microsoft/azure/hooks/data_factory.py +++ b/airflow/providers/microsoft/azure/hooks/data_factory.py @@ -1175,3 +1175,25 @@ async def get_adf_pipeline_run_status( return status except Exception as e: raise AirflowException(e) + + @provide_targeted_factory_async + async def cancel_pipeline_run( + self, + run_id: str, + resource_group_name: str | None = None, + factory_name: str | None = 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. + """ + client = await self.get_async_conn() + try: + await client.pipeline_runs.cancel(resource_group_name, factory_name, run_id) + except Exception as e: + raise AirflowException(e) diff --git a/airflow/providers/microsoft/azure/triggers/data_factory.py b/airflow/providers/microsoft/azure/triggers/data_factory.py index c68d6e96b7916..40b95559405c7 100644 --- a/airflow/providers/microsoft/azure/triggers/data_factory.py +++ b/airflow/providers/microsoft/azure/triggers/data_factory.py @@ -191,4 +191,11 @@ async def run(self) -> AsyncIterator[TriggerEvent]: } ) except Exception as e: + if self.run_id: + await hook.cancel_pipeline_run( + run_id=self.run_id, + resource_group_name=self.resource_group_name, + factory_name=self.factory_name, + ) + self.log.info("Unexpected error %s caught. Cancel pipeline run %s", str(e), self.run_id) yield TriggerEvent({"status": "error", "message": str(e), "run_id": self.run_id}) diff --git a/tests/providers/microsoft/azure/operators/test_azure_data_factory.py b/tests/providers/microsoft/azure/operators/test_azure_data_factory.py index 461c1502c6b0c..2a545cd1371f2 100644 --- a/tests/providers/microsoft/azure/operators/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/operators/test_azure_data_factory.py @@ -363,7 +363,7 @@ class CreateRunResponse: ) @mock.patch("airflow.providers.microsoft.azure.hooks.data_factory.AzureDataFactoryHook.run_pipeline") def test_azure_data_factory_run_pipeline_operator_async(self, mock_run_pipeline, mock_get_status, status): - """Assert that AzureDataFactoryRunPipelineOperatorAsync deferred""" + """Assert that AzureDataFactoryRunPipelineOperator(..., deferrable=True) deferred""" class CreateRunResponse: pass diff --git a/tests/providers/microsoft/azure/triggers/test_azure_data_factory.py b/tests/providers/microsoft/azure/triggers/test_azure_data_factory.py index c7c31d2818eeb..9df51bae71a5e 100644 --- a/tests/providers/microsoft/azure/triggers/test_azure_data_factory.py +++ b/tests/providers/microsoft/azure/triggers/test_azure_data_factory.py @@ -316,8 +316,11 @@ async def test_azure_data_factory_trigger_run_cancelled(self, mock_pipeline_run_ assert expected == actual @pytest.mark.asyncio + @mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.cancel_pipeline_run") @mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_adf_pipeline_run_status") - async def test_azure_data_factory_trigger_run_exception(self, mock_pipeline_run_status): + async def test_azure_data_factory_trigger_run_exception( + self, mock_pipeline_run_status, mock_cancel_pipeline_run + ): """Assert that run catch exception if Azure API throw exception""" mock_pipeline_run_status.side_effect = Exception("Test exception") @@ -331,6 +334,7 @@ async def test_azure_data_factory_trigger_run_exception(self, mock_pipeline_run_ ) assert len(task) == 1 assert response in task + mock_cancel_pipeline_run.assert_called_once() @pytest.mark.asyncio @mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_adf_pipeline_run_status")