ktx/python/ktx-sl/tests/test_predefined_measures_only.py

149 lines
4.6 KiB
Python
Raw Normal View History

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 01:54:17 -07:00
"""predefined_measures_only rejects runtime-composed measures while leaving
predefined measures, predefined derived chains, dimensions, and filters usable."""
from __future__ import annotations
import pytest
from semantic_layer.engine import SemanticEngine
from semantic_layer.models import SourceDefinition
def _engine() -> SemanticEngine:
orders = SourceDefinition(
name="orders",
table="public.orders",
grain=["id"],
columns=[
{"name": "id", "type": "number"},
{"name": "amount", "type": "number"},
{"name": "status", "type": "string"},
],
measures=[
{"name": "revenue", "expr": "sum(amount)"},
{"name": "order_count", "expr": "count(*)"},
{"name": "aov", "expr": "revenue / order_count"},
],
)
return SemanticEngine.from_sources({"orders": orders})
def test_rejects_composed_string_measure() -> None:
with pytest.raises(ValueError, match="composed measure") as excinfo:
_engine().query(
{
"measures": ["sum(orders.amount)"],
"predefined_measures_only": True,
}
)
assert "sum(orders.amount)" in str(excinfo.value)
assert "query_policy: semantic-layer-only" in str(excinfo.value)
def test_rejects_composed_dict_measure() -> None:
with pytest.raises(ValueError, match="composed measure"):
_engine().query(
{
"measures": [{"expr": "avg(orders.amount)", "name": "avg_amount"}],
"predefined_measures_only": True,
}
)
def test_rejects_query_time_derivation_over_predefined_measures() -> None:
with pytest.raises(ValueError, match="composed measure"):
_engine().query(
{
"measures": [
{"expr": "orders.revenue / orders.order_count", "name": "ratio"}
],
"predefined_measures_only": True,
}
)
def test_rejects_composed_aggregate_in_filter() -> None:
# A HAVING-classified filter must not smuggle a runtime aggregate the
# measures guard would reject (threshold-probing bypass).
with pytest.raises(ValueError, match="compose aggregate expressions") as excinfo:
_engine().query(
{
"measures": ["orders.revenue"],
"dimensions": ["orders.status"],
"filters": ["avg(orders.amount) > 100"],
"predefined_measures_only": True,
}
)
assert "avg(orders.amount) > 100" in str(excinfo.value)
assert "query_policy: semantic-layer-only" in str(excinfo.value)
def test_rejects_composed_aggregate_in_compound_filter() -> None:
with pytest.raises(ValueError, match="compose aggregate expressions"):
_engine().query(
{
"measures": ["orders.revenue"],
"filters": ["orders.status = 'active' AND sum(orders.amount) > 5000"],
"predefined_measures_only": True,
}
)
def test_allows_predefined_measure_having_filter() -> None:
result = _engine().query(
{
"measures": ["orders.revenue"],
"dimensions": ["orders.status"],
"filters": ["orders.revenue > 100"],
"predefined_measures_only": True,
}
)
assert "having" in result.sql.lower()
def test_composed_aggregate_filter_allowed_when_flag_absent() -> None:
result = _engine().query(
{
"measures": ["orders.revenue"],
"filters": ["avg(orders.amount) > 100"],
}
)
assert "having" in result.sql.lower()
def test_allows_predefined_measure_with_dimensions_and_filters() -> None:
result = _engine().query(
{
"measures": ["orders.revenue"],
"dimensions": ["orders.status"],
"filters": ["orders.status != 'cancelled'"],
"predefined_measures_only": True,
}
)
assert "sum" in result.sql.lower()
def test_allows_unqualified_predefined_measure() -> None:
result = _engine().query(
{
"measures": ["revenue"],
"predefined_measures_only": True,
}
)
assert "sum" in result.sql.lower()
def test_allows_predefined_derived_measure_chain() -> None:
result = _engine().query(
{
"measures": ["orders.aov"],
"predefined_measures_only": True,
}
)
assert "sum" in result.sql.lower()
def test_composed_measures_allowed_when_flag_absent() -> None:
result = _engine().query({"measures": ["sum(orders.amount)"]})
assert "sum" in result.sql.lower()