From c11a1e5e5b2beb49ad9c2151df1bf838875d2aea Mon Sep 17 00:00:00 2001 From: Nico Ritschel Date: Sun, 19 Jul 2026 13:26:49 -0700 Subject: [PATCH] Keep registration contexts task-local --- sidemantic/core/registry.py | 15 +++++-- sidemantic/core/semantic_layer.py | 37 +++++++++++++---- tests/test_metric_auto_registration.py | 56 ++++++++++++++++++++++++++ 3 files changed, 97 insertions(+), 11 deletions(-) diff --git a/sidemantic/core/registry.py b/sidemantic/core/registry.py index 60009a96a..11ddd91f2 100644 --- a/sidemantic/core/registry.py +++ b/sidemantic/core/registry.py @@ -1,8 +1,10 @@ """Global registry for auto-registration of models and measures.""" -from contextvars import ContextVar +from contextvars import ContextVar, Token from typing import TYPE_CHECKING +from sidemantic.validation import MetricValidationError + if TYPE_CHECKING: from .semantic_layer import SemanticLayer @@ -15,9 +17,14 @@ def get_current_layer() -> "SemanticLayer | None": return _current_layer.get() -def set_current_layer(layer: "SemanticLayer | None"): +def set_current_layer(layer: "SemanticLayer | None") -> Token: """Set the current semantic layer context.""" - _current_layer.set(layer) + return _current_layer.set(layer) + + +def reset_current_layer(token: Token) -> None: + """Restore the layer that was active before ``set_current_layer``.""" + _current_layer.reset(token) def auto_register_model(model): @@ -48,7 +55,7 @@ def auto_register_metric(metric): if not metric.agg or metric.type in ("derived", "ratio"): try: layer.add_metric(metric) - except Exception: + except (ValueError, MetricValidationError): # Best-effort: metric might already exist or validation might fail during init import logging diff --git a/sidemantic/core/semantic_layer.py b/sidemantic/core/semantic_layer.py index 79a932176..84c05d521 100644 --- a/sidemantic/core/semantic_layer.py +++ b/sidemantic/core/semantic_layer.py @@ -4,6 +4,7 @@ import os from collections.abc import Callable +from contextvars import ContextVar, Token from pathlib import Path import yaml @@ -269,24 +270,46 @@ def __init__( else: raise TypeError(f"connection must be a string URL or BaseDatabaseAdapter instance, got {type(connection)}") - # Set as current layer for auto-registration + self._registration_token = None + self._context_registration_tokens: ContextVar[tuple[Token, ...]] = ContextVar( + f"semantic_layer_registration_tokens_{id(self)}", default=() + ) if auto_register: from .registry import set_current_layer - set_current_layer(self) + self._registration_token = set_current_layer(self) def __enter__(self): """Context manager entry - set as current layer.""" from .registry import set_current_layer - set_current_layer(self) + token = set_current_layer(self) + tokens = self._context_registration_tokens.get() + self._context_registration_tokens.set((*tokens, token)) return self def __exit__(self, exc_type, exc_val, exc_tb): - """Context manager exit - clear current layer and close adapter.""" - from .registry import set_current_layer - - set_current_layer(None) + """Context manager exit - restore the prior layer and close adapter.""" + from .registry import get_current_layer, reset_current_layer + + tokens = self._context_registration_tokens.get() + if tokens: + token = tokens[-1] + remaining_tokens = tokens[:-1] + self._context_registration_tokens.set(remaining_tokens) + reset_current_layer(token) + + # An inline ``with SemanticLayer(auto_register=True)`` should still + # restore the layer that preceded construction. A copied async + # context cannot reset the constructor's token, so leave that token + # for its owning context instead of raising from ``__exit__``. + if not remaining_tokens and self._registration_token is not None and get_current_layer() is self: + try: + reset_current_layer(self._registration_token) + except ValueError: + pass + else: + self._registration_token = None if hasattr(self.adapter, "close"): self.adapter.close() diff --git a/tests/test_metric_auto_registration.py b/tests/test_metric_auto_registration.py index 64034d833..c43340d5f 100644 --- a/tests/test_metric_auto_registration.py +++ b/tests/test_metric_auto_registration.py @@ -1,6 +1,62 @@ """Tests for metric auto-registration.""" +import asyncio + +import pytest + from sidemantic import Dimension, Metric, Model, SemanticLayer +from sidemantic.core.registry import get_current_layer + + +def test_explicit_registration_mode_has_no_ambient_side_effect(): + layer = SemanticLayer(auto_register=False) + Model(name="orders", table="orders") + + assert layer.list_models() == [] + assert get_current_layer() is None + + +def test_nested_registration_context_restores_outer_layer(): + outer = SemanticLayer(auto_register=False) + inner = SemanticLayer(auto_register=False) + + with outer: + assert get_current_layer() is outer + with inner: + assert get_current_layer() is inner + assert get_current_layer() is outer + + assert get_current_layer() is None + + +def test_auto_registered_layer_context_is_safe_in_copied_async_context(): + layer = SemanticLayer(auto_register=True) + + async def use_layer_in_child_context(): + assert get_current_layer() is layer + with layer: + assert get_current_layer() is layer + await asyncio.sleep(0) + assert get_current_layer() is layer + + asyncio.run(use_layer_in_child_context()) + + assert get_current_layer() is layer + with layer: + assert get_current_layer() is layer + assert get_current_layer() is None + + +def test_auto_registration_does_not_hide_unexpected_failures(monkeypatch): + layer = SemanticLayer(auto_register=True) + + def fail_unexpectedly(_metric): + raise RuntimeError("implementation bug") + + monkeypatch.setattr(layer, "add_metric", fail_unexpectedly) + + with pytest.raises(RuntimeError, match="implementation bug"): + Metric(name="broken", type="derived", sql="orders.revenue") def test_standalone_metric_auto_registers_with_context_manager():