fix: adding axon SQL to sdk (#767)

Author: dines-rl

src/runloop_api_client/sdk/__init__.py

@@ -5,7 +5,7 @@
 
 from __future__ import annotations
 
-from .axon import Axon
+from .axon import Axon, AxonSqlOps
 from .sync import (
     AxonOps,
     AgentOps,
@@ -48,7 +48,7 @@
 from .benchmark import Benchmark
 from .blueprint import Blueprint
 from .execution import Execution
-from .async_axon import AsyncAxon
+from .async_axon import AsyncAxon, AsyncAxonSqlOps
 from .mcp_config import McpConfig
 from .async_agent import AsyncAgent
 from .async_devbox import AsyncDevbox, AsyncNamedShell
@@ -111,6 +111,8 @@
     "AsyncAgent",
     "Axon",
     "AsyncAxon",
+    "AxonSqlOps",
+    "AsyncAxonSqlOps",
     "AsyncSecret",
     "Benchmark",
     "AsyncBenchmark",

src/runloop_api_client/sdk/_types.py

@@ -42,6 +42,8 @@
 from ..types.devboxes import DiskSnapshotListParams, DiskSnapshotUpdateParams
 from ..types.scenarios import ScorerListParams, ScorerCreateParams, ScorerUpdateParams
 from ..types.devbox_create_params import DevboxBaseCreateParams
+from ..types.axons.sql_batch_params import SqlBatchParams
+from ..types.axons.sql_query_params import SqlQueryParams
 from ..types.scenario_start_run_params import ScenarioStartRunBaseParams
 from ..types.benchmark_start_run_params import BenchmarkSelfStartRunParams
 from ..types.devbox_execute_async_params import DevboxNiceExecuteAsyncParams
@@ -196,6 +198,14 @@ class SDKAxonPublishParams(AxonPublishParams, LongRequestOptions):
     pass
 
 
+class SDKAxonSqlQueryParams(SqlQueryParams, LongRequestOptions):
+    pass
+
+
+class SDKAxonSqlBatchParams(SqlBatchParams, LongRequestOptions):
+    pass
+
+
 class SDKScenarioListParams(ScenarioListParams, BaseRequestOptions):
     pass
 

src/runloop_api_client/sdk/async_axon.py

@@ -7,19 +7,48 @@
 from ._types import (
     BaseRequestOptions,
     SDKAxonPublishParams,
+    SDKAxonSqlBatchParams,
+    SDKAxonSqlQueryParams,
 )
 from .._client import AsyncRunloop
 from .._streaming import AsyncStream
 from ..types.axon_view import AxonView
 from ..types.axon_event_view import AxonEventView
 from ..types.publish_result_view import PublishResultView
+from ..types.axons.sql_batch_result_view import SqlBatchResultView
+from ..types.axons.sql_query_result_view import SqlQueryResultView
+
+
+class AsyncAxonSqlOps:
+    """[Beta] Async SQL operations for an axon's SQLite database.
+
+    Access via ``axon.sql``.
+
+    Example:
+        >>> axon = await runloop.axon.create()
+        >>> await axon.sql.query(sql="CREATE TABLE tasks (id INTEGER PRIMARY KEY, name TEXT)")
+        >>> result = await axon.sql.query(sql="SELECT * FROM tasks WHERE id = ?", params=[1])
+    """
+
+    def __init__(self, client: AsyncRunloop, axon_id: str) -> None:
+        self._client = client
+        self._axon_id = axon_id
+
+    async def query(self, **params: Unpack[SDKAxonSqlQueryParams]) -> SqlQueryResultView:
+        """[Beta] Execute a single parameterized SQL statement against this axon's SQLite database."""
+        return await self._client.axons.sql.query(self._axon_id, **params)
+
+    async def batch(self, **params: Unpack[SDKAxonSqlBatchParams]) -> SqlBatchResultView:
+        """[Beta] Execute multiple SQL statements atomically within a single transaction."""
+        return await self._client.axons.sql.batch(self._axon_id, **params)
 
 
 class AsyncAxon:
     """[Beta] Wrapper around asynchronous axon operations.
 
-    Axons are event communication channels that support publishing events
-    and subscribing to event streams via server-sent events (SSE).
+    Axons are event communication channels that support publishing events,
+    subscribing to event streams via server-sent events (SSE), and executing
+    SQL queries against an embedded SQLite database.
     Obtain instances via ``runloop.axon.create()`` or ``runloop.axon.from_id()``.
 
     Example:
@@ -29,11 +58,17 @@ class AsyncAxon:
         >>> async with await axon.subscribe_sse() as stream:
         ...     async for event in stream:
         ...         print(event.event_type, event.payload)
+        >>> await axon.sql.query(sql="CREATE TABLE tasks (id INTEGER PRIMARY KEY, name TEXT)")
     """
 
     def __init__(self, client: AsyncRunloop, axon_id: str) -> None:
         self._client = client
         self._id = axon_id
+        self._sql = AsyncAxonSqlOps(client, axon_id)
+
+    @property
+    def sql(self) -> AsyncAxonSqlOps:
+        return self._sql
 
     @override
     def __repr__(self) -> str:

src/runloop_api_client/sdk/axon.py

@@ -7,19 +7,48 @@
 from ._types import (
     BaseRequestOptions,
     SDKAxonPublishParams,
+    SDKAxonSqlBatchParams,
+    SDKAxonSqlQueryParams,
 )
 from .._client import Runloop
 from .._streaming import Stream
 from ..types.axon_view import AxonView
 from ..types.axon_event_view import AxonEventView
 from ..types.publish_result_view import PublishResultView
+from ..types.axons.sql_batch_result_view import SqlBatchResultView
+from ..types.axons.sql_query_result_view import SqlQueryResultView
+
+
+class AxonSqlOps:
+    """[Beta] SQL operations for an axon's SQLite database.
+
+    Access via ``axon.sql``.
+
+    Example:
+        >>> axon = runloop.axon.create()
+        >>> axon.sql.query(sql="CREATE TABLE tasks (id INTEGER PRIMARY KEY, name TEXT)")
+        >>> result = axon.sql.query(sql="SELECT * FROM tasks WHERE id = ?", params=[1])
+    """
+
+    def __init__(self, client: Runloop, axon_id: str) -> None:
+        self._client = client
+        self._axon_id = axon_id
+
+    def query(self, **params: Unpack[SDKAxonSqlQueryParams]) -> SqlQueryResultView:
+        """[Beta] Execute a single parameterized SQL statement against this axon's SQLite database."""
+        return self._client.axons.sql.query(self._axon_id, **params)
+
+    def batch(self, **params: Unpack[SDKAxonSqlBatchParams]) -> SqlBatchResultView:
+        """[Beta] Execute multiple SQL statements atomically within a single transaction."""
+        return self._client.axons.sql.batch(self._axon_id, **params)
 
 
 class Axon:
     """[Beta] Wrapper around synchronous axon operations.
 
-    Axons are event communication channels that support publishing events
-    and subscribing to event streams via server-sent events (SSE).
+    Axons are event communication channels that support publishing events,
+    subscribing to event streams via server-sent events (SSE), and executing
+    SQL queries against an embedded SQLite database.
     Obtain instances via ``runloop.axon.create()`` or ``runloop.axon.from_id()``.
 
     Example:
@@ -29,11 +58,17 @@ class Axon:
         >>> with axon.subscribe_sse() as stream:
         ...     for event in stream:
         ...         print(event.event_type, event.payload)
+        >>> axon.sql.query(sql="CREATE TABLE tasks (id INTEGER PRIMARY KEY, name TEXT)")
     """
 
     def __init__(self, client: Runloop, axon_id: str) -> None:
         self._client = client
         self._id = axon_id
+        self._sql = AxonSqlOps(client, axon_id)
+
+    @property
+    def sql(self) -> AxonSqlOps:
+        return self._sql
 
     @override
     def __repr__(self) -> str:

tests/sdk/conftest.py

@@ -133,6 +133,47 @@ class MockPublishResultView:
     timestamp_ms: int = 1234567890000
 
 
+@dataclass
+class MockSqlColumnMetaView:
+    """Mock SqlColumnMetaView for testing."""
+
+    name: str = "id"
+    type: str = "INTEGER"
+
+
+@dataclass
+class MockSqlResultMetaView:
+    """Mock SqlResultMetaView for testing."""
+
+    changes: int = 0
+    duration_ms: float = 1.5
+    rows_read_limit_reached: bool = False
+
+
+@dataclass
+class MockSqlQueryResultView:
+    """Mock SqlQueryResultView for testing."""
+
+    columns: list[Any] = field(default_factory=lambda: [MockSqlColumnMetaView()])
+    meta: Any = field(default_factory=MockSqlResultMetaView)
+    rows: list[Any] = field(default_factory=lambda: [[1, "hello"]])
+
+
+@dataclass
+class MockSqlStepResultView:
+    """Mock SqlStepResultView for testing."""
+
+    success: Any = field(default_factory=lambda: MockSqlQueryResultView())
+    error: Any = None
+
+
+@dataclass
+class MockSqlBatchResultView:
+    """Mock SqlBatchResultView for testing."""
+
+    results: list[Any] = field(default_factory=lambda: [MockSqlStepResultView()])
+
+
 @dataclass
 class MockScenarioView:
     """Mock ScenarioView for testing."""

tests/sdk/test_async_axon.py

@@ -6,8 +6,9 @@
 
 import pytest
 
-from tests.sdk.conftest import MockAxonView, MockPublishResultView
+from tests.sdk.conftest import MockAxonView, MockPublishResultView, MockSqlBatchResultView, MockSqlQueryResultView
 from runloop_api_client.sdk import AsyncAxon
+from runloop_api_client.types.axons.sql_statement_params import SqlStatementParams
 
 
 class TestAsyncAxon:
@@ -75,3 +76,38 @@ async def test_subscribe_sse(self, mock_async_client: AsyncMock) -> None:
 
         assert result == mock_stream
         mock_async_client.axons.subscribe_sse.assert_awaited_once_with("axn_123")
+
+    @pytest.mark.asyncio
+    async def test_sql_query(self, mock_async_client: AsyncMock) -> None:
+        """Test sql.query method delegates to client.axons.sql.query."""
+        mock_result = MockSqlQueryResultView()
+        mock_async_client.axons.sql.query = AsyncMock(return_value=mock_result)
+
+        axon = AsyncAxon(mock_async_client, "axn_123")
+        result = await axon.sql.query(sql="SELECT * FROM test WHERE id = ?", params=[1])
+
+        assert result == mock_result
+        mock_async_client.axons.sql.query.assert_awaited_once_with(
+            "axn_123",
+            sql="SELECT * FROM test WHERE id = ?",
+            params=[1],
+        )
+
+    @pytest.mark.asyncio
+    async def test_sql_batch(self, mock_async_client: AsyncMock) -> None:
+        """Test sql.batch method delegates to client.axons.sql.batch."""
+        mock_result = MockSqlBatchResultView()
+        mock_async_client.axons.sql.batch = AsyncMock(return_value=mock_result)
+
+        statements: list[SqlStatementParams] = [
+            {"sql": "CREATE TABLE t (id INTEGER PRIMARY KEY)"},
+            {"sql": "INSERT INTO t (id) VALUES (?)", "params": [1]},
+        ]
+        axon = AsyncAxon(mock_async_client, "axn_123")
+        result = await axon.sql.batch(statements=statements)
+
+        assert result == mock_result
+        mock_async_client.axons.sql.batch.assert_awaited_once_with(
+            "axn_123",
+            statements=statements,
+        )

tests/sdk/test_axon.py

@@ -4,8 +4,9 @@
 
 from unittest.mock import Mock
 
-from tests.sdk.conftest import MockAxonView, MockPublishResultView
+from tests.sdk.conftest import MockAxonView, MockPublishResultView, MockSqlBatchResultView, MockSqlQueryResultView
 from runloop_api_client.sdk import Axon
+from runloop_api_client.types.axons.sql_statement_params import SqlStatementParams
 
 
 class TestAxon:
@@ -70,3 +71,36 @@ def test_subscribe_sse(self, mock_client: Mock) -> None:
 
         assert result == mock_stream
         mock_client.axons.subscribe_sse.assert_called_once_with("axn_123")
+
+    def test_sql_query(self, mock_client: Mock) -> None:
+        """Test sql.query method delegates to client.axons.sql.query."""
+        mock_result = MockSqlQueryResultView()
+        mock_client.axons.sql.query.return_value = mock_result
+
+        axon = Axon(mock_client, "axn_123")
+        result = axon.sql.query(sql="SELECT * FROM test WHERE id = ?", params=[1])
+
+        assert result == mock_result
+        mock_client.axons.sql.query.assert_called_once_with(
+            "axn_123",
+            sql="SELECT * FROM test WHERE id = ?",
+            params=[1],
+        )
+
+    def test_sql_batch(self, mock_client: Mock) -> None:
+        """Test sql.batch method delegates to client.axons.sql.batch."""
+        mock_result = MockSqlBatchResultView()
+        mock_client.axons.sql.batch.return_value = mock_result
+
+        statements: list[SqlStatementParams] = [
+            {"sql": "CREATE TABLE t (id INTEGER PRIMARY KEY)"},
+            {"sql": "INSERT INTO t (id) VALUES (?)", "params": [1]},
+        ]
+        axon = Axon(mock_client, "axn_123")
+        result = axon.sql.batch(statements=statements)
+
+        assert result == mock_result
+        mock_client.axons.sql.batch.assert_called_once_with(
+            "axn_123",
+            statements=statements,
+        )