Skip to content

Commit f3e62bb

Browse files
ClaudeClaude
Claude
authored and
Claude
committed
Merge remote-tracking branch 'origin/trunk' into feat/runtime-status
# Conflicts: # README.md
2 parents 098b383 + d1cf846 commit f3e62bb

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
### Runtime health and status
144172

145173
`is_ready()` is a single boolean for the whole runtime. When you need to know *which*

‎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},
@@ -275,7 +275,125 @@ impl SpiceClient {
275275
.to_string(),
276276
})?;
277277

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

@@ -1032,6 +1150,83 @@ mod tests {
10321150
));
10331151
}
10341152

1153+
#[tokio::test]
1154+
async fn test_query_with_bindings_submits_parameters() {
1155+
let server = MockServer::start().await;
1156+
let query_id = "qry_params";
1157+
let sql = "SELECT * FROM t WHERE status = $1 AND n > $2";
1158+
1159+
Mock::given(method("POST"))
1160+
.and(path("/v1/queries"))
1161+
.and(body_json(
1162+
json!({ "sql": sql, "parameters": ["active", 5] }),
1163+
))
1164+
.respond_with(ResponseTemplate::new(202).set_body_json(json!({
1165+
"query_id": query_id,
1166+
"status": "PENDING",
1167+
"status_url": format!("{}/v1/queries/{query_id}/status", server.uri()),
1168+
"results_url": format!("{}/v1/queries/{query_id}/results", server.uri())
1169+
})))
1170+
.mount(&server)
1171+
.await;
1172+
1173+
let client = test_client(Some(&server.uri()));
1174+
let job = client
1175+
.query_with_bindings(sql, QueryParameters::new().push("active").push(5_i64))
1176+
.await
1177+
.expect("submit parameterized async query");
1178+
assert_eq!(job.id(), query_id);
1179+
}
1180+
1181+
#[tokio::test]
1182+
async fn test_query_with_options_submits_all_fields() {
1183+
let server = MockServer::start().await;
1184+
let query_id = "qry_opts";
1185+
let sql = "SELECT * FROM t WHERE status = $1";
1186+
1187+
Mock::given(method("POST"))
1188+
.and(path("/v1/queries"))
1189+
.and(body_json(json!({
1190+
"sql": sql,
1191+
"parameters": ["active"],
1192+
"timeout_seconds": 300,
1193+
"maximum_size": 100_000_000
1194+
})))
1195+
.respond_with(ResponseTemplate::new(202).set_body_json(json!({
1196+
"query_id": query_id,
1197+
"status": "PENDING",
1198+
"status_url": format!("{}/v1/queries/{query_id}/status", server.uri()),
1199+
"results_url": format!("{}/v1/queries/{query_id}/results", server.uri())
1200+
})))
1201+
.mount(&server)
1202+
.await;
1203+
1204+
let client = test_client(Some(&server.uri()));
1205+
let job = client
1206+
.query_with_options(
1207+
sql,
1208+
QuerySubmitOptions::new()
1209+
.bindings(QueryParameters::new().push("active"))
1210+
.timeout_seconds(300)
1211+
.maximum_size(100_000_000),
1212+
)
1213+
.await
1214+
.expect("submit async query with options");
1215+
assert_eq!(job.id(), query_id);
1216+
}
1217+
1218+
#[tokio::test]
1219+
async fn test_query_with_bindings_rejects_unsupported_param() {
1220+
// A binary bind value has no JSON scalar form; the client must reject it
1221+
// locally, before any HTTP request is attempted.
1222+
let client = test_client(Some("http://127.0.0.1:1"));
1223+
let err = client
1224+
.query_with_bindings("SELECT $1", QueryParameters::new().push(vec![1_u8, 2, 3]))
1225+
.await
1226+
.expect_err("binary bind parameter should be rejected before submit");
1227+
assert!(matches!(err, QueryError::InvalidParameter { .. }));
1228+
}
1229+
10351230
#[tokio::test]
10361231
async fn test_async_query_wrappers_and_job_lifecycle() {
10371232
let server = MockServer::start().await;

‎src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ pub use dataset::{
2020
pub use params::{QueryParameter, QueryParameterError, QueryParameters};
2121
pub use query::{
2222
QueryError, QueryInfo, QueryJob, QueryListResponse, QueryResult, QueryResultStream,
23-
QueryStatus, QuerySummary,
23+
QueryStatus, QuerySubmitOptions, QuerySummary,
2424
};
2525
pub use status::{ComponentStatus, ConnectionDetails, StatusError};
2626

0 commit comments

Comments
 (0)