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 2cdc403cd19 Convert BigQueryAsyncHook.get_records to async method
(#72286)
2cdc403cd19 is described below
commit 2cdc403cd19ebc19dc9c808f22cd34b76b70f48f
Author: sundeep8967 <[email protected]>
AuthorDate: Sat Oct 31 13:17:26 2026 +0530
Convert BigQueryAsyncHook.get_records to async method (#72286)
Co-authored-by: Shahar Epstein <[email protected]>
---
.../providers/google/cloud/hooks/bigquery.py | 46 ++++++++++-
.../providers/google/cloud/triggers/bigquery.py | 12 +--
.../tests/unit/google/cloud/hooks/test_bigquery.py | 89 +++++++++++++++++++++-
.../unit/google/cloud/triggers/test_bigquery.py | 24 +++---
4 files changed, 150 insertions(+), 21 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/hooks/bigquery.py
b/providers/google/src/airflow/providers/google/cloud/hooks/bigquery.py
index b5f63f5812d..1b3471db3ae 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/bigquery.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/bigquery.py
@@ -20,6 +20,7 @@
from __future__ import annotations
+import asyncio
import json
import logging
import re
@@ -70,6 +71,7 @@ from airflow.providers.google.cloud.utils.bigquery import
bq_cast
from airflow.providers.google.cloud.utils.credentials_provider import
_get_scopes
from airflow.providers.google.cloud.utils.lineage import
send_hook_lineage_for_bq_job
from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.deprecated import deprecated
from airflow.providers.google.common.hooks.base_google import (
_UNSET,
PROVIDE_PROJECT_ID,
@@ -2448,7 +2450,11 @@ class BigQueryAsyncHook(GoogleBaseAsyncHook):
self.log.error("Failed to cancel BigQuery job %s: %s", job_id,
str(e))
raise
- # TODO: Convert get_records into an async method
+ @deprecated(
+ planned_removal_date="April 30, 2027",
+
use_instead="airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.aget_records",
+ category=AirflowProviderDeprecationWarning,
+ )
def get_records(
self,
query_results: dict[str, Any],
@@ -2479,6 +2485,44 @@ class BigQueryAsyncHook(GoogleBaseAsyncHook):
buffer.append(typed_row)
return buffer
+ async def aget_records(
+ self,
+ query_results: dict[str, Any],
+ as_dict: bool = False,
+ selected_fields: str | list[str] | None = None,
+ yield_frequency: int = 1000,
+ ) -> list[Any]:
+ """
+ Convert a response from BigQuery to records asynchronously.
+
+ Periodically yields control back to the event loop when processing
large
+ result sets to prevent blocking the Triggerer event loop.
+
+ :param query_results: the results from a SQL query
+ :param as_dict: if True returns the result as a list of dictionaries,
otherwise as list of lists.
+ :param selected_fields:
+ :param yield_frequency: number of rows processed before yielding to
the asyncio event loop.
+ Set to 0 to disable yielding.
+ """
+ if isinstance(selected_fields, str):
+ selected_fields = selected_fields.split(",")
+ buffer: list[Any] = []
+ if rows := query_results.get("rows"):
+ fields = query_results["schema"]["fields"]
+ fields = [field for field in fields if not selected_fields or
field["name"] in selected_fields]
+ fields_names = [field["name"] for field in fields]
+ col_types = [field["type"] for field in fields]
+ for i, dict_row in enumerate(rows):
+ if yield_frequency > 0 and i > 0 and i % yield_frequency == 0:
+ await asyncio.sleep(0)
+ typed_row = [bq_cast(vs["v"], col_type) for vs, col_type in
zip(dict_row["f"], col_types)]
+ if as_dict:
+ typed_row_dict = dict(zip(fields_names, typed_row))
+ buffer.append(typed_row_dict)
+ else:
+ buffer.append(typed_row)
+ return buffer
+
def value_check(
self,
sql: str,
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 25b5fc529f2..74fc923eafa 100644
--- a/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
+++ b/providers/google/src/airflow/providers/google/cloud/triggers/bigquery.py
@@ -270,7 +270,7 @@ class BigQueryCheckTrigger(BigQueryInsertJobTrigger):
job_status = await hook.get_job_status(job_id=self.job_id,
project_id=self.project_id)
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)
+ records = await hook.aget_records(query_results)
# If empty list, then no records are available
if not records:
@@ -356,7 +356,7 @@ class BigQueryGetDataTrigger(BigQueryInsertJobTrigger):
query_results = await hook.get_job_output(
job_id=self.job_id, project_id=self.project_id
)
- records = hook.get_records(
+ records = await hook.aget_records(
query_results=query_results,
as_dict=self.as_dict,
selected_fields=self.selected_fields,
@@ -524,9 +524,9 @@ class
BigQueryIntervalCheckTrigger(BigQueryInsertJobTrigger):
job_id=self.second_job_id,
project_id=self.project_id
)
- first_records = hook.get_records(first_query_results)
+ first_records = await
hook.aget_records(first_query_results)
- second_records = hook.get_records(second_query_results)
+ second_records = await
hook.aget_records(second_query_results)
# If empty list, then no records are available
if not first_records:
@@ -686,7 +686,7 @@ class BigQueryValueCheckTrigger(BigQueryInsertJobTrigger):
response_from_hook = await
hook.get_job_status(job_id=self.job_id, project_id=self.project_id)
if response_from_hook["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)
+ records = await hook.aget_records(query_results)
_records = records.pop(0) if records else None
hook.value_check(self.sql, self.pass_value, _records,
self.tolerance)
yield TriggerEvent({"status": "success", "message": "Job
completed", "records": _records})
@@ -882,7 +882,7 @@ class
BigQueryTablePartitionExistenceTrigger(BigQueryTableExistenceTrigger):
async def _partition_exists(self, hook: BigQueryAsyncHook, job_id: str |
None, project_id: str):
query_results = await hook.get_job_output(job_id=job_id,
project_id=project_id)
- records = hook.get_records(query_results)
+ records = await hook.aget_records(query_results)
if records:
records = [row[0] for row in records]
return self.partition_id in records
diff --git a/providers/google/tests/unit/google/cloud/hooks/test_bigquery.py
b/providers/google/tests/unit/google/cloud/hooks/test_bigquery.py
index 1e7349aba50..2f106138a45 100644
--- a/providers/google/tests/unit/google/cloud/hooks/test_bigquery.py
+++ b/providers/google/tests/unit/google/cloud/hooks/test_bigquery.py
@@ -2140,7 +2140,8 @@ class TestBigQueryAsyncHookMethods:
"cacheHit": False,
}
hook = BigQueryAsyncHook(use_legacy_sql=True)
- result = hook.get_records(query_result)
+ with pytest.deprecated_call(match="get_records.*deprecated"):
+ result = hook.get_records(query_result)
assert isinstance(result[0][0], int)
assert isinstance(result[0][1], float)
assert isinstance(result[0][2], str)
@@ -2168,9 +2169,93 @@ class TestBigQueryAsyncHookMethods:
"cacheHit": False,
}
hook = BigQueryAsyncHook(use_legacy_sql=True)
- result = hook.get_records(query_result, as_dict=True)
+ with pytest.deprecated_call(match="get_records.*deprecated"):
+ result = hook.get_records(query_result, as_dict=True)
assert result == [{"f0_": 22, "f1_": 3.14, "f2_": "PI"}]
+ @pytest.mark.asyncio
+ async def test_aget_records_return_type(self):
+ query_result = {
+ "kind": "bigquery#getQueryResultsResponse",
+ "etag": "test_etag",
+ "schema": {
+ "fields": [
+ {"name": "f0_", "type": "INTEGER", "mode": "NULLABLE"},
+ {"name": "f1_", "type": "FLOAT", "mode": "NULLABLE"},
+ {"name": "f2_", "type": "STRING", "mode": "NULLABLE"},
+ ]
+ },
+ "jobReference": {
+ "projectId": "test_airflow-providers",
+ "jobId": "test_jobid",
+ "location": "US",
+ },
+ "totalRows": "1",
+ "rows": [{"f": [{"v": "22"}, {"v": "3.14"}, {"v": "PI"}]}],
+ "totalBytesProcessed": "0",
+ "jobComplete": True,
+ "cacheHit": False,
+ }
+ hook = BigQueryAsyncHook(use_legacy_sql=True)
+ result = await hook.aget_records(query_result)
+ assert isinstance(result[0][0], int)
+ assert isinstance(result[0][1], float)
+ assert isinstance(result[0][2], str)
+
+ @pytest.mark.asyncio
+ async def test_aget_records_as_dict(self):
+ query_result = {
+ "kind": "bigquery#getQueryResultsResponse",
+ "etag": "test_etag",
+ "schema": {
+ "fields": [
+ {"name": "f0_", "type": "INTEGER", "mode": "NULLABLE"},
+ {"name": "f1_", "type": "FLOAT", "mode": "NULLABLE"},
+ {"name": "f2_", "type": "STRING", "mode": "NULLABLE"},
+ ]
+ },
+ "jobReference": {
+ "projectId": "test_airflow-providers",
+ "jobId": "test_jobid",
+ "location": "US",
+ },
+ "totalRows": "1",
+ "rows": [{"f": [{"v": "22"}, {"v": "3.14"}, {"v": "PI"}]}],
+ "totalBytesProcessed": "0",
+ "jobComplete": True,
+ "cacheHit": False,
+ }
+ hook = BigQueryAsyncHook(use_legacy_sql=True)
+ result = await hook.aget_records(query_result, as_dict=True)
+ assert result == [{"f0_": 22, "f1_": 3.14, "f2_": "PI"}]
+
+ @pytest.mark.asyncio
+ @mock.patch("asyncio.sleep", new_callable=mock.AsyncMock)
+ async def test_aget_records_yields_control(self, mock_sleep):
+ rows = [{"f": [{"v": str(i)}]} for i in range(2500)]
+ query_result = {
+ "schema": {"fields": [{"name": "col", "type": "INTEGER"}]},
+ "rows": rows,
+ }
+ hook = BigQueryAsyncHook(use_legacy_sql=True)
+ result = await hook.aget_records(query_result, yield_frequency=1000)
+ assert len(result) == 2500
+ assert mock_sleep.call_count == 2
+ mock_sleep.assert_awaited_with(0)
+
+ @pytest.mark.asyncio
+ @mock.patch("asyncio.sleep", new_callable=mock.AsyncMock)
+ async def test_aget_records_no_yield_when_disabled(self, mock_sleep):
+ rows = [{"f": [{"v": str(i)}]} for i in range(500)]
+ query_result = {
+ "schema": {"fields": [{"name": "col", "type": "INTEGER"}]},
+ "rows": rows,
+ }
+ hook = BigQueryAsyncHook(use_legacy_sql=True)
+ result = await hook.aget_records(query_result, yield_frequency=0)
+ assert len(result) == 500
+ mock_sleep.assert_not_called()
+
@pytest.mark.db_test
class TestHookLevelLineage(_BigQueryBaseTestClass):
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 27e93522ed9..a60a7b5daf8 100644
--- a/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
+++ b/providers/google/tests/unit/google/cloud/triggers/test_bigquery.py
@@ -641,14 +641,14 @@ class TestBigQueryIntervalCheckTrigger:
@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")
+
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.aget_records")
async def test_interval_check_trigger_success(
- self, mock_get_records, mock_get_job_output, mock_job_status,
mock_sync_hook, interval_check_trigger
+ self, mock_aget_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.
"""
- mock_get_records.return_value = {}
+ mock_aget_records.return_value = []
mock_job_status.return_value = {"status": "success", "message": "Job
completed"}
mock_get_job_output.return_value = ["0"]
@@ -767,7 +767,7 @@ class TestBigQueryIntervalCheckTrigger:
)
@mock.patch("asyncio.sleep", new_callable=AsyncMock)
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.interval_check")
-
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_records")
+
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.aget_records")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_output")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_sync_hook")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_status")
@@ -776,7 +776,7 @@ class TestBigQueryIntervalCheckTrigger:
mock_job_status,
mock_sync_hook,
mock_get_job_output,
- mock_get_records,
+ mock_aget_records,
mock_interval_check,
mock_sleep,
first_status,
@@ -791,7 +791,7 @@ class TestBigQueryIntervalCheckTrigger:
{"status": "success", "message": "Job completed"},
]
mock_sync_hook.return_value.is_default_universe.return_value = True
- mock_get_records.side_effect = lambda *_, **__: [[100]]
+ mock_aget_records.side_effect = lambda *_, **__: [[100]]
actual = await interval_check_trigger.run().asend(None)
@@ -884,18 +884,18 @@ class TestBigQueryValueCheckTrigger:
}
@pytest.mark.asyncio
-
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_records")
+
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.aget_records")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_output")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_status")
async def test_value_check_op_trigger_success(
- self, mock_job_status, get_job_output, get_records, value_check_trigger
+ self, mock_job_status, get_job_output, aget_records,
value_check_trigger
):
"""
Tests BigQueryValueCheckTrigger only fires once the query execution
reaches a successful state.
"""
mock_job_status.return_value = {"status": "success", "message": "Job
completed"}
get_job_output.return_value = {}
- get_records.return_value = [[2], [4]]
+ aget_records.return_value = [[2], [4]]
await value_check_trigger.run().__anext__()
await asyncio.sleep(0.5)
@@ -929,14 +929,14 @@ class TestBigQueryValueCheckTrigger:
@pytest.mark.asyncio
@mock.patch("asyncio.sleep", new_callable=AsyncMock)
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.value_check")
-
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_records")
+
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.aget_records")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_output")
@mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryAsyncHook.get_job_status")
async def test_value_check_op_trigger_keeps_polling_while_job_is_running(
self,
mock_job_status,
mock_get_job_output,
- mock_get_records,
+ mock_aget_records,
mock_value_check,
mock_sleep,
value_check_trigger,
@@ -946,7 +946,7 @@ class TestBigQueryValueCheckTrigger:
{"status": "running", "message": "Job running"},
{"status": "success", "message": "Job completed"},
]
- mock_get_records.side_effect = lambda *_, **__: [[4]]
+ mock_aget_records.side_effect = lambda *_, **__: [[4]]
actual = await value_check_trigger.run().asend(None)
