Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix BigQueryInsertJobOperatorAsync failure after the google provider upgrade #471

Merged
merged 7 commits into from
Jun 27, 2022
Merged

Fix BigQueryInsertJobOperatorAsync failure after the google provider upgrade #471

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
a0363fb
Fix BigQueryInsertJobOperatorAsync bug
rajaths010494 Jun 24, 2022
24ba015
Remove unused line
rajaths010494 Jun 24, 2022
e6c81ad
fix mypy
rajaths010494 Jun 24, 2022
7071dff
fix mypy
rajaths010494 Jun 24, 2022
190e52c
Fix test case
rajaths010494 Jun 24, 2022
3992d3a
Update version for google provider
rajaths010494 Jun 24, 2022
9756e9e
Add test case for generate_job_id
rajaths010494 Jun 24, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion astronomer/providers/google/cloud/operators/bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,14 @@ def execute(self, context: Context) -> None: # noqa: D102
hook = _BigQueryHook(gcp_conn_id=self.gcp_conn_id)

self.hook = hook
job_id = self._job_id(context)
job_id = self.hook.generate_job_id(
job_id=self.job_id,
dag_id=self.dag_id,
task_id=self.task_id,
logical_date=context["logical_date"],
configuration=self.configuration,
force_rerun=self.force_rerun,
)
Comment on lines +89 to +96
Copy link
Collaborator

@kaxil kaxil Jun 27, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you wanted to make this backwards and forwards-compatible without bumping the dependency, you could do the following:

try:
    # for google provider>=8.1.0
    job_id = self.hook.generate_job_id(
        job_id=self.job_id,
        dag_id=self.dag_id,
        task_id=self.task_id,
        logical_date=context["logical_date"],
        configuration=self.configuration,
        force_rerun=self.force_rerun,
    )
except AttributeError:
    # for google provider<8.1.0
    job_id = self._job_id(context)


try:
job = self._submit_job(hook, job_id)
Expand Down
4 changes: 2 additions & 2 deletions setup.cfg
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ cncf.kubernetes =
databricks =
apache-airflow-providers-databricks>=2.2.0
google =
apache-airflow-providers-google>=8.0.0
apache-airflow-providers-google>=8.1.0
gcloud-aio-storage
gcloud-aio-bigquery
snowflake =
Expand Down Expand Up @@ -117,7 +117,7 @@ all =
apache-airflow-providers-amazon>=3.0.0
apache-airflow-providers-cncf-kubernetes>=4
apache-airflow-providers-databricks>=2.2.0
apache-airflow-providers-google>=8.0.0
apache-airflow-providers-google>=8.1.0
apache-airflow-providers-apache-livy
apache-airflow-providers-apache-hive
apache-airflow-providers-snowflake
Expand Down
7 changes: 4 additions & 3 deletions tests/google/cloud/extractors/test_bigquery_async_extractor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,11 +75,11 @@ def context():

def create_context(task):
dag = DAG(dag_id="dag")
execution_date = datetime(2022, 1, 1, 0, 0, 0)
logical_date = datetime(2022, 1, 1, 0, 0, 0)
dag_run = DagRun(
dag_id=dag.dag_id,
execution_date=execution_date,
run_id=DagRun.generate_run_id(DagRunType.MANUAL, execution_date),
execution_date=logical_date,
run_id=DagRun.generate_run_id(DagRunType.MANUAL, logical_date),
)
task_instance = TaskInstance(task=task)
task_instance.dag_run = dag_run
Expand All @@ -91,6 +91,7 @@ def create_context(task):
"task": task,
"ti": task_instance,
"task_instance": task_instance,
"logical_date": logical_date,
}


Expand Down
69 changes: 44 additions & 25 deletions tests/google/cloud/operators/test_bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,11 +98,11 @@ def test_bigquery_insert_job_operator_execute_failure(context):

def create_context(task):
dag = DAG(dag_id="dag")
execution_date = datetime(2022, 1, 1, 0, 0, 0)
logical_date = datetime(2022, 1, 1, 0, 0, 0)
dag_run = DagRun(
dag_id=dag.dag_id,
execution_date=execution_date,
run_id=DagRun.generate_run_id(DagRunType.MANUAL, execution_date),
execution_date=logical_date,
run_id=DagRun.generate_run_id(DagRunType.MANUAL, logical_date),
)
task_instance = TaskInstance(task=task)
task_instance.dag_run = dag_run
Expand All @@ -114,6 +114,7 @@ def create_context(task):
"task": task,
"ti": task_instance,
"task_instance": task_instance,
"logical_date": logical_date,
}


