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..58c23cab4bd9d 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 = True, + 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({})