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 70419e5b08d3f..711531d9e572a 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 @@ 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 @@ -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 @@ -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) + self.log.debug("Response from hook: %s", job_status["status"]) yield TriggerEvent( { @@ -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, 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 22f4a17cf54bd..9221c91d6cb8a 100644 --- a/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py +++ b/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py @@ -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. @@ -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 @@ -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. @@ -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):