AI-Cognitive-Leap/agnes-the-ai-analyst
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add BQ Storage API fallback to REST when readsessions permission missing

Browse source
This commit is contained in:
Petr 2026-03-11 13:59:09 +01:00
parent d438438e33
commit e26e47a071
1 changed files with 14 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
connectors/bigquery/client.py
View file

@ -274,7 +274,20 @@ class BigQueryClient:
logger.debug(f"Executing BQ query: {sql[:200]}...") logger.debug(f"Executing BQ query: {sql[:200]}...")
query_job = self.client.query(sql, job_config=job_config) query_job = self.client.query(sql, job_config=job_config)
arrow_table = query_job.to_arrow()
# Try BQ Storage API first (faster for large results).
# Fall back to REST API if SA lacks bigquery.readsessions.create permission.
try:
arrow_table = query_job.to_arrow()
except Exception as storage_err:
if "readsessions" in str(storage_err) or "PERMISSION_DENIED" in str(storage_err):
logger.warning(
"BQ Storage API unavailable (missing readsessions permission), "
"falling back to REST API"
)
arrow_table = query_job.to_arrow(create_bqstorage_client=False)
else:
raise
logger.debug(f"Query returned {arrow_table.num_rows} rows, {arrow_table.num_columns} columns") logger.debug(f"Query returned {arrow_table.num_rows} rows, {arrow_table.num_columns} columns")
return arrow_table return arrow_table