diff --git a/datajunction-server/datajunction_server/api/semantic_layer.py b/datajunction-server/datajunction_server/api/semantic_layer.py index dc9737f4a..363a4f0f8 100644 --- a/datajunction-server/datajunction_server/api/semantic_layer.py +++ b/datajunction-server/datajunction_server/api/semantic_layer.py @@ -22,6 +22,7 @@ from datajunction_server.errors import DJException from datajunction_server.internal.access.authentication.http import SecureAPIRouter from datajunction_server.internal.sql import ( + build_row_count_sql, generate_dimensions_sql, generate_metrics_sql, ) @@ -185,11 +186,25 @@ def _view_payload(cube: NodeRevision) -> "ViewDetail": # Filter / order translation (spec QueryPayload -> DJ build_metrics_sql args) # --------------------------------------------------------------------------- -# Comparison operators accepted for filters; anything else (IN/LIKE/…) is rejected. +# Comparison operators accepted for filters. These cover the portable operators +# used by the semantic-layer spec; pattern matching remains unsupported. _ALLOWED_OPERATORS = frozenset( - {"=", "!=", ">", "<", ">=", "<=", "IS NULL", "IS NOT NULL"}, + { + "=", + "!=", + ">", + "<", + ">=", + "<=", + "IN", + "NOT IN", + "BETWEEN", + "IS NULL", + "IS NOT NULL", + }, ) _NULLARY_OPERATORS = frozenset({"IS NULL", "IS NOT NULL"}) +_COLLECTION_OPERATORS = frozenset({"IN", "NOT IN"}) def _quote_value(value: Any) -> str: @@ -220,6 +235,22 @@ def _filter_to_sql(flt: "FilterPayload") -> str: ) if op in _NULLARY_OPERATORS: return f"{flt.column} {op}" + if op in _COLLECTION_OPERATORS: + if not isinstance(flt.value, (list, tuple, set, frozenset)) or not flt.value: + raise DJException( + message=f"Filter operator {flt.operator} requires a non-empty list", + http_status_code=400, + ) + values = ", ".join(_quote_value(value) for value in flt.value) + return f"{flt.column} {op} ({values})" + if op == "BETWEEN": + if not isinstance(flt.value, (list, tuple)) or len(flt.value) != 2: + raise DJException( + message="Filter operator BETWEEN requires exactly two values", + http_status_code=400, + ) + start, end = (_quote_value(value) for value in flt.value) + return f"{flt.column} BETWEEN {start} AND {end}" return f"{flt.column} {op} {_quote_value(flt.value)}" @@ -260,9 +291,18 @@ class QueryPayload(BaseModel): class QueryRequest(BaseModel): """Body for POST /views/{view_name}/sql (backs the spec's /query).""" + additional_configuration: dict[str, Any] = Field(default_factory=dict) query: QueryPayload +class ValuesRequest(BaseModel): + """Body for POST /views/{view_name}/values.""" + + additional_configuration: dict[str, Any] = Field(default_factory=dict) + dimension: str + filters: list[FilterPayload] = Field(default_factory=list) + + # --------------------------------------------------------------------------- # Response schemas # --------------------------------------------------------------------------- @@ -387,6 +427,7 @@ async def _generate_sql( session: AsyncSession, payload: QueryPayload, cube_rev: NodeRevision, + row_count: bool = False, ) -> GeneratedSQLResponse: """Generate physical SQL pinned to this cube (``matched_cube=cube_rev``), so we don't let ``find_matching_cube`` pick a different/differently-filtered @@ -422,6 +463,8 @@ async def _generate_sql( limit=limit, endpoint="/semantic-layer/views/sql", ) + if row_count: + generated_sql = build_row_count_sql(generated_sql) # ``generated_sql.sql`` renders via ``to_sql(query, dialect)``, which already # applies DJ's dialect rules and transpiles to the resolved dialect, so it is # execution-ready for the caller (which runs it directly). We return the @@ -438,15 +481,13 @@ async def _generate_sql( ) -@router.post("/views/{view_name}/sql", response_model=GeneratedSQLResponse) -async def generate_query_sql( +async def _generate_view_sql( view_name: str, - body: QueryRequest, - session: AsyncSession = Depends(get_session), - current_user: User = Depends(get_current_user), + payload: QueryPayload, + session: AsyncSession, + row_count: bool = False, ) -> GeneratedSQLResponse | JSONResponse: - """Generate physical SQL for a view query.""" - payload = body.query + """Validate and generate SQL for a semantic view query.""" if payload.offset: return _problem( 400, @@ -483,7 +524,72 @@ async def generate_query_sql( f"metrics={bad_metrics} dimensions={bad_dims} " f"filter_columns={bad_filters}", ) - result = await _generate_sql(session, payload, cube_rev) + result = await _generate_sql(session, payload, cube_rev, row_count=row_count) except DJException as exc: return _problem(exc.http_status_code or 400, exc.message) return result + + +@router.post("/views/{view_name}/sql", response_model=GeneratedSQLResponse) +async def generate_query_sql( + view_name: str, + body: QueryRequest, + session: AsyncSession = Depends(get_session), + current_user: User = Depends(get_current_user), +) -> GeneratedSQLResponse | JSONResponse: + """Generate physical SQL for a view query.""" + return await _generate_view_sql(view_name, body.query, session) + + +@router.post("/views/{view_name}/row-count", response_model=GeneratedSQLResponse) +async def generate_row_count_sql( + view_name: str, + body: QueryRequest, + session: AsyncSession = Depends(get_session), + current_user: User = Depends(get_current_user), +) -> GeneratedSQLResponse | JSONResponse: + """Generate SQL that counts the rows returned by a semantic query.""" + return await _generate_view_sql(view_name, body.query, session, row_count=True) + + +@router.post("/views/{view_name}/values", response_model=GeneratedSQLResponse) +async def generate_values_sql( + view_name: str, + body: ValuesRequest, + session: AsyncSession = Depends(get_session), + current_user: User = Depends(get_current_user), +) -> GeneratedSQLResponse | JSONResponse: + """Generate SQL for a dimension's distinct values, optionally filtered.""" + try: + cube_node = await Node.get_cube_by_name(session, view_name) + except DJException as exc: + return _problem(exc.http_status_code or 400, exc.message) + if cube_node is None or cube_node.current is None: + return _problem(404, f"View `{view_name}` does not exist.") + + cube_rev = cube_node.current + if body.dimension not in cube_rev.cube_node_dimensions: + return _problem( + 400, + f"View `{view_name}` does not contain dimension `{body.dimension}`.", + ) + + bad_filters = [ + flt.column + for flt in body.filters + if flt.column and flt.column not in cube_rev.cube_node_dimensions + ] + if bad_filters: + return _problem( + 400, + f"View `{view_name}` does not contain filter dimensions: {bad_filters}", + ) + + query = QueryPayload( + dimensions=[body.dimension], + filters=body.filters, + ) + try: + return await _generate_sql(session, query, cube_rev) + except DJException as exc: + return _problem(exc.http_status_code or 400, exc.message) diff --git a/datajunction-server/datajunction_server/internal/sql.py b/datajunction-server/datajunction_server/internal/sql.py index 3b9efa723..fd20ec0ca 100644 --- a/datajunction-server/datajunction_server/internal/sql.py +++ b/datajunction-server/datajunction_server/internal/sql.py @@ -369,6 +369,55 @@ async def generate_dimensions_sql( return project_dimension_values_sql(generated, dimensions, orderby, limit) +def build_row_count_sql( + generated: "BuildV3GeneratedSQL", +) -> "BuildV3GeneratedSQL": + """Wrap generated SQL in ``COUNT(*)`` before dialect rendering.""" + from copy import deepcopy + + from datajunction_server.construction.build_v3.types import ( + ColumnMetadata as V3ColumnMetadata, + GeneratedSQL as V3GeneratedSQL, + ) + + inner = deepcopy(generated.query) + inner.parenthesized = True + inner.alias = ast.Name("semantic_query") + inner.as_ = True + + count = ast.Function( + name=ast.Name("COUNT"), + args=[ast.Column(name=ast.Name("*"))], + ).set_alias( + ast.Name( + "COUNT", + quote_style="`", + ), + ) + count.set_as(True) + query = ast.Query( + select=ast.Select( + projection=[count], + from_=ast.From(relations=[ast.Relation(primary=inner)]), + ), + ) + return V3GeneratedSQL( + query=query, + columns=[ + V3ColumnMetadata( + name="COUNT", + semantic_name="COUNT", + type="bigint", + semantic_type="metric", + ), + ], + dialect=generated.dialect, + cube_name=generated.cube_name, + scan_estimate=generated.scan_estimate, + warnings=generated.warnings, + ) + + async def build_sql_for_multiple_metrics( session: AsyncSession, metrics: list[str], diff --git a/datajunction-server/datajunction_server/sql/parsing/ast.py b/datajunction-server/datajunction_server/sql/parsing/ast.py index 324bb2abb..0faa92790 100644 --- a/datajunction-server/datajunction_server/sql/parsing/ast.py +++ b/datajunction-server/datajunction_server/sql/parsing/ast.py @@ -18,6 +18,7 @@ from functools import reduce from itertools import chain, zip_longest import re +from sqlglot import Dialect as SQLGlotDialect from sqlglot import exp as sqlglot_exp from typing import ( TYPE_CHECKING, @@ -111,6 +112,10 @@ "_render_dialect", default=None, ) +_render_identifier_dialect: ContextVar[Dialect | None] = ContextVar( + "_render_identifier_dialect", + default=None, +) def get_render_dialect() -> Dialect | None: @@ -236,6 +241,16 @@ def render_for_dialect(dialect: Dialect): _render_dialect.reset(token) +@contextmanager +def render_identifiers_for_dialect(dialect: Dialect): + """Render quoted identifiers using a dialect's native quote character.""" + token = _render_identifier_dialect.set(dialect) + try: + yield + finally: + _render_identifier_dialect.reset(token) + + def to_sql(query: Query, dialect: Dialect | None = None) -> str: """ Render a query AST to SQL for a specific dialect. @@ -258,15 +273,21 @@ def to_sql(query: Query, dialect: Dialect | None = None) -> str: with render_for_dialect(dialect): rendered = str(query) + + def native_render() -> str: + with render_for_dialect(dialect), render_identifiers_for_dialect(dialect): + return str(query) + try: - return transpile_sql(rendered, dialect, schema=_sqlglot_schema(query)) + transpiled = transpile_sql(rendered, dialect, schema=_sqlglot_schema(query)) + return native_render() if transpiled == rendered else transpiled except Exception: # pragma: no cover - fall back to native render logger.warning( "Transpilation to %s failed; falling back to native render", dialect, exc_info=True, ) - return rendered + return native_render() # When True, skip parent-pointer wiring in __post_init__ and __setattr__. @@ -878,6 +899,12 @@ def identifier(self, quotes: bool = True) -> str: the name with or without quotes """ quote_style = "" if not quotes else self.quote_style + if quote_style and (dialect := _render_identifier_dialect.get()): + dialect_class = SQLGlotDialect.classes.get( + str(dialect), + SQLGlotDialect.classes["spark"], + ) + quote_style = dialect_class().tokenizer_class.IDENTIFIERS[0] namespace = str(self.namespace) + "." if self.namespace else "" return f"{namespace}{quote_style}{self.name}{quote_style}" diff --git a/datajunction-server/tests/api/semantic_layer_test.py b/datajunction-server/tests/api/semantic_layer_test.py index b0dba2401..bf6672ce2 100644 --- a/datajunction-server/tests/api/semantic_layer_test.py +++ b/datajunction-server/tests/api/semantic_layer_test.py @@ -22,7 +22,14 @@ _metrics_payload, _quote_value, ) +from datajunction_server.construction.build_v3.types import ( + GeneratedSQL as V3GeneratedSQL, +) from datajunction_server.errors import DJException +from datajunction_server.internal.sql import build_row_count_sql +from datajunction_server.models.dialect import Dialect, DialectRegistry +from datajunction_server.sql.parsing.backends.antlr4 import parse +from datajunction_server.transpilation import SQLTranspilationPlugin class TestFilterToSql: @@ -50,6 +57,34 @@ def test_missing_column_rejected_400(self): _filter_to_sql(flt) assert exc.value.http_status_code == 400 + def test_in_quotes_each_value(self): + flt = FilterPayload( + column="ns.dim", + operator="IN", + value=["North", "O'Brien"], + ) + assert _filter_to_sql(flt) == "ns.dim IN ('North', 'O\\'Brien')" + + @pytest.mark.parametrize("value", [[], "North"]) + def test_collection_operator_requires_non_empty_collection(self, value): + flt = FilterPayload(column="ns.dim", operator="IN", value=value) + + with pytest.raises(DJException) as exc: + _filter_to_sql(flt) + + assert exc.value.http_status_code == 400 + assert "requires a non-empty list" in exc.value.message + + def test_between_requires_two_values(self): + flt = FilterPayload(column="ns.dim", operator="between", value=[1, 10]) + assert _filter_to_sql(flt) == "ns.dim BETWEEN 1 AND 10" + + with pytest.raises(DJException) as exc: + _filter_to_sql( + FilterPayload(column="ns.dim", operator="between", value=[1]), + ) + assert exc.value.http_status_code == 400 + class TestQuoteValue: def test_none_is_null(self): @@ -199,6 +234,43 @@ def test_generated_sql_metric_column_type_fallback(self): assert _generated_column_arrow_type_name(column) == "floating" +@pytest.mark.parametrize( + ("dialect", "quoted_alias"), + [ + (Dialect.SPARK, "`COUNT`"), + (Dialect.TRINO, '"COUNT"'), + (Dialect.DRUID, '"COUNT"'), + (Dialect.POSTGRES, '"COUNT"'), + (Dialect.DUCKDB, '"COUNT"'), + ], +) +def test_row_count_uses_native_ast_and_quotes_alias( + dialect, + quoted_alias, + monkeypatch, +): + monkeypatch.setitem( + DialectRegistry._registry, + dialect.value, + SQLTranspilationPlugin, + ) + generated = V3GeneratedSQL( + query=parse("SELECT region_name FROM sales"), + columns=[], + dialect=dialect, + cube_name="sem.sales_cube", + ) + + result = build_row_count_sql(generated) + + assert f"COUNT(*) AS {quoted_alias}" in result.sql + assert ") AS semantic_query" in result.sql + assert [(column.name, column.type) for column in result.columns] == [ + ("COUNT", "bigint"), + ] + assert result.cube_name == "sem.sales_cube" + + # --------------------------------------------------------------------------- # DB-backed integration tests (require the testcontainers Postgres harness) # --------------------------------------------------------------------------- @@ -381,6 +453,41 @@ async def test_semantic_endpoints_end_to_end(client: AsyncClient): assert "SELECT" in sql_body["sql"].upper() assert sql_body["columns"] + # /row-count wraps the same semantic query in COUNT(*). + resp = await _expect( + await client.post(f"/semantic/views/{view}/row-count", json=query), + 200, + ) + count_body = resp.json() + assert 'COUNT(*) AS "COUNT"' in count_body["sql"] + assert count_body["columns"] == [{"name": "COUNT", "type": "int"}] + assert count_body["dialect"] == "trino" + + # /values returns a filtered distinct-dimension query. + resp = await _expect( + await client.post( + f"/semantic/views/{view}/values", + json={ + "additional_configuration": {}, + "dimension": "sem.region.region_name", + "filters": [ + { + "column": "sem.region.region_name", + "operator": "IN", + "value": ["North", "South"], + }, + ], + }, + ), + 200, + ) + values_body = resp.json() + assert "DISTINCT" in values_body["sql"] + assert "total_amount" in values_body["sql"] + assert " IN ('North', 'South')" in values_body["sql"] + assert values_body["columns"] == [{"name": "region_name", "type": "utf8"}] + assert values_body["dialect"] == "trino" + # Unknown metric/dimension ids are rejected (400). resp = await client.post( f"/semantic/views/{view}/sql", @@ -502,6 +609,71 @@ async def test_get_view_unknown_returns_404(client: AsyncClient): assert "does not exist" in resp.json()["detail"] +@pytest.mark.asyncio +async def test_values_rejects_unknown_dimension(client: AsyncClient): + view = await _setup_cube(client) + resp = await client.post( + f"/semantic/views/{view}/values", + json={"dimension": "sem.region.unknown", "filters": []}, + ) + assert resp.status_code == 400, resp.text + assert "does not contain dimension" in resp.json()["detail"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cube", [None, SimpleNamespace(current=None)]) +async def test_values_unknown_view_returns_404( + client: AsyncClient, + monkeypatch, + cube, +): + monkeypatch.setattr( + "datajunction_server.api.semantic_layer.Node.get_cube_by_name", + AsyncMock(return_value=cube), + ) + + resp = await client.post( + "/semantic/views/sem.no_such_view/values", + json={"dimension": "sem.region.region_name"}, + ) + + assert resp.status_code == 404, resp.text + assert "does not exist" in resp.json()["detail"] + + +@pytest.mark.asyncio +async def test_values_rejects_filter_outside_view( + client: AsyncClient, + monkeypatch, +): + fake_cube = SimpleNamespace( + current=SimpleNamespace( + cube_node_dimensions=["sem.region.region_name"], + ), + ) + monkeypatch.setattr( + "datajunction_server.api.semantic_layer.Node.get_cube_by_name", + AsyncMock(return_value=fake_cube), + ) + + resp = await client.post( + "/semantic/views/sem.sales_cube/values", + json={ + "dimension": "sem.region.region_name", + "filters": [ + { + "column": "sem.customer.country", + "operator": "=", + "value": "US", + }, + ], + }, + ) + + assert resp.status_code == 400, resp.text + assert "does not contain filter dimensions" in resp.json()["detail"] + + # --------------------------------------------------------------------------- # DJException handlers (the ``except DJException`` branches) — forced via # monkeypatch so the underlying call raises a DJException with a specific @@ -546,6 +718,27 @@ async def test_get_view_djexception_returns_problem( assert resp.json() == {"status_code": 422, "detail": "cube blew up"} +@pytest.mark.asyncio +async def test_values_cube_lookup_djexception_returns_problem( + client: AsyncClient, + monkeypatch, +): + monkeypatch.setattr( + "datajunction_server.api.semantic_layer.Node.get_cube_by_name", + AsyncMock( + side_effect=DJException(message="cube blew up", http_status_code=418), + ), + ) + + resp = await client.post( + "/semantic/views/some_view/values", + json={"dimension": "sem.region.region_name"}, + ) + + assert resp.status_code == 418, resp.text + assert resp.json() == {"status_code": 418, "detail": "cube blew up"} + + @pytest.mark.asyncio async def test_generate_sql_djexception_returns_problem( client: AsyncClient, @@ -584,3 +777,33 @@ async def test_generate_sql_djexception_returns_problem( ) assert resp.status_code == 400, resp.text assert resp.json() == {"status_code": 400, "detail": "sql gen blew up"} + + +@pytest.mark.asyncio +async def test_values_sql_generation_djexception_returns_problem( + client: AsyncClient, + monkeypatch, +): + fake_cube = SimpleNamespace( + current=SimpleNamespace( + cube_node_dimensions=["sem.region.region_name"], + ), + ) + monkeypatch.setattr( + "datajunction_server.api.semantic_layer.Node.get_cube_by_name", + AsyncMock(return_value=fake_cube), + ) + monkeypatch.setattr( + "datajunction_server.api.semantic_layer._generate_sql", + AsyncMock( + side_effect=DJException(message="values SQL blew up", http_status_code=422), + ), + ) + + resp = await client.post( + "/semantic/views/sem.sales_cube/values", + json={"dimension": "sem.region.region_name"}, + ) + + assert resp.status_code == 422, resp.text + assert resp.json() == {"status_code": 422, "detail": "values SQL blew up"} diff --git a/datajunction-server/tests/test_route_coverage.py b/datajunction-server/tests/test_route_coverage.py index b098453fa..3855acfce 100644 --- a/datajunction-server/tests/test_route_coverage.py +++ b/datajunction-server/tests/test_route_coverage.py @@ -111,6 +111,8 @@ def flatten(buckets: dict[str, list[tuple[str, str]]]) -> set[tuple[str, str]]: ("POST", "/semantic/views/list"), ("POST", "/semantic/views/{view_name}"), ("POST", "/semantic/views/{view_name}/sql"), + ("POST", "/semantic/views/{view_name}/row-count"), + ("POST", "/semantic/views/{view_name}/values"), ], "GraphQL; access enforced within resolvers, not the route": [ ("POST", "/graphql"),