diff --git a/src/runloop_api_client/sdk/__init__.py b/src/runloop_api_client/sdk/__init__.py index 136080fe8..38b3d0e41 100644 --- a/src/runloop_api_client/sdk/__init__.py +++ b/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", diff --git a/src/runloop_api_client/sdk/_types.py b/src/runloop_api_client/sdk/_types.py index ac678f126..aa7e0208a 100644 --- a/src/runloop_api_client/sdk/_types.py +++ b/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 diff --git a/src/runloop_api_client/sdk/async_axon.py b/src/runloop_api_client/sdk/async_axon.py index dda6a3037..f4410126f 100644 --- a/src/runloop_api_client/sdk/async_axon.py +++ b/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: diff --git a/src/runloop_api_client/sdk/axon.py b/src/runloop_api_client/sdk/axon.py index eb310a88f..c86a31efa 100644 --- a/src/runloop_api_client/sdk/axon.py +++ b/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: diff --git a/tests/sdk/conftest.py b/tests/sdk/conftest.py index ab8d0d048..92ab25f42 100644 --- a/tests/sdk/conftest.py +++ b/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.""" diff --git a/tests/sdk/test_async_axon.py b/tests/sdk/test_async_axon.py index f5ce9761d..f24101673 100644 --- a/tests/sdk/test_async_axon.py +++ b/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, + ) diff --git a/tests/sdk/test_axon.py b/tests/sdk/test_axon.py index 58871d339..16cf565c3 100644 --- a/tests/sdk/test_axon.py +++ b/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, + ) diff --git a/tests/smoketests/sdk/test_async_axon.py b/tests/smoketests/sdk/test_async_axon.py index ced85820d..065535ad2 100644 --- a/tests/smoketests/sdk/test_async_axon.py +++ b/tests/smoketests/sdk/test_async_axon.py @@ -3,11 +3,17 @@ from __future__ import annotations import json +import uuid import pytest from runloop_api_client.sdk import AsyncRunloopSDK + +def _unique_table() -> str: + return f"t_{uuid.uuid4().hex[:12]}" + + pytestmark = [pytest.mark.smoketest, pytest.mark.asyncio] THIRTY_SECOND_TIMEOUT = 30 @@ -69,6 +75,48 @@ async def test_axon_publish(self, async_sdk_client: AsyncRunloopSDK) -> None: pass +class TestAsyncAxonSql: + """Test async axon SQL operations.""" + + @pytest.mark.timeout(THIRTY_SECOND_TIMEOUT) + async def test_sql_query_create_and_select(self, async_sdk_client: AsyncRunloopSDK) -> None: + """Test creating a table and querying it via sql.query.""" + axon = await async_sdk_client.axon.create() + table = _unique_table() + + await axon.sql.query(sql=f"CREATE TABLE {table} (id INTEGER PRIMARY KEY, value TEXT)") + + await axon.sql.query(sql=f"INSERT INTO {table} (id, value) VALUES (?, ?)", params=[1, "hello"]) + + result = await axon.sql.query(sql=f"SELECT * FROM {table} WHERE id = ?", params=[1]) + + assert result.columns is not None + assert len(result.columns) > 0 + assert len(result.rows) == 1 + assert result.meta.duration_ms >= 0 + + @pytest.mark.timeout(THIRTY_SECOND_TIMEOUT) + async def test_sql_batch(self, async_sdk_client: AsyncRunloopSDK) -> None: + """Test executing multiple statements atomically via sql.batch.""" + axon = await async_sdk_client.axon.create() + table = _unique_table() + + result = await axon.sql.batch( + statements=[ + {"sql": f"CREATE TABLE {table} (id INTEGER PRIMARY KEY, name TEXT)"}, + {"sql": f"INSERT INTO {table} (id, name) VALUES (?, ?)", "params": [1, "alice"]}, + {"sql": f"INSERT INTO {table} (id, name) VALUES (?, ?)", "params": [2, "bob"]}, + {"sql": f"SELECT * FROM {table} ORDER BY id"}, + ], + ) + + assert result.results is not None + assert len(result.results) == 4 + select_result = result.results[3] + assert select_result.success is not None + assert len(select_result.success.rows) == 2 + + class TestAsyncAxonListing: """Test axon listing operations.""" diff --git a/tests/smoketests/sdk/test_axon.py b/tests/smoketests/sdk/test_axon.py index c8947d006..ffdd2dd7e 100644 --- a/tests/smoketests/sdk/test_axon.py +++ b/tests/smoketests/sdk/test_axon.py @@ -3,11 +3,17 @@ from __future__ import annotations import json +import uuid import pytest from runloop_api_client.sdk import RunloopSDK + +def _unique_table() -> str: + return f"t_{uuid.uuid4().hex[:12]}" + + pytestmark = [pytest.mark.smoketest] THIRTY_SECOND_TIMEOUT = 30 @@ -69,6 +75,48 @@ def test_axon_publish(self, sdk_client: RunloopSDK) -> None: pass +class TestAxonSql: + """Test axon SQL operations.""" + + @pytest.mark.timeout(THIRTY_SECOND_TIMEOUT) + def test_sql_query_create_and_select(self, sdk_client: RunloopSDK) -> None: + """Test creating a table and querying it via sql.query.""" + axon = sdk_client.axon.create() + table = _unique_table() + + axon.sql.query(sql=f"CREATE TABLE {table} (id INTEGER PRIMARY KEY, value TEXT)") + + axon.sql.query(sql=f"INSERT INTO {table} (id, value) VALUES (?, ?)", params=[1, "hello"]) + + result = axon.sql.query(sql=f"SELECT * FROM {table} WHERE id = ?", params=[1]) + + assert result.columns is not None + assert len(result.columns) > 0 + assert len(result.rows) == 1 + assert result.meta.duration_ms >= 0 + + @pytest.mark.timeout(THIRTY_SECOND_TIMEOUT) + def test_sql_batch(self, sdk_client: RunloopSDK) -> None: + """Test executing multiple statements atomically via sql.batch.""" + axon = sdk_client.axon.create() + table = _unique_table() + + result = axon.sql.batch( + statements=[ + {"sql": f"CREATE TABLE {table} (id INTEGER PRIMARY KEY, name TEXT)"}, + {"sql": f"INSERT INTO {table} (id, name) VALUES (?, ?)", "params": [1, "alice"]}, + {"sql": f"INSERT INTO {table} (id, name) VALUES (?, ?)", "params": [2, "bob"]}, + {"sql": f"SELECT * FROM {table} ORDER BY id"}, + ], + ) + + assert result.results is not None + assert len(result.results) == 4 + select_result = result.results[3] + assert select_result.success is not None + assert len(select_result.success.rows) == 2 + + class TestAxonListing: """Test axon listing operations.""" diff --git a/uv.lock b/uv.lock index a43d415ef..970c26cc8 100644 --- a/uv.lock +++ b/uv.lock @@ -2386,7 +2386,7 @@ wheels = [ [[package]] name = "runloop-api-client" -version = "1.13.0" +version = "1.13.1" source = { editable = "." } dependencies = [ { name = "anyio" },