cp dines

Author: dines-rl

src/runloop_api_client/sdk/_types.py

@@ -1,4 +1,4 @@
-from typing import Union, Literal, Callable, Iterable, Optional
+from typing import Union, Literal, Callable, Optional
 from typing_extensions import TypedDict
 
 from ..types import (
@@ -42,8 +42,9 @@
 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.axons.sql_statement_params import SqlStatementParams
 from ..types.benchmark_start_run_params import BenchmarkSelfStartRunParams
 from ..types.devbox_execute_async_params import DevboxNiceExecuteAsyncParams
 
@@ -197,17 +198,12 @@ class SDKAxonPublishParams(AxonPublishParams, LongRequestOptions):
     pass
 
 
-class SDKAxonSqlQueryParams(LongRequestOptions, total=False):
-    sql: str
-    """SQL query with ?-style positional placeholders."""
-
-    params: list[object]
-    """Positional parameter bindings for ? placeholders."""
+class SDKAxonSqlQueryParams(SqlQueryParams, LongRequestOptions):
+    pass
 
 
-class SDKAxonSqlBatchParams(LongRequestOptions, total=False):
-    statements: Iterable[SqlStatementParams]
-    """The SQL statements to execute atomically within a transaction."""
+class SDKAxonSqlBatchParams(SqlBatchParams, LongRequestOptions):
+    pass
 
 
 class SDKScenarioListParams(ScenarioListParams, BaseRequestOptions):

tests/smoketests/sdk/test_axon.py

@@ -3,11 +3,16 @@
 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
@@ -76,12 +81,13 @@ class TestAxonSql:
     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="CREATE TABLE IF NOT EXISTS smoke_test (id INTEGER PRIMARY KEY, value TEXT)")
+        axon.sql.query(sql=f"CREATE TABLE {table} (id INTEGER PRIMARY KEY, value TEXT)")
 
-        axon.sql.query(sql="INSERT INTO smoke_test (id, value) VALUES (?, ?)", params=[1, "hello"])
+        axon.sql.query(sql=f"INSERT INTO {table} (id, value) VALUES (?, ?)", params=[1, "hello"])
 
-        result = axon.sql.query(sql="SELECT * FROM smoke_test WHERE id = ?", params=[1])
+        result = axon.sql.query(sql=f"SELECT * FROM {table} WHERE id = ?", params=[1])
 
         assert result.columns is not None
         assert len(result.columns) > 0