This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new e335fadee91 Adjust GCP BigQuery triggers for the private cloud (#66404)
e335fadee91 is described below
commit e335fadee912a820164f391f24b8884d56634661
Author: olegkachur-e <[email protected]>
AuthorDate: Wed May 6 04:05:49 2026 +0000
Adjust GCP BigQuery triggers for the private cloud (#66404)
Co-authored-by: Maksim Yermakou <[email protected]>
---
.../providers/google/cloud/triggers/bigquery.py | 111 +++++++++++++++------
.../unit/google/cloud/triggers/test_bigquery.py | 105 ++++++++++++++++++-
2 files changed, 184 insertions(+), 32 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
b/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
index 1eba05425b1..4a6d1a7780c 100644
--- a/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
+++ b/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
@@ -18,7 +18,7 @@ from __future__ import annotations
import asyncio
from collections.abc import AsyncIterator, Sequence
-from typing import TYPE_CHECKING, Any, SupportsAbs
+from typing import TYPE_CHECKING, Any, SupportsAbs, cast
from aiohttp import ClientSession
from aiohttp.client_exceptions import ClientResponseError
@@ -300,6 +300,7 @@ class BigQueryGetDataTrigger(BigQueryInsertJobTrigger):
def __init__(self, as_dict: bool = False, selected_fields: str | None =
None, **kwargs):
super().__init__(**kwargs)
+
self.as_dict = as_dict
self.selected_fields = selected_fields
@@ -324,17 +325,38 @@ class BigQueryGetDataTrigger(BigQueryInsertJobTrigger):
async def run(self) -> AsyncIterator[TriggerEvent]:
"""Get current job execution status and yields a TriggerEvent with
response data."""
hook = self._get_async_hook()
+
try:
while True:
- # Poll for job execution status
- job_status = await hook.get_job_status(job_id=self.job_id,
project_id=self.project_id)
+ job_status = await hook.get_job_status(
+ job_id=self.job_id, project_id=self.project_id,
location=self.location
+ )
+
if job_status["status"] == "success":
- query_results = await
hook.get_job_output(job_id=self.job_id, project_id=self.project_id)
- records = hook.get_records(
- query_results=query_results,
- as_dict=self.as_dict,
- selected_fields=self.selected_fields,
- )
+ sync_hook = await hook.get_sync_hook()
+
+ if sync_hook.is_default_universe():
+ query_results = await hook.get_job_output(
+ job_id=self.job_id, project_id=self.project_id
+ )
+ records = hook.get_records(
+ query_results=query_results,
+ as_dict=self.as_dict,
+ selected_fields=self.selected_fields,
+ )
+ else:
+ if not self.location:
+ raise ValueError(
+ "Location cannot be empty or None. Please
provide a valid location."
+ )
+ query_results_args = {
+ "job_id": self.job_id,
+ "location": self.location,
+ "selected_fields": self.selected_fields,
+ "project_id": self.project_id,
+ }
+ records = await
sync_to_async(sync_hook.get_query_results)(**query_results_args)
+
self.log.debug("Response from hook: %s",
job_status["status"])
yield TriggerEvent(
{
@@ -457,41 +479,70 @@ class
BigQueryIntervalCheckTrigger(BigQueryInsertJobTrigger):
try:
while True:
first_job_response_from_hook = await hook.get_job_status(
- job_id=self.first_job_id, project_id=self.project_id
+ job_id=self.first_job_id,
+ project_id=self.project_id,
+ location=self.location,
)
second_job_response_from_hook = await hook.get_job_status(
- job_id=self.second_job_id, project_id=self.project_id
+ job_id=self.second_job_id,
+ project_id=self.project_id,
+ location=self.location,
)
if (
first_job_response_from_hook["status"] == "success"
and second_job_response_from_hook["status"] == "success"
):
- first_query_results = await hook.get_job_output(
- job_id=self.first_job_id, project_id=self.project_id
- )
+ sync_hook = await hook.get_sync_hook()
- second_query_results = await hook.get_job_output(
- job_id=self.second_job_id, project_id=self.project_id
- )
+ if sync_hook.is_default_universe():
+ first_query_results = await hook.get_job_output(
+ job_id=self.first_job_id,
project_id=self.project_id
+ )
- first_records = hook.get_records(first_query_results)
+ second_query_results = await hook.get_job_output(
+ job_id=self.second_job_id,
project_id=self.project_id
+ )
- second_records = hook.get_records(second_query_results)
+ first_records = hook.get_records(first_query_results)
- # If empty list, then no records are available
- if not first_records:
- first_job_row: str | None = None
- else:
- # Extract only first record from the query results
- first_job_row = first_records.pop(0)
+ second_records = hook.get_records(second_query_results)
- # If empty list, then no records are available
- if not second_records:
- second_job_row: str | None = None
+ # If empty list, then no records are available
+ if not first_records:
+ first_job_row: str | None = None
+ else:
+ # Extract only first record from the query results
+ first_job_row = first_records.pop(0)
+
+ # If empty list, then no records are available
+ if not second_records:
+ second_job_row: str | None = None
+ else:
+ # Extract only first record from the query results
+ second_job_row = second_records.pop(0)
else:
- # Extract only first record from the query results
- second_job_row = second_records.pop(0)
+ if not self.location:
+ raise ValueError(
+ "Location cannot be empty or None. Please
provide a valid location."
+ )
+ query_args_base = {
+ "location": self.location,
+ "max_results": 1,
+ "project_id": self.project_id,
+ }
+ first_job_result = await
sync_to_async(sync_hook.get_query_results)(
+ **(query_args_base | {"job_id": self.first_job_id})
+ )
+ second_job_result = await
sync_to_async(sync_hook.get_query_results)(
+ **(query_args_base | {"job_id":
self.second_job_id})
+ )
+ first_job_row = (
+ cast("Any", list(first_job_result[0].values())) if
first_job_result else None
+ )
+ second_job_row = (
+ cast("Any", list(second_job_result[0].values()))
if second_job_result else None
+ )
hook.interval_check(
first_job_row,
diff --git a/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
b/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
index 7b06ae83196..720a9a0d806 100644
--- a/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
+++ b/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
@@ -368,10 +368,11 @@ class TestBigQueryGetDataTrigger:
assert TriggerEvent({"status": "error", "message": "Test exception"})
== actual
@pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_sync_hook")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_status")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_output")
async def test_bigquery_get_data_trigger_success_with_data(
- self, mock_job_output, mock_job_status, get_data_trigger
+ self, mock_job_output, mock_job_status, mock_get_sync_hook,
get_data_trigger
):
"""
Tests that BigQueryGetDataTrigger only fires once the query execution
reaches a successful state.
@@ -416,6 +417,46 @@ class TestBigQueryGetDataTrigger:
# Prevents error when task is destroyed while in "pending" state
asyncio.get_event_loop().stop()
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.triggers.bigquery.sync_to_async")
+
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_sync_hook")
+
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_status")
+ async def test_bigquery_get_data_trigger_success_with_data_custom_universe(
+ self, mock_job_status, mock_get_sync_hook, mock_sync_to_async,
get_data_trigger
+ ):
+ """
+ Tests that when a custom universe is detected, the trigger uses
sync_to_async
+ to call the sync hook's get_query_results with the correct arguments.
+ """
+ TEST_LOCATION = "custom_private_loc"
+ get_data_trigger.location = TEST_LOCATION
+
+ mock_job_status.return_value = {"status": "success", "message": "Job
completed"}
+
+ mock_sync_hook = mock.MagicMock()
+ mock_sync_hook.is_default_universe.return_value = False # Force the
"else" branch
+ mock_get_sync_hook.return_value = mock_sync_hook
+
+ mock_wrapped_func = mock.AsyncMock()
+ mock_sync_to_async.return_value = mock_wrapped_func
+ mock_wrapped_func.return_value = [[1, "data"]]
+
+ generator = get_data_trigger.run()
+ actual_event = await generator.asend(None)
+
+
mock_sync_to_async.assert_called_once_with(mock_sync_hook.get_query_results)
+
+ expected_args = {
+ "job_id": TEST_JOB_ID,
+ "location": TEST_LOCATION,
+ "selected_fields": TEST_SELECTED_FIELDS,
+ "project_id": TEST_GCP_PROJECT_ID,
+ }
+ mock_wrapped_func.assert_called_once_with(**expected_args)
+
+ assert actual_event.payload["status"] == "success"
+ assert actual_event.payload["records"] == [[1, "data"]]
+
class TestBigQueryCheckTrigger:
@pytest.mark.asyncio
@@ -567,11 +608,12 @@ class TestBigQueryIntervalCheckTrigger:
}
@pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_sync_hook")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_status")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_output")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_records")
async def test_interval_check_trigger_success(
- self, mock_get_records, mock_get_job_output, mock_job_status,
interval_check_trigger
+ self, mock_get_records, mock_get_job_output, mock_job_status,
mock_sync_hook, interval_check_trigger
):
"""
Tests the BigQueryIntervalCheckTrigger only fires once the query
execution reaches a successful state.
@@ -584,6 +626,65 @@ class TestBigQueryIntervalCheckTrigger:
actual = await generator.asend(None)
assert actual == TriggerEvent({"status": "error", "message": "The
second SQL query returned None"})
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.google.cloud.triggers.bigquery.sync_to_async")
+
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_sync_hook")
+
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_status")
+
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.interval_check")
+ async def test_interval_check_trigger_success_non_default_universe(
+ self,
+ mock_interval_check,
+ mock_job_status,
+ mock_get_sync_hook,
+ mock_sync_to_async,
+ interval_check_trigger,
+ ):
+ TEST_LOCATION = "custom_private_loc"
+ interval_check_trigger.location = TEST_LOCATION
+ mock_job_status.return_value = {"status": "success", "message": "Job
completed"}
+
+ mock_sync_hook = mock.MagicMock()
+ mock_sync_hook.is_default_universe.return_value = False
+ mock_get_sync_hook.return_value = mock_sync_hook
+
+ mock_wrapper = mock.AsyncMock()
+ mock_sync_to_async.return_value = mock_wrapper
+
+ mock_row_1 = {"f0_": 100}
+ mock_row_2 = {"f0_": 150}
+ mock_wrapper.side_effect = [[mock_row_1], [mock_row_2]]
+
+ generator = interval_check_trigger.run()
+ actual_event = await generator.asend(None)
+
+ mock_sync_to_async.assert_called_with(mock_sync_hook.get_query_results)
+
+ expected_args_base = {
+ "location": interval_check_trigger.location,
+ "max_results": 1,
+ "project_id": interval_check_trigger.project_id,
+ }
+
+ mock_wrapper.assert_any_call(**(expected_args_base | {"job_id":
interval_check_trigger.first_job_id}))
+ mock_wrapper.assert_any_call(
+ **(expected_args_base | {"job_id":
interval_check_trigger.second_job_id})
+ )
+
+ mock_interval_check.assert_called_once_with(
+ [100],
+ [150],
+ interval_check_trigger.metrics_thresholds,
+ interval_check_trigger.ignore_zero,
+ interval_check_trigger.ratio_formula,
+ )
+
+ assert actual_event.payload == {
+ "status": "success",
+ "message": "Job completed",
+ "first_row_data": [100],
+ "second_row_data": [150],
+ }
+
@pytest.mark.asyncio
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_status")
async def test_interval_check_trigger_pending(self, mock_job_status,
caplog, interval_check_trigger):
