opensource/WrenAI
3
0
Fork 0
mirror of https://github.com/Canner/WrenAI.git synced 2026-08-30 18:00:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(memory): avoid identifier columns in aggregation seed queries (#2358)

Browse Source
This commit is contained in:
Peter
2026-06-11 10:04:53 +08:00
committed by GitHub
parent 322bd798da
commit d5e78790b9
2 changed files with 463 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/wren/src/wren/memory/seed_queries.py
+76 -6
View File
@@ -7,6 +7,10 @@ from __future__ import annotations
from collections.abc import Sequence
import sqlglot
import sqlglot.errors
from sqlglot import exp
_NUMERIC_TYPES = {
"int",
"integer",
@@ -31,11 +35,16 @@ def generate_seed_queries(manifest: dict) -> list[dict]:
model["name"]: _prop_value(model, "dbtLayer", "dbt_layer")
for model in manifest.get("models", [])
}
relationship_keys = _relationship_key_columns(manifest)
for model in manifest.get("models", []):
if model_layers.get(model["name"]) == "raw":
continue
pairs.extend(_model_seeds(model))
pairs.extend(
_model_seeds(
model, relationship_keys.get(_norm_ident(model["name"]), frozenset())
)
)
for rel in manifest.get("relationships", []):
pair = _relationship_seed(rel, model_layers)
@@ -45,9 +54,12 @@ def generate_seed_queries(manifest: dict) -> list[dict]:
return pairs
def _model_seeds(model: dict) -> list[dict]:
def _model_seeds(
model: dict, relationship_keys: frozenset[str] = frozenset()
) -> list[dict]:
name = model["name"]
columns = model.get("columns", [])
primary_keys = _primary_key_columns(model)
pairs = []
# Template 1: basic listing
@@ -62,24 +74,29 @@ def _model_seeds(model: dict) -> list[dict]:
numeric_col = None
group_col = None
for col in columns:
col_name = col["name"]
norm_name = _norm_ident(col_name)
col_type = (col.get("type") or "").split("(")[0].lower().strip()
is_calc = col.get("isCalculated", False)
is_pk = col["name"] == model.get("primaryKey")
is_pk = norm_name in primary_keys
# Identifiers are numeric by storage but not measures: summing a join
# key (e.g. SUM(customer_id)) is semantically meaningless.
is_identifier = is_pk or norm_name in relationship_keys or _is_id_like(col_name)
if (
col_type in _NUMERIC_TYPES
and not is_calc
and not is_pk
and not is_identifier
and numeric_col is None
):
numeric_col = col["name"]
numeric_col = col_name
elif (
col_type not in _NUMERIC_TYPES
and not is_pk
and not is_calc
and group_col is None
):
group_col = col["name"]
group_col = col_name
# Template 2a: simple aggregation
if numeric_col:
@@ -132,6 +149,59 @@ def _relationship_seed(rel: dict, model_layers: dict[str, str]) -> dict | None:
}
def _relationship_key_columns(manifest: dict) -> dict[str, frozenset[str]]:
"""Map each model to the set of columns it exposes as a relationship key.
Relationship conditions (e.g. ``orders.customer_id = customers.customer_id``)
are the manifest's own declaration of join keys, so we keep both sides out
of aggregation seeds.
"""
accum: dict[str, set[str]] = {}
for rel in manifest.get("relationships", []):
condition = rel.get("condition") or ""
try:
tree = sqlglot.parse_one(condition)
except sqlglot.errors.SqlglotError:
continue
for col in tree.find_all(exp.Column):
if col.table and col.name:
accum.setdefault(_norm_ident(col.table), set()).add(
_norm_ident(col.name)
)
return {model: frozenset(cols) for model, cols in accum.items()}
def _primary_key_columns(model: dict) -> frozenset[str]:
"""Return primary key column names for string and composite-list PKs."""
primary_key = model.get("primaryKey")
if isinstance(primary_key, str):
return frozenset([_norm_ident(primary_key)])
if isinstance(primary_key, list):
return frozenset(_norm_ident(str(part)) for part in primary_key if part)
return frozenset()
def _norm_ident(name: str) -> str:
"""Canonicalize an identifier for case-insensitive membership checks.
Primary-key, relationship-key and ``*_id`` matching all compare against a
column name; normalizing keeps them consistent when a manifest mixes cases
(e.g. a condition referencing ``ORDERS.CUSTKEY`` while the column is
declared ``Custkey``). The original name is always used for generated SQL.
"""
return name.strip().lower()
def _is_id_like(col_name: str) -> bool:
"""Cheap heuristic for identifier columns not declared as relationships.
Case-insensitive: warehouses such as Snowflake/Oracle fold identifiers to
upper case, so an undeclared ``CUSTOMER_ID`` must be caught too.
"""
lowered = _norm_ident(col_name)
return lowered == "id" or lowered.endswith("_id")
def _prop_value(obj: dict, *keys: str) -> str:
props = obj.get("properties") or {}
if not isinstance(props, dict):
core/wren/tests/unit/test_seed_queries.py
+387 -2
View File
@@ -4,12 +4,16 @@ from __future__ import annotations
import pytest
from wren.memory.seed_queries import SEED_TAG, generate_seed_queries
from wren.memory.seed_queries import (
SEED_TAG,
_relationship_key_columns,
generate_seed_queries,
)
# ── Fixtures ──────────────────────────────────────────────────────────────
def _model(name: str, pk: str, columns: list[dict]) -> dict:
def _model(name: str, pk: str | list[str], columns: list[dict]) -> dict:
return {"name": name, "primaryKey": pk, "columns": columns}
@@ -381,3 +385,384 @@ class TestGenerateSeedQueries:
manifest = {"models": [_model("t", "id", [_col("id", "varchar")])]}
pairs = generate_seed_queries(manifest)
assert "LIMIT 100" in pairs[0]["sql"]
# ── relationship-key / identifier exclusion ───────────────────────────────
@pytest.mark.unit
class TestIdentifierExclusion:
"""Numeric identifier columns must not be picked as aggregation targets:
SUM(customer_id) is semantically meaningless noise (MEM-001)."""
def test_relationship_key_not_aggregated(self):
# customer_id is numeric and declared as a relationship key on orders.
manifest = {
"models": [
_model(
"orders",
"order_id",
[
_col("order_id", "int"),
_col("customer_id", "int"),
_col("amount", "double"),
_col("status", "varchar"),
],
),
_model(
"customers",
"customer_id",
[_col("customer_id", "int"), _col("name", "varchar")],
),
],
"relationships": [
{
"name": "orders_customers",
"models": ["orders", "customers"],
"condition": "orders.customer_id = customers.customer_id",
}
],
}
pairs = generate_seed_queries(manifest)
sqls = [p["sql"] for p in pairs]
# No relationship-key aggregation on either side of the join.
assert not any("SUM(customer_id)" in s for s in sqls)
# The real metric is aggregated instead
assert "SELECT SUM(amount) FROM orders" in sqls
assert "SELECT status, SUM(amount) FROM orders GROUP BY 1" in sqls
def test_composite_primary_key_columns_not_aggregated(self):
# Composite PKs are list-shaped in schema v4 MDL. Every PK member must
# be treated as an identifier even when it is numeric and not *_id-like.
manifest = {
"models": [
_model(
"store_sales",
["ss_item_sk", "ss_ticket_number"],
[
_col("ss_item_sk", "int"),
_col("ss_ticket_number", "int"),
_col("sales_price", "decimal"),
_col("status", "varchar"),
],
)
]
}
pairs = generate_seed_queries(manifest)
sqls = [p["sql"] for p in pairs]
assert not any("SUM(ss_item_sk)" in s for s in sqls)
assert not any("SUM(ss_ticket_number)" in s for s in sqls)
assert "SELECT SUM(sales_price) FROM store_sales" in sqls
assert "SELECT status, SUM(sales_price) FROM store_sales GROUP BY 1" in sqls
def test_id_like_column_not_aggregated_without_relationship(self):
# user_id is numeric, not a PK, and not declared in any relationship —
# the *_id naming heuristic should still exclude it.
manifest = {
"models": [
_model(
"raw_orders",
"id",
[
_col("id", "int"),
_col("user_id", "int"),
_col("amount", "double"),
],
)
]
}
pairs = generate_seed_queries(manifest)
sqls = [p["sql"] for p in pairs]
assert not any("SUM(user_id)" in s for s in sqls)
assert "SELECT SUM(amount) FROM raw_orders" in sqls
def test_column_named_id_not_aggregated(self):
# A bare numeric "id" that is not the declared primaryKey.
manifest = {
"models": [
_model(
"events",
"event_pk",
[
_col("event_pk", "varchar"),
_col("id", "bigint"),
_col("duration", "int"),
],
)
]
}
pairs = generate_seed_queries(manifest)
sqls = [p["sql"] for p in pairs]
assert not any("SUM(id)" in s for s in sqls)
assert "SELECT SUM(duration) FROM events" in sqls
def test_relationship_key_only_model_has_no_aggregation(self):
# When the sole numeric column is a relationship key, fall back to
# listing only.
manifest = {
"models": [
_model(
"order_items",
"line_id",
[
_col("line_id", "varchar"),
_col("order_id", "int"),
_col("label", "varchar"),
],
),
_model("orders", "order_id", [_col("order_id", "int")]),
],
"relationships": [
{
"name": "order_items_orders",
"models": ["order_items", "orders"],
"condition": "order_items.order_id = orders.order_id",
}
],
}
pairs = generate_seed_queries(manifest)
sqls = [p["sql"] for p in pairs]
assert not any("SUM(" in s for s in sqls)
def test_quoted_relationship_condition_parsed(self):
# Conditions may quote identifiers; relationship-key detection must
# still work.
manifest = {
"models": [
_model(
"orders",
"order_id",
[
_col("order_id", "int"),
_col("customer_id", "int"),
_col("amount", "double"),
],
),
_model("customers", "customer_id", [_col("customer_id", "int")]),
],
"relationships": [
{
"name": "orders_customers",
"models": ["orders", "customers"],
"condition": '"orders"."customer_id" = "customers"."customer_id"',
}
],
}
pairs = generate_seed_queries(manifest)
sqls = [p["sql"] for p in pairs]
assert not any("SUM(customer_id)" in s for s in sqls)
assert "SELECT SUM(amount) FROM orders" in sqls
def test_quoted_relationship_key_with_space_excluded(self):
# sqlglot parses quoted identifiers with spaces; the old regex-based
# extraction missed these relationship keys.
manifest = {
"models": [
_model(
"orders",
"order_pk",
[
_col("order_pk", "varchar"),
_col("Customer Key", "int"),
_col("amount", "double"),
],
),
_model(
"customers",
"customer_pk",
[
_col("customer_pk", "varchar"),
_col("Customer Key", "int"),
],
),
],
"relationships": [
{
"name": "orders_customers",
"models": ["orders", "customers"],
"condition": '"orders"."Customer Key" = "customers"."Customer Key"',
}
],
}
sqls = [p["sql"] for p in generate_seed_queries(manifest)]
assert not any("SUM(Customer Key)" in s for s in sqls)
assert "SELECT SUM(amount) FROM orders" in sqls
def test_invalid_relationship_condition_skipped_for_key_extraction(self):
manifest = {
"relationships": [
{
"name": "broken_relationship",
"models": ["orders", "customers"],
"condition": "orders.customer_id =",
},
{
"name": "valid_relationship",
"models": ["orders", "customers"],
"condition": "orders.customer_id = customers.customer_id",
},
]
}
assert _relationship_key_columns(manifest) == {
"orders": frozenset({"customer_id"}),
"customers": frozenset({"customer_id"}),
}
def test_legitimate_metric_named_with_id_suffix_is_still_excluded(self):
# Documented trade-off: the *_id heuristic also drops a would-be metric
# like "household_id". Real metrics should avoid the _id suffix; the
# noise reduction is worth this rare false-positive.
manifest = {
"models": [
_model(
"households",
"pk",
[_col("pk", "varchar"), _col("household_id", "int")],
)
]
}
pairs = generate_seed_queries(manifest)
assert not any("SUM(household_id)" in p["sql"] for p in pairs)
def test_uppercase_id_like_column_excluded(self):
# Warehouses such as Snowflake/Oracle fold identifiers to upper case.
# An undeclared CUSTOMER_ID (not a PK, not a relationship key) must
# still be caught by the case-insensitive *_id heuristic.
manifest = {
"models": [
_model(
"orders",
"ORDER_PK",
[
_col("ORDER_PK", "varchar"),
_col("CUSTOMER_ID", "int"),
_col("AMOUNT", "double"),
],
)
]
}
sqls = [p["sql"] for p in generate_seed_queries(manifest)]
assert not any("SUM(CUSTOMER_ID)" in s for s in sqls)
assert "SELECT SUM(AMOUNT) FROM orders" in sqls
def test_relationship_key_without_id_suffix_excluded(self):
# A foreign key whose name does NOT end in _id (e.g. created_by) is only
# caught via the relationship condition, not the naming heuristic — this
# is the value the relationship-key parsing adds on top of *_id.
manifest = {
"models": [
_model(
"documents",
"doc_pk",
[
_col("doc_pk", "varchar"),
_col("created_by", "int"),
_col("word_count", "int"),
],
),
_model("users", "id", [_col("id", "int")]),
],
"relationships": [
{
"name": "users_documents",
"models": ["users", "documents"],
"condition": "users.id = documents.created_by",
}
],
}
sqls = [p["sql"] for p in generate_seed_queries(manifest)]
assert not any("SUM(created_by)" in s for s in sqls)
assert "SELECT SUM(word_count) FROM documents" in sqls
def test_composite_and_condition_parsed(self):
# Composite join keys are emitted as an `AND`-joined condition; every
# referenced column on both sides must be excluded from aggregation.
manifest = {
"models": [
_model(
"a",
"akey",
[
_col("akey", "varchar"),
_col("x", "int"),
_col("y", "int"),
_col("val", "double"),
],
),
_model("b", ["x", "y"], [_col("x", "int"), _col("y", "int")]),
],
"relationships": [
{
"name": "a_b",
"models": ["a", "b"],
"condition": "a.x = b.x AND a.y = b.y",
}
],
}
sqls = [p["sql"] for p in generate_seed_queries(manifest)]
assert not any("SUM(x)" in s or "SUM(y)" in s for s in sqls)
assert "SELECT SUM(val) FROM a" in sqls
def test_identifier_matching_is_case_insensitive(self):
# A non-*_id identifier (e.g. "Custkey") is only excluded via the PK /
# relationship-key path. Those checks must be case-insensitive so a
# manifest that mixes cases — PK "custkey", column "Custkey", condition
# "ORDERS.CUSTKEY = ..." — still keeps it out of aggregation. The
# generated SQL must preserve the original column case.
manifest = {
"models": [
_model(
"orders",
"custkey", # PK declared lower-case
[
_col("Custkey", "int"), # column defined Title-case
_col("totalprice", "double"),
],
),
_model(
"customer",
"CUSTKEY",
[_col("CUSTKEY", "int"), _col("name", "varchar")],
),
],
"relationships": [
{
"name": "orders_customer",
"models": ["orders", "customer"],
"condition": "ORDERS.CUSTKEY = CUSTOMER.CUSTKEY",
}
],
}
sqls = [p["sql"] for p in generate_seed_queries(manifest)]
assert not any("SUM(Custkey)" in s for s in sqls)
assert "SELECT SUM(totalprice) FROM orders" in sqls
def test_relationship_key_case_insensitive_isolated(self):
# Isolates the relationship-key normalization path: "CreatedBy" is NOT a
# PK and NOT *_id-like, so only the (case-insensitive) relationship-key
# match can exclude it. The column is "CreatedBy" but the condition says
# "DOCUMENTS.CREATEDBY" — without normalization this leaks SUM(CreatedBy).
manifest = {
"models": [
_model(
"documents",
"doc_pk",
[
_col("doc_pk", "varchar"),
_col("CreatedBy", "int"),
_col("word_count", "int"),
],
),
_model("users", "user_pk", [_col("user_pk", "int")]),
],
"relationships": [
{
"name": "users_documents",
"models": ["users", "documents"],
"condition": "DOCUMENTS.CREATEDBY = USERS.USER_PK",
}
],
}
sqls = [p["sql"] for p in generate_seed_queries(manifest)]
assert not any("SUM(CreatedBy)" in s for s in sqls)
assert "SELECT SUM(word_count) FROM documents" in sqls