diff --git a/src/boring_semantic_layer/calc_compiler.py b/src/boring_semantic_layer/calc_compiler.py index 0c99882f..2900652f 100644 --- a/src/boring_semantic_layer/calc_compiler.py +++ b/src/boring_semantic_layer/calc_compiler.py @@ -281,6 +281,30 @@ def __getitem__(self, name: str): return self._virtual_agg_tbl[resolved] return self._base_tbl[name] + def _rebind_to_totals(self, x: Any): + """Re-point every measure reference inside *x* at the totals table. + + Returns ``None`` when *x* contains no measure references, in which + case the caller falls back to the reduction/windowed handling — + ``t.all(t.raw_column)`` has no measure formula to re-apply. + """ + try: + op = _to_op(x) + vt_op = _to_op(self._virtual_agg_tbl) + totals_op = _to_op(self._totals_virtual_agg_tbl) + referenced = { + node.name + for node in _walk(op) + if isinstance(node, Field) and id(node.rel) == id(vt_op) + } + if not referenced: + return None + subs = {Field(vt_op, name): Field(totals_op, name) for name in referenced} + return op.replace(subs).to_expr() + except Exception as exc: + logger.debug("IbisCalcScope.all() totals rebind swallowed: %s", exc) + return None + def all(self, x: Any): """Resolve a measure reference to its totals-table column. @@ -329,6 +353,16 @@ def all(self, x: Any): if isinstance(op, Field) and id(op.rel) == id(_to_op(self._virtual_agg_tbl)): return self._totals_virtual_agg_tbl[op.name] + # An expression *built from* measure references is the same + # request, one level up: evaluate it against the totals table + # rather than the per-group one. Falling through to the + # windowed-sum shape below made ``t.all(m)`` and + # ``t.all(m * 1)`` mean different things — for a mean measure + # the second silently summed the per-group means. + over_totals = self._rebind_to_totals(x) + if over_totals is not None: + return over_totals + Reduction = getattr(ibis_ops, "Reduction", None) if Reduction is not None: if isinstance(op, Reduction): diff --git a/src/boring_semantic_layer/expr.py b/src/boring_semantic_layer/expr.py index f28b675f..a1d5c5b1 100644 --- a/src/boring_semantic_layer/expr.py +++ b/src/boring_semantic_layer/expr.py @@ -263,12 +263,23 @@ def mutate(self, **post): chain order and materialized, preserving the historical ``.mutate()`` semantics without a dedicated operator node. """ - from .ops import SemanticJoinOp, SemanticTableOp, _has_prior_aggregate, _resolve_expr + from .ops import ( + SemanticJoinOp, + SemanticTableOp, + _has_prior_aggregate, + _non_additive_result_columns, + _resolve_expr, + ) if _has_prior_aggregate(self.op()): tbl = self.op().to_untagged() + # Only the aggregated rows are available here, so t.all() can only + # window-sum them; pass which columns that would misrepresent. + non_additive = _non_additive_result_columns(self.op()) for name, fn in post.items(): - proxy = MeasureScope(_tbl=tbl, _known=[], _post_agg=True) + proxy = MeasureScope( + _tbl=tbl, _known=[], _post_agg=True, _non_additive=non_additive + ) resolved = _resolve_expr(fn, proxy) tbl = tbl.mutate(resolved.name(name)) return _build_post_aggregate_model(self.op(), tbl) diff --git a/src/boring_semantic_layer/measure_scope.py b/src/boring_semantic_layer/measure_scope.py index b8f0e08d..ef9161c4 100644 --- a/src/boring_semantic_layer/measure_scope.py +++ b/src/boring_semantic_layer/measure_scope.py @@ -117,6 +117,15 @@ class MeasureScope: converter=tuple, alias="_prefer_known", ) + #: Result columns whose window sum is *not* their overall value (means, + #: medians, distinct counts, calc measures). ``t.all()`` refuses these + #: rather than returning a sum with no meaning. Empty means "not + #: classified" — the historical behaviour applies. + non_additive: frozenset[str] = field( + factory=frozenset, + converter=frozenset, + alias="_non_additive", + ) def __attrs_post_init__(self): object.__setattr__(self, "known_set", frozenset(self.known)) @@ -166,17 +175,68 @@ def all(self, ref): from ._xorq import ibis as ibis_mod if isinstance(ref, str): - return self.tbl[ref].sum().over(ibis_mod.window()) + self._reject_non_additive(ref) + return _float_total(self.tbl[ref].sum().over(ibis_mod.window())) if hasattr(ref, "__class__") and "ibis" in str(type(ref).__module__): if "Scalar" in type(ref).__name__: return ref.over(ibis_mod.window()) - return ref.sum().over(ibis_mod.window()) + self._reject_non_additive(_column_name_of(ref)) + return _float_total(ref.sum().over(ibis_mod.window())) raise TypeError( "t.all(...) expects a string column name or an ibis expression", ) + def _reject_non_additive(self, name: str | None) -> None: + """Refuse a total this scope cannot compute correctly. + + Post-aggregation, the only rows available are the grouped ones, so + the total can only be a window sum over them. For a mean, median, + distinct count or ratio that sum is not the overall value, and + returning it silently produced answers that disagreed with the same + formula written as a calc measure. + """ + if name is None or name not in self.non_additive: + return + raise NonAdditiveTotalError( + f"t.all({name!r}) cannot be computed after aggregation: {name!r} is " + "not additive, so summing its per-group values is not its overall " + "value. Define the derivation on the model instead, where the total " + "is computed from the underlying rows:\n" + f" model.with_measures(share=lambda t: t.{name} / t.all(t.{name}))\n" + "or place the .mutate() directly on the aggregate (before " + "filter/order_by/limit), which routes through the same path." + ) + + +class NonAdditiveTotalError(ValueError): + """``t.all()`` was asked for a total that post-aggregation rows can't give.""" + + +def _column_name_of(ref) -> str | None: + """Best-effort column name for an ibis value expression.""" + try: + name = ref.get_name() + except Exception: + return None + return name if isinstance(name, str) else None + + +def _float_total(total): + """Cast an integral total to float so ``measure / total`` is a true ratio. + + The calc-measure path gives its virtual columns a float64 schema for the + same reason: with two integer operands some engines (xorq's DataFusion + among them) do integer division, and ``30 / 160`` came back as ``0``. + """ + try: + if total.type().is_integer(): + return total.cast("float64") + except Exception: + pass + return total + @frozen(kw_only=True, slots=True) class ColumnScope: @@ -206,12 +266,12 @@ def all(self, ref): from ._xorq import ibis as ibis_mod if isinstance(ref, str): - return self.tbl[ref].sum().over(ibis_mod.window()) + return _float_total(self.tbl[ref].sum().over(ibis_mod.window())) if hasattr(ref, "__class__") and "ibis" in str(type(ref).__module__): if "Scalar" in type(ref).__name__: return ref.over(ibis_mod.window()) - return ref.sum().over(ibis_mod.window()) + return _float_total(ref.sum().over(ibis_mod.window())) raise TypeError( "t.all(...) expects a string column name or an ibis expression", diff --git a/src/boring_semantic_layer/ops.py b/src/boring_semantic_layer/ops.py index f32f0873..2d898ebd 100644 --- a/src/boring_semantic_layer/ops.py +++ b/src/boring_semantic_layer/ops.py @@ -338,18 +338,43 @@ def _ensure_xorq_table(table): return table +def _connection_identity(backend) -> tuple: + """Identify the physical connection a backend reads from. + + ``from_ibis()`` mints a fresh ``Backend`` wrapper per call but reuses the + caller's DBAPI connection object, so wrapper identity says nothing about + which database a table lives in while connection identity does. Backends + that expose no ``con`` fall back to their own identity, which makes + rebinding a no-op for them rather than a guess. + """ + con = getattr(backend, "con", None) + if con is None: + return ("backend", id(backend)) + return ("con", id(con)) + + def _rebind_to_backend(expr, target_backend): - """Rebind every ``DatabaseTable`` op in *expr* to *target_backend*. + """Rebind ``DatabaseTable`` ops in *expr* that share *target_backend*'s connection. Low-level primitive shared with ``serialization.reconstruct``. No-op on plain ibis expressions or when xorq is unavailable for any reason; callers must pass a xorq-vendored ``target_backend``. + + Tables belonging to a *different* connection are left alone. Rebinding + those used to repoint them at the canonical backend without checking, + so joining two same-schema databases (prod and staging, or two shards) + silently read every column from whichever one happened to be first. + Leaving them untouched lets the engine raise its own multiple-backends + error instead of returning plausible numbers from the wrong database. """ try: from ._xorq import relations as xorq_rel except Exception: return expr + target_identity = _connection_identity(target_backend) + foreign: set[str] = set() + def _recreate(op, _kwargs, **overrides): kwargs = dict(zip(op.__argnames__, op.__args__, strict=False)) if _kwargs: @@ -359,12 +384,23 @@ def _recreate(op, _kwargs, **overrides): def replacer(op, _kwargs): if isinstance(op, xorq_rel.DatabaseTable) and op.source is not target_backend: - return _recreate(op, _kwargs, source=target_backend) + if _connection_identity(op.source) == target_identity: + return _recreate(op, _kwargs, source=target_backend) + foreign.add(op.name) if _kwargs: return _recreate(op, _kwargs) return op - return expr.op().replace(replacer).to_expr() + rebound = expr.op().replace(replacer).to_expr() + if foreign: + logger.warning( + "Expression spans more than one database connection; tables %s were " + "left bound to their own backend. A query mixing them will fail in " + "the engine — read them through a single connection (or ATTACH one " + "database to the other) if they are meant to be joined.", + sorted(foreign), + ) + return rebound def _rebind_to_canonical_backend(expr): @@ -5008,9 +5044,39 @@ def _to_untagged_with_preagg( empty_count_measures = tuple(_empty_count_measures) if not preagg_results and not _deferred_count_distincts: - if tbl is not None: - return tbl.aggregate({n: f(tbl) for n, f in plan.agg_specs.items()}) - raise ValueError("No aggregation results and full join unavailable") + if tbl is None: + raise ValueError("No aggregation results and full join unavailable") + # Nothing could be pre-aggregated at a source grain. This fallback + # used to aggregate the flattened join while ignoring both the + # group keys and every calc spec: with only calc measures + # requested it returned ``tbl.aggregate({})`` — an Aggregate with + # no columns at all, which surfaces much later as an unrelated + # arrow/schema error instead of naming the problem. + if plan.calc_specs: + raise ValueError( + "Pre-aggregation cannot compute calculated measure(s) " + f"{sorted(plan.calc_specs)} on this joined model: none of " + "the requested measures aggregate at a single source's " + "grain, so there is no fan-out-safe base to calculate " + "from. This happens when a calc measure builds its " + "reduction inline — e.g. " + "t.distance.sum() / t.all(t.distance.sum()) — rather than " + "referencing a declared measure. Declare the reduction as " + "a measure and reference it by name:\n" + " .with_measures(total=lambda t: t.distance.sum())\n" + " .with_measures(share=lambda t: t.total / t.all(t.total))" + ) + if not plan.agg_specs: + raise ValueError( + f"Pre-aggregation produced no measures for {sorted(self.aggs)} " + f"with group keys {list(plan.group_by_cols)}; aggregating the " + "joined table here would ignore the request entirely." + ) + specs = {n: f(tbl) for n, f in plan.agg_specs.items()} + group_cols = [c for c in plan.group_by_cols if c in tbl.columns] + if group_cols: + return tbl.group_by(group_cols).aggregate(**specs) + return tbl.aggregate(specs) # --- 5. Combine pre-agg results --- result = None @@ -6558,25 +6624,13 @@ def _rebind_join_backends(left_tbl, right_tbl): if canonical is None: return left_tbl, right_tbl - def _recreate(op, _kwargs, **overrides): - kwargs = dict(zip(op.__argnames__, op.__args__, strict=False)) - if _kwargs: - kwargs.update(_kwargs) - kwargs.update(overrides) - return op.__recreate__(kwargs) - - def replacer(op, _kwargs): - if isinstance(op, xorq_rel.DatabaseTable) and op.source is not canonical: - return _recreate(op, _kwargs, source=canonical) - # Propagate rewritten children (e.g. SelfReference wrapping - # a replaced DatabaseTable). - if _kwargs: - return _recreate(op, _kwargs) - return op - - new_left = left_tbl.op().replace(replacer).to_expr() - new_right = right_tbl.op().replace(replacer).to_expr() - return new_left, new_right + # Shared primitive: only tables on the same physical connection are + # rebound, so a join across two distinct databases fails in the + # engine rather than silently reading both sides from one of them. + return ( + _rebind_to_backend(left_tbl, canonical), + _rebind_to_backend(right_tbl, canonical), + ) def execute(self): return _rebind_to_canonical_backend(self.to_untagged()).execute() @@ -7047,6 +7101,77 @@ def _find_all_root_models(node: Any) -> tuple[SemanticTableOp, ...]: return roots +def _non_additive_result_columns(node: Any) -> frozenset[str]: + """Result columns of a prior aggregate that must not be summed to get a total. + + A post-aggregation ``.mutate()`` only sees the aggregated rows, so its + ``t.all(x)`` can only be a window sum over those rows. That equals the + true overall value for SUM/COUNT measures and nothing else: summing + per-group means, medians, min/max or distinct counts gives a number with + no meaning, which is what ``t.all()`` used to return silently. + + Classification resolves each measure against its root's raw table, which + builds an expression but compiles nothing. Measures that cannot be + classified are omitted rather than assumed non-additive — callers keep + their historical behaviour for those instead of failing on a guess. + """ + current = node + agg_op = None + while current is not None: + if isinstance(current, SemanticAggregateOp): + agg_op = current + break + current = getattr(current, "source", None) + if agg_op is None: + return frozenset() + + try: + roots = _find_all_root_models(agg_op.source) + if not roots: + return frozenset() + merged_base = _get_merged_fields(roots, "measures") + merged_calc = _get_merged_fields(roots, "calc_measures") + probes = [] + for root in roots: + raw = getattr(root, "table", None) + if raw is None: + continue + probes.append(raw.to_expr() if hasattr(raw, "to_expr") else raw) + except Exception as exc: + logger.debug("additivity classification unavailable: %s", exc) + return frozenset() + + non_additive: set[str] = set() + for name in agg_op.aggs: + resolved = _resolve_short_name(name, merged_base, merged_calc) + if resolved is None: + continue + if resolved in merged_calc: + # A calculated measure is a ratio/window expression; summing it + # across groups is never the overall value. + non_additive.add(name) + continue + measure = merged_base.get(resolved) + expr = None + for probe in probes: + try: + expr = _resolve_expr(getattr(measure, "expr", measure), probe) + break + except Exception: + continue + if expr is None: + continue + try: + if _is_mean_expr(expr) or _reagg_op_for_expr(expr) != "sum": + non_additive.add(name) + except Exception as exc: + # _reagg_op_for_expr raises "this is a bug" for undecomposed + # mean / undeferred count-distinct: both are non-additive. + logger.debug("treating %r as non-additive: %s", name, exc) + non_additive.add(name) + return frozenset(non_additive) + + def _has_prior_aggregate(node: Any) -> bool: """True when a SemanticAggregateOp sits beneath ``node`` in the chain. diff --git a/src/boring_semantic_layer/serialization/context.py b/src/boring_semantic_layer/serialization/context.py index 7b69ee6a..a66be067 100644 --- a/src/boring_semantic_layer/serialization/context.py +++ b/src/boring_semantic_layer/serialization/context.py @@ -57,7 +57,7 @@ def parse_field(self, metadata: dict, field: str) -> dict | list: value = metadata.get(field) if not value: return {} if field != "order_keys" else [] - return thaw(value) + return thaw(value, key=field) def parse_structured_dict(self, raw: Any) -> dict: """Convert a FrozenOrderedDict-encoded tuple-of-pairs to a dict (one level). diff --git a/src/boring_semantic_layer/serialization/extract.py b/src/boring_semantic_layer/serialization/extract.py index 6f501802..9016ee83 100644 --- a/src/boring_semantic_layer/serialization/extract.py +++ b/src/boring_semantic_layer/serialization/extract.py @@ -146,6 +146,7 @@ def _extract_group_by(op, context: BSLSerializationContext) -> dict[str, Any]: @_register_lazy("SemanticAggregateOp") def _extract_aggregate(op, context: BSLSerializationContext) -> dict[str, Any]: + from ..ops import _detect_bare_name_lambda, _unwrap from ..utils import expr_to_structured metadata: dict[str, Any] = {} @@ -155,6 +156,21 @@ def _extract_aggregate(op, context: BSLSerializationContext) -> dict[str, Any]: metadata["aggs_struct"] = { name: expr_to_structured(fn).value_or(None) for name, fn in op.aggs.items() } + # ``aggregate("revenue")`` and ``aggregate(revenue=lambda t: ...)`` + # both land in ``aggs`` and can serialize to similar-looking trees, + # but they mean different things: the first must replay through + # measure resolution (which is what makes it fan-out safe on a + # joined model), the second is a query-local expression that must + # be rebuilt verbatim. Record which is which instead of guessing + # from the name on the way back in. + # Always emitted, including empty: its absence is what tells the + # reader this payload predates the marker and needs the structural + # fallback in ``_bare_ref_names``. + metadata["agg_bare_refs"] = sorted( + name + for name, fn in op.aggs.items() + if _detect_bare_name_lambda(_unwrap(fn)) is not None + ) return metadata diff --git a/src/boring_semantic_layer/serialization/freeze.py b/src/boring_semantic_layer/serialization/freeze.py index 738a98da..b650569b 100644 --- a/src/boring_semantic_layer/serialization/freeze.py +++ b/src/boring_semantic_layer/serialization/freeze.py @@ -3,34 +3,80 @@ xorq tag metadata stores dicts as tuples-of-pairs and lists as tuples (FrozenOrderedDict). These utilities convert between mutable Python types and the frozen representation. + +Two invariants keep the round-trip lossless: + +1. ``freeze`` never coerces a value it cannot represent. Anything that is + not a scalar/list/dict raises, because the previous ``str(obj)`` + fallback silently turned dates, ``Decimal``s and ``bytes`` into strings + that came back as the wrong type (or failed a type check much later, in + the query compiler). Non-scalar constants inside expressions are + encoded by ``utils._encode_scalar`` *before* they reach ``freeze``. + +2. ``thaw`` never recurses into a serialized resolver tree. Resolver nodes + such as ``("just", 0)`` are indistinguishable from the tuple-of-pairs + encoding of a dict, so thawing them collapsed multi-argument calls into + their last argument (``substr(0, 2)`` became ``substr(2)``). Values + stored under a ``*_struct`` key are therefore passed through verbatim. """ from __future__ import annotations from typing import Any +#: Metadata keys whose values are serialized resolver trees. ``thaw`` must +#: hand these back exactly as written — see invariant 2 above. +OPAQUE_STRUCT_KEYS = frozenset( + { + "expr_struct", + "predicate_struct", + "aggs_struct", + "on_struct", + "value_struct", + "post_struct", + } +) + + +class FreezeError(TypeError): + """A value in tag metadata cannot be frozen without losing information.""" -def freeze(obj: Any) -> Any: + +def freeze(obj: Any, *, path: str = "metadata") -> Any: """Recursively convert dicts to tuples-of-pairs and lists to tuples. Scalar types (str, int, float, bool, None) pass through unchanged. - Anything else is converted to ``str(obj)``. + + Raises: + FreezeError: If *obj* contains a value with no lossless frozen + representation. """ - if isinstance(obj, str | int | float | bool | type(None)): + if isinstance(obj, str | bool | int | float | type(None)): return obj if isinstance(obj, dict): - return tuple((k, freeze(v)) for k, v in obj.items()) + return tuple((k, freeze(v, path=f"{path}.{k}")) for k, v in obj.items()) if isinstance(obj, list | tuple): - return tuple(freeze(item) for item in obj) - return str(obj) + return tuple(freeze(item, path=f"{path}[{i}]") for i, item in enumerate(obj)) + raise FreezeError( + f"Cannot serialize {path}: {type(obj).__name__} has no lossless " + f"representation in xorq tag metadata (value: {obj!r}). Expression " + "constants of this type must be encoded by " + "boring_semantic_layer.utils._encode_scalar before reaching freeze()." + ) -def thaw(obj: Any) -> Any: +def thaw(obj: Any, *, key: str | None = None) -> Any: """Recursively convert frozen tuples back to mutable dicts/lists. A tuple is treated as a dict if every element is a 2-tuple with a str key. Otherwise it is treated as a list. + + Values reached under one of :data:`OPAQUE_STRUCT_KEYS` are returned + verbatim: they are resolver trees, whose ``("just", x)`` nodes would + otherwise be misread as dict entries. """ + if key in OPAQUE_STRUCT_KEYS: + return obj if isinstance(obj, tuple): if len(obj) == 0: return {} @@ -38,7 +84,7 @@ def thaw(obj: Any) -> Any: isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str) for item in obj ): - return {k: thaw(v) for k, v in obj} + return {k: thaw(v, key=k) for k, v in obj} return [thaw(item) for item in obj] return obj @@ -59,7 +105,7 @@ def thaw_shallow(obj: Any) -> dict: for item in obj ) ): - return {k: v for k, v in obj} + return dict(obj) return {} @@ -67,7 +113,9 @@ def list_to_tuple(obj: Any) -> Any: """Recursively convert lists back to tuples. Reverses ``thaw`` for structured expression data that needs to stay - as tuples for the resolver deserialization layer. + as tuples for the resolver deserialization layer. Still needed for + payloads written before ``OPAQUE_STRUCT_KEYS`` existed, and for + struct data that arrives through other paths (e.g. YAML). """ if isinstance(obj, list): return tuple(list_to_tuple(item) for item in obj) diff --git a/src/boring_semantic_layer/serialization/reconstruct.py b/src/boring_semantic_layer/serialization/reconstruct.py index 1e8b51ee..de688dec 100644 --- a/src/boring_semantic_layer/serialization/reconstruct.py +++ b/src/boring_semantic_layer/serialization/reconstruct.py @@ -15,6 +15,7 @@ from .context import BSLSerializationContext from .extract import deserialize_calc_measures +from .freeze import thaw # --------------------------------------------------------------------------- # Registry @@ -54,7 +55,16 @@ def _create_dimension(name: str, dim_data: dict) -> ops.Dimension: elif expr_struct is not None: expr = context.deserialize_expr(expr_struct, f"Dimension '{name}'") else: - expr = lambda t, n=name: t[n] # noqa: E731 + # Every serialized dimension carries either a column name or a + # resolver tree. Falling back to ``t[name]`` here turned an + # unreadable payload (a v1.0 pickle field, a future encoding) + # into a raw column silently: ``amount = _.amount * 1.1`` came + # back as ``amount``, with no error and plausible numbers. + raise ValueError( + f"Dimension {name!r} has no readable expression in this payload " + f"(keys: {sorted(dim_data)}). It was written by an incompatible " + "version of boring-semantic-layer — re-serialize the model." + ) return ops.Dimension( expr=expr, description=dim_data.get("description"), @@ -219,32 +229,70 @@ def _reconstruct_aggregate( if not aggs_struct: raise ValueError("SemanticAggregateOp has no aggs_struct") - # Model-declared measures replay by name; query-local entries (e.g. - # derivations folded in by ``.mutate()``) are not on the model, so - # rebuild their expressions from the serialized resolver structs. - known: set[str] = set() - source_op = source.op() - for getter in ("get_measures", "get_calculated_measures"): - with suppress(Exception): - known |= set(getattr(source_op, getter)().keys()) + # Entries the query referenced by bare measure name replay by name, so + # they route back through measure resolution (fan-out-safe pre-aggregation + # on joined models). Everything else is a query-local expression and is + # rebuilt from its serialized resolver tree — replaying *those* by name + # silently substitutes a same-named model measure for the user's + # expression, e.g. ``aggregate(n=lambda t: t.a.max())`` returning the + # model's ``n = a.sum()``. + bare_refs = _bare_ref_names(metadata, aggs_struct, source) names: list[str] = [] aliased: dict = {} for name, data in aggs_struct.items(): - # Mirror _resolve_short_name semantics: a bare name replays only - # when it matches a model measure exactly or by a UNIQUE suffix. - # Ambiguous suffixes fall through to struct deserialization, - # which rebuilds the exact expression. - is_known = ( - name in known - or sum(1 for k in known if k.endswith(f".{name}")) == 1 - ) - if is_known or data is None: + if name in bare_refs or data is None: names.append(name) else: aliased[name] = context.deserialize_expr(data, f"Aggregate({name})") return source.aggregate(*names, **aliased) + +def _bare_ref_names(metadata: dict, aggs_struct: dict, source) -> set[str]: + """Names in ``aggs_struct`` that were written as bare measure references. + + Payloads written by current BSL carry ``agg_bare_refs`` explicitly. For + older payloads, recover the distinction structurally: a bare reference + serializes to exactly the tree of ``make_bare_ref_lambda(name)``, so + comparing against a freshly built one separates the two cases without + consulting the model's measure names. + """ + if "agg_bare_refs" in metadata: + declared = thaw(metadata["agg_bare_refs"]) + return {n for n in declared if isinstance(n, str)} + + from ..ops import make_bare_ref_lambda + from ..utils import expr_to_structured + + from .freeze import list_to_tuple + + known: set[str] = set() + + def _looks_like_model_measure(name: str) -> bool: + # Historical heuristic, kept only for entries this function cannot + # classify structurally (see below). + nonlocal known + if not known: + source_op = source.op() + for getter in ("get_measures", "get_calculated_measures"): + with suppress(Exception): + known |= set(getattr(source_op, getter)().keys()) + return name in known or sum(1 for k in known if k.endswith(f".{name}")) == 1 + + recovered: set[str] = set() + for name, data in aggs_struct.items(): + if data is None: + # Nothing to rebuild from — name replay is the only option. + recovered.add(name) + continue + canonical = expr_to_structured(make_bare_ref_lambda(name)).value_or(None) + if canonical is not None: + if canonical == list_to_tuple(data): + recovered.add(name) + elif _looks_like_model_measure(name): + recovered.add(name) + return recovered + @register_reconstructor("SemanticProjectOp") def _reconstruct_project( metadata: dict, xorq_expr, source, context: BSLSerializationContext @@ -494,6 +542,29 @@ def is_bsl_tag(op) -> bool: # --------------------------------------------------------------------------- +#: Payload format versions this build knows how to read. ``bsl_version`` was +#: written from the start but never checked, so a v1.0 tag (whose expressions +#: were pickled — a format no longer read at all) used to load as a model with +#: silently degraded fields instead of failing. +SUPPORTED_PAYLOAD_MAJORS = frozenset({2}) + + +def _check_payload_version(metadata: dict[str, Any]) -> None: + """Refuse a payload written by an incompatible serializer version.""" + version = metadata.get("bsl_version") + if version is None: + # Metadata assembled in-process (tests, direct reconstructor calls) + # carries no version; only tagged payloads are gated. + return + major = str(version).split(".", 1)[0] + if not major.isdigit() or int(major) not in SUPPORTED_PAYLOAD_MAJORS: + raise ValueError( + f"Cannot read BSL payload version {version!r}: this build reads " + f"major version(s) {sorted(SUPPORTED_PAYLOAD_MAJORS)}. Re-serialize " + "the model with a matching boring-semantic-layer version." + ) + + def reconstruct_bsl_operation( metadata: dict[str, Any], xorq_expr, @@ -504,6 +575,7 @@ def reconstruct_bsl_operation( Walks the metadata tree recursively, dispatching to registered reconstructors by ``bsl_op_type``. """ + _check_payload_version(metadata) op_type = metadata.get("bsl_op_type") source = None source_metadata = context.parse_field(metadata, "source") diff --git a/src/boring_semantic_layer/tests/test_serialization_trust_boundary.py b/src/boring_semantic_layer/tests/test_serialization_trust_boundary.py new file mode 100644 index 00000000..09404e61 --- /dev/null +++ b/src/boring_semantic_layer/tests/test_serialization_trust_boundary.py @@ -0,0 +1,272 @@ +"""Regression tests for the serialization trust boundary and lossy encodings. + +A serialized model is data that travels: xorq's ``from_tag_node`` entry point +(pyproject.toml) routes any ``bsl``-tagged expression here automatically, and +git catalogs store the resolver trees as editable YAML. Everything in this +file is about a payload that is not necessarily written by whoever reads it, +or about an encoding that used to lose information on the way through. + +Each test below corresponds to a defect that produced either arbitrary code +execution or a silently wrong number through the public +``to_tagged``/``from_tagged`` API. +""" + +from __future__ import annotations + +import datetime +import decimal +import os + +import ibis +import pytest +from returns.result import Failure + +from boring_semantic_layer import to_semantic_table +from boring_semantic_layer.serialization import from_tagged, to_tagged +from boring_semantic_layer.serialization.context import BSLSerializationContext +from boring_semantic_layer.serialization.reconstruct import reconstruct_bsl_operation +from boring_semantic_layer.utils import ( + UntrustedCallableError, + serialize_resolver, + structured_to_expr, +) + +xorq = pytest.importorskip("xorq", reason="xorq not installed") + + +@pytest.fixture +def table(): + return ibis.memtable( + { + "a": [1, 2, 3], + "b": ["abcd", "efgh", "zzzz"], + "d": [ + datetime.date(2024, 1, 1), + datetime.date(2024, 6, 1), + datetime.date(2025, 1, 1), + ], + } + ) + + +def _refusal(struct): + """Deserialize *struct* and return the error it was refused with.""" + result = structured_to_expr(struct) + assert isinstance(result, Failure), f"payload was accepted: {result}" + return result.failure() + + +# --------------------------------------------------------------------------- +# Arbitrary code execution +# --------------------------------------------------------------------------- + + +def test_call_gadget_is_refused(tmp_path): + """``Call.resolve()`` invokes its func, so a Just(callable) is a call gadget.""" + marker = tmp_path / "pwned" + payload = ( + "call", + ("fn", "builtins", "eval"), + (("just", f"__import__('pathlib').Path({str(marker)!r}).touch()"),), + (), + ) + assert isinstance(_refusal(payload), UntrustedCallableError) + assert not marker.exists() + + +def test_import_side_effects_do_not_run(): + """Importing is itself the side effect, so it must not happen at all.""" + assert isinstance(_refusal(("fn", "antigravity", "x")), UntrustedCallableError) + assert "antigravity" not in os.sys.modules + + +@pytest.mark.parametrize( + "module,qualname", + [ + ("ibis.util", "os.system"), # ibis.util imports os + ("ibis.expr.api", "builtins.eval"), # ibis.expr.api imports builtins + ], +) +def test_attribute_chain_cannot_escape_a_trusted_module(module, qualname): + """A qualname is a getattr chain: a trusted root is not enough on its own.""" + err = _refusal(("fn", module, qualname)) + assert isinstance(err, UntrustedCallableError) + assert "outside the trusted" in str(err) + + +def test_untrusted_callable_is_refused_at_write_time(): + """Authors find out when they serialize, not readers when they load.""" + from boring_semantic_layer._xorq import Just + + def user_fn(x): + return x + + with pytest.raises(UntrustedCallableError, match="not trusted"): + serialize_resolver(Just(user_fn)) + + +def test_expression_functions_still_round_trip(table): + """The allowlist must not cost real expressions anything.""" + model = ( + to_semantic_table(table, "m") + .with_dimensions(bucket=lambda t: (t.a > 1).ifelse("hi", "lo")) + .with_measures(n=lambda t: t.count()) + ) + df = from_tagged(to_tagged(model)).group_by("bucket").aggregate("n").execute() + assert dict(zip(df["bucket"], df["n"], strict=True)) == {"lo": 1, "hi": 2} + + +# --------------------------------------------------------------------------- +# Lossy tag encodings +# --------------------------------------------------------------------------- + + +def test_multi_argument_calls_survive_the_tag_round_trip(table): + """``thaw`` read ``(("just", 0), ("just", 2))`` as a dict and kept the last. + + ``substr(0, 2)`` came back as ``substr(2)`` — still valid SQL, different + answer, no error. + """ + model = ( + to_semantic_table(table, "m") + .with_dimensions(pre=lambda t: t.b.substr(0, 2)) + .with_measures(n=lambda t: t.count()) + ) + df = from_tagged(to_tagged(model)).group_by("pre").aggregate("n").execute() + assert sorted(df["pre"]) == ["ab", "ef", "zz"] + + +def test_multi_element_isin_survives_the_tag_round_trip(table): + model = to_semantic_table(table, "m").with_measures(n=lambda t: t.count()) + filtered = model.filter(lambda t: t.b.isin(["abcd", "efgh"])) + df = from_tagged(to_tagged(filtered)).aggregate("n").execute() + assert int(df["n"][0]) == 2 + + +@pytest.mark.parametrize( + "value,expected", + [ + (datetime.date(2024, 3, 1), 2), + (datetime.datetime(2024, 3, 1), 2), + ], +) +def test_non_scalar_literals_keep_their_type(table, value, expected): + """``freeze`` used to ``str()`` these, so the predicate compared to a string.""" + model = to_semantic_table(table, "m").with_measures(n=lambda t: t.count()) + filtered = model.filter(lambda t: t.d > value) + df = from_tagged(to_tagged(filtered)).aggregate("n").execute() + assert int(df["n"][0]) == expected + + +def test_decimal_literal_round_trips(): + from boring_semantic_layer._xorq import Just + + struct = serialize_resolver(Just(decimal.Decimal("1.5"))) + from boring_semantic_layer.utils import deserialize_resolver + + assert deserialize_resolver(struct).value == decimal.Decimal("1.5") + + +def test_unrepresentable_constant_fails_at_write_time(): + """Better a loud failure than a value silently replaced by its repr.""" + from boring_semantic_layer._xorq import Just + + with pytest.raises(ValueError, match="Cannot serialize constant"): + serialize_resolver(Just(object())) + + +# --------------------------------------------------------------------------- +# Aggregate replay +# --------------------------------------------------------------------------- + + +def test_query_local_agg_beats_a_same_named_model_measure(table): + """``aggregate(n=...)`` must not be replaced by the model's ``n``.""" + model = to_semantic_table(table, "m").with_measures(n=lambda t: t.a.sum()) + query = model.aggregate(n=lambda t: t.a.max()) + assert int(from_tagged(to_tagged(query)).execute()["n"][0]) == 3 + + +def test_bare_measure_name_still_replays_by_name(table): + """Bare names must keep routing through measure resolution (fan-out safety).""" + model = to_semantic_table(table, "m").with_measures(n=lambda t: t.a.sum()) + assert int(from_tagged(to_tagged(model.aggregate("n"))).execute()["n"][0]) == 6 + + +# --------------------------------------------------------------------------- +# Payload versioning +# --------------------------------------------------------------------------- + + +def test_unsupported_payload_version_is_refused(): + metadata = { + "bsl_op_type": "SemanticTableOp", + "bsl_version": "1.0", + "dimensions": {"amount": {"expr_pickle": "gAWV"}}, + "measures": {}, + } + with pytest.raises(ValueError, match="1.0"): + reconstruct_bsl_operation(metadata, None, BSLSerializationContext()) + + +def test_dimension_without_an_expression_is_refused(): + metadata = { + "bsl_op_type": "SemanticTableOp", + "bsl_version": "2.0", + "dimensions": {"amount": {"description": "no expression here"}}, + "measures": {}, + } + with pytest.raises(ValueError, match="no readable expression"): + reconstruct_bsl_operation(metadata, None, BSLSerializationContext()) + + +# --------------------------------------------------------------------------- +# Backend rebinding +# --------------------------------------------------------------------------- + + +def test_rebinding_does_not_repoint_a_table_at_another_database(tmp_path): + """Rebinding every DatabaseTable made a join read both sides from one db.""" + prod = ibis.duckdb.connect(str(tmp_path / "prod.ddb")) + prod.create_table("t", ibis.memtable({"k": [1, 2], "v": [10, 20]}).execute()) + staging = ibis.duckdb.connect(str(tmp_path / "staging.ddb")) + staging.create_table("t", ibis.memtable({"k": [1, 2], "v": [999, 888]}).execute()) + + left = ( + to_semantic_table(prod.table("t"), "p") + .with_dimensions(k=lambda t: t.k) + .with_measures(total=lambda t: t.v.sum()) + ) + right = ( + to_semantic_table(staging.table("t"), "s") + .with_dimensions(k=lambda t: t.k) + .with_measures(other=lambda t: t.v.sum()) + ) + + assert int(left.aggregate("total").execute()["total"][0]) == 30 + + joined = left.join_one(right, on=lambda x, y: x.k == y.k) + try: + got = int(joined.aggregate("p.total").execute()["p.total"][0]) + except Exception: + return # a loud cross-database failure is the acceptable outcome + assert got == 30, "prod's measure was computed from staging's rows" + + +def test_duplicate_wrappers_of_one_connection_still_unify(tmp_path): + """The case rebinding exists for: from_ibis() mints a Backend per call.""" + con = ibis.duckdb.connect(str(tmp_path / "one.ddb")) + con.create_table("t", ibis.memtable({"k": [1, 2], "v": [10, 20]}).execute()) + + left = ( + to_semantic_table(con.table("t"), "a") + .with_dimensions(k=lambda t: t.k) + .with_measures(total=lambda t: t.v.sum()) + ) + right = ( + to_semantic_table(con.table("t"), "b") + .with_dimensions(k=lambda t: t.k) + .with_measures(n=lambda t: t.count()) + ) + joined = left.join_one(right, on=lambda x, y: x.k == y.k) + assert int(joined.aggregate("a.total").execute()["a.total"][0]) == 30 diff --git a/src/boring_semantic_layer/tests/test_totals_semantics.py b/src/boring_semantic_layer/tests/test_totals_semantics.py new file mode 100644 index 00000000..fb135f05 --- /dev/null +++ b/src/boring_semantic_layer/tests/test_totals_semantics.py @@ -0,0 +1,192 @@ +"""Regression tests for ``t.all()`` / percent-of-total semantics. + +``t.all(x)`` means "x over the whole filtered dataset, ignoring the group +by". Getting that wrong is silent: the query still runs and the shares still +look like shares. The data below is chosen so the correct answer and the +sum-of-group-values answer never coincide: + + carrier A: distance [10, 20, 30] -> sum 60, mean 20 + carrier B: distance [100] -> sum 100, mean 100 + overall: sum 160, mean 40 (sum of the two group means is 120) +""" + +from __future__ import annotations + +import ibis +import pytest + +from boring_semantic_layer import to_semantic_table +from boring_semantic_layer.measure_scope import NonAdditiveTotalError + + +@pytest.fixture +def model(): + flights = ibis.memtable( + { + "carrier": ["A", "A", "A", "B"], + "distance": [10, 20, 30, 100], + "origin": ["x", "y", "x", "y"], + } + ) + return ( + to_semantic_table(flights, "flights") + .with_dimensions(carrier=lambda t: t.carrier, origin=lambda t: t.origin) + .with_measures( + total=lambda t: t.distance.sum(), + avg=lambda t: t.distance.mean(), + ) + ) + + +@pytest.fixture +def airports(): + table = ibis.memtable({"code": ["x", "y"], "region": ["west", "east"]}) + return ( + to_semantic_table(table, "airports") + .with_dimensions(code=lambda t: t.code, region=lambda t: t.region) + .with_measures(cnt=lambda t: t.count()) + ) + + +MEAN_SHARE = {"A": 20 / 40, "B": 100 / 40} +SUM_SHARE = {"A": 60 / 160, "B": 100 / 160} +SUM_OF_GROUP_MEANS = {"A": 20 / 120, "B": 100 / 120} + + +def _shares(model, col="share", key="carrier"): + df = model.group_by(key).aggregate(col).execute() + return {k: pytest.approx(float(v)) for k, v in zip(df[key], df[col], strict=True)} + + +# --------------------------------------------------------------------------- +# t.all() over an expression, not just a bare measure reference +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "share_fn", + [ + pytest.param(lambda t: t.avg / t.all(t.avg), id="bare-reference"), + pytest.param(lambda t: t.avg / t.all(t.avg * 1), id="scaled-inside"), + pytest.param(lambda t: t.avg / (t.all(t.avg) + 0), id="arithmetic-outside"), + pytest.param(lambda t: t.avg / t.all(t.avg + 0), id="arithmetic-inside"), + ], +) +def test_all_of_a_measure_expression_uses_the_real_total(model, share_fn): + """``t.all(m)`` and ``t.all(m * 1)`` must mean the same thing. + + Anything that was not an exact Field reference fell through to + ``x.sum().over(window())``, which for a mean measure sums the per-group + means — here 120 instead of 40, a 3x error with no warning. + """ + assert _shares(model.with_measures(share=share_fn)) == MEAN_SHARE + assert _shares(model.with_measures(share=share_fn)) != SUM_OF_GROUP_MEANS + + +def test_all_of_an_additive_measure_expression(model): + share = model.with_measures(share=lambda t: t.total / t.all(t.total * 1)) + assert _shares(share) == SUM_SHARE + + +# --------------------------------------------------------------------------- +# Integer division +# --------------------------------------------------------------------------- + + +def test_integer_measure_over_integer_total_is_a_ratio(model): + """Two integer operands must not truncate to 0 (xorq's DataFusion does).""" + assert _shares(model.with_measures(share=lambda t: t.total / t.all(t.total))) == SUM_SHARE + + +def test_percent_of_total_is_order_independent(): + """Declaration order must not change the answer.""" + data = ibis.memtable({"carrier": ["AA", "UA", "DL", "WN", "B6"] * 10}) + expected = dict.fromkeys(["AA", "UA", "DL", "WN", "B6"], pytest.approx(0.2)) + + dims_after = ( + to_semantic_table(data, "f") + .with_measures(flight_count=lambda t: t.count()) + .with_measures(ratio=lambda t: t.flight_count / t.all(t.flight_count)) + .with_dimensions(carrier=lambda t: t.carrier) + ) + dims_before = ( + to_semantic_table(data, "f") + .with_dimensions(carrier=lambda t: t.carrier) + .with_measures(flight_count=lambda t: t.count()) + .with_measures(ratio=lambda t: t.flight_count / t.all(t.flight_count)) + ) + assert _shares(dims_after, "ratio") == expected + assert _shares(dims_before, "ratio") == expected + + +# --------------------------------------------------------------------------- +# Post-aggregation chain .mutate() +# --------------------------------------------------------------------------- + + +def test_chain_mutate_total_of_an_additive_measure(model): + """``aggregate().order_by().mutate()`` sees only grouped rows. + + A window sum over them is the true total for SUM/COUNT, so this must + agree with the same formula written as a calc measure. + """ + df = ( + model.group_by("carrier") + .aggregate("total") + .order_by("carrier") + .mutate(share=lambda t: t.total / t.all(t.total)) + .execute() + ) + got = {k: pytest.approx(float(v)) for k, v in zip(df["carrier"], df["share"], strict=True)} + assert got == SUM_SHARE + assert got == _shares(model.with_measures(share=lambda t: t.total / t.all(t.total))) + + +def test_chain_mutate_total_of_a_mean_measure_is_refused(model): + """Summing per-group means is not the overall mean — refuse, don't guess. + + Returning the window sum made this spelling disagree with the identical + calc-measure formula (0.167 vs 0.5) with nothing to indicate which was + right. + """ + with pytest.raises(NonAdditiveTotalError, match="not additive"): + ( + model.group_by("carrier") + .aggregate("avg") + .order_by("carrier") + .mutate(share=lambda t: t.avg / t.all(t.avg)) + .execute() + ) + + +# --------------------------------------------------------------------------- +# Totals under a fan-out join +# --------------------------------------------------------------------------- + + +def test_totals_stay_correct_under_join_many(model, airports): + """The denominator must not inflate with the join's fan-out.""" + joined = model.with_measures(share=lambda t: t.total / t.all(t.total)).join_many( + airports, on=lambda left, right: left.origin == right.code + ) + assert _shares(joined, "flights.share", "flights.carrier") == SUM_SHARE + + +def test_inline_reduction_in_totals_under_join_many_is_refused(model, airports): + """A calc measure that builds its own reduction has no fan-out-safe base. + + This used to reach the engine as an aggregate with no columns at all and + surface as an unrelated arrow error ("Schema and number of arrays + unequal"); the message now names the problem and the fix. + """ + joined = model.with_measures( + pot=lambda t: t.distance.sum() / t.all(t.distance.sum()) + ).join_many(airports, on=lambda left, right: left.origin == right.code) + with pytest.raises(ValueError, match="fan-out-safe"): + joined.group_by("flights.carrier").aggregate("flights.pot").execute() + + # ...and the suggested spelling works. + fixed = model.with_measures(share=lambda t: t.total / t.all(t.total)).join_many( + airports, on=lambda left, right: left.origin == right.code + ) + assert _shares(fixed, "flights.share", "flights.carrier") == SUM_SHARE diff --git a/src/boring_semantic_layer/tests/test_xorq_convert.py b/src/boring_semantic_layer/tests/test_xorq_convert.py index 8e7d5730..a458d668 100644 --- a/src/boring_semantic_layer/tests/test_xorq_convert.py +++ b/src/boring_semantic_layer/tests/test_xorq_convert.py @@ -203,18 +203,36 @@ def test_from_xorq_returns_bsl_expr(): @pytest.mark.skipif(not xorq, reason="xorq not available") def test_from_xorq_with_tagged_table(): - from xorq.api import memtable + """A tagged table reconstructs — and unreadable payloads are refused. - # Use nested tuples format (following xorq sklearn pipeline pattern) - xorq_table = memtable({"a": [1, 2, 3]}).tag( - tag="bsl_test", - bsl_op_type="SemanticTableOp", - bsl_version="1.0", - dimensions=(("a", (("description", "Column A"),)),), - measures=(), - ) + ``bsl_version`` used to be written and never checked, and a dimension + with no readable expression silently became a reference to the raw + column of the same name: ``amount = _.amount * 1.1`` came back as + ``amount``, with no error and plausible numbers. + """ + from xorq.api import memtable - bsl_expr = from_tagged(xorq_table) + def tagged(**overrides): + payload = { + "tag": "bsl_test", + "bsl_op_type": "SemanticTableOp", + "bsl_version": "2.0", + "dimensions": (("a", (("expr", "a"), ("description", "Column A"))),), + "measures": (), + } + payload.update(overrides) + return memtable({"a": [1, 2, 3]}).tag(**payload) + + # v1.0 stored expressions as pickles, a format no longer read at all. + with pytest.raises(ValueError, match="1.0"): + from_tagged(tagged(bsl_version="1.0")) + + # Current version, but the dimension carries neither a column name nor + # a serialized expression. + with pytest.raises(ValueError, match="no readable expression"): + from_tagged(tagged(dimensions=(("a", (("description", "Column A"),)),))) + + bsl_expr = from_tagged(tagged()) assert bsl_expr is not None assert hasattr(bsl_expr, "dimensions") diff --git a/src/boring_semantic_layer/utils.py b/src/boring_semantic_layer/utils.py index 6dc2df23..f23612c4 100644 --- a/src/boring_semantic_layer/utils.py +++ b/src/boring_semantic_layer/utils.py @@ -18,6 +18,84 @@ class SafeEvalError(Exception): pass +class UntrustedCallableError(ValueError): + """A serialized expression names a callable outside the trusted set. + + Serialized models are data, not code: a tag payload can travel through + a xorq catalog, a git repo or any other artifact store, and is not + necessarily written by whoever reads it. Restoring an arbitrary + ``(module, qualname)`` pair means importing an attacker-chosen module + and handing the result to ``Call.resolve()``, which calls it — i.e. + arbitrary code execution. Only functions from the expression libraries + BSL builds on can be restored. + """ + + +#: Module roots whose callables may be named in a serialized expression. +#: These are the libraries that actually appear in ibis resolver trees: +#: deferrable API functions (``ifelse``, ``coalesce``, ``_finish_searched_case``) +#: and the ``operator`` functions behind binary/unary nodes. +_TRUSTED_CALLABLE_ROOTS: frozenset[str] = frozenset( + { + "ibis", + "xorq", + "operator", + "_operator", + "boring_semantic_layer", + } +) + +_EXTRA_TRUSTED_CALLABLE_ROOTS: set[str] = set() + + +def trust_callable_module(root: str) -> None: + """Allow callables from an additional top-level module in serialized models. + + Only do this for modules you control, and only when every model you + deserialize comes from a source you trust as much as your own code: + a serialized expression naming a callable is equivalent to a function + call, so widening this set widens what a malicious payload can invoke. + """ + _EXTRA_TRUSTED_CALLABLE_ROOTS.add(root.split(".", 1)[0]) + + +def _trusted_roots() -> frozenset[str]: + return _TRUSTED_CALLABLE_ROOTS | frozenset(_EXTRA_TRUSTED_CALLABLE_ROOTS) + + +def _module_root(module_name: str | None) -> str: + return (module_name or "").split(".", 1)[0] + + +def _check_callable_ref(module_name: str | None, qualname: str | None) -> None: + """Reject a ``(module, qualname)`` pair that must not cross the wire. + + Applied on *both* sides: serialization refuses to emit a reference that + deserialization would refuse to load, so the failure surfaces where the + model is authored rather than in someone else's process. + """ + if not module_name or not qualname: + raise UntrustedCallableError( + f"Callable reference is incomplete: module={module_name!r} qualname={qualname!r}" + ) + root = _module_root(module_name) + if root not in _trusted_roots(): + raise UntrustedCallableError( + f"Refusing to (de)serialize callable {module_name}.{qualname}: " + f"module root {root!r} is not trusted. Serialized expressions may " + f"only reference {sorted(_trusted_roots())}. Express the logic with " + "ibis operations, or call utils.trust_callable_module() if you own " + "the module and trust every model you load." + ) + for part in qualname.split("."): + if part.startswith("__") or "<" in part or not part.isidentifier(): + raise UntrustedCallableError( + f"Refusing to (de)serialize callable {module_name}.{qualname}: " + f"qualname component {part!r} is not a plain public identifier " + "(lambdas, closures and dunder attributes cannot be restored)." + ) + + SAFE_NODES = { ast.Expression, ast.Load, @@ -564,6 +642,94 @@ def _is_ibis_literal_node(value) -> bool: return False +#: Marker for a constant that is not one of xorq's native tag scalar types. +#: Tag metadata can only hold str/int/float/bool/None (see +#: ``serialization.freeze``), so anything else — dates, ``Decimal``, +#: ``bytes`` — is carried as ``(_SCALAR_TAG, kind, payload)`` and rebuilt on +#: read. Previously these reached ``freeze()`` and were flattened with +#: ``str()``: a ``date`` predicate came back comparing against a string, and +#: a ``Decimal`` came back as a type error from the query compiler. +_SCALAR_TAG = "__bsl_scalar__" + + +def _encode_scalar(value: Any) -> Any: + """Represent a constant in a form tag metadata can hold losslessly.""" + import datetime + import decimal + import uuid + + if isinstance(value, str | bool | int | float | type(None)): + # numpy scalars subclass int/float; normalize so they survive as + # native Python values rather than as repr strings. + if type(value) is not bool and isinstance(value, int) and type(value) is not int: + return int(value) + if isinstance(value, float) and type(value) is not float: + return float(value) + return value + # datetime before date: datetime is a date subclass. + if isinstance(value, datetime.datetime): + return (_SCALAR_TAG, "datetime", value.isoformat()) + if isinstance(value, datetime.date): + return (_SCALAR_TAG, "date", value.isoformat()) + if isinstance(value, datetime.time): + return (_SCALAR_TAG, "time", value.isoformat()) + if isinstance(value, datetime.timedelta): + return (_SCALAR_TAG, "timedelta", repr(value.total_seconds())) + if isinstance(value, decimal.Decimal): + return (_SCALAR_TAG, "decimal", str(value)) + if isinstance(value, uuid.UUID): + return (_SCALAR_TAG, "uuid", str(value)) + if isinstance(value, bytes): + import base64 + + return (_SCALAR_TAG, "bytes", base64.b64encode(value).decode("ascii")) + if isinstance(value, list | tuple): + kind = "list" if isinstance(value, list) else "tuple" + return (_SCALAR_TAG, kind, tuple(_encode_scalar(item) for item in value)) + # numpy scalars that subclass nothing familiar (e.g. np.datetime64) + if hasattr(value, "item") and type(value).__module__.startswith("numpy"): + return _encode_scalar(value.item()) + raise ValueError( + f"Cannot serialize constant of type {type(value).__name__} ({value!r}): " + "tag metadata holds only scalars, dates, Decimal, UUID and bytes. " + "Previously such values were silently stringified." + ) + + +def _decode_scalar(value: Any) -> Any: + """Inverse of :func:`_encode_scalar`; untagged values pass through.""" + if not (isinstance(value, tuple | list) and len(value) == 3 and value[0] == _SCALAR_TAG): + return value + import datetime + import decimal + import uuid + + _, kind, payload = value + match kind: + case "datetime": + return datetime.datetime.fromisoformat(payload) + case "date": + return datetime.date.fromisoformat(payload) + case "time": + return datetime.time.fromisoformat(payload) + case "timedelta": + return datetime.timedelta(seconds=float(payload)) + case "decimal": + return decimal.Decimal(payload) + case "uuid": + return uuid.UUID(payload) + case "bytes": + import base64 + + return base64.b64decode(payload.encode("ascii")) + case "list": + return [_decode_scalar(item) for item in payload] + case "tuple": + return tuple(_decode_scalar(item) for item in payload) + case _: + raise ValueError(f"Unknown encoded-constant kind: {kind!r}") + + def serialize_resolver(resolver) -> tuple: """Walk a Resolver tree and produce a hashable nested-tuple representation.""" from ._xorq import ( @@ -588,23 +754,22 @@ def serialize_resolver(resolver) -> tuple: if _is_ibis_literal_node(value): py_value = value.args[0] dtype_str = str(value.args[1]) - return ("ibis_literal", py_value, dtype_str) + return ("ibis_literal", _encode_scalar(py_value), dtype_str) # callable (operator functions, deferrable functions like ifelse, _finish_searched_case) if callable(value): module = getattr(value, "__module__", None) qualname = getattr(value, "__qualname__", None) - if module and qualname: - return ("fn", module, qualname) - raise ValueError(f"Cannot serialize callable without __module__/__qualname__: {value!r}") - # primitive value (int, float, str, bool, None) - return ("just", value) + _check_callable_ref(module, qualname) + return ("fn", module, qualname) + # primitive value (int, float, str, bool, None) or an encodable constant + return ("just", _encode_scalar(value)) if isinstance(resolver, JustUnhashable): value = resolver.value.obj if _is_ibis_literal_node(value): py_value = value.args[0] dtype_str = str(value.args[1]) - return ("ibis_literal", py_value, dtype_str) + return ("ibis_literal", _encode_scalar(py_value), dtype_str) raise ValueError(f"Cannot serialize unhashable value: {value!r}") if isinstance(resolver, Attr): @@ -678,6 +843,33 @@ def _resolve_qualname(module_obj, qualname: str): return obj +def _load_trusted_callable(module_name: str, qualname: str): + """Import and return a callable named by a serialized expression. + + The pair is validated before the import — an unimportable module is a + side effect in itself, so an untrusted name must never reach + ``import_module``. After resolution the *result* is checked too: a + qualname is a ``getattr`` chain, so ``("fn", "ibis", "os.system")`` + would otherwise walk out of a trusted module into an untrusted one. + """ + _check_callable_ref(module_name, qualname) + mod = importlib.import_module(module_name) + func = _resolve_qualname(mod, qualname) + if not callable(func): + raise UntrustedCallableError( + f"{module_name}.{qualname} resolved to a non-callable " + f"{type(func).__name__}; refusing to use it as an expression function." + ) + origin = getattr(func, "__module__", None) + if _module_root(origin) not in _trusted_roots(): + raise UntrustedCallableError( + f"{module_name}.{qualname} resolves to an object defined in " + f"{origin!r}, which is outside the trusted module set. This is how " + "an attribute chain escapes a trusted module — refusing to load it." + ) + return func + + def _finalize_frozen_slotted(obj, *fields) -> None: """Set ``__precomputed_hash__`` on a FrozenSlotted built via ``object.__new__``. @@ -713,16 +905,14 @@ def deserialize_resolver(data: tuple): return Variable(name) case ("just", value): - return Just(value) + return Just(_decode_scalar(value)) case ("fn", module_name, qualname): - mod = importlib.import_module(module_name) - func = _resolve_qualname(mod, qualname) - return Just(func) + return Just(_load_trusted_callable(module_name, qualname)) case ("ibis_literal", py_value, dtype_str): from ._xorq import ibis - lit_expr = ibis.literal(py_value, type=ibis.dtype(dtype_str)) + lit_expr = ibis.literal(_decode_scalar(py_value), type=ibis.dtype(dtype_str)) return Just(lit_expr.op()) case ("attr", obj_data, name_data):