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) <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

View file

@ -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:

View file

@ -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")