From ddf4fb1b4d4f0170555942c5bd5ec9f3a4a0a143 Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Mon, 22 Sep 2025 13:17:06 +0200 Subject: [PATCH 1/5] Prefer operator's project_id over hook's project_id --- .../google/cloud/openlineage/mixins.py | 4 ++- .../google/cloud/openlineage/test_mixins.py | 35 ++++++++++++++++++- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py b/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py index 5eba61b188649..539c24f65332f 100644 --- a/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py +++ b/providers/google/src/airflow/providers/google/cloud/openlineage/mixins.py @@ -97,7 +97,9 @@ def get_openlineage_facets_on_complete(self, _): run_facets: dict[str, RunFacet] = { "externalQuery": ExternalQueryRunFacet(externalQueryId=self.job_id, source="bigquery") } - self._client = self.hook.get_client(project_id=self.hook.project_id, location=self.location) + self._client = self.hook.get_client( + project_id=self.project_id or self.hook.project_id, location=self.location + ) try: job_properties = self._client.get_job(job_id=self.job_id)._properties diff --git a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py index 7006d2f9130c4..08e09c8a10349 100644 --- a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py +++ b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py @@ -20,7 +20,7 @@ import json import logging import os -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, patch, ANY import pytest from google.cloud.bigquery.table import Table @@ -1013,3 +1013,36 @@ def test_generate_column_lineage_facet(self): ), } ) + + def test_project_id_selection(self): + """ + Check if project_id set via an argument to the operator takes prevalence over project_id set in a + connection. + """ + from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator + class TestOperator(GoogleCloudBaseOperator, _BigQueryInsertJobOperatorOpenLineageMixin): + def __init__(self, project_id: str = None, **kwargs): + self.project_id = project_id + self.job_id = "foobar" + self.location = "foobar" + self.sql = "foobar" + + # First test task where project_id is set explicitly + test = TestOperator(task_id="foo", job_id="bar", project_id="project_a") + test.hook = MagicMock() + test.hook.project_id = "project_b" + test._client = MagicMock() + + test.get_openlineage_facets_on_complete(None) + _, kwargs = test.hook.get_client.call_args + assert kwargs["project_id"] == "project_a" + + # Then test task where project_id is inherited from the hook + test = TestOperator(task_id="foo", job_id="bar") + test.hook = MagicMock() + test.hook.project_id = "project_b" + test._client = MagicMock() + + test.get_openlineage_facets_on_complete(None) + _, kwargs = test.hook.get_client.call_args + assert kwargs["project_id"] == "project_b" From 0803b7553bbe5b77e7c0d22189fbd0e0d2c0d83f Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Mon, 22 Sep 2025 13:22:48 +0200 Subject: [PATCH 2/5] Remove unused import --- .../google/tests/unit/google/cloud/openlineage/test_mixins.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py index 08e09c8a10349..e2cec49a5f86d 100644 --- a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py +++ b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py @@ -20,7 +20,7 @@ import json import logging import os -from unittest.mock import MagicMock, patch, ANY +from unittest.mock import MagicMock, patch import pytest from google.cloud.bigquery.table import Table From d79c9c7611e979247c99b5fe81ca77103c949528 Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Mon, 22 Sep 2025 13:25:05 +0200 Subject: [PATCH 3/5] Formatting --- .../google/tests/unit/google/cloud/openlineage/test_mixins.py | 1 + 1 file changed, 1 insertion(+) diff --git a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py index e2cec49a5f86d..02b6ea7ab2eef 100644 --- a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py +++ b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py @@ -1020,6 +1020,7 @@ def test_project_id_selection(self): connection. """ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator + class TestOperator(GoogleCloudBaseOperator, _BigQueryInsertJobOperatorOpenLineageMixin): def __init__(self, project_id: str = None, **kwargs): self.project_id = project_id From da083d615b4a2a042f9dbec6ec212b9192eca4d2 Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Mon, 22 Sep 2025 13:31:01 +0200 Subject: [PATCH 4/5] Remove unnecessary arguments --- .../tests/unit/google/cloud/openlineage/test_mixins.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py index 02b6ea7ab2eef..5703e9fe96bdc 100644 --- a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py +++ b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py @@ -1022,14 +1022,14 @@ def test_project_id_selection(self): from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator class TestOperator(GoogleCloudBaseOperator, _BigQueryInsertJobOperatorOpenLineageMixin): - def __init__(self, project_id: str = None, **kwargs): + def __init__(self, project_id: str = None, **_): self.project_id = project_id self.job_id = "foobar" self.location = "foobar" self.sql = "foobar" # First test task where project_id is set explicitly - test = TestOperator(task_id="foo", job_id="bar", project_id="project_a") + test = TestOperator(project_id="project_a") test.hook = MagicMock() test.hook.project_id = "project_b" test._client = MagicMock() @@ -1039,7 +1039,7 @@ def __init__(self, project_id: str = None, **kwargs): assert kwargs["project_id"] == "project_a" # Then test task where project_id is inherited from the hook - test = TestOperator(task_id="foo", job_id="bar") + test = TestOperator() test.hook = MagicMock() test.hook.project_id = "project_b" test._client = MagicMock() From 2de536c3dc755276684bca3ef8017c021359e673 Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Mon, 22 Sep 2025 13:51:47 +0200 Subject: [PATCH 5/5] Make mypy happy --- .../google/tests/unit/google/cloud/openlineage/test_mixins.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py index 5703e9fe96bdc..e7204368fd500 100644 --- a/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py +++ b/providers/google/tests/unit/google/cloud/openlineage/test_mixins.py @@ -1022,7 +1022,7 @@ def test_project_id_selection(self): from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator class TestOperator(GoogleCloudBaseOperator, _BigQueryInsertJobOperatorOpenLineageMixin): - def __init__(self, project_id: str = None, **_): + def __init__(self, project_id: str | None = None, **_): self.project_id = project_id self.job_id = "foobar" self.location = "foobar"