Allow setting the pooling time in DLPHook (#8824)

xuan616 · mik-laj · ashb · web-flow · commit 841d81664737 · 2020-05-19T04:55:41.000+02:00
Co-authored-by: Kamil Breguła &lt;mik-laj@users.noreply.github.com&gt;
Co-authored-by: Ash Berlin-Taylor &lt;ash_github@firemirror.com&gt;
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(
