Skip to content

Adjust GCP BigQuery triggers for the private cloud #66404

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

Merged
shahar1 merged 4 commits into from
May 6, 2026
Merged

Adjust GCP BigQuery triggers for the private cloud #66404

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
202b8db
Add universe_domain support for GoogleBaseHook
MaksYermak Apr 23, 2026
5b08696
Adjust bigquery triggers for the private cloud
May 4, 2026
2ff2581
Merge branch 'main' into bq_private_cloud_deff_fix
shahar1 May 5, 2026
93379a9
*resolve mypy and comments
May 5, 2026
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
111 changes: 81 additions & 30 deletions providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@

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
Expand Down Expand Up @@ -299,6 +299,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

Expand All @@ -323,17 +324,38 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
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)
Comment thread
shahar1 marked this conversation as resolved.

self.log.debug("Response from hook: %s", job_status["status"])
yield TriggerEvent(
{
Expand Down Expand Up @@ -456,41 +478,70 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
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,
Expand Down
105 changes: 103 additions & 2 deletions providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -372,10 +372,11 @@ async def test_bigquery_get_data_trigger_exception(self, mock_job_status, caplog
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.
Expand Down Expand Up @@ -420,6 +421,46 @@ async def test_bigquery_get_data_trigger_success_with_data(
# 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
Expand Down Expand Up @@ -571,11 +612,12 @@ def test_interval_check_trigger_serialization(self, interval_check_trigger):
}

@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.
Expand All @@ -588,6 +630,65 @@ async def test_interval_check_trigger_success(
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):
Expand Down
Loading