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))]