def poke(self, context: "Context") -> 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
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.")