Skip to content

Commit 7727041

Browse files
claudespiceclaudespice
committed
Merge remote-tracking branch 'origin/trunk' into feat/search
# Conflicts: # README.md # src/client.rs
2 parents 852a9e0 + 97add5e commit 7727041

6 files changed

Lines changed: 855 additions & 44 deletions

File tree

‎README.md‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,48 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
168168
}
169169
```
170170

171+
### List and cancel running queries
172+
173+
`active_queries()` reports the synchronous queries currently running in the caller's scope — the ones started by `sql()`, FlightSQL, `/v1/sql`, NSQL, and search — and `cancel_active_query()` stops one by id.
174+
175+
The runtime does not hand a query's id back to the client that submitted it, so the two are used together: list to find the query, then cancel it.
176+
177+
Two boundaries apply, and a query is reachable only inside both.
178+
179+
**One runtime instance.** The runtime tracks active synchronous queries in memory, per instance, and these endpoints report only what the instance answering them knows. A `Client` configures its Flight and HTTP endpoints independently, so behind a load balancer the query submitted over Flight may be running on a different instance than the one answering here — it will not be listed, and its id reports as not found. Point `http_url()` at the instance running the query.
180+
181+
**One authenticated principal**, not a `Client` instance. The principal is whatever credential the runtime authenticates — an API key or a client certificate — so every client presenting the same credential lists and cancels the same queries. Only requests for which the runtime establishes no principal at all share the `public` scope. A query outside the caller's scope is reported as if it did not exist.
182+
183+
> **Runtime version.** Principal scoping on these two endpoints landed in [spiceai/spiceai#12841](https://github.com/spiceai/spiceai/pull/12841) and is in no runtime release up to and including `v2.1.5`. Against an earlier runtime both calls operate on every active query the instance holds, for any caller with write access. Check your runtime version before relying on the scope described above.
184+
185+
```rust,no_run
186+
use spiceai::ClientBuilder;
187+
188+
#[tokio::main]
189+
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
190+
let client = ClientBuilder::new()
191+
.http_url("http://localhost:8090")
192+
.build()
193+
.await?;
194+
195+
let active = client.active_queries().await?;
196+
println!("{} queries running", active.total_count);
197+
198+
for query in &active.queries {
199+
println!("{} [{}] {}", query.query_id, query.protocol, query.sql_preview);
200+
}
201+
202+
if let Some(query) = active.queries.first() {
203+
let cancelled = client.cancel_active_query(&query.query_id).await?;
204+
println!("{} is now {}", cancelled.query_id, cancelled.status);
205+
}
206+
207+
Ok(())
208+
}
209+
```
210+
211+
To cancel an *async query job* instead, use `cancel_query()` — see [Async query jobs](#async-query-jobs-and-dataset-refresh) above.
212+
171213
### Search
172214

173215
`search` finds documents similar to a piece of text using the runtime's `/v1/search` endpoint. It runs against datasets that have an embedding column and a loaded embedding model — see [Search & Retrieval](https://docs.spice.ai/features/search-and-retrieval) for how to configure them. Like dataset refresh, it uses the HTTP API, so `http_url()` must be configured.

‎src/active_query.rs‎

Lines changed: 255 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,255 @@
1+
//! Listing and cancelling *synchronous* queries running on the runtime.
2+
//!
3+
//! This is distinct from the async query jobs in [`crate::query`]. A synchronous
4+
//! query is one started by [`Client::sql()`](crate::Client::sql), FlightSQL,
5+
//! `/v1/sql`, NSQL, or `/v1/search` — it streams results back on the connection
6+
//! that started it. Async jobs are submitted with
7+
//! [`Client::query()`](crate::Client::query), are polled for completion, and
8+
//! require the runtime to be running in cluster mode.
9+
//!
10+
//! The runtime assigns a `query_id` to every synchronous query but does not
11+
//! return it to the client that submitted it, so the two operations here are
12+
//! used together: list the active queries to find the one you want, then cancel
13+
//! it by id.
14+
//!
15+
//! # Scope
16+
//!
17+
//! Two boundaries apply, and a query is reachable only inside both.
18+
//!
19+
//! **One runtime instance.** The runtime tracks active synchronous queries in
20+
//! memory, per instance, and these endpoints report only what the instance
21+
//! answering them knows. This client configures its Flight and HTTP endpoints
22+
//! independently, so behind a load balancer the query submitted over Flight may
23+
//! be running on a different instance than the one answering here — it will not
24+
//! be listed, and its id reports as not found. Point
25+
//! [`http_url()`](crate::ClientBuilder::http_url) at the instance running the
26+
//! query.
27+
//!
28+
//! **One authenticated principal**, not a [`Client`](crate::Client) instance.
29+
//! The principal is whatever credential the runtime authenticates — an API key
30+
//! or a client certificate — so every client presenting the same credential
31+
//! shares one scope and can list and cancel the others' queries. Only requests
32+
//! for which the runtime establishes no principal at all share the `public`
33+
//! scope. A query outside the caller's scope is reported as if it did not
34+
//! exist.
35+
//!
36+
//! **Runtime version.** Principal scoping on these two endpoints landed in
37+
//! [spiceai/spiceai#12841](https://github.com/spiceai/spiceai/pull/12841) and is
38+
//! in no runtime release up to and including `v2.1.5`. Against an earlier
39+
//! runtime both calls operate on every active query the instance holds, for any
40+
//! caller with write access.
41+
//!
42+
//! # Example
43+
//!
44+
//! ```no_run
45+
//! use spiceai::ClientBuilder;
46+
//!
47+
//! #[tokio::main]
48+
//! async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
49+
//! let client = ClientBuilder::new()
50+
//! .http_url("http://localhost:8090")
51+
//! .build()
52+
//! .await?;
53+
//!
54+
//! let active = client.active_queries().await?;
55+
//! for query in &active.queries {
56+
//! println!("{} [{}] {}", query.query_id, query.protocol, query.sql_preview);
57+
//! }
58+
//!
59+
//! if let Some(query) = active.queries.first() {
60+
//! client.cancel_active_query(&query.query_id).await?;
61+
//! }
62+
//!
63+
//! Ok(())
64+
//! }
65+
//! ```
66+
67+
use serde::Deserialize;
68+
use snafu::Snafu;
69+
70+
/// Errors that can occur while listing or cancelling active synchronous queries.
71+
#[derive(Debug, Snafu)]
72+
pub enum ActiveQueryError {
73+
/// No active synchronous query with this id is in the caller's scope.
74+
///
75+
/// The runtime reports a query submitted under another principal the same
76+
/// way it reports one that does not exist, so a caller cannot probe for
77+
/// other principals' query ids.
78+
#[snafu(display(
79+
"No active query '{query_id}' found. It may have already finished, it was submitted under a different principal, or it is running on another runtime instance."
80+
))]
81+
NotFound { query_id: String },
82+
83+
/// The supplied id is not a UUID, so it cannot name a query.
84+
#[snafu(display(
85+
"Query id '{query_id}' is not a valid UUID. Use the query_id from active_queries()."
86+
))]
87+
InvalidQueryId { query_id: String },
88+
89+
/// The configured API key does not grant write access.
90+
#[snafu(display(
91+
"The configured API key does not allow cancelling queries. Use a key with write access."
92+
))]
93+
WriteAccessRequired,
94+
95+
/// The request failed with an unexpected status code.
96+
#[snafu(display("Request failed (HTTP {status_code}): {response_body}"))]
97+
RequestFailed {
98+
/// HTTP status code returned by the runtime.
99+
status_code: u16,
100+
/// Response body returned by the runtime.
101+
response_body: String,
102+
},
103+
104+
/// HTTP transport error.
105+
#[snafu(display("Request failed: {message}"))]
106+
HttpError { message: String },
107+
108+
/// The response could not be parsed.
109+
#[snafu(display("Failed to parse response: {message}"))]
110+
ParseError { message: String },
111+
}
112+
113+
/// A synchronous query currently executing on the runtime.
114+
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
115+
pub struct ActiveQuery {
116+
/// Server-assigned id, used to cancel the query.
117+
pub query_id: String,
118+
/// The protocol the query arrived on: `http`, `flight`, `flightsql`, or
119+
/// `internal`.
120+
pub protocol: String,
121+
/// The query's SQL, truncated by the runtime for display.
122+
pub sql_preview: String,
123+
/// When the query started, in milliseconds since the Unix epoch.
124+
pub started_at_ms: u64,
125+
}
126+
127+
/// The set of synchronous queries the caller currently has running.
128+
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
129+
pub struct ActiveQueryList {
130+
/// The active queries, most recently started first.
131+
pub queries: Vec<ActiveQuery>,
132+
/// Number of active queries reported by the runtime.
133+
pub total_count: usize,
134+
}
135+
136+
/// Whether `query_id` has the shape the runtime parses as a UUID.
137+
///
138+
/// The ids this SDK cancels always come from [`ActiveQuery::query_id`], so a
139+
/// value that is not a UUID cannot name a running query. Checking locally keeps
140+
/// such a value out of the request path entirely: `.` and `..` are unreserved,
141+
/// so percent-encoding leaves them intact and the URL parser then resolves them
142+
/// away — `..` turns `/v1/sql/{id}/cancel` into a POST at a route the caller
143+
/// never asked for.
144+
pub(crate) fn is_uuid(query_id: &str) -> bool {
145+
let bytes = query_id.as_bytes();
146+
if bytes.len() != 36 {
147+
return false;
148+
}
149+
bytes.iter().enumerate().all(|(index, byte)| {
150+
if matches!(index, 8 | 13 | 18 | 23) {
151+
*byte == b'-'
152+
} else {
153+
byte.is_ascii_hexdigit()
154+
}
155+
})
156+
}
157+
158+
/// The runtime's response to a successful cancellation.
159+
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
160+
pub struct CancelActiveQueryResponse {
161+
/// The id of the query that was cancelled.
162+
pub query_id: String,
163+
/// The query's state after cancellation, such as `cancelled`.
164+
pub status: String,
165+
}
166+
167+
#[cfg(test)]
168+
mod tests {
169+
use super::*;
170+
171+
#[test]
172+
fn deserializes_an_active_query_list() {
173+
let body = r#"{
174+
"queries": [
175+
{
176+
"query_id": "0198f0a1-9c3d-7c4e-8a11-2b3c4d5e6f70",
177+
"protocol": "flight",
178+
"sql_preview": "SELECT * FROM taxi_trips",
179+
"started_at_ms": 1750000000000
180+
}
181+
],
182+
"total_count": 1
183+
}"#;
184+
185+
let list: ActiveQueryList = serde_json::from_str(body).expect("deserialize list");
186+
assert_eq!(list.total_count, 1);
187+
assert_eq!(list.queries.len(), 1);
188+
assert_eq!(list.queries[0].protocol, "flight");
189+
assert_eq!(list.queries[0].sql_preview, "SELECT * FROM taxi_trips");
190+
assert_eq!(list.queries[0].started_at_ms, 1_750_000_000_000);
191+
}
192+
193+
#[test]
194+
fn deserializes_an_empty_active_query_list() {
195+
let list: ActiveQueryList =
196+
serde_json::from_str(r#"{"queries":[],"total_count":0}"#).expect("deserialize list");
197+
assert_eq!(list.total_count, 0);
198+
assert!(list.queries.is_empty());
199+
}
200+
201+
#[test]
202+
fn deserializes_a_cancel_response() {
203+
let response: CancelActiveQueryResponse = serde_json::from_str(
204+
r#"{"query_id":"0198f0a1-9c3d-7c4e-8a11-2b3c4d5e6f70","status":"cancelled"}"#,
205+
)
206+
.expect("deserialize cancel response");
207+
assert_eq!(response.status, "cancelled");
208+
}
209+
210+
#[test]
211+
fn not_found_error_explains_both_causes() {
212+
let message = ActiveQueryError::NotFound {
213+
query_id: "abc".to_string(),
214+
}
215+
.to_string();
216+
assert!(message.contains("abc"));
217+
assert!(message.contains("already finished"));
218+
}
219+
220+
#[test]
221+
fn is_uuid_accepts_the_ids_the_runtime_hands_out() {
222+
assert!(is_uuid("0198f0a1-9c3d-7c4e-8a11-2b3c4d5e6f70"));
223+
// The runtime parses either case.
224+
assert!(is_uuid("0198F0A1-9C3D-7C4E-8A11-2B3C4D5E6F70"));
225+
}
226+
227+
#[test]
228+
fn is_uuid_rejects_anything_that_could_reroute_a_request() {
229+
for id in [
230+
"",
231+
".",
232+
"..",
233+
"../queries/escape",
234+
"not-a-uuid",
235+
// Right length, wrong shape: hyphens off their positions.
236+
"0198f0a19c3d-7c4e-8a11-2b3c4d5e6f70-",
237+
// Right shape, a non-hex digit.
238+
"0198f0a1-9c3d-7c4e-8a11-2b3c4d5e6f7g",
239+
// A trailing segment appended to a valid id.
240+
"0198f0a1-9c3d-7c4e-8a11-2b3c4d5e6f70/cancel",
241+
] {
242+
assert!(!is_uuid(id), "{id:?} should not be accepted as a UUID");
243+
}
244+
}
245+
246+
#[test]
247+
fn invalid_query_id_error_points_at_active_queries() {
248+
let message = ActiveQueryError::InvalidQueryId {
249+
query_id: "not-a-uuid".to_string(),
250+
}
251+
.to_string();
252+
assert!(message.contains("not-a-uuid"));
253+
assert!(message.contains("active_queries()"));
254+
}
255+
}

0 commit comments

Comments
 (0)