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:
parent
618385e7e4
commit
35df940e5c
2 changed files with 11 additions and 6 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue