feat(pydantic): forward session properties to the engine (#2684)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

Author: lucifer726

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
 

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
+    )