shahar1 commented on code in PR #72286:
URL: https://github.com/apache/airflow/pull/72286#discussion_r3999607813
##########
providers/google/src/airflow/providers/google/cloud/hooks/bigquery.py:
##########
@@ -2459,7 +2459,9 @@ def get_records(
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 dict_row in rows:
+ for i, dict_row in enumerate(rows):
+ if i > 0 and i % 1000 == 0:
+ await asyncio.sleep(0)
Review Comment:
1. This `if` must be unit tested
2. What's the reasoning for choosing 1000? Do we want to allow
parametrization?
##########
providers/google/src/airflow/providers/google/cloud/hooks/bigquery.py:
##########
@@ -2437,8 +2438,7 @@ async def cancel_job(self, job_id: str, project_id: str |
None, location: str |
self.log.error("Failed to cancel BigQuery job %s: %s", job_id,
str(e))
raise
- # TODO: Convert get_records into an async method
- def get_records(
+ async def get_records(
Review Comment:
It breaks API for sync calls - if someone customized a trigger/operator that
calls `get_records`, converting to `async` will just break it.
##########
providers/google/tests/unit/google/cloud/triggers/test_bigquery.py:
##########
@@ -641,14 +641,14 @@ def test_interval_check_trigger_round_trip(self,
interval_check_trigger):
@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.get_records",
new_callable=mock.AsyncMock)
Review Comment:
I expect it to break static checks
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
