Repository: apache/airflow
Pull request: Improve idempotency of BigQueryInsertJobOperator
Closes #8903 

This operator work in the following way:
- it calculates unique hash of the job using job's configuration or uuid if ``force_rerun`` is True
- creates ``job_id`` in form of 
``[provide_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 wihtout attaching to already existing ones.

---
Make sure to mark the boxes below before creating PR: [x]

- [x] Description above provides context of the change
- [x] Unit tests coverage for changes (not needed for documentation changes)
- [x] Target Github ISSUE in description if exists
- [x] Commits follow "[How to write a good git commit message](http://chris.beams.io/posts/git-commit/)"
- [ ] Relevant documentation is updated including usage instructions.
- [x] I will engage committers as explained in [Contribution Workflow Example](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#contribution-workflow-example).

---
In case of fundamental code change, Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvements+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in [UPDATING.md](https://github.com/apache/airflow/blob/master/UPDATING.md).
Read the [Pull Request Guidelines](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#pull-request-guidelines) for more information.


Problem statement:
BigQueryHook refactor + deterministic BQ Job ID
<!--

Welcome to Apache Airflow!  For a smooth issue process, try to answer the following questions.
Don't worry if they're not all applicable; just try to include what you can :-)

If you need to include code snippets or logs, please put them in fenced code
blocks.  If they're super-long, please use the details tag like
<details><summary>super-long log</summary> lots of stuff </details>

Please delete these comment blocks before submitting the issue.

-->

**Description**

Looking at the code it seems like  a lot of the logic in the BQ Hook is already implemented in the google API python library. This includes job polling, a nicer way to use job config and of course las the validations that we now do manually.
It would be ice to make use of these and simplify the code.

My idea is then to refactor the run_<job> methods to take the google job config and a deterministic job id.
This would help in case of a pod dies because of any given reason, we’d restart polling for the async job previously started (I apologize for the crappy explanation).

See my hacky spike below:

This is the job id definition for reference
`job_id = re.sub(r"[^0-9a-zA-Z_\-]+", "-", f"{self.dag_id}_{self.task_id}_{context['execution_date'].isoformat()}__try_0")`


and here roughly how un query would work
```
def run_query(self, job_id: str, job_config: QueryJobConfig, sql: str, destination_dataset_table: str = None) -> str:
        def _recurse(job_id: str):
            [j, try_num] = job_id.split("__try_")
            new_job_id = f"{j}__try_{int(try_num) + 1}"
            return run_query(new_job_id, job_config, sql)

        def run_query(job_id: str, job_config: QueryJobConfig, sql: str):
            if not self.project_id:
                raise ValueError("The project_id should be set")

            if destination_dataset_table is not None:
                job_config.destination = TableReference.from_string(destination_dataset_table, self.project_id)

            try:
                job: QueryJob = self.client.get_job(job_id, self.project_id)
                if job.state == 'RUNNING':
                    if job.query != sql:
                        job.cancel()
                        self.log.info(f"Job {job_id} found, but sql is different. "
                                      f"Cancelling the current job and starting a new one")
                        return _recurse(job_id)
                    self.log.info(f"Job {job_id} still running, re-starting to poll.")
                    return job.result()
                else:
                    self.log.info(f"Job {job_id} already executed once. Restarting")
                    return _recurse(job_id)
            except NotFound:
                self.log.info(f"Job {job_id} not found, starting a new job.")
            job: QueryJob = self.client.query(sql, job_config, job_id, project=self.project_id)
            self.log.info(f"Running Job {job_id}...")
            return job.result()

        return run_query(job_id, job_config, sql)
```

**the encoded __try_<try_num> is not the airflow but a secondary try in case the task is cleared since BQ Job Ids are a unique key and can't be re-used.**

**Use case / motivation**
Trying to use the functionalities in the google cloud library rather than re-implementing them ourselves.
This would allow us to pass through a Deterministic Job ID too, useful for picking up jobs which are still running in case a pod dies.
**Related Issues**

<!-- Is there currently another issue associated with this? -->

Hints from issue discussion:
Thanks for opening your first issue here! Be sure to follow the issue template!

Hi @albertocalderari, I'm curently workining on refactor of BQ integration. I decided to abandon custom "run" and use `insert_job` method which will accept `job_id`:

https://github.com/PolideaInternal/airflow/blob/71328b19be73ff0e4820be1ee25867eb6a578820/airflow/providers/google/cloud/hooks/bigquery.py#L1432-L1441
My idea was yo use the methods exposed from google.cloud.bigquery.Client along with  google.cloud.bigquery.{JobType}JobConfig rather than using dictionaries.
Doing this we can ditch most of the custom validation we have in the hook and make the code a bit slimmer.

I'd build the job configs in the operator itself and mark run_with_configuration as deprecated.

I like the Idea of having only 1 method to execute any kind of job though, we could add the recursive logic in there.

What do you think @turbaszek?
> My idea was yo use the methods exposed from google.cloud.bigquery.Client along with google.cloud.bigquery.{JobType}JobConfig rather than using dictionaries.

I see your point, however this will make the migration harder. On the bright side, we are ditching the custom validation. If you will take a look at `inser_job` method you will see that I'm constructing a `XJobConfig`:
https://github.com/apache/airflow/blob/375d1ca229464617780623c61c6e8a1bf570c87f/airflow/providers/google/cloud/hooks/bigquery.py#L1418-L1434

@turbaszek Hey, sorry work was hell this past week and did not find the time or the energy to reply :(
I have noticed that inner job does that, and that we'd be building the config from the dictionary that is passed in. Though going full on with using {JobType}JobConfig would make method signatures a lot more friendly to the eye. I do agree also that the migration would be harder - that could be done in 2 steps:

- 1 create a new method used only by the operators and mark the old ones as deprecated
- 2 re move the deprecated methods

If I can't sell this to you with this, I'll go your way and see how it comes :)
Summoning @edejong to hear his opinion :)
I should really check GitHub more often, I only saw the notification now.

Let me know if I understand the question correctly: should the BigQueryHook's interface rely on classes from the Google API client library, or should all data be passed in as dictionaries?

I think it's one thing to have the Airflow hooks/operators coupled to the BigQuery REST interface which I guess is what you get passing in the config in a Python dict. This allows you to translate online examples very easily to a DAG.

But it's a much bigger step to rely on the Google client library in the API because that introduces a tight coupling to this specific library. It would only look good if the Airflow code can stay 100% agnostic about what is passed to the library. Can we guarantee that, even for the future? And does that align with other GCP products?

So my personal opinion is stick with the dict :)

As for generating a job id for all job types, I agree that would be a very good move. Without it you would have to wait for a response to even have something to check up on after. That works most of the time, but in cases where it goes wrong it makes it harder to troubleshoot.

I love the suggested job ids string. One small thing I would change is to add a prefix such as `airflow_` or even just `af_` to make it even easier to spot these in Stackdriver for example. I would generate some well defined job id string every time one wasn't provided by the user.

See https://cloud.google.com/bigquery/docs/running-jobs#generate-jobid for recommendations.
> So my personal opinion is stick with the dict :)

