Improve idempotency of BigQueryInsertJobOperator

airflow/providers/google/cloud/hooks/bigquery.py

@@ -49,6 +49,8 @@
 
 log = logging.getLogger(__name__)
 
+BigQueryJob = Union[CopyJob, QueryJob, LoadJob, ExtractJob]
+
 
 # pylint: disable=too-many-public-methods
 class BigQueryHook(GoogleBaseHook, DbApiHook):
@@ -1435,7 +1437,7 @@ def insert_job(
         job_id: Optional[str] = None,
         project_id: Optional[str] = None,
         location: Optional[str] = None,
-    ) -> Union[CopyJob, QueryJob, LoadJob, ExtractJob]:
+    ) -> BigQueryJob:
         """
         Executes a BigQuery job. Waits for the job to complete and returns job id.
         See here:

airflow/providers/google/cloud/operators/bigquery.py

@@ -21,21 +21,22 @@
 This module contains Google BigQuery operators.
 """
 import enum
+import hashlib
 import json
+import re
+import uuid
 import warnings
-from time import sleep, time
-from typing import Any, Dict, Iterable, List, Optional, SupportsAbs, Union
+from typing import Any, Dict, Iterable, List, Optional, Set, SupportsAbs, Union
 
 import attr
 from google.api_core.exceptions import Conflict
-from google.api_core.retry import exponential_sleep_generator
 from google.cloud.bigquery import TableReference
 
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator, BaseOperatorLink
 from airflow.models.taskinstance import TaskInstance
 from airflow.operators.check_operator import CheckOperator, IntervalCheckOperator, ValueCheckOperator
-from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
+from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
 from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url
 from airflow.utils.decorators import apply_defaults
 
@@ -1631,7 +1632,18 @@ def execute(self, context):
 class BigQueryInsertJobOperator(BaseOperator):
     """
     Executes a BigQuery job. Waits for the job to complete and returns job id.
-    See here:
+    This operator work in the following way:
+
+    - it calculates a unique hash of the job using job's configuration or uuid if ``force_rerun`` is True
+    - creates ``job_id`` in form of
+        ``[provided_job_id | airflow_{dag_id}_{task_id}_{exec_date}]_{uniqueness_suffix}``
+    - submits a BigQuery job using the ``job_id``
+    - if job with given id already exists then it tries to reattach to the job if its not done and its
+        state is in ``reattach_states``. If the job is done the operator will raise ``AirflowException``.
+
+    Using ``force_rerun`` will submit a new job everytime without attaching to already existing ones.
+
+    For job definition see here:
 
         https://cloud.google.com/bigquery/docs/reference/v2/jobs
 
@@ -1644,15 +1656,21 @@ class BigQueryInsertJobOperator(BaseOperator):
         configuration field in the job  object. For more details see
         https://cloud.google.com/bigquery/docs/reference/v2/jobs
     :type configuration: Dict[str, Any]
-    :param job_id: The ID of the job. The ID must contain only letters (a-z, A-Z),
-        numbers (0-9), underscores (_), or dashes (-). The maximum length is 1,024
-        characters. If not provided then uuid will be generated.
+    :param job_id: The ID of the job. It will be suffixed with hash of job configuration
+        unless ``force_rerun`` is True.
+        The ID must contain only letters (a-z, A-Z), numbers (0-9), underscores (_), or
+        dashes (-). The maximum length is 1,024 characters. If not provided then uuid will
+        be generated.
     :type job_id: str
+    :param force_rerun: If True then operator will use hash of uuid as job id suffix
+    :type force_rerun: bool
+    :param reattach_states: Set of BigQuery job's states in case of which we should reattach
+        to the job. Should be other than final states.
     :param project_id: Google Cloud Project where the job is running
     :type project_id: str
     :param location: location the job is running
     :type location: str
-    :param gcp_conn_id: (Optional) The connection ID used to connect to Google Cloud Platform.
+    :param gcp_conn_id: The connection ID used to connect to Google Cloud Platform.
     :type gcp_conn_id: str
     """
 
@@ -1666,6 +1684,8 @@ def __init__(
         project_id: Optional[str] = None,
         location: Optional[str] = None,
         job_id: Optional[str] = None,
+        force_rerun: bool = True,
+        reattach_states: Optional[Set[str]] = None,
         gcp_conn_id: str = 'google_cloud_default',
         delegate_to: Optional[str] = None,
         **kwargs,
@@ -1677,39 +1697,78 @@ def __init__(
         self.project_id = project_id
         self.gcp_conn_id = gcp_conn_id
         self.delegate_to = delegate_to
+        self.force_rerun = force_rerun
+        self.reattach_states: Set[str] = reattach_states or set()
 
     def prepare_template(self) -> None:
         # If .json is passed then we have to read the file
         if isinstance(self.configuration, str) and self.configuration.endswith('.json'):
             with open(self.configuration, 'r') as file:
                 self.configuration = json.loads(file.read())
 
+    def _submit_job(
+        self,
+        hook: BigQueryHook,
+        job_id: str,
+    ) -> BigQueryJob:
+        # Submit a new job
+        job = hook.insert_job(
+            configuration=self.configuration,
+            project_id=self.project_id,
+            location=self.location,
+            job_id=job_id,
+        )
+        # Start the job and wait for it to complete and get the result.
+        job.result()
+        return job
+
+    @staticmethod
+    def _handle_job_error(job: BigQueryJob) -> None:
+        if job.error_result:
+            raise AirflowException(f"BigQuery job {job.job_id} failed: {job.error_result}")
+
+    def _job_id(self, context):
+        if self.force_rerun:
+            hash_base = str(uuid.uuid4())
+        else:
+            hash_base = json.dumps(self.configuration, sort_keys=True)
+
+        uniqueness_suffix = hashlib.md5(hash_base.encode()).hexdigest()
+
+        if self.job_id:
+            return f"{self.job_id}_{uniqueness_suffix}"
+
+        exec_date = re.sub(r"\:|-|\+", "_", context['execution_date'].isoformat())
+        return f"airflow_{self.dag_id}_{self.task_id}_{exec_date}_{uniqueness_suffix}"
+
     def execute(self, context: Any):
         hook = BigQueryHook(
             gcp_conn_id=self.gcp_conn_id,
             delegate_to=self.delegate_to,
         )
 
-        job_id = self.job_id or f"airflow_{self.task_id}_{int(time())}"
+        job_id = self._job_id(context)
+
         try:
-            job = hook.insert_job(
-                configuration=self.configuration,
-                project_id=self.project_id,
-                location=self.location,
-                job_id=job_id,
-            )
-            # Start the job and wait for it to complete and get the result.
-            job.result()
+            job = self._submit_job(hook, job_id)
+            self._handle_job_error(job)
         except Conflict:
+            # If the job already exists retrieve it
             job = hook.get_job(
                 project_id=self.project_id,
                 location=self.location,
                 job_id=job_id,
             )
-            # Get existing job and wait for it to be ready
-            for time_to_wait in exponential_sleep_generator(initial=10, maximum=120):
-                sleep(time_to_wait)
-                job.reload()
-                if job.done():
-                    break
+            if job.state in self.reattach_states:
+                # We are reattaching to a job
+                job.result()
+                self._handle_job_error(job)
+            else:
+                # Same job configuration so we need force_rerun
+                raise AirflowException(
+                    f"Job with id: {job_id} already exists and is in {job.state} state. If you "
+                    f"want to force rerun it consider setting `force_rerun=True`."
+                    f"Or, if you want to reattach in this scenario add {job.state} to `reattach_states`"
+                )
+
         return job.job_id