Skip to content

Commit d1cf846

Browse files
lukekimclaudephillipleblanc
authored
feat(query): parameterized async /v1/queries submit + options (#80)
The async /v1/queries API (added in #68/#70) could only submit a bare SQL string. Add parameterized submit and submit options, matching the Spice v2 API (verified against spiceai/spiceai@v2.1.0): - SpiceClient::query_with_bindings(sql, QueryParameters) — mirrors the sync sql_with_bindings; positional $1, $2, ... scalar bindings. - SpiceClient::query_with_options(sql, QuerySubmitOptions) — bindings plus timeout_seconds and maximum_size. - QueryParameters::to_json_values() encodes scalar bindings as the JSON array the HTTP API expects; non-scalar/binary/non-finite params fail fast with QueryError::InvalidParameter. - Export QuerySubmitOptions; add README example. Tests: +9 (6 param-encoding unit tests, 3 client wiremock tests). All lib tests and doctests pass. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com>
1 parent 6be5cb3 commit d1cf846

5 files changed

Lines changed: 439 additions & 7 deletions

File tree

‎README.md‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,34 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
140140
}
141141
```
142142

143+
Async queries also accept positional bindings (`$1`, `$2`, ...) and submit options. Use `query_with_bindings` for the common parameterized case, or `query_with_options` to also set an execution `timeout_seconds` or a `maximum_size` cap on the materialized result.
144+
145+
```rust,no_run
146+
use spiceai::{ClientBuilder, QueryParameters, QuerySubmitOptions};
147+
148+
#[tokio::main]
149+
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
150+
let client = ClientBuilder::new()
151+
.http_url("http://localhost:8090")
152+
.build()
153+
.await?;
154+
155+
let job = client
156+
.query_with_options(
157+
"SELECT * FROM large_table WHERE status = $1 AND created_at > $2",
158+
QuerySubmitOptions::new()
159+
.bindings(QueryParameters::new().push("active").push("2025-01-01"))
160+
.timeout_seconds(300)
161+
.maximum_size(100_000_000),
162+
)
163+
.await?;
164+
165+
let result = job.wait().await?;
166+
println!("completed with {} rows", result.total_rows);
167+
Ok(())
168+
}
169+
```
170+
143171
## Documentation
144172

145173
Check out our [Documentation](https://docs.spice.ai/sdks/rust-sdk) to learn more about how to use the Rust SDK.

‎src/client.rs‎

Lines changed: 197 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::flight::RetryableQueryStream;
22
use crate::params::{QueryParameterError, QueryParameters};
3-
use crate::query::{QueryError, QueryHttpClient, QueryJob};
3+
use crate::query::{QueryError, QueryHttpClient, QueryJob, QuerySubmitOptions};
44
use crate::util::{FibonacciBackoffBuilder, RetryError, retry};
55
use crate::{
66
config::{GenericError, SPICE_CLOUD_FLIGHT_ADDR, SPICE_LOCAL_FLIGHT_ADDR},
@@ -274,7 +274,125 @@ impl SpiceClient {
274274
.to_string(),
275275
})?;
276276

277-
let response = http_client.submit(sql).await?;
277+
let response = http_client.submit(sql, None, None, None).await?;
278+
Ok(QueryJob::new(response.query_id, Arc::clone(http_client)))
279+
}
280+
281+
/// Submits an async parameterized SQL query using scalar bindings.
282+
///
283+
/// This is the async (`/v1/queries`) counterpart to
284+
/// [`sql_with_bindings()`](Self::sql_with_bindings): it binds a single row of
285+
/// positional scalar values (`$1`, `$2`, ...) and returns a [`QueryJob`]
286+
/// handle. To also set a timeout or result-size cap, use
287+
/// [`query_with_options()`](Self::query_with_options).
288+
///
289+
/// **Note:** Requires [`http_url()`](SpiceClientBuilder::http_url) to be configured.
290+
///
291+
/// # Example
292+
///
293+
/// ```no_run
294+
/// # use spiceai::{ClientBuilder, QueryParameters};
295+
/// # #[tokio::main]
296+
/// # async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
297+
/// let client = ClientBuilder::new()
298+
/// .http_url("http://localhost:8090")
299+
/// .build()
300+
/// .await?;
301+
///
302+
/// let job = client
303+
/// .query_with_bindings(
304+
/// "SELECT * FROM large_table WHERE status = $1",
305+
/// QueryParameters::new().push("active"),
306+
/// )
307+
/// .await?;
308+
/// println!("Query submitted: {}", job.id());
309+
/// # Ok(())
310+
/// # }
311+
/// ```
312+
///
313+
/// # Errors
314+
///
315+
/// - [`QueryError::InvalidParameter`] if a binding cannot be encoded as JSON
316+
/// - [`QueryError::ClusterModeRequired`] if async queries are not enabled
317+
/// - [`QueryError::SubmitFailed`] if the query submission fails
318+
/// - [`QueryError::HttpError`] if the HTTP endpoint is not configured or unreachable
319+
pub async fn query_with_bindings(
320+
&self,
321+
sql: &str,
322+
params: QueryParameters,
323+
) -> Result<QueryJob, QueryError> {
324+
self.query_with_options(sql, QuerySubmitOptions::new().bindings(params))
325+
.await
326+
}
327+
328+
/// Submits an async SQL query with explicit submit options.
329+
///
330+
/// [`QuerySubmitOptions`] carries the optional bind parameters, a server-side
331+
/// execution `timeout_seconds`, and a `maximum_size` cap on the materialized
332+
/// result. Any option left unset is omitted from the request so the server
333+
/// applies its default.
334+
///
335+
/// **Note:** Requires [`http_url()`](SpiceClientBuilder::http_url) to be configured.
336+
///
337+
/// # Example
338+
///
339+
/// ```no_run
340+
/// # use spiceai::{ClientBuilder, QueryParameters, QuerySubmitOptions};
341+
/// # #[tokio::main]
342+
/// # async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
343+
/// let client = ClientBuilder::new()
344+
/// .http_url("http://localhost:8090")
345+
/// .build()
346+
/// .await?;
347+
///
348+
/// let job = client
349+
/// .query_with_options(
350+
/// "SELECT * FROM large_table WHERE status = $1",
351+
/// QuerySubmitOptions::new()
352+
/// .bindings(QueryParameters::new().push("active"))
353+
/// .timeout_seconds(300)
354+
/// .maximum_size(100_000_000),
355+
/// )
356+
/// .await?;
357+
/// let result = job.wait().await?;
358+
/// println!("Completed with {} rows", result.total_rows);
359+
/// # Ok(())
360+
/// # }
361+
/// ```
362+
///
363+
/// # Errors
364+
///
365+
/// - [`QueryError::InvalidParameter`] if a binding cannot be encoded as JSON
366+
/// - [`QueryError::ClusterModeRequired`] if async queries are not enabled
367+
/// - [`QueryError::SubmitFailed`] if the query submission fails
368+
/// - [`QueryError::HttpError`] if the HTTP endpoint is not configured or unreachable
369+
pub async fn query_with_options(
370+
&self,
371+
sql: &str,
372+
options: QuerySubmitOptions,
373+
) -> Result<QueryJob, QueryError> {
374+
let http_client = self.http_client.as_ref().ok_or(QueryError::HttpError {
375+
message: "HTTP endpoint not configured. Use ClientBuilder::http_url() to set it."
376+
.to_string(),
377+
})?;
378+
379+
let parameters = match &options.bindings {
380+
Some(params) if !params.is_empty() => Some(
381+
params
382+
.to_json_values()
383+
.map_err(|source| QueryError::InvalidParameter { source })?,
384+
),
385+
_ => None,
386+
};
387+
388+
let response = http_client
389+
.submit(
390+
sql,
391+
parameters,
392+
options.timeout_seconds,
393+
options.maximum_size,
394+
)
395+
.await?;
278396
Ok(QueryJob::new(response.query_id, Arc::clone(http_client)))
279397
}
280398

@@ -984,6 +1102,83 @@ mod tests {
9841102
));
9851103
}
9861104

1105+
#[tokio::test]
1106+
async fn test_query_with_bindings_submits_parameters() {
1107+
let server = MockServer::start().await;
1108+
let query_id = "qry_params";
1109+
let sql = "SELECT * FROM t WHERE status = $1 AND n > $2";
1110+
1111+
Mock::given(method("POST"))
1112+
.and(path("/v1/queries"))
1113+
.and(body_json(
1114+
json!({ "sql": sql, "parameters": ["active", 5] }),
1115+
))
1116+
.respond_with(ResponseTemplate::new(202).set_body_json(json!({
1117+
"query_id": query_id,
1118+
"status": "PENDING",
1119+
"status_url": format!("{}/v1/queries/{query_id}/status", server.uri()),
1120+
"results_url": format!("{}/v1/queries/{query_id}/results", server.uri())
1121+
})))
1122+
.mount(&server)
1123+
.await;
1124+
1125+
let client = test_client(Some(&server.uri()));
1126+
let job = client
1127+
.query_with_bindings(sql, QueryParameters::new().push("active").push(5_i64))
1128+
.await
1129+
.expect("submit parameterized async query");
1130+
assert_eq!(job.id(), query_id);
1131+
}
1132+
1133+
#[tokio::test]
1134+
async fn test_query_with_options_submits_all_fields() {
1135+
let server = MockServer::start().await;
1136+
let query_id = "qry_opts";
1137+
let sql = "SELECT * FROM t WHERE status = $1";
1138+
1139+
Mock::given(method("POST"))
1140+
.and(path("/v1/queries"))
1141+
.and(body_json(json!({
1142+
"sql": sql,
1143+
"parameters": ["active"],
1144+
"timeout_seconds": 300,
1145+
"maximum_size": 100_000_000
1146+
})))
1147+
.respond_with(ResponseTemplate::new(202).set_body_json(json!({
1148+
"query_id": query_id,
1149+
"status": "PENDING",
1150+
"status_url": format!("{}/v1/queries/{query_id}/status", server.uri()),
1151+
"results_url": format!("{}/v1/queries/{query_id}/results", server.uri())
1152+
})))
1153+
.mount(&server)
1154+
.await;
1155+
1156+
let client = test_client(Some(&server.uri()));
1157+
let job = client
1158+
.query_with_options(
1159+
sql,
1160+
QuerySubmitOptions::new()
1161+
.bindings(QueryParameters::new().push("active"))
1162+
.timeout_seconds(300)
1163+
.maximum_size(100_000_000),
1164+
)
1165+
.await
1166+
.expect("submit async query with options");
1167+
assert_eq!(job.id(), query_id);
1168+
}
1169+
1170+
#[tokio::test]
1171+
async fn test_query_with_bindings_rejects_unsupported_param() {
1172+
// A binary bind value has no JSON scalar form; the client must reject it
1173+
// locally, before any HTTP request is attempted.
1174+
let client = test_client(Some("http://127.0.0.1:1"));
1175+
let err = client
1176+
.query_with_bindings("SELECT $1", QueryParameters::new().push(vec![1_u8, 2, 3]))
1177+
.await
1178+
.expect_err("binary bind parameter should be rejected before submit");
1179+
assert!(matches!(err, QueryError::InvalidParameter { .. }));
1180+
}
1181+
9871182
#[tokio::test]
9881183
async fn test_async_query_wrappers_and_job_lifecycle() {
9891184
let server = MockServer::start().await;

‎src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ pub use dataset::{
1919
pub use params::{QueryParameter, QueryParameterError, QueryParameters};
2020
pub use query::{
2121
QueryError, QueryInfo, QueryJob, QueryListResponse, QueryResult, QueryResultStream,
22-
QueryStatus, QuerySummary,
22+
QueryStatus, QuerySubmitOptions, QuerySummary,
2323
};
2424

2525
// Further public exports and integrations

0 commit comments

Comments
 (0)