Same is mine. And what is more... dict is JSON serializable so it can be used as template field! 

@albertocalderari should we consider this issue as resolved? 
@turbaszek That's only half of the issue, the other half is implementing the deterministic job id.
So as far as I understand, Operators will interact with the hook using json config and then we'll build the job from it.
As soon as U finish you refactor I'll start implementing the feature :)

Commit message: Update airflow/providers/google/cloud/operators/bigquery.py

Co-authored-by: Jacob Ferriero <jferriero@google.com>

Patch to review:
diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py
index 1ec27fe3114a5..6a9650c7b4603 100644
--- a/airflow/providers/google/cloud/hooks/bigquery.py
+++ b/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:
diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py
index c729302976e20..c6d3b2db0710e 100644
--- a/airflow/providers/google/cloud/operators/bigquery.py
+++ b/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 = False,
+        reattach_states: Optional[Set[str]] = None,
         gcp_conn_id: str = 'google_cloud_default',
         delegate_to: Optional[str] = None,
         **kwargs,
@@ -1677,6 +1697,8 @@ 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
@@ -1684,32 +1706,69 @@ def prepare_template(self) -> None:
             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
diff --git a/tests/providers/google/cloud/operators/test_bigquery.py b/tests/providers/google/cloud/operators/test_bigquery.py
index 54d9919363cc7..246c8ec282a6a 100644
--- a/tests/providers/google/cloud/operators/test_bigquery.py
+++ b/tests/providers/google/cloud/operators/test_bigquery.py
@@ -21,6 +21,7 @@
 from unittest.mock import MagicMock
 
 import mock
