respect soft_fail argument when exception is raised for google sensors (

#34501)
apache · Sep 26, 2023 · 20b7cfc · 20b7cfc
1 parent f99e65b
commit 20b7cfc
Show file tree

Hide file tree

Showing 25 changed files with 461 additions and 163 deletions.
diff --git a/airflow/providers/google/cloud/sensors/bigquery.py b/airflow/providers/google/cloud/sensors/bigquery.py
@@ -23,7 +23,7 @@
 from typing import TYPE_CHECKING, Any, Sequence
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, AirflowSkipException
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
 from airflow.providers.google.cloud.triggers.bigquery import (
  BigQueryTableExistenceTrigger,
@@ -141,8 +141,16 @@ def execute_complete(self, context: dict[str, Any], event: dict[str, str] | None
  if event:
  if event["status"] == "success":
  return event["message"]
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ if self.soft_fail:
+ raise AirflowSkipException(event["message"])
  raise AirflowException(event["message"])
- raise AirflowException("No event received in trigger callback")
+
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = "No event received in trigger callback"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
 
 
 class BigQueryTablePartitionExistenceSensor(BaseSensorOperator):
@@ -248,8 +256,17 @@ def execute_complete(self, context: dict[str, Any], event: dict[str, str] | None
  if event:
  if event["status"] == "success":
  return event["message"]
+
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ if self.soft_fail:
+ raise AirflowSkipException(event["message"])
  raise AirflowException(event["message"])
- raise AirflowException("No event received in trigger callback")
+
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = "No event received in trigger callback"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
 
 
 class BigQueryTableExistenceAsyncSensor(BigQueryTableExistenceSensor):

diff --git a/airflow/providers/google/cloud/sensors/bigquery_dts.py b/airflow/providers/google/cloud/sensors/bigquery_dts.py
@@ -23,7 +23,7 @@
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.bigquery_datatransfer_v1 import TransferState
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.providers.google.cloud.hooks.bigquery_dts import BiqQueryDataTransferServiceHook
 from airflow.sensors.base import BaseSensorOperator
 
@@ -140,5 +140,9 @@ def poke(self, context: Context) -> bool:
  self.log.info("Status of %s run: %s", self.run_id, str(run.state))
 
  if run.state in (TransferState.FAILED, TransferState.CANCELLED):
- raise AirflowException(f"Transfer {self.run_id} did not succeed")
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = f"Transfer {self.run_id} did not succeed"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
  return run.state in self.expected_statuses
diff --git a/airflow/providers/google/cloud/sensors/cloud_composer.py b/airflow/providers/google/cloud/sensors/cloud_composer.py
@@ -21,7 +21,7 @@
 
 from typing import TYPE_CHECKING, Any, Sequence
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.providers.google.cloud.triggers.cloud_composer import CloudComposerExecutionTrigger
 from airflow.sensors.base import BaseSensorOperator
 
@@ -90,5 +90,14 @@ def execute_complete(self, context: dict[str, Any], event: dict[str, str] | None
  if event:
  if event.get("operation_done"):
  return event["operation_done"]
+
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ if self.soft_fail:
+ raise AirflowSkipException(event["message"])
  raise AirflowException(event["message"])
- raise AirflowException("No event received in trigger callback")
+
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = "No event received in trigger callback"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
diff --git a/airflow/providers/google/cloud/sensors/dataflow.py b/airflow/providers/google/cloud/sensors/dataflow.py
@@ -20,7 +20,7 @@
 
 from typing import TYPE_CHECKING, Callable, Sequence
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.providers.google.cloud.hooks.dataflow import (
  DEFAULT_DATAFLOW_LOCATION,
  DataflowHook,
@@ -106,7 +106,11 @@ def poke(self, context: Context) -> bool:
  if job_status in self.expected_statuses:
  return True
  elif job_status in DataflowJobStatus.TERMINAL_STATES:
- raise AirflowException(f"Job with id '{self.job_id}' is already in terminal state: {job_status}")
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = f"Job with id '{self.job_id}' is already in terminal state: {job_status}"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
 
  return False
 
@@ -178,9 +182,11 @@ def poke(self, context: Context) -> bool:
  )
  job_status = job["currentState"]
  if job_status in DataflowJobStatus.TERMINAL_STATES:
- raise AirflowException(
- f"Job with id '{self.job_id}' is already in terminal state: {job_status}"
- )
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = f"Job with id '{self.job_id}' is already in terminal state: {job_status}"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
 
  result = self.hook.fetch_job_metrics_by_id(
  job_id=self.job_id,
@@ -257,9 +263,11 @@ def poke(self, context: Context) -> bool:
  )
  job_status = job["currentState"]
  if job_status in DataflowJobStatus.TERMINAL_STATES:
- raise AirflowException(
- f"Job with id '{self.job_id}' is already in terminal state: {job_status}"
- )
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = f"Job with id '{self.job_id}' is already in terminal state: {job_status}"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
 
  result = self.hook.fetch_job_messages_by_id(
  job_id=self.job_id,
@@ -336,9 +344,11 @@ def poke(self, context: Context) -> bool:
  )
  job_status = job["currentState"]
  if job_status in DataflowJobStatus.TERMINAL_STATES:
- raise AirflowException(
- f"Job with id '{self.job_id}' is already in terminal state: {job_status}"
- )
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = f"Job with id '{self.job_id}' is already in terminal state: {job_status}"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
 
  result = self.hook.fetch_job_autoscaling_events_by_id(
  job_id=self.job_id,

diff --git a/airflow/providers/google/cloud/sensors/dataform.py b/airflow/providers/google/cloud/sensors/dataform.py
@@ -20,7 +20,7 @@
 
 from typing import TYPE_CHECKING, Iterable, Sequence
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.providers.google.cloud.hooks.dataform import DataformHook
 from airflow.sensors.base import BaseSensorOperator
 
@@ -95,9 +95,13 @@ def poke(self, context: Context) -> bool:
  workflow_status = workflow_invocation.state
  if workflow_status is not None:
  if self.failure_statuses and workflow_status in self.failure_statuses:
- raise AirflowException(
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = (
  f"Workflow Invocation with id '{self.workflow_invocation_id}' "
  f"state is: {workflow_status}. Terminating sensor..."
  )
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
 
  return workflow_status in self.expected_statuses
diff --git a/airflow/providers/google/cloud/sensors/datafusion.py b/airflow/providers/google/cloud/sensors/datafusion.py
@@ -20,7 +20,7 @@
 
 from typing import TYPE_CHECKING, Iterable, Sequence
 
-from airflow.exceptions import AirflowException, AirflowNotFoundException
+from airflow.exceptions import AirflowException, AirflowNotFoundException, AirflowSkipException
 from airflow.providers.google.cloud.hooks.datafusion import DataFusionHook
 from airflow.sensors.base import BaseSensorOperator
 
@@ -109,15 +109,23 @@ def poke(self, context: Context) -> bool:
  )
  pipeline_status = pipeline_workflow["status"]
  except AirflowNotFoundException:
- raise AirflowException("Specified Pipeline ID was not found.")
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = "Specified Pipeline ID was not found."
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
  except AirflowException:
  pass # Because the pipeline may not be visible in system yet
  if pipeline_status is not None:
  if self.failure_statuses and pipeline_status in self.failure_statuses:
- raise AirflowException(
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = (
  f"Pipeline with id '{self.pipeline_id}' state is: {pipeline_status}. "
  f"Terminating sensor..."
  )
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
 
  self.log.debug(
  "Current status of the pipeline workflow for %s: %s.", self.pipeline_id, pipeline_status

diff --git a/airflow/providers/google/cloud/sensors/dataplex.py b/airflow/providers/google/cloud/sensors/dataplex.py
@@ -24,11 +24,12 @@
  from google.api_core.retry import Retry
 
  from airflow.utils.context import Context
+
 from google.api_core.exceptions import GoogleAPICallError
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.dataplex_v1.types import DataScanJob
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.providers.google.cloud.hooks.dataplex import (
  AirflowDataQualityScanException,
  AirflowDataQualityScanResultTimeoutException,
@@ -116,7 +117,11 @@ def poke(self, context: Context) -> bool:
  task_status = task.state
 
  if task_status == TaskState.DELETING:
- raise AirflowException(f"Task is going to be deleted {self.dataplex_task_id}")
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = f"Task is going to be deleted {self.dataplex_task_id}"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
 
  self.log.info("Current status of the Dataplex task %s => %s", self.dataplex_task_id, task_status)
 
@@ -196,9 +201,13 @@ def poke(self, context: Context) -> bool:
  if self.result_timeout:
  duration = self._duration()
  if duration > self.result_timeout:
- raise AirflowDataQualityScanResultTimeoutException(
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = (
  f"Timeout: Data Quality scan {self.job_id} is not ready after {self.result_timeout}s"
  )
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowDataQualityScanResultTimeoutException(message)
 
  hook = DataplexHook(
  gcp_conn_id=self.gcp_conn_id,
@@ -217,22 +226,36 @@ def poke(self, context: Context) -> bool:
  metadata=self.metadata,
  )
  except GoogleAPICallError as e:
- raise AirflowException(
- f"Error occurred when trying to retrieve Data Quality scan job: {self.data_scan_id}", e
- )
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = f"Error occurred when trying to retrieve Data Quality scan job: {self.data_scan_id}"
+ if self.soft_fail:
+ raise AirflowSkipException(message, e)
+ raise AirflowException(message, e)
 
  job_status = job.state
  self.log.info(
  "Current status of the Dataplex Data Quality scan job %s => %s", self.job_id, job_status
  )
  if job_status == DataScanJob.State.FAILED:
- raise AirflowException(f"Data Quality scan job failed: {self.job_id}")
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = f"Data Quality scan job failed: {self.job_id}"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
  if job_status == DataScanJob.State.CANCELLED:
- raise AirflowException(f"Data Quality scan job cancelled: {self.job_id}")
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = f"Data Quality scan job cancelled: {self.job_id}"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
  if self.fail_on_dq_failure:
  if job_status == DataScanJob.State.SUCCEEDED and not job.data_quality_result.passed:
- raise AirflowDataQualityScanException(
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = (
  f"Data Quality job {self.job_id} execution failed due to failure of its scanning "
  f"rules: {self.data_scan_id}"
  )
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowDataQualityScanException(message)
  return job_status == DataScanJob.State.SUCCEEDED
diff --git a/airflow/providers/google/cloud/sensors/dataproc.py b/airflow/providers/google/cloud/sensors/dataproc.py
@@ -24,7 +24,7 @@
 from google.api_core.exceptions import ServerError
 from google.cloud.dataproc_v1.types import Batch, JobStatus
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.providers.google.cloud.hooks.dataproc import DataprocHook
 from airflow.sensors.base import BaseSensorOperator
 
@@ -83,24 +83,36 @@ def poke(self, context: Context) -> bool:
  duration = self._duration()
  self.log.info("DURATION RUN: %f", duration)
  if duration > self.wait_timeout:
- raise AirflowException(
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = (
  f"Timeout: dataproc job {self.dataproc_job_id} "
  f"is not ready after {self.wait_timeout}s"
  )
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
  self.log.info("Retrying. Dataproc API returned server error when waiting for job: %s", err)
  return False
  else:
  job = hook.get_job(job_id=self.dataproc_job_id, region=self.region, project_id=self.project_id)
 
  state = job.status.state
  if state == JobStatus.State.ERROR:
- raise AirflowException(f"Job failed:\n{job}")
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = f"Job failed:\n{job}"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
  elif state in {
  JobStatus.State.CANCELLED,
  JobStatus.State.CANCEL_PENDING,
  JobStatus.State.CANCEL_STARTED,
  }:
- raise AirflowException(f"Job was cancelled:\n{job}")
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = f"Job was cancelled:\n{job}"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
  elif JobStatus.State.DONE == state:
  self.log.debug("Job %s completed successfully.", self.dataproc_job_id)
  return True
@@ -171,12 +183,20 @@ def poke(self, context: Context) -> bool:
 
  state = batch.state
  if state == Batch.State.FAILED:
- raise AirflowException("Batch failed")
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = "Batch failed"
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
  elif state in {
  Batch.State.CANCELLED,
  Batch.State.CANCELLING,
  }:
- raise AirflowException("Batch was cancelled.")
+ # TODO: remove this if check when min_airflow_version is set to higher than 2.7.1
+ message = "Batch was cancelled."
+ if self.soft_fail:
+ raise AirflowSkipException(message)
+ raise AirflowException(message)
  elif state == Batch.State.SUCCEEDED:
  self.log.debug("Batch %s completed successfully.", self.batch_id)
  return True