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

fix: BQ COUNT subquery alias, wrap ImportError in RemoteQueryError

Browse source 
- Add AS _cnt alias to COUNT(*) subquery (BQ Standard SQL requires it)
- Catch ImportError in _get_bq_client() and raise RemoteQueryError
  so API endpoint returns proper 400 instead of 500

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
ZdenekSrotyr 2026-04-11 20:29:03 +02:00
parent 618385e7e4
commit 35df940e5c
2 changed files with 11 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
src/remote_query.py
View file

@ -267,7 +267,7 @@ class RemoteQueryEngine:
client = self._get_bq_client() client = self._get_bq_client()
# --- Phase 1a: COUNT(*) pre-check --- # --- Phase 1a: COUNT(*) pre-check ---
count_sql = f"SELECT COUNT(*) FROM ({bq_sql})" count_sql = f"SELECT COUNT(*) FROM ({bq_sql}) AS _cnt"
try: try:
count_job = client.query(count_sql) count_job = client.query(count_sql)
count_arrow = count_job.to_arrow() count_arrow = count_job.to_arrow()
@ -413,9 +413,14 @@ class RemoteQueryEngine:
project = os.environ.get("BIGQUERY_PROJECT", "unknown") project = os.environ.get("BIGQUERY_PROJECT", "unknown")
return self._bq_client_factory(project) return self._bq_client_factory(project)
# Trigger ImportError early if the package is missing. # Lazy import so the module stays usable without BQ installed.
# This is a lazy import so the module stays usable without BQ installed. try:
import google.cloud.bigquery as _bq_module # noqa: PLC0415, F401 import google.cloud.bigquery as _bq_module # noqa: PLC0415, F401
except ImportError:
raise RemoteQueryError(
"google-cloud-bigquery is not installed. Install with: pip install google-cloud-bigquery",
error_type="bq_error",
)
project = os.environ.get("BIGQUERY_PROJECT") project = os.environ.get("BIGQUERY_PROJECT")
if not project: if not project:

4
tests/test_remote_query.py
View file

@ -124,7 +124,7 @@ class TestRemoteQueryEngineRegister:
pass # Expected — no BQ package in test env pass # Expected — no BQ package in test env
def test_register_bq_missing_package(self, analytics_conn): def test_register_bq_missing_package(self, analytics_conn):
"""When google-cloud-bigquery is not installed, engine must raise ImportError.""" """When google-cloud-bigquery is not installed, engine must raise RemoteQueryError."""
engine = RemoteQueryEngine( engine = RemoteQueryEngine(
analytics_conn, analytics_conn,
# No factory — will try to import google.cloud.bigquery # No factory — will try to import google.cloud.bigquery
@ -133,7 +133,7 @@ class TestRemoteQueryEngineRegister:
) )
with patch.dict(sys.modules, {"google": None, "google.cloud": None, "google.cloud.bigquery": None}): with patch.dict(sys.modules, {"google": None, "google.cloud": None, "google.cloud.bigquery": None}):
with pytest.raises((ImportError, ModuleNotFoundError)): with pytest.raises(RemoteQueryError, match="google-cloud-bigquery"):
engine.register_bq("bq_alias", "SELECT 1") engine.register_bq("bq_alias", "SELECT 1")