Allow setting the pooling time in DLPHook (#8824)

Co-authored-by: Kamil Breguła <[email protected]> Co-authored-by: Ash Berlin-Taylor <[email protected]>
apache · May 19, 2020 · 841d816 · 841d816
1 parent 2121f49
commit 841d816
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/airflow/providers/google/cloud/hooks/dlp.py b/airflow/providers/google/cloud/hooks/dlp.py
@@ -38,8 +38,6 @@
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 DLP_JOB_PATH_PATTERN = "^projects/[^/]+/dlpJobs/(?P<job>.*?)$"
-# Time to sleep between active checks of the operation results
-TIME_TO_SLEEP_IN_SECONDS = 1
 
 
 # pylint: disable=R0904, C0302
@@ -180,6 +178,7 @@ def create_dlp_job(
  timeout: Optional[float] = None,
  metadata: Optional[Sequence[Tuple[str, str]]] = None,
  wait_until_finished: bool = True,
+ time_to_sleep_in_seconds: int = 60
  ) -> DlpJob:
  """
  Creates a new job to inspect storage or calculate risk metrics.
@@ -207,6 +206,9 @@ def create_dlp_job(
  until it is set to DONE.
  :type wait_until_finished: bool
  :rtype: google.cloud.dlp_v2.types.DlpJob
+ :param time_to_sleep_in_seconds: (Optional) Time to sleep, in seconds, between active checks
+ of the operation results. Defaults to 60.
+ :type time_to_sleep_in_seconds: int
  """
 
  client = self.get_conn()
@@ -242,7 +244,7 @@ def create_dlp_job(
  DlpJob.JobState.RUNNING,
  DlpJob.JobState.JOB_STATE_UNSPECIFIED,
  ]:
- time.sleep(TIME_TO_SLEEP_IN_SECONDS)
+ time.sleep(time_to_sleep_in_seconds)
  else:
  raise AirflowException(
  "Stopped polling DLP job state. DLP job {} state: {}.".format(