+import pytest
 from google.cloud.exceptions import Conflict
 from parameterized import parameterized
 
@@ -792,16 +793,21 @@ def test_execute(self, mock_hook):
 
 
 class TestBigQueryInsertJobOperator:
+    @mock.patch('airflow.providers.google.cloud.operators.bigquery.hashlib.md5')
     @mock.patch('airflow.providers.google.cloud.operators.bigquery.BigQueryHook')
-    def test_execute(self, mock_hook):
+    def test_execute_success(self, mock_hook, mock_md5):
         job_id = "123456"
+        hash_ = "hash"
+        real_job_id = f"{job_id}_{hash_}"
+        mock_md5.return_value.hexdigest.return_value = hash_
+
         configuration = {
             "query": {
                 "query": "SELECT * FROM any",
                 "useLegacySql": False,
             }
         }
-        mock_hook.return_value.insert_job.return_value = MagicMock(job_id=job_id)
+        mock_hook.return_value.insert_job.return_value = MagicMock(job_id=real_job_id, error_result=False)
 
         op = BigQueryInsertJobOperator(
             task_id="insert_query_job",
@@ -815,16 +821,19 @@ def test_execute(self, mock_hook):
         mock_hook.return_value.insert_job.assert_called_once_with(
             configuration=configuration,
             location=TEST_DATASET_LOCATION,
-            job_id=job_id,
+            job_id=real_job_id,
             project_id=TEST_GCP_PROJECT_ID,
         )
 
-        assert result == job_id
+        assert result == real_job_id
 
-    @mock.patch('airflow.providers.google.cloud.operators.bigquery.exponential_sleep_generator')
+    @mock.patch('airflow.providers.google.cloud.operators.bigquery.hashlib.md5')
     @mock.patch('airflow.providers.google.cloud.operators.bigquery.BigQueryHook')
-    def test_execute_idempotency(self, mock_hook, mock_sleep_generator):
+    def test_execute_failure(self, mock_hook, mock_md5):
         job_id = "123456"
+        hash_ = "hash"
+        real_job_id = f"{job_id}_{hash_}"
+        mock_md5.return_value.hexdigest.return_value = hash_
 
         configuration = {
             "query": {
@@ -832,46 +841,128 @@ def test_execute_idempotency(self, mock_hook, mock_sleep_generator):
                 "useLegacySql": False,
             }
         }
+        mock_hook.return_value.insert_job.return_value = MagicMock(job_id=real_job_id, error_result=True)
 
-        class MockJob:
-            _call_no = 0
-            _done = False
-
-            def __init__(self):
-                pass
-
-            def reload(self):
-                if MockJob._call_no == 3:
-                    MockJob._done = True
-                else:
-                    MockJob._call_no += 1
+        op = BigQueryInsertJobOperator(
+            task_id="insert_query_job",
+            configuration=configuration,
+            location=TEST_DATASET_LOCATION,
+            job_id=job_id,
+            project_id=TEST_GCP_PROJECT_ID
+        )
+        with pytest.raises(AirflowException):
+            op.execute({})
 
-            def done(self):
-                return MockJob._done
+    @mock.patch('airflow.providers.google.cloud.operators.bigquery.hashlib.md5')
+    @mock.patch('airflow.providers.google.cloud.operators.bigquery.BigQueryHook')
+    def test_execute_reattach(self, mock_hook, mock_md5):
+        job_id = "123456"
+        hash_ = "hash"
+        real_job_id = f"{job_id}_{hash_}"
+        mock_md5.return_value.hexdigest.return_value = hash_
 
-            @property
-            def job_id(self):
-                return job_id
+        configuration = {
+            "query": {
+                "query": "SELECT * FROM any",
+                "useLegacySql": False,
+            }
+        }
 
         mock_hook.return_value.insert_job.return_value.result.side_effect = Conflict("any")
-        mock_sleep_generator.return_value = [0, 0, 0, 0, 0]
-        mock_hook.return_value.get_job.return_value = MockJob()
+        job = MagicMock(
+            job_id=real_job_id, error_result=False, state="PENDING", done=lambda: False,
+        )
+        mock_hook.return_value.get_job.return_value = job
 
         op = BigQueryInsertJobOperator(
             task_id="insert_query_job",
             configuration=configuration,
             location=TEST_DATASET_LOCATION,
             job_id=job_id,
-            project_id=TEST_GCP_PROJECT_ID
+            project_id=TEST_GCP_PROJECT_ID,
+            reattach_states={"PENDING"}
         )
         result = op.execute({})
 
-        assert MockJob._call_no == 3
-
         mock_hook.return_value.get_job.assert_called_once_with(
+            location=TEST_DATASET_LOCATION,
+            job_id=real_job_id,
+            project_id=TEST_GCP_PROJECT_ID,
+        )
+
+        job.result.assert_called_once_with()
+
+        assert result == real_job_id
+
+    @mock.patch('airflow.providers.google.cloud.operators.bigquery.hashlib.md5')
+    @mock.patch('airflow.providers.google.cloud.operators.bigquery.uuid')
+    @mock.patch('airflow.providers.google.cloud.operators.bigquery.BigQueryHook')
+    def test_execute_force_rerun(self, mock_hook, mock_uuid, mock_md5):
+        job_id = "123456"
+        hash_ = mock_uuid.uuid4.return_value.encode.return_value
+        real_job_id = f"{job_id}_{hash_}"
+        mock_md5.return_value.hexdigest.return_value = hash_
+
+        configuration = {
+            "query": {
+                "query": "SELECT * FROM any",
+                "useLegacySql": False,
+            }
+        }
+
+        job = MagicMock(
+            job_id=real_job_id, error_result=False,
+        )
+        mock_hook.return_value.insert_job.return_value = job
+
+        op = BigQueryInsertJobOperator(
+            task_id="insert_query_job",
+            configuration=configuration,
             location=TEST_DATASET_LOCATION,
             job_id=job_id,
             project_id=TEST_GCP_PROJECT_ID,
+            force_rerun=True,
+        )
+        result = op.execute({})
+
+        mock_hook.return_value.insert_job.assert_called_once_with(
+            configuration=configuration,
+            location=TEST_DATASET_LOCATION,
+            job_id=real_job_id,
+            project_id=TEST_GCP_PROJECT_ID,
+        )
+
+        assert result == real_job_id
+
+    @mock.patch('airflow.providers.google.cloud.operators.bigquery.hashlib.md5')
+    @mock.patch('airflow.providers.google.cloud.operators.bigquery.BigQueryHook')
+    def test_execute_no_force_rerun(self, mock_hook, mock_md5):
+        job_id = "123456"
+        hash_ = "hash"
+        real_job_id = f"{job_id}_{hash_}"
+        mock_md5.return_value.hexdigest.return_value = hash_
+
+        configuration = {
+            "query": {
+                "query": "SELECT * FROM any",
+                "useLegacySql": False,
+            }
+        }
+
+        mock_hook.return_value.insert_job.return_value.result.side_effect = Conflict("any")
+        job = MagicMock(
+            job_id=real_job_id, error_result=False, state="DONE", done=lambda: True,
         )
+        mock_hook.return_value.get_job.return_value = job
 
-        assert result == job_id
+        op = BigQueryInsertJobOperator(
+            task_id="insert_query_job",
+            configuration=configuration,
+            location=TEST_DATASET_LOCATION,
+            job_id=job_id,
+            project_id=TEST_GCP_PROJECT_ID,
+            reattach_states={"PENDING"}
+        )
+        # No force rerun
+        with pytest.raises(AirflowException):
+            op.execute({})


Code review comments:
diff_hunk: @@ -1666,6 +1684,8 @@ def __init__(
         project_id: Optional[str] = None,
         location: Optional[str] = None,
         job_id: Optional[str] = None,
+        force_rerun: bool = False,
path: airflow/providers/google/cloud/operators/bigquery.py
text: @author:
```suggestion
        force_rerun: bool = True,
```