e5fb913cec
Five hottest BQ-touching endpoints were `async def` but invoked synchronous DuckDB / BQ-extension calls inside the body. Under uvicorn's single event loop that meant a single heavy `agnes query --remote` (waiting up to ~200 s for BQ's jobs.query) froze EVERY other request — /api/health, dashboard, auth, even another query — for the full BQ wait. Operators saw "VM idle, app frozen" during PR #188's testing. Convert to plain `def` so FastAPI auto-offloads the body to the anyio thread pool. Event loop stays free for non-BQ requests. - app/api/query.py:execute_query - app/api/v2_scan.py:scan_estimate_endpoint, scan_endpoint - app/api/v2_sample.py:sample - app/api/v2_schema.py:schema Audit: 0 `await` statements in any converted handler (verified file-by- file), so the rename is safe. Tests in tests/test_v2_*.py called the handlers via `asyncio.run(...)` which now fails on a non-coroutine return; swapped for direct calls (asyncio.run( -> ( ) — keeps paren balance). Plus AGNES_THREADPOOL_SIZE env var (default 200, was anyio's stock 40) in app/main.py:lifespan. Set via anyio.to_thread.current_default_thread_limiter().total_tokens. 200 is comfortable headroom for <50 concurrent analysts; bump for more. 480/480 impacted tests pass (the 2 remaining errors are a pre-existing fixture setup issue in test_reader_smoke_matrix.py unrelated to this change).
333 lines
14 KiB
Python
333 lines
14 KiB
Python
# tests/test_v2_schema.py
|
|
import asyncio
|
|
import importlib
|
|
from unittest.mock import patch, MagicMock
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
|
|
@pytest.fixture
|
|
def reload_db(tmp_path, monkeypatch):
|
|
monkeypatch.setenv("DATA_DIR", str(tmp_path))
|
|
import src.db as db_module
|
|
importlib.reload(db_module)
|
|
yield db_module
|
|
|
|
|
|
def _seed_bq_table(conn):
|
|
_ensure_admin1(conn)
|
|
from src.repositories.table_registry import TableRegistryRepository
|
|
TableRegistryRepository(conn).register(
|
|
id="bq_view", name="bq_view", source_type="bigquery",
|
|
bucket="ds", source_table="bq_view", query_mode="remote",
|
|
)
|
|
|
|
|
|
def _ensure_admin1(conn):
|
|
"""Seed an admin user with id='admin1' + Admin group membership so
|
|
{"id": "admin1", ...} dicts pass the can_access admin shortcut."""
|
|
from src.db import SYSTEM_ADMIN_GROUP
|
|
from src.repositories.users import UserRepository
|
|
from src.repositories.user_group_members import UserGroupMembersRepository
|
|
if UserRepository(conn).get_by_id('admin1') is None:
|
|
UserRepository(conn).create(id='admin1', email='admin1@test.com', name='Admin')
|
|
admin_gid = conn.execute(
|
|
'SELECT id FROM user_groups WHERE name = ?', [SYSTEM_ADMIN_GROUP]
|
|
).fetchone()
|
|
if admin_gid:
|
|
UserGroupMembersRepository(conn).add_member(
|
|
'admin1', admin_gid[0], source='system_seed',
|
|
)
|
|
|
|
|
|
def _bq(billing="billing-proj", data="data-proj"):
|
|
"""Build a BqAccess wired to default factories. For tests that monkeypatch
|
|
`_fetch_bq_schema` / `_fetch_bq_table_options` whole, the inner factories
|
|
are never called."""
|
|
from connectors.bigquery.access import BqAccess, BqProjects
|
|
return BqAccess(BqProjects(billing=billing, data=data))
|
|
|
|
|
|
class TestSchemaEndpoint:
|
|
def test_bq_table_returns_columns_and_dialect_hints(self, reload_db, monkeypatch):
|
|
from app.api import v2_schema
|
|
# Stub the BQ schema fetch to avoid hitting real BQ
|
|
monkeypatch.setattr(
|
|
v2_schema, "_fetch_bq_schema",
|
|
lambda bq, dataset, table: [
|
|
{"name": "event_date", "type": "DATE", "nullable": False, "description": ""},
|
|
{"name": "country_code", "type": "STRING", "nullable": True, "description": ""},
|
|
],
|
|
)
|
|
monkeypatch.setattr(
|
|
v2_schema, "_fetch_bq_table_options",
|
|
lambda bq, dataset, table: {"partition_by": "event_date", "clustered_by": []},
|
|
)
|
|
|
|
conn = reload_db.get_system_db()
|
|
try:
|
|
_seed_bq_table(conn)
|
|
user = {"id": "admin1", "email": "a@x.com"}
|
|
data = v2_schema.build_schema(conn, user, "bq_view", bq=_bq())
|
|
finally:
|
|
conn.close()
|
|
assert data["table_id"] == "bq_view"
|
|
assert data["sql_flavor"] == "bigquery"
|
|
assert {c["name"] for c in data["columns"]} == {"event_date", "country_code"}
|
|
assert "where_dialect_hints" in data
|
|
assert data["partition_by"] == "event_date"
|
|
|
|
def test_unknown_table_raises_404(self, reload_db):
|
|
from app.api.v2_schema import build_schema, NotFound
|
|
conn = reload_db.get_system_db()
|
|
try:
|
|
user = {"id": "admin1", "email": "a@x.com"}
|
|
with pytest.raises(NotFound):
|
|
build_schema(conn, user, "missing", bq=_bq())
|
|
finally:
|
|
conn.close()
|
|
|
|
def test_rbac_check_runs_before_cache(self, reload_db, monkeypatch):
|
|
"""Regression: cache lookup used to happen before the RBAC check, and the
|
|
cache key had no user component — so an unauthorized user could read
|
|
cached schema fetched by an authorized one. The fix moves RBAC ahead."""
|
|
from app.api import v2_schema
|
|
monkeypatch.setattr(
|
|
v2_schema, "_fetch_bq_schema",
|
|
lambda *a, **kw: [{"name": "x", "type": "STRING", "nullable": True, "description": ""}],
|
|
)
|
|
monkeypatch.setattr(v2_schema, "_fetch_bq_table_options", lambda *a, **kw: {})
|
|
# Stub can_access_table — admin1 always allowed (warms cache),
|
|
# everyone else denied (must hit RBAC check, not cached payload).
|
|
monkeypatch.setattr(
|
|
"app.api.v2_schema.can_access_table",
|
|
lambda user, tid, conn: user.get("id") == "admin1",
|
|
)
|
|
|
|
conn = reload_db.get_system_db()
|
|
try:
|
|
_seed_bq_table(conn)
|
|
# Admin warms the cache.
|
|
admin = {"id": "admin1", "email": "admin@x.com"}
|
|
v2_schema.build_schema(conn, admin, "bq_view", bq=_bq())
|
|
# Non-admin must hit RBAC denial — cache must NOT short-circuit.
|
|
other = {"id": "viewer1", "email": "viewer@x.com"}
|
|
with pytest.raises(PermissionError):
|
|
v2_schema.build_schema(conn, other, "bq_view", bq=_bq())
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
class TestBqAccessErrors:
|
|
"""Issue #134: structured 502 translation on BQ errors in the strict /schema path.
|
|
|
|
These tests exercise the REAL translation path through `BqAccess` +
|
|
`translate_bq_error` by injecting a duckdb_session whose execute() raises
|
|
the Google API exception. That's the production path — Phase 1
|
|
monkeypatches of `_fetch_bq_schema` whole would skip the translation logic
|
|
and only test the outer wrap (which has been removed in Phase 2).
|
|
|
|
`/schema` differs from `/scan` and `/scan/estimate`: the SQL is
|
|
server-constructed (queries `INFORMATION_SCHEMA.COLUMNS` with validated
|
|
identifiers, no user-derived fragments), so `BadRequest` → 502
|
|
(`bq_upstream_error`) rather than 400. Same as `/sample`.
|
|
"""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _clear_schema_cache(self):
|
|
# Prevent payloads from earlier tests from short-circuiting the fetch.
|
|
from app.api import v2_schema
|
|
v2_schema._schema_cache.clear()
|
|
yield
|
|
v2_schema._schema_cache.clear()
|
|
|
|
def test_schema_returns_502_on_bq_forbidden_serviceusage(self, reload_db, bq_access):
|
|
"""When the BQ extension raises Forbidden mentioning serviceusage, the
|
|
endpoint must translate to HTTP 502 with `cross_project_forbidden`
|
|
and a hint that mentions billing_project."""
|
|
from app.api import v2_schema
|
|
from google.api_core.exceptions import Forbidden
|
|
|
|
mock_conn = MagicMock()
|
|
mock_conn.execute.side_effect = Forbidden(
|
|
"Permission denied: serviceusage.services.use on project foo"
|
|
)
|
|
bq = bq_access(duckdb_conn=mock_conn, billing="billing-proj", data="data-proj")
|
|
|
|
conn = reload_db.get_system_db()
|
|
try:
|
|
_seed_bq_table(conn)
|
|
user = {"id": "admin1", "email": "a@x.com"}
|
|
# Endpoint is async — drive it directly. dependency_overrides only
|
|
# fires through TestClient/HTTP, so pass `bq=bq` explicitly.
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
(v2_schema.schema(
|
|
table_id="bq_view", user=user, conn=conn, bq=bq,
|
|
))
|
|
finally:
|
|
conn.close()
|
|
|
|
assert exc_info.value.status_code == 502
|
|
detail = exc_info.value.detail
|
|
assert isinstance(detail, dict)
|
|
assert detail["error"] == "cross_project_forbidden"
|
|
assert "hint" in detail["details"]
|
|
assert "billing_project" in detail["details"]["hint"].lower()
|
|
|
|
def test_schema_returns_502_on_bq_forbidden_non_serviceusage(self, reload_db, bq_access):
|
|
"""A Forbidden that is NOT about serviceusage (e.g. dataset-level ACL)
|
|
still becomes a 502, but with `bq_forbidden`."""
|
|
from app.api import v2_schema
|
|
from google.api_core.exceptions import Forbidden
|
|
|
|
mock_conn = MagicMock()
|
|
mock_conn.execute.side_effect = Forbidden(
|
|
"Access Denied: Dataset foo.bar: User does not have permission"
|
|
)
|
|
bq = bq_access(duckdb_conn=mock_conn, billing="billing-proj", data="data-proj")
|
|
|
|
conn = reload_db.get_system_db()
|
|
try:
|
|
_seed_bq_table(conn)
|
|
user = {"id": "admin1", "email": "a@x.com"}
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
(v2_schema.schema(
|
|
table_id="bq_view", user=user, conn=conn, bq=bq,
|
|
))
|
|
finally:
|
|
conn.close()
|
|
|
|
assert exc_info.value.status_code == 502
|
|
assert exc_info.value.detail["error"] == "bq_forbidden"
|
|
|
|
def test_schema_returns_502_on_bq_bad_request(self, reload_db, bq_access):
|
|
"""`/schema` SQL is server-constructed (INFORMATION_SCHEMA.COLUMNS with
|
|
validated identifiers); a BadRequest here means registry corruption →
|
|
upstream error, not user fault. Translate to HTTP 502 (`bq_upstream_error`),
|
|
same as `/sample`, opposite of `/scan*`."""
|
|
from app.api import v2_schema
|
|
from google.api_core.exceptions import BadRequest
|
|
|
|
mock_conn = MagicMock()
|
|
mock_conn.execute.side_effect = BadRequest("Syntax error at line 1, column 5")
|
|
bq = bq_access(duckdb_conn=mock_conn, billing="billing-proj", data="data-proj")
|
|
|
|
conn = reload_db.get_system_db()
|
|
try:
|
|
_seed_bq_table(conn)
|
|
user = {"id": "admin1", "email": "a@x.com"}
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
(v2_schema.schema(
|
|
table_id="bq_view", user=user, conn=conn, bq=bq,
|
|
))
|
|
finally:
|
|
conn.close()
|
|
|
|
assert exc_info.value.status_code == 502
|
|
detail = exc_info.value.detail
|
|
assert isinstance(detail, dict)
|
|
assert detail["error"] == "bq_upstream_error"
|
|
assert "Syntax error" in detail["message"]
|
|
|
|
def test_schema_returns_200_with_empty_partition_on_table_options_failure(
|
|
self, reload_db, monkeypatch,
|
|
):
|
|
"""REGRESSION GUARD: `_fetch_bq_table_options` is best-effort
|
|
(`try/except Exception → return {}` with a logger.warning). On internal
|
|
failure the helper returns `{}`; the endpoint MUST still return 200 with
|
|
the column list and `partition_by=None, clustered_by=[]`.
|
|
|
|
Phase 2 preserved the swallow-all contract — this test verifies it
|
|
end-to-end with the helper patched to return the empty dict it would
|
|
yield on failure, AND a Phase-2-specific regression guard that exercises
|
|
the real helper raising an upstream error and confirms it's swallowed
|
|
(not surfaced as 502)."""
|
|
from app.api import v2_schema
|
|
|
|
# Strict path returns real schema.
|
|
monkeypatch.setattr(
|
|
v2_schema, "_fetch_bq_schema",
|
|
lambda bq, dataset, table: [
|
|
{"name": "event_date", "type": "DATE", "nullable": False, "description": ""},
|
|
{"name": "country_code", "type": "STRING", "nullable": True, "description": ""},
|
|
],
|
|
)
|
|
# Best-effort path returns `{}` (the documented swallow-all output).
|
|
monkeypatch.setattr(v2_schema, "_fetch_bq_table_options", lambda bq, dataset, table: {})
|
|
|
|
conn = reload_db.get_system_db()
|
|
try:
|
|
_seed_bq_table(conn)
|
|
user = {"id": "admin1", "email": "a@x.com"}
|
|
data = (v2_schema.schema(
|
|
table_id="bq_view", user=user, conn=conn, bq=_bq(),
|
|
))
|
|
finally:
|
|
conn.close()
|
|
|
|
# Endpoint returns 200 (no HTTPException raised) with columns present.
|
|
assert data["table_id"] == "bq_view"
|
|
assert {c["name"] for c in data["columns"]} == {"event_date", "country_code"}
|
|
# Partition info absent or None (the swallow-all returns {} on failure).
|
|
assert data.get("partition_by") is None
|
|
assert data.get("clustered_by") == []
|
|
|
|
def test_table_options_swallows_bq_errors_returns_empty_dict(self, reload_db, bq_access):
|
|
"""REGRESSION GUARD (Phase 2 production path): `_fetch_bq_table_options`
|
|
must swallow ANY exception from BQ and return `{}`. Calling it directly
|
|
with a bq whose duckdb_session().execute() raises Forbidden / BadRequest /
|
|
unknown — all must produce `{}`, NOT a BqAccessError that the endpoint
|
|
would 502 on. This is the load-bearing contract for `/schema` to keep
|
|
returning 200 on permissioned tables / cross-project misconfigurations."""
|
|
from app.api import v2_schema
|
|
from google.api_core.exceptions import Forbidden, BadRequest
|
|
|
|
for exc in [
|
|
Forbidden("Permission denied: serviceusage.services.use on project foo"),
|
|
Forbidden("Access Denied: Dataset foo.bar"),
|
|
BadRequest("Syntax error"),
|
|
RuntimeError("totally unexpected"),
|
|
]:
|
|
mock_conn = MagicMock()
|
|
mock_conn.execute.side_effect = exc
|
|
bq = bq_access(duckdb_conn=mock_conn, billing="billing-proj", data="data-proj")
|
|
assert v2_schema._fetch_bq_table_options(bq, "ds", "bq_view") == {}, (
|
|
f"swallow-all contract violated for {type(exc).__name__}: {exc}"
|
|
)
|
|
|
|
def test_schema_passes_billing_project_to_bigquery_query(self, reload_db, bq_access):
|
|
"""Regression guard: bq.projects.billing must be passed to bigquery_query()
|
|
as the billing project (positional arg 0), and the FROM clause must
|
|
reference bq.projects.data (NOT billing). Verifies the migration didn't
|
|
regress the cross-project bug fix from Phase 1."""
|
|
from app.api import v2_schema
|
|
|
|
captured = {}
|
|
|
|
def _fake_execute(sql, params):
|
|
if "bigquery_query" in sql:
|
|
# _fetch_bq_schema executes first; capture once.
|
|
if "bq_sql" not in captured:
|
|
captured["billing_project"] = params[0]
|
|
captured["bq_sql"] = params[1]
|
|
result = MagicMock()
|
|
result.fetchall.return_value = []
|
|
return result
|
|
|
|
mock_conn = MagicMock()
|
|
mock_conn.execute.side_effect = _fake_execute
|
|
bq = bq_access(duckdb_conn=mock_conn, billing="billing-proj", data="data-proj")
|
|
|
|
conn = reload_db.get_system_db()
|
|
try:
|
|
_seed_bq_table(conn)
|
|
user = {"id": "admin1", "email": "a@x.com"}
|
|
(v2_schema.schema(
|
|
table_id="bq_view", user=user, conn=conn, bq=bq,
|
|
))
|
|
finally:
|
|
conn.close()
|
|
|
|
assert captured["billing_project"] == "billing-proj"
|
|
# FROM clause uses data project (where INFORMATION_SCHEMA.COLUMNS lives)
|
|
assert "`data-proj.ds.INFORMATION_SCHEMA.COLUMNS`" in captured["bq_sql"]
|