From 69d581d577c28b8fc542ce665e7918dd2a2900d3 Mon Sep 17 00:00:00 2001 From: Tomek Urbaszek Date: Tue, 30 Jun 2020 15:42:42 +0200 Subject: [PATCH 1/6] Improve idempotency of BigQueryInsertJobOperator --- .../google/cloud/operators/bigquery.py | 50 +++++++++++++------ 1 file changed, 35 insertions(+), 15 deletions(-) diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py index c729302976e20..3005d5ed89f85 100644 --- a/airflow/providers/google/cloud/operators/bigquery.py +++ b/airflow/providers/google/cloud/operators/bigquery.py @@ -1684,32 +1684,52 @@ 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): + # 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 + 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())}" + exec_date = context['execution_date'].isoformat() + job_id = self.job_id or f"airflow_{self.dag_id}_{self.task_id}_{exec_date}" + 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() + # Submit a new job + job = self._submit_job(hook, job_id) 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.done() and job.error_result: + # The job exists and finished with an error and we are probably reruning it + # So we have to make a new job_id because it has to be unique + job_id = f"{self.job_id}_{int(time())}" + job = self._submit_job(hook, job_id) + elif not job.done(): + # The job is still running so 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.error_result: + raise AirflowException(f"BigQuery job {job_id} failed: {job.error_result}") + return job.job_id From 256d093afb6fd522d425532e29f6a7d96ddc66ca Mon Sep 17 00:00:00 2001 From: Tomek Urbaszek Date: Sat, 4 Jul 2020 11:52:40 +0200 Subject: [PATCH 2/6] fixup! Improve idempotency of BigQueryInsertJobOperator --- airflow/providers/google/cloud/operators/bigquery.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py index 3005d5ed89f85..79f0bf4b636c4 100644 --- a/airflow/providers/google/cloud/operators/bigquery.py +++ b/airflow/providers/google/cloud/operators/bigquery.py @@ -22,6 +22,7 @@ """ import enum import json +import re import warnings from time import sleep, time from typing import Any, Dict, Iterable, List, Optional, SupportsAbs, Union @@ -1702,8 +1703,8 @@ def execute(self, context: Any): delegate_to=self.delegate_to, ) - exec_date = context['execution_date'].isoformat() - job_id = self.job_id or f"airflow_{self.dag_id}_{self.task_id}_{exec_date}" + exec_date = re.sub("\:|\-|\+", "_", context['execution_date'].isoformat()) + job_id = self.job_id or f"airflow_{self.dag_id}_{self.task_id}_{exec_date}_" try: # Submit a new job @@ -1723,11 +1724,7 @@ def execute(self, context: Any): job = self._submit_job(hook, job_id) elif not job.done(): # The job is still running so 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 + job.result() if job.error_result: raise AirflowException(f"BigQuery job {job_id} failed: {job.error_result}") From f704528169ab9ce9a75e251f3f48a3316326cc81 Mon Sep 17 00:00:00 2001 From: Tomek Urbaszek Date: Mon, 3 Aug 2020 14:55:02 +0200 Subject: [PATCH 3/6] fixup! fixup! Improve idempotency of BigQueryInsertJobOperator --- .../providers/google/cloud/hooks/bigquery.py | 4 +- .../google/cloud/operators/bigquery.py | 87 +++++++--- .../google/cloud/operators/test_bigquery.py | 149 ++++++++++++++---- 3 files changed, 187 insertions(+), 53 deletions(-) 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 79f0bf4b636c4..cf0bf011eb969 100644 --- a/airflow/providers/google/cloud/operators/bigquery.py +++ b/airflow/providers/google/cloud/operators/bigquery.py @@ -21,22 +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 @@ -1632,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 @@ -1645,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 """ @@ -1667,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, @@ -1678,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 @@ -1685,7 +1706,11 @@ 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): + def _submit_job( + self, + hook: BigQueryHook, + job_id: str, + ) -> BigQueryJob: # Submit a new job job = hook.insert_job( configuration=self.configuration, @@ -1697,18 +1722,36 @@ def _submit_job(self, hook: BigQueryHook, job_id: str): 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, ) - exec_date = re.sub("\:|\-|\+", "_", context['execution_date'].isoformat()) - job_id = self.job_id or f"airflow_{self.dag_id}_{self.task_id}_{exec_date}_" + job_id = self._job_id(context) try: - # Submit a new job 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( @@ -1716,17 +1759,15 @@ def execute(self, context: Any): location=self.location, job_id=job_id, ) - - if job.done() and job.error_result: - # The job exists and finished with an error and we are probably reruning it - # So we have to make a new job_id because it has to be unique - job_id = f"{self.job_id}_{int(time())}" - job = self._submit_job(hook, job_id) - elif not job.done(): + if job.state in self.reattach_states and not job.done(): # The job is still running so wait for it to be ready job.result() - - if job.error_result: - raise AirflowException(f"BigQuery job {job_id} failed: {job.error_result}") + self._handle_job_error(job) + elif job.done(): + # 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`." + ) 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({}) From ce440b93d8737bcfd6814b5d9a33f0def5eae894 Mon Sep 17 00:00:00 2001 From: Tomek Urbaszek Date: Thu, 6 Aug 2020 11:07:26 +0200 Subject: [PATCH 4/6] fixup! fixup! fixup! Improve idempotency of BigQueryInsertJobOperator --- airflow/providers/google/cloud/operators/bigquery.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py index cf0bf011eb969..af0e38e68f58b 100644 --- a/airflow/providers/google/cloud/operators/bigquery.py +++ b/airflow/providers/google/cloud/operators/bigquery.py @@ -1759,11 +1759,11 @@ def execute(self, context: Any): location=self.location, job_id=job_id, ) - if job.state in self.reattach_states and not job.done(): - # The job is still running so wait for it to be ready + if job.state in self.reattach_states: + # We are reattaching to a job job.result() self._handle_job_error(job) - elif job.done(): + 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 " From 147aaadddbd2ab8c873060ea9254006b6784dc88 Mon Sep 17 00:00:00 2001 From: Tomek Urbaszek Date: Thu, 6 Aug 2020 13:24:47 +0200 Subject: [PATCH 5/6] Update airflow/providers/google/cloud/operators/bigquery.py Co-authored-by: Jacob Ferriero --- airflow/providers/google/cloud/operators/bigquery.py | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py index af0e38e68f58b..c6d3b2db0710e 100644 --- a/airflow/providers/google/cloud/operators/bigquery.py +++ b/airflow/providers/google/cloud/operators/bigquery.py @@ -1768,6 +1768,7 @@ def execute(self, context: Any): 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 From 7dd582dd7cc824221d60ffaad2887827fb84e003 Mon Sep 17 00:00:00 2001 From: Tomek Urbaszek Date: Mon, 10 Aug 2020 11:51:29 +0200 Subject: [PATCH 6/6] Apply suggestions from code review --- airflow/providers/google/cloud/operators/bigquery.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py index c6d3b2db0710e..58c23cab4bd9d 100644 --- a/airflow/providers/google/cloud/operators/bigquery.py +++ b/airflow/providers/google/cloud/operators/bigquery.py @@ -1684,7 +1684,7 @@ def __init__( project_id: Optional[str] = None, location: Optional[str] = None, job_id: Optional[str] = None, - force_rerun: bool = False, + force_rerun: bool = True, reattach_states: Optional[Set[str]] = None, gcp_conn_id: str = 'google_cloud_default', delegate_to: Optional[str] = None,