diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index f2d5285fa..801035391 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -501,6 +501,8 @@ class JobDescription: class SqlQuery: @property def id(self) -> UUID: ... + @property + def job_id(self) -> Optional[str]: ... async def describe(self) -> QueryDescription: ... async def reader(self) -> RecordBatchStream: ... async def cancel(self) -> None: ... @@ -509,6 +511,8 @@ class QueryDescription: @property def id(self) -> UUID: ... @property + def job_id(self) -> Optional[str]: ... + @property def status(self) -> str: ... @property def progress(self) -> Optional[float]: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index e73a44aec..a5d7aeb41 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -1515,7 +1515,10 @@ class DBConnection(EnforceOverrides): This submits through :meth:`execute_query_async` and waits until the initial result stream is readable. It does not wait for the full query - to finish. + to finish. A statement that ran as a server-side SQL job exposes the + durable Job Registry id in the reader schema metadata under + ``x-lancedb-sql-job-id``. Use :meth:`execute_query_async` when the + query handle itself is needed. """ return self.execute_query_async( query, @@ -1554,6 +1557,9 @@ class DBConnection(EnforceOverrides): stops a directly streamed query; deferred execution requires explicit cancellation. Open ``reader()`` directly: a streamed query needs its reader to make progress and reports ``running`` until the stream ends. + Once a server-side SQL job exists, ``describe().job_id`` exposes its + durable Job Registry id while the query is still running. This is + separate from the connection-scoped query ``id``. LanceDB Embedded connections do not support SQL. """ @@ -4128,7 +4134,10 @@ class AsyncConnection(object): This submits through :meth:`execute_query_async` and waits until the initial result stream is readable. It does not wait for the full query - to finish. + to finish. A statement that ran as a server-side SQL job exposes the + durable Job Registry id in the reader schema metadata under + ``x-lancedb-sql-job-id``. Use :meth:`execute_query_async` when the + query handle itself is needed. """ submitted = await self.execute_query_async( query, diff --git a/python/python/lancedb/sql.py b/python/python/lancedb/sql.py index 7498cddcc..5da4d35dd 100644 --- a/python/python/lancedb/sql.py +++ b/python/python/lancedb/sql.py @@ -104,6 +104,15 @@ class AsyncQuery: """The stable identifier scoped to the connection that submitted it.""" return self._inner.id + @property + def job_id(self) -> Optional[str]: + """The durable Job Registry id once the server reports this query as a job. + + A running query reports it after :meth:`describe` observes the latest + server status. This is separate from the connection-scoped ``id``. + """ + return self._inner.job_id + async def describe(self) -> QueryDescription: """Get a point-in-time description of the query.""" return await self._inner.describe() @@ -118,7 +127,12 @@ class AsyncQuery: return AsyncRecordBatchReader(await self._inner.reader()) async def cancel(self) -> None: - """Request cancellation of the query.""" + """Request cancellation of this query's Flight execution. + + A durable server-side job may reject this operation because its Flight + waiter cannot cancel the registry job. Use ``connection.cancel_job()`` + with ``job_id`` to cancel that job. + """ await self._inner.cancel() @@ -133,6 +147,15 @@ class Query: """The stable identifier scoped to the connection that submitted it.""" return self._inner.id + @property + def job_id(self) -> Optional[str]: + """The durable Job Registry id once the server reports this query as a job. + + A running query reports it after :meth:`describe` observes the latest + server status. This is separate from the connection-scoped ``id``. + """ + return self._inner.job_id + def describe(self) -> QueryDescription: """Get a point-in-time description of the query.""" return LOOP.run(self._inner.describe()) @@ -159,7 +182,12 @@ class Query: return pa.RecordBatchReader.from_batches(reader.schema, batches()) def cancel(self) -> None: - """Request cancellation of the query.""" + """Request cancellation of this query's Flight execution. + + A durable server-side job may reject this operation because its Flight + waiter cannot cancel the registry job. Use ``connection.cancel_job()`` + with ``job_id`` to cancel that job. + """ LOOP.run(self._inner.cancel()) diff --git a/python/python/tests/test_sql.py b/python/python/tests/test_sql.py index 0ed382d93..b4c555635 100644 --- a/python/python/tests/test_sql.py +++ b/python/python/tests/test_sql.py @@ -20,6 +20,7 @@ NIL_QUERY_ID = UUID(int=0) class FakeNativeQuery: id = UUID("0198f1b2-c3d4-7e5f-8123-456789abcdef") + job_id = None async def reader(self): return pa.table({"value": [1, 2]}) @@ -68,6 +69,15 @@ def test_query_id_is_uuid(): assert Query(query).id == query.id +def test_query_job_id_is_optional_and_separate_from_query_id(): + native = FakeNativeQuery() + native.job_id = "j1-durable" + query = AsyncQuery(native) + assert query.id != query.job_id + assert query.job_id == "j1-durable" + assert Query(query).job_id == "j1-durable" + + def test_connection_serializes_sql_host_override(): endpoint = "grpc+tls://sql.example.com:10026" restored = lancedb.deserialize_conn( diff --git a/python/src/sql.rs b/python/src/sql.rs index 612521207..166cf0df5 100644 --- a/python/src/sql.rs +++ b/python/src/sql.rs @@ -31,6 +31,11 @@ impl Query { self.inner.id() } + #[getter] + pub fn job_id(&self) -> Option { + self.inner.job_id() + } + pub fn describe(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.inner.clone(); future_into_py(self_.py(), async move { @@ -63,6 +68,7 @@ impl Query { #[derive(Clone)] pub struct QueryDescription { id: Uuid, + job_id: Option, status: String, progress: Option, expires_at: Option>, @@ -72,8 +78,8 @@ pub struct QueryDescription { impl QueryDescription { fn __repr__(&self) -> String { format!( - "QueryDescription(id={:?}, status={:?}, progress={:?}, expires_at={:?})", - self.id, self.status, self.progress, self.expires_at + "QueryDescription(id={:?}, job_id={:?}, status={:?}, progress={:?}, expires_at={:?})", + self.id, self.job_id, self.status, self.progress, self.expires_at ) } } @@ -82,6 +88,7 @@ impl From for QueryDescription { fn from(description: lancedb::sql::QueryDescription) -> Self { Self { id: description.id, + job_id: description.job_id, status: description.status.to_string(), progress: description.progress, expires_at: description.expires_at, diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index 724c813b3..0b87ff5a5 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -19,11 +19,12 @@ use arrow_flight::{ FlightDescriptor, FlightEndpoint, FlightInfo, PollInfo, }; use arrow_schema::{Schema, SchemaRef}; -use futures::TryStreamExt; +use futures::{StreamExt, TryStreamExt}; use http::header::{HeaderMap, HeaderName, HeaderValue}; use prost::Message; use serde::{Serialize, de::DeserializeOwned}; use tokio::sync::{Mutex, Notify, OnceCell, mpsc}; +use tonic::Request; use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity}; use uuid::Uuid; @@ -87,6 +88,12 @@ struct ResultEndpointStream { stream: FlightRecordBatchStream, request_id: String, read_timeout: Duration, + job_id: Option, +} + +struct PollResult { + info: PollInfo, + job_id: Option, } struct PreparedSqlResult { @@ -233,10 +240,19 @@ impl SqlClient { materialized_view_job_handle, ) .await?; - let query = match opened.submission { + let stream::OpenedSubmission { + submission, + materialized_view_job_handle_supported, + job_id, + } = opened; + let query = match submission { Submission::Stream(response) => { - let query = - Arc::new(StreamingQuery::new(query_id, self.inner.clone(), response)); + let query = Arc::new(StreamingQuery::new( + query_id, + self.inner.clone(), + response, + job_id.clone(), + )); self.queries .insert(query_id, RegisteredQuery::Streamed(query.clone())); Query::new(Arc::new(StreamingQueryHandle::new(query))) @@ -247,6 +263,7 @@ impl SqlClient { self.inner.clone(), request.default_namespace_path, poll_info, + job_id, )?); self.queries .insert(query_id, RegisteredQuery::Polled(query.clone())); @@ -255,8 +272,7 @@ impl SqlClient { }; Ok(SubmittedQuery { query, - materialized_view_job_handle_supported: opened - .materialized_view_job_handle_supported, + materialized_view_job_handle_supported, }) }) .await @@ -283,6 +299,7 @@ impl SqlClient { self.inner.clone(), default_namespace_path.to_vec(), poll_info, + None, )?); self.queries .insert(query_id, RegisteredQuery::Polled(query.clone())); @@ -383,16 +400,29 @@ impl SqlClientInner { &self, descriptor: FlightDescriptor, default_namespace_path: &[String], - ) -> Result> { + ) -> Result> { let request_id = uuid::Uuid::new_v4().to_string(); - let mut client = self - .client_with_headers(default_namespace_path, &request_id) - .await - .map_err(|err| sql_error(&request_id, err))?; - match tokio::time::timeout(STATUS_POLL_TIMEOUT, client.poll_flight_info(descriptor)).await { - Ok(result) => result - .map(Some) - .map_err(|err| sql_flight_error(&request_id, &err)), + let connection = self.connection(&request_id).await?; + let headers = self.headers(default_namespace_path, &request_id).await?; + let metadata = client_with_headers(connection.client.clone(), &headers)? + .metadata() + .clone(); + let mut request = Request::new(descriptor); + *request.metadata_mut() = metadata; + match tokio::time::timeout( + STATUS_POLL_TIMEOUT, + connection.client.clone().poll_flight_info(request), + ) + .await + { + Ok(Ok(response)) => { + let job_id = sql_job_id_from_metadata(response.metadata()); + Ok(Some(PollResult { + info: response.into_inner(), + job_id, + })) + } + Ok(Err(status)) => Err(sql_status_error(&request_id, &status)), Err(_) => Ok(None), } } @@ -401,7 +431,7 @@ impl SqlClientInner { &self, descriptor: FlightDescriptor, default_namespace_path: &[String], - ) -> Result { + ) -> Result { let read_timeout = resolve_timeout( self.client_config.timeout_config.read_timeout, "LANCE_CLIENT_READ_TIMEOUT", @@ -413,12 +443,18 @@ impl SqlClientInner { loop { let started = Instant::now(); let request_id = uuid::Uuid::new_v4().to_string(); - let mut client = self - .client_with_headers(default_namespace_path, &request_id) - .await?; - let result = - tokio::time::timeout(read_timeout, client.poll_flight_info(descriptor.clone())) - .await; + let connection = self.connection(&request_id).await?; + let headers = self.headers(default_namespace_path, &request_id).await?; + let metadata = client_with_headers(connection.client.clone(), &headers)? + .metadata() + .clone(); + let mut request = Request::new(descriptor.clone()); + *request.metadata_mut() = metadata; + let result = tokio::time::timeout( + read_timeout, + connection.client.clone().poll_flight_info(request), + ) + .await; let poll_info = match result { Err(_) if retry_count < retry_config.read_retries => { retry_count += 1; @@ -426,7 +462,7 @@ impl SqlClientInner { continue; } Err(_) => return Err(sql_error(&request_id, "SQL query poll timed out")), - Ok(Err(FlightError::Tonic(status))) + Ok(Err(status)) if matches!( status.code(), tonic::Code::DeadlineExceeded | tonic::Code::Unavailable @@ -436,8 +472,14 @@ impl SqlClientInner { tokio::time::sleep(poll_retry_delay(&retry_config, retry_count)).await; continue; } - Ok(Err(error)) => return Err(sql_flight_error(&request_id, &error)), - Ok(Ok(poll_info)) => poll_info, + Ok(Err(status)) => return Err(sql_status_error(&request_id, &status)), + Ok(Ok(response)) => { + let job_id = sql_job_id_from_metadata(response.metadata()); + PollResult { + info: response.into_inner(), + job_id, + } + } }; if let Some(delay) = MIN_POLL_INTERVAL.checked_sub(started.elapsed()) { tokio::time::sleep(delay).await; @@ -461,17 +503,27 @@ impl SqlClientInner { let ticket = endpoint.ticket.ok_or_else(|| { sql_error(&request_id, "SQL result endpoint did not include a ticket") })?; - let mut endpoint_client = self - .client_with_headers(default_namespace_path, &request_id) - .await?; - let stream = tokio::time::timeout(read_timeout, endpoint_client.do_get(ticket)) - .await - .map_err(|_| sql_error(&request_id, "SQL result fetch timed out"))? - .map_err(|err| sql_flight_error(&request_id, &err))?; + let connection = self.connection(&request_id).await?; + let headers = self.headers(default_namespace_path, &request_id).await?; + let metadata = client_with_headers(connection.client.clone(), &headers)? + .metadata() + .clone(); + let mut request = Request::new(ticket); + *request.metadata_mut() = metadata; + let response = + tokio::time::timeout(read_timeout, connection.client.clone().do_get(request)) + .await + .map_err(|_| sql_error(&request_id, "SQL result fetch timed out"))? + .map_err(|status| sql_status_error(&request_id, &status))?; + let job_id = sql_job_id_from_metadata(response.metadata()); + let stream = response + .into_inner() + .map_err(|status| FlightError::Tonic(Box::new(status))); Ok(ResultEndpointStream { - stream, + stream: FlightRecordBatchStream::new_from_flight_data(stream), request_id, read_timeout, + job_id, }) } @@ -547,6 +599,7 @@ impl SqlClientInner { Ok(CancelOutcome::Status(status)) } + #[cfg(test)] async fn client_with_headers( &self, default_namespace_path: &[String], @@ -684,6 +737,7 @@ impl QueryRegistry { struct RemoteQuery { id: Uuid, + job_id: StdMutex>, client: Arc, default_namespace_path: Vec, state: Mutex, @@ -713,6 +767,7 @@ impl RemoteQuery { client: Arc, default_namespace_path: Vec, poll_info: PollInfo, + job_id: Option, ) -> Result { let expires_at = query_expiration(&poll_info)?; let terminal_at = OnceLock::new(); @@ -724,6 +779,7 @@ impl RemoteQuery { }; Ok(Self { id, + job_id: StdMutex::new(job_id.filter(|job_id| !job_id.is_empty())), client, default_namespace_path, state: Mutex::new(poll_info), @@ -850,6 +906,44 @@ impl RemoteQuery { *self.last_accessed.lock().unwrap() = Instant::now(); } + fn job_id(&self) -> Option { + self.job_id.lock().unwrap().clone() + } + + fn observe_job_id(&self, job_id: Option<&str>) -> Result<()> { + let Some(job_id) = job_id.filter(|job_id| !job_id.is_empty()) else { + return Ok(()); + }; + let mut current = self.job_id.lock().unwrap(); + match current.as_deref() { + None => *current = Some(job_id.to_string()), + Some(existing) if existing == job_id => {} + Some(existing) => { + return Err(Error::Runtime { + message: format!( + "SQL query changed durable job identity from {existing} to {job_id}" + ), + }); + } + } + Ok(()) + } + + fn observe_error_job_id(&self, error: &Error) -> Result<()> { + self.observe_job_id( + error + .remote_error_metadata() + .and_then(|metadata| metadata.job_id.as_deref()), + ) + } + + fn record_error_job_id(&self, result: Result) -> Result { + if let Err(error) = &result { + self.observe_error_job_id(error)?; + } + result + } + async fn poll_next_state(&self, descriptor: FlightDescriptor) -> Result { self.touch(); if self.is_cancellation_requested() { @@ -874,10 +968,11 @@ impl RemoteQuery { Err(_) if self.is_cancellation_requested() => { return Err(self.cancelled_error()); } - result => result?, + result => self.record_error_job_id(result)?, }, }; - self.update_state(&descriptor, updated).await + self.observe_job_id(updated.job_id.as_deref())?; + self.update_state(&descriptor, updated.info).await } async fn prepare_result(self: &Arc) -> Result { @@ -894,8 +989,9 @@ impl RemoteQuery { result = self.client.open_result_endpoint( endpoint, &self.default_namespace_path, - ) => result?, + ) => self.record_error_job_id(result)?, }; + self.observe_job_id(endpoint_stream.job_id.as_deref())?; let buffered_batch = tokio::select! { biased; _ = self.wait_for_cancellation() => return Err(self.cancelled_error()), @@ -908,6 +1004,7 @@ impl RemoteQuery { .ok_or_else(|| Error::Runtime { message: "SQL result endpoint did not include a schema".to_string(), })?; + self.observe_job_id(sql_job_id_from_schema(&schema).as_deref())?; return Ok(PreparedSqlResult { schema, next_endpoint: 1, @@ -925,6 +1022,7 @@ impl RemoteQuery { .map_err(|err| sql_error(&request_id, err))?, ) }; + self.observe_job_id(sql_job_id_from_schema(&schema).as_deref())?; return Ok(PreparedSqlResult { schema, next_endpoint: 0, @@ -1001,7 +1099,7 @@ impl RemoteQuery { result = self.client.open_result_endpoint( endpoint, &self.default_namespace_path, - ) => result?, + ) => self.record_error_job_id(result)?, }); continue; } @@ -1067,24 +1165,22 @@ impl RemoteQuery { self.touch(); if self.is_cancellation_requested() { let state = self.state.lock().await.clone(); - return query_description(self.id, &state, self.lifecycle()); + return query_description(self.id, self.job_id(), &state, self.lifecycle()); } let state = self.state.lock().await.clone(); let state = if let Some(descriptor) = state.flight_descriptor.clone() { let poll_guard = tokio::select! { biased; - _ = self.wait_for_cancellation() => return query_description( - self.id, - &state, - self.lifecycle(), - ), + _ = self.wait_for_cancellation() => { + return query_description(self.id, self.job_id(), &state, self.lifecycle()) + } poll_guard = tokio::time::timeout( STATUS_POLL_TIMEOUT, self.poll_gate.lock(), ) => poll_guard, }; let Ok(_poll_guard) = poll_guard else { - return query_description(self.id, &state, self.lifecycle()); + return query_description(self.id, self.job_id(), &state, self.lifecycle()); }; let latest = self.state.lock().await.clone(); if latest.flight_descriptor.as_ref() != Some(&descriptor) { @@ -1092,25 +1188,28 @@ impl RemoteQuery { } else { let updated = tokio::select! { biased; - _ = self.wait_for_cancellation() => return query_description( - self.id, - &latest, - self.lifecycle(), - ), + _ = self.wait_for_cancellation() => { + return query_description(self.id, self.job_id(), &latest, self.lifecycle()) + } result = self.client.poll_status( descriptor.clone(), &self.default_namespace_path, ) => match result { - Err(_) if self.is_cancellation_requested() => return query_description( - self.id, - &latest, - self.lifecycle(), - ), - result => result?, + Err(error) if self.is_cancellation_requested() => { + self.observe_error_job_id(&error)?; + return query_description( + self.id, + self.job_id(), + &latest, + self.lifecycle(), + ); + } + result => self.record_error_job_id(result)?, }, }; if let Some(updated) = updated { - self.update_state(&descriptor, updated).await? + self.observe_job_id(updated.job_id.as_deref())?; + self.update_state(&descriptor, updated.info).await? } else { latest } @@ -1118,7 +1217,7 @@ impl RemoteQuery { } else { state }; - query_description(self.id, &state, self.lifecycle()) + query_description(self.id, self.job_id(), &state, self.lifecycle()) } async fn cancel(&self) -> Result<()> { @@ -1214,11 +1313,15 @@ impl RemoteQuery { if self.state.lock().await.flight_descriptor.as_ref() != Some(&descriptor) { continue; } - let updated = self.client.poll_continuation( + let updated = self.record_error_job_id( + self.client.poll_continuation( descriptor.clone(), &self.default_namespace_path, - ).await?; - self.update_state(&descriptor, updated).await?; + ) + .await, + )?; + self.observe_job_id(updated.job_id.as_deref())?; + self.update_state(&descriptor, updated.info).await?; } _ = notified => {} } @@ -1247,6 +1350,10 @@ impl QueryHandle for RemoteQueryHandle { self.query.id } + fn job_id(&self) -> Option { + self.query.job_id() + } + async fn describe(&self) -> Result { self.query.describe().await } @@ -1286,7 +1393,12 @@ impl QueryHandle for RemoteQueryHandle { receiver.recv().await.map(|item| (item, receiver)) }); result_start.commit(); - Ok(Box::pin(SimpleRecordBatchStream::new(stream, schema))) + let stream: SendableRecordBatchStream = + Box::pin(SimpleRecordBatchStream::new(stream, schema)); + Ok(match self.job_id() { + Some(job_id) => attach_sql_job_id(stream, &job_id), + None => stream, + }) } async fn cancel(&self) -> Result<()> { @@ -1296,12 +1408,14 @@ impl QueryHandle for RemoteQueryHandle { fn query_description( id: Uuid, + job_id: Option, poll_info: &PollInfo, lifecycle: QueryLifecycle, ) -> Result { let expires_at = query_expiration(poll_info)?; Ok(QueryDescription { id, + job_id, status: match lifecycle { QueryLifecycle::Cancelling => QueryStatus::Cancelling, QueryLifecycle::Cancelled => QueryStatus::Cancelled, @@ -1317,6 +1431,36 @@ fn query_description( }) } +fn sql_job_id_from_metadata(metadata: &tonic::metadata::MetadataMap) -> Option { + metadata + .get(SQL_JOB_ID_METADATA_KEY) + .and_then(|value| value.to_str().ok()) + .filter(|value| !value.is_empty()) + .map(str::to_owned) +} + +fn sql_job_id_from_schema(schema: &Schema) -> Option { + schema + .metadata() + .get(SQL_JOB_ID_METADATA_KEY) + .filter(|value| !value.is_empty()) + .cloned() +} + +pub(super) fn attach_sql_job_id( + stream: SendableRecordBatchStream, + job_id: &str, +) -> SendableRecordBatchStream { + let mut metadata = stream.schema().metadata().clone(); + metadata.insert(SQL_JOB_ID_METADATA_KEY.to_string(), job_id.to_string()); + let schema = Arc::new(stream.schema().as_ref().clone().with_metadata(metadata)); + let batch_schema = schema.clone(); + let stream = stream.map(move |batch| { + batch.and_then(|batch| batch.with_schema(batch_schema.clone()).map_err(Error::from)) + }); + Box::pin(SimpleRecordBatchStream::new(stream, schema)) +} + fn query_expiration(poll_info: &PollInfo) -> Result>> { poll_info .expiration_time diff --git a/rust/lancedb/src/remote/sql/stream.rs b/rust/lancedb/src/remote/sql/stream.rs index 420eecd15..16155bf7c 100644 --- a/rust/lancedb/src/remote/sql/stream.rs +++ b/rust/lancedb/src/remote/sql/stream.rs @@ -55,6 +55,7 @@ pub(super) struct StreamingQuery { cancelled: Notify, terminal_at: OnceLock, last_accessed: StdMutex, + job_id: StdMutex>, } struct PreparedStream { @@ -125,6 +126,7 @@ pub(super) enum Submission { pub(super) struct OpenedSubmission { pub(super) submission: Submission, pub(super) materialized_view_job_handle_supported: bool, + pub(super) job_id: Option, } impl StreamingQuery { @@ -190,6 +192,7 @@ impl StreamingQuery { } .map_err(|_| sql_error(&request_id, "SQL query submission timed out"))? .map_err(|error| sql_status_error(&request_id, &error))?; + let response_job_id = super::sql_job_id_from_metadata(response.metadata()); let negotiated = response .metadata() .get(PROTOCOL_HEADER) @@ -233,6 +236,7 @@ impl StreamingQuery { return Ok(OpenedSubmission { submission: Submission::Deferred(info), materialized_view_job_handle_supported, + job_id: response_job_id, }); } } @@ -245,8 +249,10 @@ impl StreamingQuery { stream, request_id, read_timeout, + job_id: None, }), materialized_view_job_handle_supported: false, + job_id: response_job_id, }) } @@ -254,6 +260,7 @@ impl StreamingQuery { id: Uuid, client: Arc, response: ResultEndpointStream, + job_id: Option, ) -> Self { Self { id, @@ -263,6 +270,7 @@ impl StreamingQuery { cancelled: Notify::new(), terminal_at: OnceLock::new(), last_accessed: StdMutex::new(Instant::now()), + job_id: StdMutex::new(job_id.filter(|job_id| !job_id.is_empty())), } } @@ -291,6 +299,7 @@ impl StreamingQuery { }; Ok(QueryDescription { id: self.id, + job_id: self.job_id(), status, progress: (status == QueryStatus::Finished).then_some(1.0), expires_at: None, @@ -326,6 +335,29 @@ impl StreamingQuery { *self.lifecycle.lock().unwrap() == Lifecycle::Cancelled } + fn job_id(&self) -> Option { + self.job_id.lock().unwrap().clone() + } + + fn observe_job_id(&self, job_id: Option<&str>) -> Result<()> { + let Some(job_id) = job_id.filter(|job_id| !job_id.is_empty()) else { + return Ok(()); + }; + let mut current = self.job_id.lock().unwrap(); + match current.as_deref() { + None => *current = Some(job_id.to_string()), + Some(existing) if existing == job_id => {} + Some(existing) => { + return Err(Error::Runtime { + message: format!( + "SQL query changed durable job identity from {existing} to {job_id}" + ), + }); + } + } + Ok(()) + } + fn cancelled_error(&self) -> Error { Error::JobCancelled { job_id: Some(self.id.to_string()), @@ -373,6 +405,7 @@ impl StreamingQuery { .ok_or_else(|| Error::Runtime { message: "SQL result stream did not include a schema".to_string(), })?; + self.observe_job_id(super::sql_job_id_from_schema(&schema).as_deref())?; if first.is_none() { self.settle(Lifecycle::Finished); } @@ -476,6 +509,10 @@ impl QueryHandle for StreamingQueryHandle { self.query.id } + fn job_id(&self) -> Option { + self.query.job_id() + } + async fn describe(&self) -> Result { self.query.describe() } @@ -526,7 +563,12 @@ impl QueryHandle for StreamingQueryHandle { receiver.recv().await.map(|item| (item, receiver)) }); result_start.commit(); - Ok(Box::pin(SimpleRecordBatchStream::new(stream, schema))) + let stream: SendableRecordBatchStream = + Box::pin(SimpleRecordBatchStream::new(stream, schema)); + Ok(match self.job_id() { + Some(job_id) => super::attach_sql_job_id(stream, &job_id), + None => stream, + }) } async fn cancel(&self) -> Result<()> { diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 65d08c984..8ac08f10c 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -160,8 +160,17 @@ impl TestSqlService { return Response::new(rows); } let poll_info = self.defer.then(|| { - let mut info = FlightInfo::new() - .with_endpoint(FlightEndpoint::new().with_ticket(Ticket::new(query.to_string()))); + let mut info = FlightInfo::new(); + if !query.starts_with("SELECT durable job") { + info = info.with_endpoint( + FlightEndpoint::new().with_ticket(Ticket::new(query.to_string())), + ); + } + if query.starts_with("SELECT durable job") { + info = info + .try_with_schema(&self.result_for_query(query).schema()) + .unwrap(); + } if self.materialized_view_job_handle_supported { info.app_metadata = b"early-v1".to_vec().into(); } @@ -198,6 +207,13 @@ impl TestSqlService { response .metadata_mut() .insert(PROTOCOL_HEADER, PROTOCOL_VERSION.parse().unwrap()); + if query.starts_with("SELECT direct durable job") { + let job_id = query.rsplit(' ').next().unwrap(); + response.metadata_mut().insert( + SQL_JOB_ID_METADATA_KEY, + format!("job-{job_id}").parse().unwrap(), + ); + } response } } @@ -351,7 +367,14 @@ impl FlightService for TestSqlService { .insert("x-lancedb-sql-job-id", "job-1".parse().unwrap()); return Err(status); } - let complete = if query == "SELECT no info" { + if query == "SELECT durable job failure" && stage == 1 { + let mut status = Status::internal("durable SQL job failed"); + status + .metadata_mut() + .insert("x-lancedb-sql-job-id", "job-1".parse().unwrap()); + return Err(status); + } + let complete = if query == "SELECT no info" || query.starts_with("SELECT durable job") { stage >= 2 } else { stage >= 1 @@ -362,11 +385,14 @@ impl FlightService for TestSqlService { } else { query.clone() }; - let mut info = FlightInfo::new().with_endpoint( - FlightEndpoint::new() - .with_ticket(Ticket::new(first_ticket)) - .with_location("grpc://127.0.0.1:1"), - ); + let mut info = FlightInfo::new(); + if !(query.starts_with("SELECT durable job") && !complete) { + info = info.with_endpoint( + FlightEndpoint::new() + .with_ticket(Ticket::new(first_ticket)) + .with_location("grpc://127.0.0.1:1"), + ); + } if query == "SELECT incremental" && stage > 0 { info = info.with_endpoint( FlightEndpoint::new() @@ -378,13 +404,21 @@ impl FlightService for TestSqlService { let schema = self.result_for_query(&query).schema(); info = info.try_with_schema(&schema).unwrap(); } - Ok(Response::new(PollInfo { + let mut response = Response::new(PollInfo { info: (query != "SELECT no info" || stage > 0).then_some(info), flight_descriptor: (!complete) .then(|| FlightDescriptor::new_cmd(format!("poll:{}:{query}", stage + 1))), progress: Some(if complete { 1.0 } else { 0.25 }), expiration_time: None, - })) + }); + if query.starts_with("SELECT durable job") && stage > 0 { + let job_id = query.rsplit(' ').next().unwrap(); + response.metadata_mut().insert( + SQL_JOB_ID_METADATA_KEY, + format!("job-{job_id}").parse().unwrap(), + ); + } + Ok(response) } async fn get_schema( @@ -730,6 +764,106 @@ async fn parameters_travel_on_one_exchange() { assert_eq!(headers[0].api_key, "test-key"); } +#[tokio::test] +async fn running_job_backed_queries_expose_their_own_durable_job_id() { + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + + let first = client + .execute(ExecuteQueryRequest::new("SELECT durable job first")) + .await + .unwrap(); + let second = client + .execute(ExecuteQueryRequest::new("SELECT durable job second")) + .await + .unwrap(); + assert_eq!(first.job_id(), None); + assert_eq!(second.job_id(), None); + + let first_description = first.describe().await.unwrap(); + let second_description = second.describe().await.unwrap(); + assert_eq!(first_description.status, QueryStatus::Running); + assert_eq!(first_description.job_id.as_deref(), Some("job-first")); + assert_eq!(second_description.status, QueryStatus::Running); + assert_eq!(second_description.job_id.as_deref(), Some("job-second")); + assert_ne!(first.id(), second.id()); + assert_ne!(first_description.job_id, second_description.job_id); + + let reader = first.reader().await.unwrap(); + assert_eq!( + reader + .schema() + .metadata() + .get(SQL_JOB_ID_METADATA_KEY) + .map(String::as_str), + Some("job-first") + ); + assert_eq!(first.job_id().as_deref(), Some("job-first")); + assert_eq!(second.job_id().as_deref(), Some("job-second")); + + let reader_only = client + .execute(ExecuteQueryRequest::new("SELECT durable job reader-only")) + .await + .unwrap(); + let reader_only_result = reader_only.reader().await.unwrap(); + assert_eq!( + reader_only_result + .schema() + .metadata() + .get(SQL_JOB_ID_METADATA_KEY) + .map(String::as_str), + Some("job-reader-only") + ); + + let ordinary = client + .execute(ExecuteQueryRequest::new("SELECT ordinary")) + .await + .unwrap(); + assert_eq!(ordinary.describe().await.unwrap().job_id, None); +} + +#[tokio::test] +async fn direct_job_backed_queries_expose_the_initial_response_job_id() { + let service = TestSqlService { + adaptive: true, + ..Default::default() + }; + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + let query = client + .execute(ExecuteQueryRequest::new("SELECT direct durable job direct")) + .await + .unwrap(); + + assert_eq!(query.job_id().as_deref(), Some("job-direct")); + assert_eq!( + query.describe().await.unwrap().job_id.as_deref(), + Some("job-direct") + ); + let reader = query.reader().await.unwrap(); + assert_eq!( + reader + .schema() + .metadata() + .get(SQL_JOB_ID_METADATA_KEY) + .map(String::as_str), + Some("job-direct") + ); + let batches = reader.try_collect::>().await.unwrap(); + assert_eq!(batches.len(), 1); + assert_eq!( + batches[0] + .schema() + .metadata() + .get(SQL_JOB_ID_METADATA_KEY) + .map(String::as_str), + Some("job-direct") + ); +} + /// The server plans before it answers, so a refusal is the submission's. #[tokio::test] async fn a_refused_statement_fails_its_submission() { @@ -1213,6 +1347,35 @@ async fn function_failure_preserves_the_isolated_sql_job_id() { shutdown.send(()).unwrap(); } +#[tokio::test] +async fn fast_durable_job_failure_exposes_the_same_job_id_on_the_query_handle() { + let (client, shutdown) = start_sql_service( + TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }, + ClientConfig::default(), + ) + .await; + let query = client + .execute(ExecuteQueryRequest { + query: "SELECT durable job failure".into(), + default_namespace_path: vec![], + parameters: None, + }) + .await + .unwrap(); + assert_eq!(query.job_id(), None); + + let Err(Error::JobFailed { job_id, .. }) = collect_result(&query).await else { + panic!("a failed durable query must remain a structured JobFailed error"); + }; + assert_eq!(job_id.as_deref(), Some("job-1")); + assert_eq!(query.job_id().as_deref(), Some("job-1")); + shutdown.send(()).unwrap(); +} + #[tokio::test] async fn rejected_job_cancellation_does_not_abort_sql_completion() { let service = TestSqlService { @@ -1843,6 +2006,7 @@ async fn submits_polls_fetches_cancels_and_reuses_client() { flight_descriptor: Some(FlightDescriptor::new_cmd("restored")), ..Default::default() }, + None, ) .unwrap(), ); @@ -1975,6 +2139,7 @@ async fn submits_polls_fetches_cancels_and_reuses_client() { flight_descriptor: Some(FlightDescriptor::new_cmd(descriptor)), ..Default::default() }, + None, ) .unwrap(), ); @@ -1996,6 +2161,7 @@ async fn submits_polls_fetches_cancels_and_reuses_client() { expiration_time: Some(Default::default()), ..Default::default() }, + None, ) .unwrap(), ); @@ -2012,6 +2178,7 @@ async fn submits_polls_fetches_cancels_and_reuses_client() { flight_descriptor: Some(FlightDescriptor::new_cmd("stale")), ..Default::default() }, + None, ) .unwrap(), ); diff --git a/rust/lancedb/src/sql.rs b/rust/lancedb/src/sql.rs index 5c01fe898..bda5df085 100644 --- a/rust/lancedb/src/sql.rs +++ b/rust/lancedb/src/sql.rs @@ -44,6 +44,10 @@ impl fmt::Display for QueryStatus { pub struct QueryDescription { /// The stable, connection-scoped identifier assigned when the query was submitted. pub id: Uuid, + /// The durable Job Registry identifier for a server-side SQL job, when the + /// submitted statement ran as one. This identifier can be passed to + /// `open_job`; `id` remains connection-scoped. + pub job_id: Option, /// The server-visible lifecycle state. pub status: QueryStatus, /// Server-reported completion progress, when known. Values are in `[0.0, 1.0]`, @@ -56,6 +60,9 @@ pub struct QueryDescription { #[async_trait] pub(crate) trait QueryHandle: Send + Sync { fn id(&self) -> Uuid; + fn job_id(&self) -> Option { + None + } async fn describe(&self) -> Result; async fn reader(&self) -> Result; async fn cancel(&self) -> Result<()>; @@ -68,8 +75,9 @@ pub(crate) trait QueryHandle: Send + Sync { /// 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. +/// [`Self::id`] is scoped to the connection that submitted the query. When the +/// server runs the statement as a SQL job, [`Self::job_id`] exposes the durable +/// Job Registry identifier accepted by `open_job`. pub struct Query { handle: Arc, } @@ -79,6 +87,7 @@ impl std::fmt::Debug for Query { formatter .debug_struct("Query") .field("id", &self.id()) + .field("job_id", &self.job_id()) .finish() } } @@ -94,6 +103,16 @@ impl Query { self.handle.id() } + /// Return the durable Job Registry identifier when the server has reported + /// one for this query. + /// + /// A running query may report it after [`Self::describe`] observes the + /// server's latest status. The connection-scoped [`Self::id`] remains a + /// separate identifier. + pub fn job_id(&self) -> Option { + self.handle.job_id() + } + /// Get a point-in-time description of the query. /// /// Directly streamed queries stay [`QueryStatus::Running`] until the result @@ -115,7 +134,12 @@ impl Query { self.handle.reader().await } - /// Request cancellation of the query. + /// Request cancellation of this query's Flight execution. + /// + /// For a durable server-side SQL job, the operation may return + /// [`Error::NotSupported`] because the Flight waiter cannot cancel the + /// registry job. Use [`Connection::cancel_job`](crate::Connection::cancel_job) + /// with [`Self::job_id`] to cancel that job. pub async fn cancel(&self) -> Result<()> { self.handle.cancel().await }