From 35df940e5cc1c04a92e0fde842bebc3197970ce4 Mon Sep 17 00:00:00 2001 From: ZdenekSrotyr Date: Sat, 11 Apr 2026 20:29:03 +0200 Subject: [PATCH] fix: BQ COUNT subquery alias, wrap ImportError in RemoteQueryError - 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) --- src/remote_query.py | 13 +++++++++---- tests/test_remote_query.py | 4 ++-- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/remote_query.py b/src/remote_query.py index 5ebcfba..c0df463 100644 --- a/src/remote_query.py +++ b/src/remote_query.py @@ -267,7 +267,7 @@ class RemoteQueryEngine: client = self._get_bq_client() # --- Phase 1a: COUNT(*) pre-check --- - count_sql = f"SELECT COUNT(*) FROM ({bq_sql})" + count_sql = f"SELECT COUNT(*) FROM ({bq_sql}) AS _cnt" try: count_job = client.query(count_sql) count_arrow = count_job.to_arrow() @@ -413,9 +413,14 @@ class RemoteQueryEngine: project = os.environ.get("BIGQUERY_PROJECT", "unknown") return self._bq_client_factory(project) - # Trigger ImportError early if the package is missing. - # This is a lazy import so the module stays usable without BQ installed. - import google.cloud.bigquery as _bq_module # noqa: PLC0415, F401 + # Lazy import so the module stays usable without BQ installed. + try: + 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") if not project: diff --git a/tests/test_remote_query.py b/tests/test_remote_query.py index f4de108..992462a 100644 --- a/tests/test_remote_query.py +++ b/tests/test_remote_query.py @@ -124,7 +124,7 @@ class TestRemoteQueryEngineRegister: pass # Expected — no BQ package in test env 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( analytics_conn, # 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 pytest.raises((ImportError, ModuleNotFoundError)): + with pytest.raises(RemoteQueryError, match="google-cloud-bigquery"): engine.register_bq("bq_alias", "SELECT 1")