Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 27 additions & 9 deletions sdk/wren-pydantic/src/wren_pydantic/_toolkit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
58 changes: 55 additions & 3 deletions sdk/wren-pydantic/tests/unit/test_toolkit_runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -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
)
Loading