feat(sql): follow adaptive streaming and polling responses (#4664)

SQL execution now starts with one Flight call. The server can stream
rows immediately or return a polling continuation, which the SDK follows
through the existing query handle and reader APIs. The same behavior
applies with and without parameters.

Both initial transports share direct-stream lifecycle handling. Handoffs
are explicitly negotiated; existing servers can continue returning
ordinary streams. Add coverage for streaming, deferred results,
cancellation, malformed decisions, and legacy responses.

```mermaid
flowchart TD
    A[execute_query or execute_query_async] --> B{Parameters supplied?}
    B -->|No| C[DoGet: SQL command ticket]
    B -->|Yes| D[DoExchange: SQL and Arrow parameter row]
    C --> E[Server chooses execution path]
    D --> E
    E -->|Stream| F[Initial response: streaming decision and Arrow batches]
    E -->|Defer| G[Initial response: PollInfo continuation]
    G --> H[SDK follows PollFlightInfo]
    H -->|Still running| H
    H -->|Result endpoints available| I[SDK fetches tickets with DoGet]
    F --> J[Query reader yields Arrow batches]
    I --> J
```

`execute_query` opens the reader; `execute_query_async` returns its
query handle. The caller does not select the transport or execution
path.
This commit is contained in:
Jack Ye
2026-10-25 15:47:29 +08:00
committed by GitHub
parent 6b6778e684
commit 24c2669a50
8 changed files with 414 additions and 133 deletions
+4 -6
View File
@@ -372,12 +372,10 @@ impl ExecuteQueryAsyncBuilder {
/// Do not follow a named placeholder directly with `$`: `$a$` starts a
/// dollar-quoted string.
///
/// A parameterized statement runs on the call that returns its rows rather
/// than detached from it: cancelling the query, or dropping its handle
/// before reading it, stops the statement on the server. Its status follows
/// that stream too -- it reports `Running` until the reader has received
/// every row -- so open the reader directly rather than waiting for
/// `Finished`.
/// With or without parameters, the server may stream rows on the initial
/// call or return a polling continuation. The reader follows either path
/// transparently. Open it directly instead of waiting for `Finished`:
/// directly streamed queries need their reader to make progress.
///
/// # Example
///
+38 -48
View File
@@ -34,9 +34,9 @@ use crate::remote::client::{ClientConfig, TlsConfig};
use crate::remote::retry::ResolvedRetryConfig;
use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus};
mod exchange;
mod stream;
use exchange::{ExchangeQuery, ExchangeQueryHandle};
use stream::{StreamingQuery, StreamingQueryHandle, Submission};
const DEFAULT_SQL_PORT: u16 = 10025;
const DEFAULT_SQL_TLS_PORT: u16 = 10026;
@@ -191,20 +191,39 @@ impl SqlClient {
}
pub(super) async fn execute(&self, request: ExecuteQueryRequest) -> Result<Query> {
match request.parameters {
Some(parameters) => {
self.submit_with_parameters(
&request.query,
parameters,
&request.default_namespace_path,
)
.await
let timeout = self.inner.overall_timeout()?;
with_overall_timeout(timeout, "SQL query submission", async {
validate_namespace_path(&request.default_namespace_path)?;
let query_id = Uuid::now_v7();
let submission = StreamingQuery::open(
self.inner.clone(),
&request.query,
request.parameters,
&request.default_namespace_path,
)
.await?;
match submission {
Submission::Stream(response) => {
let query =
Arc::new(StreamingQuery::new(query_id, self.inner.clone(), response));
self.queries
.insert(query_id, RegisteredQuery::Streamed(query.clone()));
Ok(Query::new(Arc::new(StreamingQueryHandle::new(query))))
}
Submission::Deferred(poll_info) => {
let query = Arc::new(RemoteQuery::new(
query_id,
self.inner.clone(),
request.default_namespace_path,
poll_info,
)?);
self.queries
.insert(query_id, RegisteredQuery::Polled(query.clone()));
Ok(Query::new(Arc::new(RemoteQueryHandle::new(query))))
}
}
None => {
self.submit(&request.query, &request.default_namespace_path)
.await
}
}
})
.await
}
pub(super) async fn submit(
@@ -271,35 +290,6 @@ impl SqlClient {
})
}
/// Values cannot ride on a poll, so a statement that has them is sent on
/// its own exchange instead.
pub(super) async fn submit_with_parameters(
&self,
query: &str,
parameters: RecordBatch,
default_namespace_path: &[String],
) -> Result<Query> {
let timeout = self.inner.overall_timeout()?;
with_overall_timeout(timeout, "SQL query submission", async {
validate_namespace_path(default_namespace_path)?;
let query_id = Uuid::now_v7();
let query = Arc::new(
ExchangeQuery::open(
query_id,
self.inner.clone(),
query,
parameters,
default_namespace_path,
)
.await?,
);
self.queries
.insert(query_id, RegisteredQuery::Exchanged(query.clone()));
Ok(Query::new(Arc::new(ExchangeQueryHandle::new(query))))
})
.await
}
pub(super) async fn describe(&self, query_id: Uuid) -> Result<QueryDescription> {
let query = self
.queries
@@ -586,14 +576,14 @@ impl SqlClientInner {
#[derive(Clone)]
enum RegisteredQuery {
Polled(Arc<RemoteQuery>),
Exchanged(Arc<ExchangeQuery>),
Streamed(Arc<StreamingQuery>),
}
impl RegisteredQuery {
fn touch(&self) {
match self {
Self::Polled(query) => query.touch(),
Self::Exchanged(query) => query.touch(),
Self::Streamed(query) => query.touch(),
}
}
@@ -602,14 +592,14 @@ impl RegisteredQuery {
fn registry_expired(&self) -> bool {
match self {
Self::Polled(query) => query.registry_expired(Arc::strong_count(query) == 1),
Self::Exchanged(query) => query.registry_expired(Arc::strong_count(query) == 1),
Self::Streamed(query) => query.registry_expired(Arc::strong_count(query) == 1),
}
}
async fn describe(&self) -> Result<QueryDescription> {
match self {
Self::Polled(query) => query.describe().await,
Self::Exchanged(query) => query.describe(),
Self::Streamed(query) => query.describe(),
}
}
}
@@ -1,26 +1,19 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! SQL statements sent with parameters.
//!
//! A statement without parameters is polled: `PollFlightInfo` starts it
//! detached from the connection and `DoGet` reads its endpoints. Neither call
//! carries Arrow data, so a statement's values would have to be spelled into
//! its text. A statement with parameters is sent on one `DoExchange` instead:
//! the command and the parameter schema, then one row, and the result rows
//! come back on the same call. The call is the statement's lifetime, so
//! cancelling it -- or dropping the handle before reading -- stops the
//! statement on the server.
//! SQL submitted on one call, with either streamed rows or a polling handoff.
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex as StdMutex, OnceLock};
use std::time::Instant;
use arrow_array::RecordBatch;
use arrow_flight::FlightDescriptor;
use arrow_flight::decode::FlightRecordBatchStream;
use arrow_flight::encode::FlightDataEncoderBuilder;
use arrow_flight::sql::{CommandStatementQuery, ProstMessageExt};
use arrow_flight::sql::{Any, CommandStatementQuery, ProstMessageExt};
use arrow_flight::{FlightDescriptor, PollInfo, Ticket};
use arrow_schema::SchemaRef;
use futures::TryStreamExt;
use prost::Message;
use tokio::sync::{Notify, mpsc};
use uuid::Uuid;
@@ -51,8 +44,8 @@ fn failure(error: &Error) -> Lifecycle {
})
}
/// A statement running on its own `DoExchange` call.
pub(super) struct ExchangeQuery {
/// A statement streaming on its initial `DoGet` or `DoExchange` call.
pub(super) struct StreamingQuery {
id: Uuid,
client: Arc<SqlClientInner>,
/// The call's rows, until a reader takes them or a cancellation drops them.
@@ -63,7 +56,7 @@ pub(super) struct ExchangeQuery {
last_accessed: StdMutex<Instant>,
}
struct PreparedExchange {
struct PreparedStream {
schema: SchemaRef,
first: Option<RecordBatch>,
rest: Option<ResultEndpointStream>,
@@ -76,7 +69,7 @@ struct PreparedExchange {
/// this one did. A query cancelled meanwhile has its rows dropped instead,
/// which ends the call.
struct Borrowed<'a> {
query: &'a ExchangeQuery,
query: &'a StreamingQuery,
rows: Option<ResultEndpointStream>,
}
@@ -104,18 +97,37 @@ impl Drop for Borrowed<'_> {
}
}
impl ExchangeQuery {
/// Send the statement and its parameters, returning once the server has
/// planned it: a statement it refuses is refused here, as a polled one is
/// on its first poll.
pub(super) const PROTOCOL_HEADER: &str = "lancedb-sql-protocol";
pub(super) const PROTOCOL_VERSION: &str = "adaptive-v1";
const RESPONSE_TYPE_URL: &str = "type.googleapis.com/lancedb.sql.QueryResponse";
#[derive(Clone, PartialEq, Message)]
pub(super) struct QueryResponse {
#[prost(message, optional, tag = "1")]
pub poll_info: Option<PollInfo>,
}
#[derive(Clone, PartialEq, Message)]
struct CommandTicket {
#[prost(bytes = "vec", tag = "2")]
command: Vec<u8>,
}
pub(super) enum Submission {
Stream(ResultEndpointStream),
Deferred(PollInfo),
}
impl StreamingQuery {
pub(super) async fn open(
id: Uuid,
client: Arc<SqlClientInner>,
query: &str,
parameters: RecordBatch,
parameters: Option<RecordBatch>,
default_namespace_path: &[String],
) -> Result<Self> {
if parameters.num_rows() != 1 {
) -> Result<Submission> {
if let Some(parameters) = &parameters
&& parameters.num_rows() != 1
{
return Err(Error::InvalidInput {
message: format!(
"SQL query parameters must be exactly one row, got {}",
@@ -130,39 +142,102 @@ impl ExchangeQuery {
Some(DEFAULT_READ_TIMEOUT),
)?
.unwrap();
let descriptor = FlightDescriptor::new_cmd(
CommandStatementQuery {
query: query.to_string(),
transaction_id: None,
let command = CommandStatementQuery {
query: query.to_string(),
transaction_id: None,
}
.as_any()
.encode_to_vec();
let connection = client.connection(&request_id).await?;
let headers = client.headers(default_namespace_path, &request_id).await?;
let mut metadata = super::client_with_headers(connection.client.clone(), &headers)?
.metadata()
.clone();
metadata.insert(PROTOCOL_HEADER, PROTOCOL_VERSION.parse().unwrap());
let mut flight = connection.client.clone();
let response = if let Some(parameters) = parameters {
let messages = FlightDataEncoderBuilder::new()
.with_flight_descriptor(Some(FlightDescriptor::new_cmd(command)))
.with_schema(parameters.schema())
.build(futures::stream::iter([Ok(parameters)]))
.try_collect::<Vec<_>>()
.await
.map_err(|error| sql_error(&request_id, error))?;
let mut request = tonic::Request::new(futures::stream::iter(messages));
*request.metadata_mut() = metadata;
tokio::time::timeout(read_timeout, flight.do_exchange(request)).await
} else {
let mut request =
tonic::Request::new(Ticket::new(CommandTicket { command }.encode_to_vec()));
*request.metadata_mut() = metadata;
tokio::time::timeout(read_timeout, flight.do_get(request)).await
}
.map_err(|_| sql_error(&request_id, "SQL query submission timed out"))?
.map_err(|error| sql_error(&request_id, error))?;
let negotiated = response
.metadata()
.get(PROTOCOL_HEADER)
.is_some_and(|value| value == PROTOCOL_VERSION);
let mut response_stream = response.into_inner();
if negotiated {
let first = tokio::time::timeout(read_timeout, response_stream.message())
.await
.map_err(|_| sql_error(&request_id, "SQL query handoff timed out"))?
.map_err(|error| sql_error(&request_id, error))?
.ok_or_else(|| {
sql_error(&request_id, "SQL response omitted its execution decision")
})?;
if !first.data_header.is_empty() || !first.data_body.is_empty() {
return Err(sql_error(
&request_id,
"SQL response sent rows before its execution decision",
));
}
.as_any()
.encode_to_vec(),
let envelope =
Any::decode(first.app_metadata).map_err(|error| sql_error(&request_id, error))?;
if envelope.type_url != RESPONSE_TYPE_URL {
return Err(sql_error(
&request_id,
"SQL response returned an unknown execution decision",
));
}
let decision = QueryResponse::decode(envelope.value)
.map_err(|error| sql_error(&request_id, error))?;
if let Some(info) = decision.poll_info {
if info.info.is_none() && info.flight_descriptor.is_none() {
return Err(sql_error(
&request_id,
"SQL handoff omitted result and continuation information",
));
}
return Ok(Submission::Deferred(info));
}
}
let stream = FlightRecordBatchStream::new_from_flight_data(
response_stream
.map_err(|status| arrow_flight::error::FlightError::Tonic(Box::new(status))),
);
// The descriptor rides on the schema message, and the row follows it.
let request = FlightDataEncoderBuilder::new()
.with_flight_descriptor(Some(descriptor))
.with_schema(parameters.schema())
.build(futures::stream::iter([Ok(parameters)]));
let mut flight = client
.client_with_headers(default_namespace_path, &request_id)
.await?;
let stream = tokio::time::timeout(read_timeout, flight.do_exchange(request))
.await
.map_err(|_| sql_error(&request_id, "SQL query submission timed out"))?
.map_err(|err| sql_error(&request_id, err))?;
Ok(Self {
Ok(Submission::Stream(ResultEndpointStream {
stream,
request_id,
read_timeout,
}))
}
pub(super) fn new(
id: Uuid,
client: Arc<SqlClientInner>,
response: ResultEndpointStream,
) -> Self {
Self {
id,
client,
response: StdMutex::new(Some(ResultEndpointStream {
stream,
request_id,
read_timeout,
})),
response: StdMutex::new(Some(response)),
lifecycle: StdMutex::new(Lifecycle::Running),
cancelled: Notify::new(),
terminal_at: OnceLock::new(),
last_accessed: StdMutex::new(Instant::now()),
})
}
}
pub(super) fn touch(&self) {
@@ -242,7 +317,7 @@ impl ExchangeQuery {
}
/// Take the rows and wait for the first batch, which carries the schema.
async fn prepare(&self) -> Result<PreparedExchange> {
async fn prepare(&self) -> Result<PreparedStream> {
let Some(rows) = self.response.lock().unwrap().take() else {
// Taken by an earlier attempt that failed, or dropped by a cancel.
return Err(match &*self.lifecycle.lock().unwrap() {
@@ -275,7 +350,7 @@ impl ExchangeQuery {
if first.is_none() {
self.settle(Lifecycle::Finished);
}
Ok(PreparedExchange {
Ok(PreparedStream {
schema,
rest: first.is_some().then_some(rows),
first,
@@ -284,7 +359,7 @@ impl ExchangeQuery {
async fn run(
self: Arc<Self>,
mut prepared: PreparedExchange,
mut prepared: PreparedStream,
sender: mpsc::Sender<Result<RecordBatch>>,
) -> Result<()> {
if let Some(batch) = prepared.first.take()
@@ -343,13 +418,13 @@ impl ExchangeQuery {
}
}
pub(super) struct ExchangeQueryHandle {
query: Arc<ExchangeQuery>,
pub(super) struct StreamingQueryHandle {
query: Arc<StreamingQuery>,
result_started: AtomicBool,
}
impl ExchangeQueryHandle {
pub(super) fn new(query: Arc<ExchangeQuery>) -> Self {
impl StreamingQueryHandle {
pub(super) fn new(query: Arc<StreamingQuery>) -> Self {
Self {
query,
result_started: AtomicBool::new(false),
@@ -357,7 +432,7 @@ impl ExchangeQueryHandle {
}
}
impl Drop for ExchangeQueryHandle {
impl Drop for StreamingQueryHandle {
// Rows nobody took are rows nobody will read. Keeping the call open would
// hold the statement running on the server, stalled on flow control,
// until the registry forgot it.
@@ -369,7 +444,7 @@ impl Drop for ExchangeQueryHandle {
}
#[async_trait::async_trait]
impl QueryHandle for ExchangeQueryHandle {
impl QueryHandle for StreamingQueryHandle {
fn id(&self) -> Uuid {
self.query.touch();
self.query.id
+209 -3
View File
@@ -19,6 +19,7 @@ use futures::{StreamExt, TryStreamExt};
use tonic::metadata::MetadataMap;
use tonic::{Request, Response, Status, Streaming};
use super::stream::{PROTOCOL_HEADER, PROTOCOL_VERSION, QueryResponse};
use super::*;
use crate::database::Database;
use crate::remote::RemoteCatalogOptions;
@@ -70,6 +71,9 @@ type ReceivedExchanges = Arc<std::sync::Mutex<Vec<(String, Option<RecordBatch>)>
#[derive(Clone)]
struct TestSqlService {
adaptive: bool,
defer: bool,
invalid_decision: bool,
query_count: Arc<AtomicUsize>,
do_get_count: Arc<AtomicUsize>,
cancel_count: Arc<AtomicUsize>,
@@ -120,6 +124,51 @@ impl TestSqlService {
.unwrap_or_default(),
});
}
fn adaptive_response(
&self,
query: &str,
rows: BoxStream<'static, std::result::Result<FlightData, Status>>,
) -> Response<BoxStream<'static, std::result::Result<FlightData, Status>>> {
if !self.adaptive {
return Response::new(rows);
}
let poll_info =
self.defer.then(|| PollInfo {
info: Some(FlightInfo::new().with_endpoint(
FlightEndpoint::new().with_ticket(Ticket::new(query.to_string())),
)),
flight_descriptor: Some(FlightDescriptor::new_cmd(format!("poll:1:{query}"))),
progress: Some(0.25),
expiration_time: None,
});
let first = FlightData {
app_metadata: if self.invalid_decision {
b"invalid protobuf".to_vec().into()
} else {
Any {
type_url: "type.googleapis.com/lancedb.sql.QueryResponse".to_string(),
value: QueryResponse { poll_info }.encode_to_vec().into(),
}
.encode_to_vec()
.into()
},
..Default::default()
};
let rest = if self.defer {
futures::stream::empty().boxed()
} else {
rows
};
let mut response = Response::new(
futures::stream::once(async { Ok(first) })
.chain(rest)
.boxed(),
);
response
.metadata_mut()
.insert(PROTOCOL_HEADER, PROTOCOL_VERSION.parse().unwrap());
response
}
}
impl Default for TestSqlService {
@@ -154,6 +203,9 @@ impl Default for TestSqlService {
let dictionary_result =
RecordBatch::try_new(dictionary_schema, vec![Arc::new(dictionary)]).unwrap();
Self {
adaptive: false,
defer: false,
invalid_decision: false,
query_count: Arc::new(AtomicUsize::new(0)),
do_get_count: Arc::new(AtomicUsize::new(0)),
cancel_count: Arc::new(AtomicUsize::new(0)),
@@ -308,7 +360,22 @@ impl FlightService for TestSqlService {
request: Request<Ticket>,
) -> std::result::Result<Response<<Self as FlightService>::DoGetStream>, Status> {
self.do_get_count.fetch_add(1, Ordering::SeqCst);
let ticket = request.get_ref().ticket.as_ref();
let initial = Any::decode(request.get_ref().ticket.clone())
.ok()
.and_then(|envelope| Any::decode(envelope.value).ok())
.and_then(|command| command.unpack::<CommandStatementQuery>().ok().flatten());
if initial.is_some() {
self.capture_headers(request.metadata());
assert_eq!(
request.metadata().get(PROTOCOL_HEADER).unwrap(),
PROTOCOL_VERSION
);
}
let query = initial
.as_ref()
.map(|command| command.query.as_bytes())
.unwrap_or(request.get_ref().ticket.as_ref());
let ticket = query;
let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\"";
let empty = ticket == b"SELECT empty";
let slow = ticket == b"SELECT slow get";
@@ -337,7 +404,11 @@ impl FlightService for TestSqlService {
encoder = encoder.with_max_flight_data_size(8 * 1024 * 1024);
}
let stream = encoder.build(input).map_err(Status::from);
Ok(Response::new(Box::pin(stream)))
if let Some(initial) = initial {
Ok(self.adaptive_response(&initial.query, stream.boxed()))
} else {
Ok(Response::new(stream.boxed()))
}
}
async fn do_put(
@@ -486,7 +557,7 @@ impl FlightService for TestSqlService {
let _dropped = &dropped;
message
});
Ok(Response::new(Box::pin(stream)))
Ok(self.adaptive_response(&query, stream.boxed()))
}
}
@@ -1609,3 +1680,138 @@ fn validates_metadata_with_header_map() {
assert!(insert_header(&mut headers, "bad header", "value").is_err());
assert!(insert_header(&mut headers, "valid-header", "bad\nvalue").is_err());
}
#[tokio::test]
async fn adaptive_streams_start_on_one_rpc_for_both_inputs() {
for parameters in [None, Some(parameter_row())] {
let service = TestSqlService {
adaptive: true,
..Default::default()
};
let polls = service.query_count.clone();
let gets = service.do_get_count.clone();
let exchanges = service.exchanges.clone();
let expected = parameters.clone().unwrap_or_else(|| service.result.clone());
let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await;
let with_parameters = parameters.is_some();
let query = client
.execute(ExecuteQueryRequest {
parameters,
..ExecuteQueryRequest::new("SELECT echo")
})
.await
.unwrap();
assert_eq!(query.describe().await.unwrap().status, QueryStatus::Running);
assert_eq!(collect_result(&query).await.unwrap(), vec![expected]);
assert_eq!(
query.describe().await.unwrap().status,
QueryStatus::Finished
);
assert!(query.reader().await.is_err());
assert_eq!(polls.load(Ordering::SeqCst), 0);
assert_eq!(gets.load(Ordering::SeqCst), usize::from(!with_parameters));
assert_eq!(
exchanges.lock().unwrap().len(),
usize::from(with_parameters)
);
}
}
#[tokio::test]
async fn adaptive_handoffs_share_polling_and_survive_handle_drop() {
for parameters in [None, Some(parameter_row())] {
let service = TestSqlService {
adaptive: true,
defer: true,
..Default::default()
};
let polls = service.query_count.clone();
let gets = service.do_get_count.clone();
let exchanges = service.exchanges.clone();
let expected = service.result.clone();
let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await;
let with_parameters = parameters.is_some();
let query = client
.execute(ExecuteQueryRequest {
parameters,
..ExecuteQueryRequest::new("SELECT adaptive")
})
.await
.unwrap();
let id = query.id();
assert_eq!(collect_result(&query).await.unwrap(), vec![expected]);
assert_eq!(
query.describe().await.unwrap().status,
QueryStatus::Finished
);
assert!(query.reader().await.is_err());
drop(query);
assert_eq!(
client.describe(id).await.unwrap().status,
QueryStatus::Finished
);
assert_eq!(
polls.load(Ordering::SeqCst),
0,
"the original SQL must not be resubmitted through polling"
);
assert_eq!(
gets.load(Ordering::SeqCst),
1 + usize::from(!with_parameters)
);
assert_eq!(
exchanges.lock().unwrap().len(),
usize::from(with_parameters)
);
}
}
#[tokio::test]
async fn adaptive_handoffs_use_server_cancellation_for_both_inputs() {
for parameters in [None, Some(parameter_row())] {
let service = TestSqlService {
adaptive: true,
defer: true,
..Default::default()
};
let cancels = service.cancel_count.clone();
let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await;
let query = client
.execute(ExecuteQueryRequest {
parameters,
..ExecuteQueryRequest::new("SELECT cancel")
})
.await
.unwrap();
query.cancel().await.unwrap();
assert_eq!(cancels.load(Ordering::SeqCst), 1);
assert_eq!(
query.describe().await.unwrap().status,
QueryStatus::Cancelled
);
assert!(query.reader().await.is_err());
}
}
#[tokio::test]
async fn malformed_adaptive_decision_never_resubmits_the_statement() {
for parameters in [None, Some(parameter_row())] {
let service = TestSqlService {
adaptive: true,
invalid_decision: true,
..Default::default()
};
let polls = service.query_count.clone();
let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await;
assert!(
client
.execute(ExecuteQueryRequest {
parameters,
..ExecuteQueryRequest::new("SELECT invalid")
})
.await
.is_err()
);
assert_eq!(polls.load(Ordering::SeqCst), 0);
}
}
+10 -10
View File
@@ -18,9 +18,9 @@ pub enum QueryStatus {
Running,
/// The server has made the complete result available.
///
/// A query submitted with parameters streams its rows on the call that
/// runs it, so it reports `Finished` only once its reader has received
/// the last batch. Read it rather than waiting for this status.
/// Directly streamed queries report `Finished` when their stream reaches
/// its end. Deferred queries report the server's execution state. Open the
/// reader directly instead of waiting for this status before reading.
Finished,
/// The server accepted cancellation but has not confirmed it yet.
Cancelling,
@@ -64,9 +64,10 @@ pub(crate) trait QueryHandle: Send + Sync {
/// A handle to a submitted SQL query.
///
/// The handle can be inspected, opened as an Arrow reader, or cancelled.
/// Dropping it does not cancel the server-side query, unless the query was
/// submitted with parameters: those run on the call that returns their rows,
/// so dropping the handle before reading, or dropping the reader, stops them.
/// The server may stream results on the initial call or defer execution and
/// return result tickets. Dropping an unread handle or its reader stops a
/// directly streamed query; deferred execution continues until explicitly
/// cancelled. This applies equally with and without parameters.
/// Identifier lookup is scoped to the connection that submitted the query and
/// is not a durable resume mechanism.
pub struct Query {
@@ -95,10 +96,9 @@ impl Query {
/// Get a point-in-time description of the query.
///
/// For a query submitted with parameters the status follows its result
/// stream: it stays [`QueryStatus::Running`] until the reader has received
/// every row, so open [`Self::reader`] directly instead of polling for
/// [`QueryStatus::Finished`] first.
/// Directly streamed queries stay [`QueryStatus::Running`] until the result
/// stream ends. Deferred queries expose the server's execution state. Open
/// [`Self::reader`] directly instead of waiting for [`QueryStatus::Finished`].
pub async fn describe(&self) -> Result<QueryDescription> {
self.handle.describe().await
}