From 089c37e169cef907e398db4baca8073cf3a40334 Mon Sep 17 00:00:00 2001 From: Peter Date: Sat, 4 Jul 2026 20:41:23 +0800 Subject: [PATCH] fix(wren): bound the session-context cache get_session_context is keyed on the per-query extracted manifest, so an unbounded functools.cache grows one SessionContext per distinct table subset for the life of the process. Use lru_cache(maxsize=32). --- core/wren/.claude/CLAUDE.md | 2 +- core/wren/src/wren/mdl/__init__.py | 10 +++- .../tests/unit/test_session_context_cache.py | 59 +++++++++++++++++++ 3 files changed, 68 insertions(+), 3 deletions(-) create mode 100644 core/wren/tests/unit/test_session_context_cache.py diff --git a/core/wren/.claude/CLAUDE.md b/core/wren/.claude/CLAUDE.md index ba80766d42..dd924e217e 100644 --- a/core/wren/.claude/CLAUDE.md +++ b/core/wren/.claude/CLAUDE.md @@ -39,7 +39,7 @@ Uses `uv` (not Poetry). `pyproject.toml` uses `hatchling` as build backend. - **WrenEngine** is the main entry point. It accepts a base64-encoded MDL JSON string, a `DataSource`, and a connection dict. - **Query flow**: `_plan()` → wren-core `SessionContext.transform_sql()` → `_transpile()` via sqlglot → connector `.query()`. - **Manifest extraction**: `_plan()` tries to extract a minimal sub-manifest scoped to the query's referenced tables before calling wren-core — this reduces planning overhead. Falls back to the full manifest on error. -- **`get_session_context` is `@cache`-decorated** — same `(manifest_str, function_path, properties, data_source)` tuple reuses the same SessionContext. Avoid mutating session state. +- **`get_session_context` is `@lru_cache(maxsize=32)`-decorated** — same `(manifest_str, function_path, properties, data_source)` tuple reuses the same SessionContext; least-recently-used entries are evicted past 32. Never mutate session state: eviction makes the next call for the same tuple rebuild a fresh SessionContext, so any mutation silently disappears at an unpredictable point. - **Write dialect mapping**: `canner` → `trino`; file sources (`local_file`, `s3_file`, `minio_file`, `gcs_file`) → `duckdb`. All others use `data_source.name` directly. - **WrenEngine is a context manager** (`__enter__` / `__exit__` call `close()`). - **Profile-based workflow**: When no explicit `--connection-*` flags are given, the CLI auto-discovers the active profile from `~/.wren/profiles.yml`. Profiles store datasource type + connection fields. diff --git a/core/wren/src/wren/mdl/__init__.py b/core/wren/src/wren/mdl/__init__.py index a492cdbb59..d911894c2a 100644 --- a/core/wren/src/wren/mdl/__init__.py +++ b/core/wren/src/wren/mdl/__init__.py @@ -1,17 +1,23 @@ """MDL processing utilities backed by wren-core-py.""" -from functools import cache +from functools import lru_cache import wren_core -@cache +@lru_cache(maxsize=32) def get_session_context( manifest_str: str | None, function_path: str | None, properties: frozenset | None = None, data_source: str | None = None, ) -> wren_core.SessionContext: + """Build (or reuse) a SessionContext for the given manifest tuple. + + Bounded LRU: the cache key includes the per-query extracted manifest + (engine.py ``extract_by``), so an unbounded cache would grow one + SessionContext per distinct table subset for the life of the process. + """ return wren_core.SessionContext( manifest_str, function_path, properties, data_source ) diff --git a/core/wren/tests/unit/test_session_context_cache.py b/core/wren/tests/unit/test_session_context_cache.py new file mode 100644 index 0000000000..9f39d36a74 --- /dev/null +++ b/core/wren/tests/unit/test_session_context_cache.py @@ -0,0 +1,59 @@ +"""Tests bounded LRU behavior for get_session_context.""" + +import pytest + +import wren.mdl +from wren.mdl import get_session_context + +pytestmark = pytest.mark.unit + +# Asserted literally, never derived from the cache under test: the tests, +# not the implementation, pin the capacity policy. +EXPECTED_MAXSIZE = 32 + + +class _FakeSessionContext: + def __init__(self, *args, **kwargs): + self.args = args + + +@pytest.fixture() +def fake_session_context(monkeypatch): + # Keep this unit test fast and isolated from the native SessionContext. + monkeypatch.setattr(wren.mdl.wren_core, "SessionContext", _FakeSessionContext) + get_session_context.cache_clear() + yield + get_session_context.cache_clear() + + +def _manifest(i: int) -> str: + # Mirrors production keying: engine.py passes the per-query extracted + # manifest as manifest_str, so distinct table subsets are distinct keys. + return f"manifest-{i}" + + +def test_cache_is_bounded_to_32(fake_session_context): + # Pin the configured cache bound as part of the public cache policy. + assert get_session_context.cache_info().maxsize == EXPECTED_MAXSIZE + + for i in range(EXPECTED_MAXSIZE + 4): + get_session_context(_manifest(i), None, None, None) + assert get_session_context.cache_info().currsize == EXPECTED_MAXSIZE + + +def test_cache_evicts_least_recently_used(fake_session_context): + initial_contexts = [ + get_session_context(_manifest(i), None, None, None) + for i in range(EXPECTED_MAXSIZE) + ] + # Holding the originals keeps the evicted instance alive: the identity + # assertions prove removal from the cache, not object destruction. + victim = initial_contexts[1] + + # Refresh key 0, then overflow by one: key 1 becomes the LRU and is + # evicted; key 0 survives. + kept = get_session_context(_manifest(0), None, None, None) + get_session_context(_manifest(EXPECTED_MAXSIZE), None, None, None) + + assert get_session_context(_manifest(0), None, None, None) is kept + assert get_session_context(_manifest(1), None, None, None) is not victim