ktx/python/ktx-daemon/tests/test_semantic_layer.py
Luca Martial a651b82e2f
feat: query_policy semantic-layer-only restricts agents to predefined semantic-layer measures (#334)
* feat(sl): add predefined_measures_only guard to semantic query planning

SemanticQuery gains a predefined_measures_only flag; the planner rejects
any measure resolved with Provenance.COMPOSED (runtime aggregate
expressions and query-time derivations) while predefined measures,
predefined derived chains, dimensions, filters, and segments pass.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* feat(config): add per-connection query_policy to warehouse connections

query_policy: semantic-layer-only | read-only-sql (default) on the
warehouse connection schema, plus a policy module with the raw-SQL
guard, federated member restriction lookup, and the project-level
predicate used to gate sql_execution registration.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* feat(cli): enforce query_policy on raw SQL through one shared executor

ktx sql and the MCP sql_execution tool now share executeProjectRawSql
(resolve, policy check, read-only validation, execute), collapsing
their duplicated validate-then-execute paths. Restricted connections
are rejected before validation; federated raw SQL is rejected when any
member is restricted. sql_execution is not registered when every SQL
connection is restricted, and connection_list marks restricted
connections so agents route to sl_query. executeProjectReadOnlySql
stays generic for ktx-internal SQL (scan, ingest, SL-generated).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* feat(sl): compile queries with predefined_measures_only from query_policy

compileLocalSlQuery injects the flag from the connection's query_policy,
never from caller input, covering both ktx sl query and the MCP
sl_query tool through the daemon compile path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* docs: document query_policy semantic-layer-only

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* fix(sl): close semantic-layer-only bypasses via filters and federated hint

The predefined_measures_only guard only inspected query.measures, so a
composed aggregate written into `filters` slipped through _classify_filters
into a HAVING clause untouched — letting a restricted agent evaluate
arbitrary aggregates (e.g. threshold-probing `sum(x) BETWEEN a AND b`).
Reject filter clauses that compose an aggregate function; a HAVING that
compares a predefined measure by name (`orders.revenue > 100`) still works.

Also make the federated sl_query error policy-aware: when a member is
restricted, raw federated SQL is disabled too, so stop directing the agent
to `ktx sql -c _ktx_federated` / sql_execution (a guaranteed failure) and
point to per-connection semantic-layer queries instead.

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Andrey Avtomonov <andreybavt@gmail.com>
2026-07-03 08:54:17 +00:00

180 lines
5.3 KiB
Python

from __future__ import annotations
import json
from pathlib import Path
import pytest
from ktx_daemon.semantic_layer import (
SemanticLayerQueryRequest,
ValidateSourcesRequest,
query_semantic_layer,
validate_semantic_layer,
)
ORDERS_SOURCE = {
"name": "orders",
"table": "public.orders",
"grain": ["id"],
"columns": [
{"name": "id", "type": "number"},
{"name": "status", "type": "string"},
{"name": "amount", "type": "number"},
],
"joins": [],
"measures": [
{"name": "order_count", "expr": "count(*)"},
{"name": "revenue", "expr": "sum(amount)"},
],
}
def test_query_semantic_layer_generates_sql_and_plan() -> None:
response = query_semantic_layer(
SemanticLayerQueryRequest(
sources=[ORDERS_SOURCE],
dialect="postgres",
query={
"measures": ["orders.order_count"],
"dimensions": ["orders.status"],
"limit": 25,
},
)
)
assert response.dialect == "postgres"
assert "public.orders" in response.sql
assert "orders.status" in response.sql
assert response.columns[0]["name"] == "orders.status"
assert response.columns[1]["name"] == "orders.order_count"
assert response.plan["sources_used"] == ["orders"]
def test_query_semantic_layer_enforces_predefined_measures_only() -> None:
with pytest.raises(ValueError, match="composed measure"):
query_semantic_layer(
SemanticLayerQueryRequest(
sources=[ORDERS_SOURCE],
dialect="postgres",
query={
"measures": ["sum(orders.amount)"],
"predefined_measures_only": True,
},
)
)
def test_query_semantic_layer_allows_predefined_measures_under_policy() -> None:
response = query_semantic_layer(
SemanticLayerQueryRequest(
sources=[ORDERS_SOURCE],
dialect="postgres",
query={
"measures": ["orders.revenue"],
"predefined_measures_only": True,
},
)
)
assert "public.orders" in response.sql
def test_query_semantic_layer_emits_plan_and_sql_debug_events(
tmp_path: Path,
monkeypatch,
capsys,
) -> None:
from ktx_daemon.telemetry.identity import reset_identity_cache
reset_identity_cache()
identity_path = tmp_path / ".ktx" / "telemetry.json"
identity_path.parent.mkdir(parents=True)
identity_path.write_text(
json.dumps(
{
"installId": "00000000-0000-4000-8000-000000000000",
"enabled": True,
"createdAt": "2026-05-22T14:33:02.000Z",
}
)
+ "\n",
encoding="utf-8",
)
monkeypatch.setenv("HOME", str(tmp_path))
monkeypatch.setenv("KTX_TELEMETRY_DEBUG", "1")
monkeypatch.delenv("CI", raising=False)
monkeypatch.delenv("KTX_TELEMETRY_DISABLED", raising=False)
monkeypatch.delenv("DO_NOT_TRACK", raising=False)
query_semantic_layer(
SemanticLayerQueryRequest(
sources=[ORDERS_SOURCE],
dialect="postgres",
projectId="a" * 64,
query={
"measures": ["orders.order_count"],
"dimensions": ["orders.status"],
"limit": 25,
},
)
)
captured = capsys.readouterr()
assert '"event": "sl_plan_completed"' in captured.err
assert '"event": "sql_gen_completed"' in captured.err
assert "public.orders" not in captured.err
def test_query_semantic_layer_reports_exception(monkeypatch) -> None:
from ktx_daemon import semantic_layer as semantic_layer_module
reports: list[dict[str, object]] = []
def fake_report(exception: BaseException, **kwargs: object) -> None:
reports.append({"exception": exception, **kwargs})
monkeypatch.setattr(semantic_layer_module, "report_exception", fake_report)
with pytest.raises(ValueError):
query_semantic_layer(
SemanticLayerQueryRequest(
sources=[ORDERS_SOURCE, ORDERS_SOURCE],
dialect="postgres",
projectId="a" * 64,
query={"measures": ["orders.order_count"]},
)
)
assert reports
assert reports[0]["source"] == "semantic-query"
assert reports[0]["handled"] is True
assert reports[0]["fatal"] is False
assert reports[0]["project_id"] == "a" * 64
def test_semantic_layer_request_rejects_project_id_field_name() -> None:
with pytest.raises(ValueError):
SemanticLayerQueryRequest(
sources=[],
dialect="postgres",
project_id="a" * 64,
query={"measures": ["orders.order_count"]},
)
def test_validate_semantic_layer_reports_duplicate_measure_names() -> None:
invalid_source = {
**ORDERS_SOURCE,
"measures": [
{"name": "revenue", "expr": "sum(amount)"},
{"name": "revenue", "expr": "sum(amount)"},
],
}
response = validate_semantic_layer(
ValidateSourcesRequest(sources=[invalid_source], dialect="postgres")
)
assert response.valid is False
assert any("Duplicate measure" in error for error in response.errors)
assert response.warnings == []