From 7b144a167fb97d0f8eea541253d643f8141252d3 Mon Sep 17 00:00:00 2001 From: lucifer726 <31751140+lucifer726@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:17:08 +0800 Subject: [PATCH] feat(pydantic): forward session properties to the engine WrenEngine.query/dry_plan/dry_run already accept `properties`, but the toolkit's direct Python API never passed them, so a model guarded by row-level access control could not be read from the SDK: planning fails when a required session property is missing and there was no way to supply one per call. The Pydantic AI tools keep their existing signatures on purpose. Their `ctx` is documented as ignored because the toolkit already captures all required state, and exposing a session property as a model-fillable tool argument would let the LLM choose the identity its own access control is keyed on. Closes #2638 Co-Authored-By: Claude Opus 5 --- .../src/wren_pydantic/_toolkit.py | 36 +++++++++--- .../tests/unit/test_toolkit_runtime.py | 58 ++++++++++++++++++- 2 files changed, 82 insertions(+), 12 deletions(-) diff --git a/sdk/wren-pydantic/src/wren_pydantic/_toolkit.py b/sdk/wren-pydantic/src/wren_pydantic/_toolkit.py index 8531bdfab7..60d925a19b 100644 --- a/sdk/wren-pydantic/src/wren_pydantic/_toolkit.py +++ b/sdk/wren-pydantic/src/wren_pydantic/_toolkit.py @@ -118,24 +118,42 @@ def instructions(self, *, toolset: object | None = None) -> str: # ── Direct Python API (sync only — see module docstring) ────────────── - def query(self, sql: str, limit: int | None = None) -> pa.Table: - """Execute SQL through the Wren context layer. Returns a pyarrow Table.""" + def query( + self, + sql: str, + limit: int | None = None, + properties: dict[str, Any] | None = None, + ) -> pa.Table: + """Execute SQL through the Wren context layer. Returns a pyarrow Table. + + ``properties`` carries MDL session properties and is forwarded to the + engine's planning path. A model guarded by row-level access control + needs the value its rule declares required — without it planning fails, + so RLAC-protected models cannot be read at all. + """ engine = self._build_engine() try: - result = engine.query(sql, limit=limit) + result = engine.query(sql, limit=limit, properties=properties) finally: self._connector_cache = engine._connector return result - def dry_plan(self, sql: str) -> str: - """Plan SQL through MDL and return the expanded SQL in target dialect.""" - return self._build_engine().dry_plan(sql) + def dry_plan(self, sql: str, properties: dict[str, Any] | None = None) -> str: + """Plan SQL through MDL and return the expanded SQL in target dialect. + + ``properties`` carries MDL session properties (see :meth:`query`); RLAC + predicates are injected during planning, so they apply here too. + """ + return self._build_engine().dry_plan(sql, properties=properties) + + def dry_run(self, sql: str, properties: dict[str, Any] | None = None) -> None: + """Validate SQL by planning and asking the DB to plan it without executing. - def dry_run(self, sql: str) -> None: - """Validate SQL by planning and asking the DB to plan it without executing.""" + ``properties`` carries MDL session properties (see :meth:`query`). + """ engine = self._build_engine() try: - engine.dry_run(sql) + engine.dry_run(sql, properties=properties) finally: self._connector_cache = engine._connector diff --git a/sdk/wren-pydantic/tests/unit/test_toolkit_runtime.py b/sdk/wren-pydantic/tests/unit/test_toolkit_runtime.py index 36db37db37..308fd9a019 100644 --- a/sdk/wren-pydantic/tests/unit/test_toolkit_runtime.py +++ b/sdk/wren-pydantic/tests/unit/test_toolkit_runtime.py @@ -26,7 +26,7 @@ def test_query_invokes_wren_engine_with_resolved_manifest( result = toolkit.query("SELECT 1", limit=10) assert result is fake_table - fake_engine.query.assert_called_once_with("SELECT 1", limit=10) + fake_engine.query.assert_called_once_with("SELECT 1", limit=10, properties=None) # Engine constructed with manifest bytes + datasource + connection_info engine_ctor.assert_called_once() kwargs = engine_ctor.call_args.kwargs @@ -106,7 +106,9 @@ def test_dry_plan_delegates_to_engine(tmp_project, fake_active_profile): result = toolkit.dry_plan("SELECT * FROM orders") assert result == "SELECT * FROM cte_orders" - fake_engine.dry_plan.assert_called_once_with("SELECT * FROM orders") + fake_engine.dry_plan.assert_called_once_with( + "SELECT * FROM orders", properties=None + ) def test_dry_run_delegates_to_engine(tmp_project, fake_active_profile): @@ -118,4 +120,54 @@ def test_dry_run_delegates_to_engine(tmp_project, fake_active_profile): with patch("wren_pydantic._toolkit.WrenEngine", return_value=fake_engine): toolkit.dry_run("SELECT 1") - fake_engine.dry_run.assert_called_once_with("SELECT 1") + fake_engine.dry_run.assert_called_once_with("SELECT 1", properties=None) + + +def test_query_forwards_session_properties(tmp_project, fake_active_profile): + """Session properties reach the engine, so RLAC-protected models are readable.""" + fake_engine = MagicMock(name="engine") + fake_engine.query.return_value = pa.table({"x": [1]}) + fake_engine._connector = MagicMock() + properties = {"session_user_id": "'u_42'"} + + toolkit = WrenToolkit.from_project(tmp_project) + + with patch("wren_pydantic._toolkit.WrenEngine", return_value=fake_engine): + toolkit.query("SELECT * FROM orders", limit=5, properties=properties) + + fake_engine.query.assert_called_once_with( + "SELECT * FROM orders", limit=5, properties=properties + ) + + +def test_dry_plan_forwards_session_properties(tmp_project, fake_active_profile): + """dry_plan forwards properties: RLAC predicates are injected while planning.""" + fake_engine = MagicMock(name="engine") + fake_engine.dry_plan.return_value = "SELECT 1" + fake_engine._connector = MagicMock() + properties = {"session_user_id": "'u_42'"} + + toolkit = WrenToolkit.from_project(tmp_project) + + with patch("wren_pydantic._toolkit.WrenEngine", return_value=fake_engine): + toolkit.dry_plan("SELECT * FROM orders", properties=properties) + + fake_engine.dry_plan.assert_called_once_with( + "SELECT * FROM orders", properties=properties + ) + + +def test_dry_run_forwards_session_properties(tmp_project, fake_active_profile): + """dry_run forwards properties so validation matches what query will run.""" + fake_engine = MagicMock(name="engine") + fake_engine._connector = MagicMock() + properties = {"session_user_id": "'u_42'"} + + toolkit = WrenToolkit.from_project(tmp_project) + + with patch("wren_pydantic._toolkit.WrenEngine", return_value=fake_engine): + toolkit.dry_run("SELECT * FROM orders", properties=properties) + + fake_engine.dry_run.assert_called_once_with( + "SELECT * FROM orders", properties=properties + )