Expand Down Expand Up @@ -142,37 +143,56 @@ def test_bigquery_insert_job_operator_execute_complete():
mock_log_info.assert_called_with("%s completed with response %s ", "insert_query_job", "Job completed")


@mock.patch("airflow.providers.google.cloud.operators.bigquery.hashlib.md5")
@pytest.mark.parametrize(
"test_dag_id, expected_job_id",
[("test-dag-id-1.1", "airflow_test_dag_id_1_1_test_job_id_2020_01_23T00_00_00_00_00_hash")],
ids=["test-dag-id-1.1"],
)
def test_job_id_validity(mock_md5, test_dag_id, expected_job_id):
rajaths010494 marked this conversation as resolved.
Show resolved Hide resolved
"""Asserts that job id is correctly generated"""
@mock.patch("astronomer.providers.google.cloud.operators.bigquery._BigQueryHook")
def test_bigquery_insert_job_operator_with_job_id_generate(mock_hook):
job_id = "123456"
hash_ = "hash"
mock_md5.return_value.hexdigest.return_value = hash_
context = {"logical_date": datetime(2020, 1, 23)}
real_job_id = f"{job_id}_{hash_}"

configuration = {
"query": {
"query": "SELECT * FROM any",
"useLegacySql": False,
}
}
with DAG(dag_id=test_dag_id, start_date=datetime(2020, 1, 23)):
op = BigQueryInsertJobOperatorAsync(
task_id="test_job_id", configuration=configuration, project_id=TEST_GCP_PROJECT_ID
)
assert op._job_id(context) == expected_job_id

mock_hook.return_value.insert_job.side_effect = Conflict("any")
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 = BigQueryInsertJobOperatorAsync(
task_id="insert_query_job",
configuration=configuration,
location=TEST_DATASET_LOCATION,
job_id=job_id,
project_id=TEST_GCP_PROJECT_ID,
reattach_states={"PENDING"},
)

with pytest.raises(TaskDeferred):
op.execute(create_context(op))

mock_hook.return_value.generate_job_id.assert_called_once_with(
job_id=job_id,
dag_id="adhoc_airflow",
task_id="insert_query_job",
logical_date=datetime(2022, 1, 1, 0, 0),
configuration=configuration,
force_rerun=True,
)


@mock.patch("airflow.providers.google.cloud.operators.bigquery.hashlib.md5")
@mock.patch("astronomer.providers.google.cloud.operators.bigquery._BigQueryHook")
def test_execute_reattach(mock_hook, mock_md5):
def test_execute_reattach(mock_hook):
job_id = "123456"
hash_ = "hash"
real_job_id = f"{job_id}_{hash_}"
mock_md5.return_value.hexdigest.return_value = hash_
mock_hook.return_value.generate_job_id.return_value = f"{job_id}_{hash_}"

configuration = {
"query": {
Expand Down Expand Up @@ -211,13 +231,12 @@ def test_execute_reattach(mock_hook, mock_md5):
job._begin.assert_called_once_with()


@mock.patch("airflow.providers.google.cloud.operators.bigquery.hashlib.md5")
@mock.patch("astronomer.providers.google.cloud.operators.bigquery._BigQueryHook")
def test_execute_force_rerun(mock_hook, mock_md5):
def test_execute_force_rerun(mock_hook):
job_id = "123456"
hash_ = "hash"
real_job_id = f"{job_id}_{hash_}"
mock_md5.return_value.hexdigest.return_value = hash_
mock_hook.return_value.generate_job_id.return_value = f"{job_id}_{hash_}"

configuration = {
"query": {
Expand Down Expand Up @@ -245,7 +264,7 @@ def test_execute_force_rerun(mock_hook, mock_md5):
)

with pytest.raises(AirflowException) as exc:
op.execute(context)
op.execute(create_context(op))

expected_exception_msg = (
f"Job with id: {real_job_id} already exists and is in {job.state} state. "
Expand Down