cp dines

Author: dines-rl

src/runloop_api_client/sdk/_types.py

@@ -1,4 +1,4 @@
-from typing import Union, Literal, Callable, Optional
+from typing import Union, Literal, Callable, Iterable, Optional
 from typing_extensions import TypedDict
 
 from ..types import (
@@ -39,6 +39,7 @@
 )
 from .._types import Body, Query, Headers, Timeout, NotGiven
 from ..lib.polling import PollingConfig
+from ..types.axons.sql_statement_params import SqlStatementParams
 from ..types.devboxes import DiskSnapshotListParams, DiskSnapshotUpdateParams
 from ..types.scenarios import ScorerListParams, ScorerCreateParams, ScorerUpdateParams
 from ..types.devbox_create_params import DevboxBaseCreateParams
@@ -205,7 +206,7 @@ class SDKAxonSqlQueryParams(LongRequestOptions, total=False):
 
 
 class SDKAxonSqlBatchParams(LongRequestOptions, total=False):
-    statements: list[dict[str, object]]
+    statements: Iterable[SqlStatementParams]
     """The SQL statements to execute atomically within a transaction."""
 
 

tests/sdk/test_async_axon.py

@@ -8,6 +8,7 @@
 
 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:
@@ -98,7 +99,7 @@ async def test_sql_batch(self, mock_async_client: AsyncMock) -> None:
         mock_result = MockSqlBatchResultView()
         mock_async_client.axons.sql.batch = AsyncMock(return_value=mock_result)
 
-        statements = [
+        statements: list[SqlStatementParams] = [
             {"sql": "CREATE TABLE t (id INTEGER PRIMARY KEY)"},
             {"sql": "INSERT INTO t (id) VALUES (?)", "params": [1]},
         ]

tests/sdk/test_axon.py

@@ -6,6 +6,7 @@
 
 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:
@@ -91,7 +92,7 @@ def test_sql_batch(self, mock_client: Mock) -> None:
         mock_result = MockSqlBatchResultView()
         mock_client.axons.sql.batch.return_value = mock_result
 
-        statements = [
+        statements: list[SqlStatementParams] = [
             {"sql": "CREATE TABLE t (id INTEGER PRIMARY KEY)"},
             {"sql": "INSERT INTO t (id) VALUES (?)", "params": [1]},
         ]