From b09735529b7720ba610ecc4f51a8f2f32cbd86d2 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 4 Nov 2026 12:04:02 +0800 Subject: [PATCH] feat(sql): return durable materialized view job handles (#4824) Materialized-view creation and refresh can outlive the submitting connection, but their SQL query IDs cannot be reopened through the Job Registry. On servers that negotiate early handles, return the durable registry job ID and use the existing job API to describe, reopen, wait for, or cancel it. Older servers retain the connection-scoped query handle. Registry refresh handles forward running and terminal progress observations and establish the committed-version read-freshness fence before wait returns. Preserve a nested SQL failure ID separately from the outer job ID through Rust and the Python native bindings, so Function failures remain traceable. Rust callers constructing `JobFailure` directly must supply the new optional `sql_job_id` field or use `..Default::default()`. Validation: seven SDK regression cases and the docs-link test pass, including plain wait and running-to-terminal progress observations; the observer regression fails before the fix. Formatting and strict Clippy pass. Earlier accepted coverage includes a real Flight test service exercising materialized-view creation followed by registry describe/reopen/cancel, and two Python tests after rebuilding the native extension. The published product tree matches the validated candidate. Refs ENT-3562. The matching Sophon server change is being validated separately before publication. --- python/python/lancedb/_lancedb.pyi | 2 + python/python/lancedb/db.py | 11 +- python/python/lancedb/exceptions.py | 19 +- python/python/lancedb/materialized_view.py | 7 +- python/python/tests/test_remote_db.py | 17 +- python/src/error.rs | 8 +- python/src/job.rs | 11 +- rust/lancedb/src/error.rs | 3 + rust/lancedb/src/job.rs | 16 +- rust/lancedb/src/remote/db.rs | 18 +- rust/lancedb/src/remote/job.rs | 30 +++ rust/lancedb/src/remote/sql.rs | 93 ++++++-- .../src/remote/sql/materialized_view.rs | 72 ++++++ rust/lancedb/src/remote/sql/stream.rs | 48 +++- rust/lancedb/src/remote/sql_test.rs | 219 +++++++++++++++++- rust/lancedb/src/remote/table.rs | 219 ++++++++++++++++-- 16 files changed, 715 insertions(+), 78 deletions(-) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index fcad61e71..02d9c9e87 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -466,6 +466,8 @@ class JobFailureInfo: def message(self) -> Optional[str]: ... @property def retryable(self) -> Optional[bool]: ... + @property + def sql_job_id(self) -> Optional[str]: ... class JobDescription: @property diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 05b0560bb..2cf438e72 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -739,11 +739,14 @@ class DBConnection(EnforceOverrides): """Submit materialized-view creation and return its job. The job may already be complete in LanceDB Embedded. In LanceDB - Enterprise, creation - submits CREATE MATERIALIZED VIEW through the SQL server. Its job has - ``job_type == "query"`` and ``id`` identifies the connection-scoped - query. Registry event history is unavailable. + Distributed, creation submits CREATE MATERIALIZED VIEW through the SQL + server and the returned ``id`` is the registry job ID. That ID can be + reopened with :meth:`open_job`, and the handle supports registry + progress, events, cancellation, and ``wait`` while the job runs. Wait for the job before opening or querying the view. + + Servers without early registry-handle support return a connection-scoped + query job instead. Its ID cannot be reopened with :meth:`open_job`. """ raise NotImplementedError( _append_doc_url( diff --git a/python/python/lancedb/exceptions.py b/python/python/lancedb/exceptions.py index 67f15cabe..516b54eed 100644 --- a/python/python/lancedb/exceptions.py +++ b/python/python/lancedb/exceptions.py @@ -3,6 +3,8 @@ """Custom exception handling""" +from typing import Optional + class MissingValueError(ValueError): """Exception raised when a required value is missing.""" @@ -26,9 +28,22 @@ class MissingColumnError(KeyError): class JobFailedError(RuntimeError): - """Exception raised when an asynchronous job reaches the failed state.""" + """Exception raised when an asynchronous job reaches the failed state. - pass + ``job_id`` identifies the failed job when available. ``sql_job_id`` preserves + the SQL execution identity reported by the server, which may differ from + the outer job's identity. + """ + + def __init__( + self, + message: str, + job_id: Optional[str] = None, + sql_job_id: Optional[str] = None, + ): + super().__init__(message) + self.job_id = job_id + self.sql_job_id = sql_job_id class JobCancelledError(RuntimeError): diff --git a/python/python/lancedb/materialized_view.py b/python/python/lancedb/materialized_view.py index 3594f5603..a1c83dd6d 100644 --- a/python/python/lancedb/materialized_view.py +++ b/python/python/lancedb/materialized_view.py @@ -208,9 +208,14 @@ class AsyncMaterializedView: """Submit a refresh and return its job without waiting. The job may already be complete in LanceDB Embedded. In LanceDB - Enterprise, refreshes are submitted through the SQL service. Options and + Distributed, refreshes are submitted through the SQL service and the + returned ``id`` is the registry job ID, usable with ``open_job`` for + progress, events, and cancellation while the refresh runs. Options and result metrics follow [refresh][lancedb.materialized_view.AsyncMaterializedView.refresh]. + + Servers without early registry-handle support return a connection-scoped + query job instead. Its ID cannot be reopened with ``open_job``. """ from ._lancedb import RefreshMaterializedViewResult diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 3a9db8812..84b973d90 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -2268,8 +2268,10 @@ def test_remote_job_wait_raises_on_failure(): with mock_lancedb_connection(handler) as db: table = db.create_table("test", [{"id": 1}]) job = table.create_index_async("id", config=BTree()) - with pytest.raises(JobFailedError, match="job-2"): + with pytest.raises(JobFailedError, match="job-2") as exc_info: job.wait() + assert exc_info.value.job_id == "job-2" + assert exc_info.value.sql_job_id is None @pytest.mark.parametrize("server_version", ["0.2.0", "0.6.0"]) @@ -5123,7 +5125,10 @@ def test_remote_connection_jobs_surface(): waiting_for_job_id="job-maintenance", paused_at_millis=1200, failure=dict( - phase="execute", message="worker died", retryable=True + phase="execute", + message="worker died", + retryable=True, + sql_job_id="sql-job-1", ), ) ).encode() @@ -5211,13 +5216,17 @@ def test_remote_connection_jobs_surface(): assert progress["waiting_for_job_id"] == "job-maintenance" assert progress["paused_at_millis"] == 1200 assert job.failure.message == "worker died" + assert job.failure.sql_job_id == "sql-job-1" assert repr(job.failure) == ( - 'JobFailureInfo(phase="execute", message="worker died", retryable=True)' + 'JobFailureInfo(phase="execute", message="worker died", retryable=True, ' + 'sql_job_id="sql-job-1")' ) assert job.status() == "failed" observed = [] - with pytest.raises(JobFailedError, match="worker died"): + with pytest.raises(JobFailedError, match="worker died") as exc_info: job.wait(timeout=timedelta(seconds=5), progress=observed.append) + assert exc_info.value.job_id == "job-1" + assert exc_info.value.sql_job_id == "sql-job-1" assert observed[0]["state"] == "failed" assert observed[0]["rows_computed"] == 4 assert finished.progress()["paused"] is False diff --git a/python/src/error.rs b/python/src/error.rs index 6528851c4..369b68c47 100644 --- a/python/src/error.rs +++ b/python/src/error.rs @@ -125,11 +125,15 @@ impl PythonErrorExt for std::result::Result { err.setattr(intern!(py, "__cause__"), cause_err)?; Err(PyErr::from_value(err)) }), - LanceError::JobFailed { .. } => Python::attach(|py| { + LanceError::JobFailed { job_id, failure } => Python::attach(|py| { let cls = py .import(intern!(py, "lancedb.exceptions"))? .getattr(intern!(py, "JobFailedError"))?; - Err(PyErr::from_value(cls.call1((err.to_string(),))?)) + Err(PyErr::from_value(cls.call1(( + err.to_string(), + job_id.clone(), + failure.sql_job_id.clone(), + ))?)) }), LanceError::JobCancelled { .. } => Python::attach(|py| { let cls = py diff --git a/python/src/job.rs b/python/src/job.rs index fd5ef64be..aa4d8890f 100644 --- a/python/src/job.rs +++ b/python/src/job.rs @@ -233,6 +233,7 @@ pub struct JobFailureInfo { phase: Option, message: Option, retryable: Option, + sql_job_id: Option, } #[pymethods] @@ -251,7 +252,14 @@ impl JobFailureInfo { Some(false) => "False", None => "None", }; - format!("JobFailureInfo(phase={phase}, message={message}, retryable={retryable})") + let sql_job_id = self + .sql_job_id + .as_ref() + .map_or_else(|| "None".to_string(), |value| format!("{value:?}")); + format!( + "JobFailureInfo(phase={phase}, message={message}, retryable={retryable}, \ + sql_job_id={sql_job_id})" + ) } } @@ -335,6 +343,7 @@ impl From for JobDescription { phase: failure.phase, message: failure.message, retryable: failure.retryable, + sql_job_id: failure.sql_job_id, }), } } diff --git a/rust/lancedb/src/error.rs b/rust/lancedb/src/error.rs index ddd04fa45..b607bd9b8 100644 --- a/rust/lancedb/src/error.rs +++ b/rust/lancedb/src/error.rs @@ -111,6 +111,9 @@ pub struct JobFailure { pub message: Option, /// Whether a retry could clear the failure, when known. pub retryable: Option, + /// Isolated SQL job that raised this failure, when the server supplies + /// one across the Flight SQL boundary. + pub sql_job_id: Option, /// The error the job failed with, when it ran in this process. pub source: Option>, } diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 7b801bc0a..823834798 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -118,6 +118,11 @@ impl TerminalResult { self.value.as_ref() } + pub(crate) fn with_value(mut self, value: Value) -> Self { + self.value = Some(value); + self + } + fn decode(self) -> Result { let value = self.value.ok_or_else(|| match &self.request_id { Some(request_id) => remote_job_error( @@ -575,10 +580,13 @@ impl Outcome { fn into_result(self) -> Result { match self { Self::Succeeded(result) => Ok(result), - Self::Failed(source) => Err(Error::JobFailed { - job_id: None, - failure: JobFailure::from_source(source), - }), + Self::Failed(source) => { + let (job_id, failure) = match source.as_ref() { + Error::JobFailed { job_id, failure } => (job_id.clone(), failure.clone()), + _ => (None, JobFailure::from_source(source)), + }; + Err(Error::JobFailed { job_id, failure }) + } Self::Cancelled => Err(Error::JobCancelled { job_id: None }), } } diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 229cdd28c..b9a3f8ae1 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -775,16 +775,26 @@ impl Database for RemoteDatabase { } else { format!(" WITH ({})", options.join(", ")) }; - client - .submit_as_job( + let submission = client + .submit_materialized_view_job( format!( "CREATE MATERIALIZED VIEW \"{name}\"{options} AS {}", request.query ), request.namespace_path, - || async { Ok(()) }, ) - .await + .await?; + match submission { + crate::remote::sql::materialized_view::MaterializedViewJobSubmission::Registry { + job_id, + } => Ok(Job::new(Box::new(RemoteJob::new( + self.client.clone(), + job_id, + )))), + crate::remote::sql::materialized_view::MaterializedViewJobSubmission::Query(query) => { + Ok(client.query_job(query, || async { Ok(()) })) + } + } } async fn drop_materialized_view_async( diff --git a/rust/lancedb/src/remote/job.rs b/rust/lancedb/src/remote/job.rs index 22ae938ed..a5baab7b0 100644 --- a/rust/lancedb/src/remote/job.rs +++ b/rust/lancedb/src/remote/job.rs @@ -75,6 +75,8 @@ pub(super) struct ReportedFailure { message: Option, #[serde(default)] retryable: Option, + #[serde(default)] + sql_job_id: Option, } /// Forward-compatible `/v1/jobs/pause` wire envelope. @@ -133,6 +135,7 @@ impl ReportedFailure { phase: self.phase, message: self.message, retryable: self.retryable, + sql_job_id: self.sql_job_id, source: None, } } @@ -481,4 +484,31 @@ mod tests { assert_eq!(snapshots[1]["state"], "finished"); assert_eq!(snapshots[1]["rows_computed"], 10); } + + #[test] + fn reported_failure_fields_are_passed_through() { + let response: DescribeJobResponse = serde_json::from_str( + r#"{ + "job_state":"FAILED", + "failure":{ + "phase":"initialize", + "message":"worker initialization timed out", + "retryable":true, + "sql_job_id":"job-1" + } + }"#, + ) + .expect("failure description"); + let failure = response + .failure + .expect("the wire response has a failure") + .into_job_failure(); + assert_eq!(failure.phase.as_deref(), Some("initialize")); + assert_eq!( + failure.message.as_deref(), + Some("worker initialization timed out") + ); + assert_eq!(failure.retryable, Some(true)); + assert_eq!(failure.sql_job_id.as_deref(), Some("job-1")); + } } diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index 7e3a42ac7..cd19f733e 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -29,7 +29,7 @@ use uuid::Uuid; use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream}; use crate::database::ExecuteQueryRequest; -use crate::error::{Error, Result}; +use crate::error::{Error, JobFailure, Result}; use crate::job::Job; use crate::remote::client::{ClientConfig, TlsConfig}; use crate::remote::retry::ResolvedRetryConfig; @@ -51,6 +51,7 @@ const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(300); const STATUS_POLL_TIMEOUT: Duration = Duration::from_secs(1); const MIN_POLL_INTERVAL: Duration = Duration::from_millis(50); const MAX_SQL_MESSAGE_SIZE: usize = 1024 * 1024 * 1024; +const SQL_JOB_ID_METADATA_KEY: &str = "x-lancedb-sql-job-id"; const TERMINAL_QUERY_RETENTION: Duration = Duration::from_secs(300); const ABANDONED_QUERY_RETENTION: Duration = Duration::from_secs(24 * 60 * 60); @@ -60,6 +61,11 @@ pub(super) struct SqlClient { queries: Arc, } +pub(super) struct SubmittedQuery { + pub(super) query: Query, + pub(super) materialized_view_job_handle_supported: bool, +} + struct SqlClientInner { database: String, database_prefix: Option, @@ -95,7 +101,7 @@ impl ResultEndpointStream { tokio::time::timeout(self.read_timeout, self.stream.try_next()) .await .map_err(|_| sql_error(&self.request_id, "SQL result read timed out"))? - .map_err(|err| sql_error(&self.request_id, err)) + .map_err(|err| sql_flight_error(&self.request_id, &err)) } } @@ -204,24 +210,36 @@ impl SqlClient { } pub(super) async fn execute(&self, request: ExecuteQueryRequest) -> Result { + Ok(self + .execute_with_materialized_view_job_handle(request, false) + .await? + .query) + } + + pub(super) async fn execute_with_materialized_view_job_handle( + &self, + request: ExecuteQueryRequest, + materialized_view_job_handle: bool, + ) -> Result { 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( + let opened = StreamingQuery::open( self.inner.clone(), &request.query, request.parameters, &request.default_namespace_path, + materialized_view_job_handle, ) .await?; - match submission { + let query = match opened.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)))) + Query::new(Arc::new(StreamingQueryHandle::new(query))) } Submission::Deferred(poll_info) => { let query = Arc::new(RemoteQuery::new( @@ -232,9 +250,14 @@ impl SqlClient { )?); self.queries .insert(query_id, RegisteredQuery::Polled(query.clone())); - Ok(Query::new(Arc::new(RemoteQueryHandle::new(query)))) + Query::new(Arc::new(RemoteQueryHandle::new(query))) } - } + }; + Ok(SubmittedQuery { + query, + materialized_view_job_handle_supported: opened + .materialized_view_job_handle_supported, + }) }) .await } @@ -287,6 +310,15 @@ impl SqlClient { }) .await?, ); + Ok(self.query_job(query, finish)) + } + + pub(super) fn query_job(&self, query: Arc, finish: F) -> Job + where + T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static, + F: FnOnce() -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + { let query_id = query.id(); let execution = query.clone(); let task = tokio::spawn(async move { @@ -294,10 +326,10 @@ impl SqlClient { while reader.try_next().await?.is_some() {} finish().await }); - Ok(Job::from_query_task(query_id, task, move || { + Job::from_query_task(query_id, task, move || { let query = query.clone(); async move { query.cancel().await } - })) + }) } pub(super) async fn describe(&self, query_id: Uuid) -> Result { @@ -344,7 +376,7 @@ impl SqlClientInner { tokio::time::timeout(read_timeout, client.poll_flight_info(descriptor)) .await .map_err(|_| sql_error(&request_id, "SQL query poll timed out"))? - .map_err(|err| sql_error(&request_id, err)) + .map_err(|err| sql_flight_error(&request_id, &err)) } async fn poll_status( @@ -358,7 +390,9 @@ impl SqlClientInner { .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_error(&request_id, err)), + Ok(result) => result + .map(Some) + .map_err(|err| sql_flight_error(&request_id, &err)), Err(_) => Ok(None), } } @@ -402,7 +436,7 @@ impl SqlClientInner { tokio::time::sleep(poll_retry_delay(&retry_config, retry_count)).await; continue; } - Ok(Err(error)) => return Err(sql_error(&request_id, error)), + Ok(Err(error)) => return Err(sql_flight_error(&request_id, &error)), Ok(Ok(poll_info)) => poll_info, }; if let Some(delay) = MIN_POLL_INTERVAL.checked_sub(started.elapsed()) { @@ -433,7 +467,7 @@ impl SqlClientInner { 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_error(&request_id, err))?; + .map_err(|err| sql_flight_error(&request_id, &err))?; Ok(ResultEndpointStream { stream, request_id, @@ -501,9 +535,9 @@ impl SqlClientInner { } Err(FlightError::Tonic(status)) if !cancellation_status_is_ambiguous(status.code()) => { attempt.resolve(); - return Err(sql_error(&request_id, status)); + return Err(sql_status_error(&request_id, &status)); } - Err(error) => return Err(sql_error(&request_id, error)), + Err(error) => return Err(sql_flight_error(&request_id, &error)), }; let status = CancelStatus::try_from(result.status) .map_err(|_| sql_error(&request_id, "SQL query returned an invalid cancel status"))?; @@ -1601,6 +1635,35 @@ fn sql_error(request_id: &str, error: impl std::fmt::Display) -> Error { } } +pub(super) fn sql_flight_error(request_id: &str, error: &FlightError) -> Error { + match error { + FlightError::Tonic(status) => sql_status_error(request_id, status), + _ => sql_error(request_id, error), + } +} + +pub(super) fn sql_status_error(request_id: &str, status: &tonic::Status) -> Error { + let sql_job_id = status + .metadata() + .get(SQL_JOB_ID_METADATA_KEY) + .and_then(|value| value.to_str().ok()) + .filter(|value| !value.is_empty()) + .map(str::to_string); + if let Some(sql_job_id) = sql_job_id { + return Error::JobFailed { + job_id: Some(sql_job_id.clone()), + failure: JobFailure { + phase: None, + message: Some(status.message().to_string()), + retryable: None, + sql_job_id: Some(sql_job_id), + source: None, + }, + }; + } + sql_error(request_id, status) +} + #[cfg(test)] #[path = "sql_test.rs"] mod tests; diff --git a/rust/lancedb/src/remote/sql/materialized_view.rs b/rust/lancedb/src/remote/sql/materialized_view.rs index 46c0a7bd2..4eaad5156 100644 --- a/rust/lancedb/src/remote/sql/materialized_view.rs +++ b/rust/lancedb/src/remote/sql/materialized_view.rs @@ -1,15 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +use std::sync::Arc; + use arrow_array::{Array, RecordBatch, StringArray}; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; use futures::TryStreamExt; use super::SqlClient; +use crate::arrow::SendableRecordBatchStream; use crate::database::ExecuteQueryRequest; use crate::error::{Error, Result}; use crate::listing::{ListingOptions, ListingPage}; use crate::materialized_view::{MaterializedViewDefinition, MaterializedViewInfo}; +use crate::sql::Query; use crate::utils::{validate_namespace, validate_table_name}; pub(in crate::remote) fn quote_identifier(name: &str) -> String { @@ -34,7 +38,75 @@ fn strings(batch: &RecordBatch, name: &str) -> Result> { .collect()) } +pub(in crate::remote) enum MaterializedViewJobSubmission { + Registry { job_id: String }, + Query(Arc), +} + impl SqlClient { + pub(in crate::remote) async fn submit_materialized_view_job( + &self, + statement: String, + namespace: Vec, + ) -> Result { + let submitted = self + .execute_with_materialized_view_job_handle( + ExecuteQueryRequest { + query: statement, + default_namespace_path: namespace, + parameters: None, + }, + true, + ) + .await?; + let query = Arc::new(submitted.query); + if !submitted.materialized_view_job_handle_supported { + return Ok(MaterializedViewJobSubmission::Query(query)); + } + let mut reader: SendableRecordBatchStream = query.reader().await?; + let mut rows = 0; + let mut registry_job_id = None; + while let Some(batch) = reader.try_next().await? { + rows += batch.num_rows(); + if rows > 1 { + return Err(Error::Runtime { + message: "materialized-view submission must return exactly one status row" + .into(), + }); + } + if let Some(column) = batch.column_by_name("job_id") { + let values = column + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::Runtime { + message: "materialized-view submission job_id is not UTF-8".into(), + })?; + if values.null_count() != 0 { + return Err(Error::Runtime { + message: "materialized-view submission contains a null job_id".into(), + }); + } + if values.is_empty() { + continue; + } + let value = values.value(0); + if !value.is_empty() { + registry_job_id = Some(value.to_string()); + } + } + } + if rows != 1 { + return Err(Error::Runtime { + message: "materialized-view submission must return exactly one status row".into(), + }); + } + let job_id = registry_job_id.ok_or_else(|| Error::Runtime { + message: "server acknowledged an early materialized-view handle but returned no job_id" + .into(), + })?; + Ok(MaterializedViewJobSubmission::Registry { job_id }) + } + pub(in crate::remote) async fn describe_materialized_view( &self, name: &str, diff --git a/rust/lancedb/src/remote/sql/stream.rs b/rust/lancedb/src/remote/sql/stream.rs index ed85c6c92..420eecd15 100644 --- a/rust/lancedb/src/remote/sql/stream.rs +++ b/rust/lancedb/src/remote/sql/stream.rs @@ -20,7 +20,8 @@ use uuid::Uuid; use super::{ ABANDONED_QUERY_RETENTION, DEFAULT_READ_TIMEOUT, ResultEndpointStream, ResultStartGuard, - SqlClientInner, TERMINAL_QUERY_RETENTION, resolve_timeout, sql_error, with_overall_timeout, + SqlClientInner, TERMINAL_QUERY_RETENTION, resolve_timeout, sql_error, sql_flight_error, + sql_status_error, with_overall_timeout, }; use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream}; use crate::error::{Error, Result}; @@ -99,6 +100,9 @@ impl Drop for Borrowed<'_> { pub(super) const PROTOCOL_HEADER: &str = "lancedb-sql-protocol"; pub(super) const PROTOCOL_VERSION: &str = "adaptive-v1"; +const MATERIALIZED_VIEW_JOB_HANDLE_HEADER: &str = "x-lancedb-mv-job-handle"; +const MATERIALIZED_VIEW_JOB_HANDLE_VALUE: &str = "early"; +const MATERIALIZED_VIEW_JOB_HANDLE_ACK: &[u8] = b"early-v1"; const RESPONSE_TYPE_URL: &str = "type.googleapis.com/lancedb.sql.QueryResponse"; #[derive(Clone, PartialEq, Message)] @@ -118,13 +122,19 @@ pub(super) enum Submission { Deferred(PollInfo), } +pub(super) struct OpenedSubmission { + pub(super) submission: Submission, + pub(super) materialized_view_job_handle_supported: bool, +} + impl StreamingQuery { pub(super) async fn open( client: Arc, query: &str, parameters: Option, default_namespace_path: &[String], - ) -> Result { + materialized_view_job_handle: bool, + ) -> Result { if let Some(parameters) = ¶meters && parameters.num_rows() != 1 { @@ -154,6 +164,12 @@ impl StreamingQuery { .metadata() .clone(); metadata.insert(PROTOCOL_HEADER, PROTOCOL_VERSION.parse().unwrap()); + if materialized_view_job_handle { + metadata.insert( + MATERIALIZED_VIEW_JOB_HANDLE_HEADER, + MATERIALIZED_VIEW_JOB_HANDLE_VALUE.parse().unwrap(), + ); + } let mut flight = connection.client.clone(); let response = if let Some(parameters) = parameters { let messages = FlightDataEncoderBuilder::new() @@ -162,7 +178,7 @@ impl StreamingQuery { .build(futures::stream::iter([Ok(parameters)])) .try_collect::>() .await - .map_err(|error| sql_error(&request_id, error))?; + .map_err(|error| sql_flight_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 @@ -173,7 +189,7 @@ impl StreamingQuery { 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))?; + .map_err(|error| sql_status_error(&request_id, &error))?; let negotiated = response .metadata() .get(PROTOCOL_HEADER) @@ -183,7 +199,7 @@ impl StreamingQuery { 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))? + .map_err(|error| sql_status_error(&request_id, &error))? .ok_or_else(|| { sql_error(&request_id, "SQL response omitted its execution decision") })?; @@ -210,18 +226,28 @@ impl StreamingQuery { "SQL handoff omitted result and continuation information", )); } - return Ok(Submission::Deferred(info)); + let materialized_view_job_handle_supported = + info.info.as_ref().is_some_and(|info| { + info.app_metadata.as_ref() == MATERIALIZED_VIEW_JOB_HANDLE_ACK + }); + return Ok(OpenedSubmission { + submission: Submission::Deferred(info), + materialized_view_job_handle_supported, + }); } } let stream = FlightRecordBatchStream::new_from_flight_data( response_stream .map_err(|status| arrow_flight::error::FlightError::Tonic(Box::new(status))), ); - Ok(Submission::Stream(ResultEndpointStream { - stream, - request_id, - read_timeout, - })) + Ok(OpenedSubmission { + submission: Submission::Stream(ResultEndpointStream { + stream, + request_id, + read_timeout, + }), + materialized_view_job_handle_supported: false, + }) } pub(super) fn new( diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 3c9edb322..7031b1422 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -68,6 +68,7 @@ struct CapturedHeaders { request_id: String, api_key: String, database_prefix: String, + materialized_view_job_handle: Option, } /// Every exchange's statement and the parameter row it arrived with. @@ -93,6 +94,7 @@ struct TestSqlService { gate_refresh_release: Arc, exchanges: ReceivedExchanges, result_overrides: Arc>>, + materialized_view_job_handle_supported: bool, /// How many exchange responses have been dropped, finished or not. exchange_drops: Arc, result: RecordBatch, @@ -128,6 +130,9 @@ impl TestSqlService { .get("x-lancedb-database-prefix") .map(|value| value.to_str().unwrap().to_string()) .unwrap_or_default(), + materialized_view_job_handle: metadata + .get("x-lancedb-mv-job-handle") + .map(|value| value.to_str().unwrap().to_string()), }); } @@ -154,15 +159,19 @@ impl TestSqlService { 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())), - )), + let poll_info = self.defer.then(|| { + let mut info = FlightInfo::new() + .with_endpoint(FlightEndpoint::new().with_ticket(Ticket::new(query.to_string()))); + if self.materialized_view_job_handle_supported { + info.app_metadata = b"early-v1".to_vec().into(); + } + PollInfo { + info: Some(info), 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() @@ -243,6 +252,7 @@ impl Default for TestSqlService { gate_refresh_release: Arc::new(Notify::new()), exchanges: Arc::new(std::sync::Mutex::new(Vec::new())), result_overrides: Arc::new(std::sync::Mutex::new(HashMap::new())), + materialized_view_job_handle_supported: false, exchange_drops: Arc::new(AtomicUsize::new(0)), result, large_result, @@ -334,6 +344,13 @@ impl FlightService for TestSqlService { { return Err(Status::unavailable("transient polling failure")); } + if query == "SELECT function failure" && stage == 1 { + let mut status = Status::failed_precondition("Function 'frames' raised ValueError"); + status + .metadata_mut() + .insert("x-lancedb-sql-job-id", "job-1".parse().unwrap()); + return Err(status); + } let complete = if query == "SELECT no info" { stage >= 2 } else { @@ -977,6 +994,166 @@ async fn remote_create_uses_sql_and_preserves_options_and_namespace() { } } +#[tokio::test] +async fn remote_create_returns_the_registry_job_from_the_early_status_row() { + let service = TestSqlService { + adaptive: true, + defer: true, + materialized_view_job_handle_supported: true, + ..Default::default() + }; + let headers = service.headers.clone(); + let schema = Arc::new(Schema::new(vec![ + Field::new("operation", DataType::Utf8, false), + Field::new("name", DataType::Utf8, false), + Field::new("job_id", DataType::Utf8, false), + Field::new("status", DataType::Utf8, false), + ])); + service.result_overrides.lock().unwrap().insert( + "CREATE MATERIALIZED VIEW \"adults\" AS SELECT * FROM people".into(), + RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(vec!["CREATE MATERIALIZED VIEW"])), + Arc::new(StringArray::from(vec!["adults"])), + Arc::new(StringArray::from(vec!["job-1"])), + Arc::new(StringArray::from(vec!["submitted"])), + ], + ) + .unwrap(), + ); + let (sql_client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let db = RemoteDatabase::new_mock(|request| -> http::Response { + assert_eq!(request.url().path(), "/v1/jobs/describe"); + http::Response::builder() + .status(200) + .body( + r#"{"job_id":"job-1","job_type":"sql","job_state":"DONE","creation_ms":1}"# + .to_string(), + ) + .unwrap() + }) + .with_sql_client(sql_client); + + let job = db + .create_materialized_view_async(CreateMaterializedViewRequest { + name: "adults".into(), + namespace_path: vec!["team".into(), "reports".into()], + query: "SELECT * FROM people".into(), + with_no_data: false, + auto_refresh_enabled: None, + auto_refresh_interval_seconds: None, + }) + .await + .unwrap(); + + assert_eq!(job.id(), Some("job-1")); + job.refresh().await.unwrap(); + assert_eq!(job.job_type().as_deref(), Some("sql")); + job.wait().await.unwrap(); + let headers = headers.lock().unwrap(); + assert!(headers.iter().any(|header| { + header.materialized_view_job_handle.as_deref() == Some("early") + && header.namespace_path == "team$reports" + })); + assert!( + headers + .iter() + .all(|header| header.namespace_path == "team$reports") + ); + shutdown.send(()).unwrap(); +} + +#[tokio::test] +async fn remote_create_registry_handle_reopens_describes_and_cancels() { + let service = TestSqlService { + adaptive: true, + defer: true, + materialized_view_job_handle_supported: true, + ..Default::default() + }; + let schema = Arc::new(Schema::new(vec![ + Field::new("operation", DataType::Utf8, false), + Field::new("name", DataType::Utf8, false), + Field::new("job_id", DataType::Utf8, false), + Field::new("status", DataType::Utf8, false), + ])); + service.result_overrides.lock().unwrap().insert( + "CREATE MATERIALIZED VIEW \"adults\" AS SELECT * FROM people".into(), + RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from(vec!["CREATE MATERIALIZED VIEW"])), + Arc::new(StringArray::from(vec!["adults"])), + Arc::new(StringArray::from(vec!["job-1"])), + Arc::new(StringArray::from(vec!["submitted"])), + ], + ) + .unwrap(), + ); + let (sql_client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let describes = Arc::new(AtomicUsize::new(0)); + let seen_describes = describes.clone(); + let describe_ids = Arc::new(std::sync::Mutex::new(Vec::new())); + let seen_describe_ids = describe_ids.clone(); + let cancel_ids = Arc::new(std::sync::Mutex::new(Vec::new())); + let seen_cancel_ids = cancel_ids.clone(); + let db = RemoteDatabase::new_mock(move |request| { + let body = request + .body() + .and_then(|body| body.as_bytes()) + .and_then(|body| serde_json::from_slice::(body).ok()) + .unwrap_or_default(); + let job_id = body["job_id"].as_str().unwrap_or_default().to_string(); + match request.url().path() { + "/v1/jobs/describe" => { + assert_eq!(job_id, "job-1"); + seen_describe_ids.lock().unwrap().push(job_id.clone()); + seen_describes.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body(format!( + r#"{{"job_id":"{job_id}","job_type":"sql","job_state":"IN_PROGRESS","creation_ms":1}}"# + )) + .unwrap() + } + "/v1/jobs/cancel" => { + assert_eq!(job_id, "job-1"); + seen_cancel_ids.lock().unwrap().push(job_id); + http::Response::builder().status(200).body(String::new()).unwrap() + } + path => panic!("unexpected request: {path}"), + } + }) + .with_sql_client(sql_client); + + let job = db + .create_materialized_view_async(CreateMaterializedViewRequest { + name: "adults".into(), + namespace_path: vec!["team".into(), "reports".into()], + query: "SELECT * FROM people".into(), + with_no_data: false, + auto_refresh_enabled: None, + auto_refresh_interval_seconds: None, + }) + .await + .unwrap(); + assert_eq!(job.id(), Some("job-1")); + + job.refresh().await.unwrap(); + assert_eq!(job.job_type().as_deref(), Some("sql")); + assert_eq!(job.state().as_deref(), Some("running")); + + let reopened = db.open_job("job-1").await.unwrap(); + assert_eq!(reopened.id(), Some("job-1")); + assert_eq!(reopened.job_type().as_deref(), Some("sql")); + reopened.cancel().await.unwrap(); + + assert_eq!(*describe_ids.lock().unwrap(), vec!["job-1", "job-1"]); + assert_eq!(*cancel_ids.lock().unwrap(), vec!["job-1"]); + shutdown.send(()).unwrap(); +} + #[tokio::test] async fn query_backed_jobs_propagate_submission_and_execution_failures() { let (client, shutdown) = start_sql_service( @@ -1008,6 +1185,34 @@ async fn query_backed_jobs_propagate_submission_and_execution_failures() { shutdown.send(()).unwrap(); } +#[tokio::test] +async fn function_failure_preserves_the_isolated_sql_job_id() { + let (client, shutdown) = start_sql_service( + TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }, + ClientConfig::default(), + ) + .await; + let job = client + .submit_as_job("SELECT function failure".into(), vec![], || async { + Ok(()) + }) + .await + .unwrap(); + + let Err(Error::JobFailed { job_id, failure }) = job.wait().await else { + panic!("a typed Function failure must stay a JobFailed error"); + }; + assert_eq!(job_id.as_deref(), Some("job-1")); + assert_eq!(failure.sql_job_id.as_deref(), Some("job-1")); + assert_eq!(failure.phase, None); + assert_eq!(failure.retryable, None); + shutdown.send(()).unwrap(); +} + #[tokio::test] async fn rejected_job_cancellation_does_not_abort_sql_completion() { let service = TestSqlService { @@ -1216,6 +1421,8 @@ async fn cancelling_remote_refresh_cancels_sql_query() { .refresh_materialized_view_async(None, false) .await .unwrap(); + assert_eq!(job.job_type().as_deref(), Some("query")); + assert!(job.id().is_some()); tokio::time::timeout(Duration::from_secs(10), async { while do_get_count.load(Ordering::SeqCst) == 0 { tokio::task::yield_now().await; diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index a46fb8a11..9e9d507b7 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -332,6 +332,82 @@ impl crate::job::JobHandle for FreshnessJob { } } +/// A registry-backed materialized-view refresh. Registry completion is only +/// the first half of `wait`; the submitting handle still has to observe the +/// committed table version before reporting the typed refresh result. +struct MaterializedViewRefreshJob { + inner: RemoteJob, + client: RestfulLanceDbClient, + db_name: String, + name: String, + namespace: Vec, + identifier: String, + server_version: ServerVersion, + sql_client: Option, + freshness: Arc>, + freshness_request: FreshnessHeaders, +} + +#[async_trait] +impl crate::job::JobHandle for MaterializedViewRefreshJob { + fn id(&self) -> Option<&str> { + crate::job::JobHandle::id(&self.inner) + } + + async fn status(&self) -> Result { + crate::job::JobHandle::status(&self.inner).await + } + + async fn describe(&self) -> Result { + crate::job::JobHandle::describe(&self.inner).await + } + + async fn events( + &self, + request: crate::job::JobEventsRequest, + ) -> Result> { + crate::job::JobHandle::events(&self.inner, request).await + } + + async fn wait(&self) -> Result { + crate::job::JobHandle::wait_with_progress(self, None).await + } + + async fn wait_with_progress( + &self, + progress: Option, + ) -> Result { + let result = crate::job::JobHandle::wait_with_progress(&self.inner, progress).await?; + let table = RemoteTable::new_with_sql_client( + self.client.clone(), + self.db_name.clone(), + self.name.clone(), + self.namespace.clone(), + self.identifier.clone(), + self.server_version.clone(), + self.sql_client.clone(), + ); + table.checkout_latest().await?; + let version = table.version().await?; + self.freshness_request + .observe_version(&self.freshness, version); + let refresh = RefreshMaterializedViewResult { + mode: None, + rows_written: None, + source_version: None, + version, + }; + let value = serde_json::to_value(refresh).map_err(|error| Error::Runtime { + message: format!("failed to encode materialized-view refresh result: {error}"), + })?; + Ok(result.with_value(value)) + } + + async fn cancel(&self) -> Result<()> { + crate::job::JobHandle::cancel(&self.inner).await + } +} + fn compute_min_timestamp( state: &FreshnessState, interval: Option, @@ -2320,30 +2396,49 @@ impl BaseTable for RemoteTable { let freshness = self.freshness.clone(); let freshness_request = self.snapshot_freshness_headers(); let job_sql_client = sql_client.clone(); - sql_client - .submit_as_job(statement, namespace.clone(), move || async move { - let table = Self::new_with_sql_client( - client, - db_name, - name, - namespace, - identifier, - server_version, - Some(job_sql_client), - ); - table.checkout_latest().await?; - let version = table.version().await?; - freshness_request.observe_version(&freshness, version); - Ok(RefreshMaterializedViewResult { - // The SQL job adapter does not expose refresh metrics. Table - // metadata cannot tell us what this refresh actually did. - mode: None, - rows_written: None, - source_version: None, - version, - }) - }) - .await + let submission = sql_client + .submit_materialized_view_job(statement, self.namespace.clone()) + .await?; + match submission { + crate::remote::sql::materialized_view::MaterializedViewJobSubmission::Registry { + job_id, + } => Ok(Job::new_typed(Box::new(MaterializedViewRefreshJob { + inner: RemoteJob::new(self.client.clone(), job_id), + client: self.client.clone(), + db_name: self.db_name.clone(), + name: self.name.clone(), + namespace: self.namespace.clone(), + identifier: self.identifier.clone(), + server_version: self.server_version.clone(), + sql_client: self.sql_client.clone(), + freshness: self.freshness.clone(), + freshness_request: self.snapshot_freshness_headers(), + }))), + crate::remote::sql::materialized_view::MaterializedViewJobSubmission::Query(query) => { + Ok(sql_client.query_job(query, move || async move { + let table = Self::new_with_sql_client( + client, + db_name, + name, + namespace, + identifier, + server_version, + Some(job_sql_client), + ); + table.checkout_latest().await?; + let version = table.version().await?; + freshness_request.observe_version(&freshness, version); + Ok(RefreshMaterializedViewResult { + // The SQL job adapter does not expose refresh metrics. Table + // metadata cannot tell us what this refresh actually did. + mode: None, + rows_written: None, + source_version: None, + version, + }) + })) + } + } } async fn query_snapshot(&self) -> Result> { let description = self.describe().await?; @@ -4243,6 +4338,7 @@ mod tests { use super::*; + use crate::remote::client::test_utils::client_with_handler; use crate::remote::client::{ClientConfig, RetryConfig}; use crate::table::{AddDataMode, FieldMetadataUpdate, FtsToken}; @@ -10567,6 +10663,81 @@ mod tests { ); } + #[rstest] + #[case(false)] + #[case(true)] + #[tokio::test] + async fn materialized_view_registry_job_waits_then_reports_the_committed_version( + #[case] with_progress: bool, + ) { + let polls = AtomicUsize::new(0); + let client = client_with_handler(move |request| -> http::Response { + match request.url().path() { + "/v1/jobs/describe" => http::Response::builder() + .status(200) + .body( + json!({ + "job_id": "job-1", + "job_type": "sql", + "job_state": if polls.fetch_add(1, Ordering::SeqCst) == 0 { + "IN_PROGRESS" + } else { + "DONE" + }, + "creation_ms": 1, + }) + .to_string(), + ) + .unwrap(), + "/v1/table/ns$view/describe/" => http::Response::builder() + .status(200) + .body(r#"{"version":42,"schema":{"fields":[]}}"#.to_string()) + .unwrap(), + path => panic!("unexpected request: {path}"), + } + }); + let freshness = Arc::new(Mutex::new(FreshnessState::default())); + let job = Job::new_typed(Box::new(MaterializedViewRefreshJob { + inner: RemoteJob::new(client.clone(), "job-1".to_string()), + client, + db_name: "analytics".to_string(), + name: "view".to_string(), + namespace: vec!["ns".to_string()], + identifier: "ns$view".to_string(), + server_version: Default::default(), + sql_client: None, + freshness: freshness.clone(), + freshness_request: FreshnessHeaders::default(), + })); + + assert_eq!(job.id(), Some("job-1")); + let observations = Arc::new(Mutex::new(Vec::new())); + let observed = observations.clone(); + let result: RefreshMaterializedViewResult = if with_progress { + job.wait_with_progress(move |description| { + observed + .lock() + .unwrap() + .push(description.progress_snapshot()); + }) + .await + .unwrap() + } else { + job.wait().await.unwrap() + }; + if with_progress { + let observations = observations.lock().unwrap(); + assert_eq!(observations.len(), 2); + assert_eq!(observations[0]["state"], "running"); + assert_eq!(observations[1]["state"], "finished"); + } + assert_eq!(result.version, 42); + assert_eq!(result.mode, None); + assert_eq!(result.rows_written, None); + assert_eq!(result.source_version, None); + assert_eq!(freshness.lock().unwrap().min_read_version, Some(42)); + } + #[rstest] #[case(None)] #[case(Some(1))]