diff --git a/python/src/error.rs b/python/src/error.rs index a86051e5d..3fd27f7be 100644 --- a/python/src/error.rs +++ b/python/src/error.rs @@ -310,6 +310,9 @@ fn new_py_sql_error<'py>( grpc_code: i32, metadata: &RemoteErrorMetadata, ) -> PyResult> { + if sql_has_resource_error(metadata, grpc_code) { + return new_py_sql_resource_error(py, message, request_id, grpc_code, metadata); + } py.import("lancedb.remote.errors")? .getattr("SqlError")? .call1(( @@ -322,6 +325,51 @@ fn new_py_sql_error<'py>( )) } +fn new_py_sql_resource_error<'py>( + py: Python<'py>, + message: &str, + request_id: &str, + grpc_code: i32, + metadata: &RemoteErrorMetadata, +) -> PyResult> { + let resource = metadata.resource.as_ref().unwrap(); + py.import("lancedb.remote.errors")? + .getattr("HttpError")? + .call1(( + message, + request_id, + grpc_code_to_http_status(grpc_code), + metadata.error_code.as_deref(), + resource.kind.as_str(), + resource.name.as_str(), + resource.namespace_path.clone(), + false, + metadata.job_id.as_deref(), + metadata.job_request_id.as_deref(), + )) +} + +fn sql_has_resource_error(metadata: &RemoteErrorMetadata, grpc_code: i32) -> bool { + metadata.resource.is_some() + && matches!( + (metadata.error_code.as_deref(), grpc_code), + // These are gRPC status codes, not HTTP status codes. + (Some("not_found"), 5) | (Some("already_exists"), 6) + ) +} + +fn grpc_code_to_http_status(grpc_code: i32) -> Option { + match grpc_code { + 3 => Some(400), + 5 => Some(404), + 6 => Some(409), + 7 => Some(403), + 9 => Some(400), + 12 => Some(501), + _ => None, + } +} + // Preserve the typed transport classification across the Python boundary. // A missing status code alone also covers malformed job responses, which // waiting longer cannot repair. @@ -335,3 +383,178 @@ fn is_transport_error(err: &(dyn std::error::Error + 'static)) -> bool { } err.source().is_some_and(is_transport_error) } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use lancedb::error::{JobFailure, RemoteResourceIdentity}; + use pyo3::types::PyAnyMethods; + + use super::*; + + #[test] + fn generic_job_failure_retains_the_named_resource_as_its_cause() { + Python::initialize(); + let source = LanceError::Sql { + message: "Function 'missing' was not found".into(), + request_id: "current-request".into(), + grpc_code: 5, + metadata: Box::new({ + let mut metadata = RemoteErrorMetadata::default(); + metadata.error_code = Some("not_found".into()); + metadata.job_id = Some("outer-job".into()); + metadata.job_request_id = Some("worker-request".into()); + metadata.resource = Some(RemoteResourceIdentity::new( + "function", + "missing", + vec!["ns".into()], + )); + metadata + }), + }; + let error = LanceError::JobFailed { + job_id: Some("outer-job".into()), + failure: Box::new(JobFailure { + message: Some("Function 'missing' was not found".into()), + traceback: None, + sql_job_id: Some("outer-job".into()), + source: Some(Arc::new(source)), + ..Default::default() + }), + }; + + let py_err = Python::attach(|_| { + let result: std::result::Result<(), LanceError> = Err(error); + result.infer_error().unwrap_err() + }); + Python::attach(|py| { + let value = py_err.value(py); + let job_failed = py + .import("lancedb.exceptions") + .unwrap() + .getattr("JobFailedError") + .unwrap(); + assert!(value.is_instance(&job_failed).unwrap(), "{py_err:?}"); + let cause = value.getattr("__cause__").unwrap(); + let function_not_found = py + .import("lancedb.remote.errors") + .unwrap() + .getattr("FunctionNotFoundError") + .unwrap(); + assert!( + cause.is_instance(&function_not_found).unwrap(), + "{py_err:?}" + ); + assert_eq!( + cause + .getattr("error_code") + .unwrap() + .extract::() + .unwrap(), + "not_found" + ); + assert_eq!( + cause + .getattr("resource_kind") + .unwrap() + .extract::() + .unwrap(), + "function" + ); + assert_eq!( + cause + .getattr("resource_name") + .unwrap() + .extract::() + .unwrap(), + "missing" + ); + assert_eq!( + cause + .getattr("namespace_path") + .unwrap() + .extract::>() + .unwrap(), + ["ns"] + ); + assert_eq!( + cause + .getattr("status_code") + .unwrap() + .extract::() + .unwrap(), + 404 + ); + assert_eq!( + cause + .getattr("request_id") + .unwrap() + .extract::() + .unwrap(), + "current-request" + ); + assert_eq!( + cause + .getattr("job_id") + .unwrap() + .extract::() + .unwrap(), + "outer-job" + ); + assert_eq!( + cause + .getattr("job_request_id") + .unwrap() + .extract::() + .unwrap(), + "worker-request" + ); + }); + } + + #[test] + fn direct_catalog_lookup_is_raised_as_the_named_resource_error() { + Python::initialize(); + for (kind, name, namespace_path, exception) in [ + ( + "function", + "missing", + vec!["ns".into()], + "FunctionNotFoundError", + ), + ("namespace", "ns", Vec::new(), "NamespaceNotFoundError"), + ] { + let mut metadata = RemoteErrorMetadata::default(); + metadata.error_code = Some("not_found".into()); + metadata.resource = Some(RemoteResourceIdentity::new(kind, name, namespace_path)); + let error = LanceError::Sql { + message: format!("{kind} '{name}' was not found"), + request_id: "current-request".into(), + grpc_code: 5, + metadata: Box::new(metadata), + }; + let py_err = Python::attach(|_| { + let result: std::result::Result<(), LanceError> = Err(error); + result.infer_error().unwrap_err() + }); + Python::attach(|py| { + let value = py_err.value(py); + let class = py + .import("lancedb.remote.errors") + .unwrap() + .getattr(exception) + .unwrap(); + assert!(value.is_instance(&class).unwrap(), "{py_err:?}"); + assert_eq!( + value + .getattr("resource_name") + .unwrap() + .extract::() + .unwrap(), + name + ); + }); + } + } +} diff --git a/rust/lancedb/src/error.rs b/rust/lancedb/src/error.rs index 9c79fcf95..cc5bd1f4e 100644 --- a/rust/lancedb/src/error.rs +++ b/rust/lancedb/src/error.rs @@ -27,6 +27,35 @@ pub struct HttpErrorDetails { pub namespace_path: Vec, } +/// Stable resource identity reported with a remote Flight SQL failure. +#[cfg(feature = "remote")] +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub struct RemoteResourceIdentity { + /// The kind of resource named by the error. + pub kind: String, + /// The leaf name of the resource. + pub name: String, + /// The parent namespace path of the resource. + pub namespace_path: Vec, +} + +#[cfg(feature = "remote")] +impl RemoteResourceIdentity { + /// Construct a resource identity as reported by a remote service. + pub fn new( + kind: impl Into, + name: impl Into, + namespace_path: Vec, + ) -> Self { + Self { + kind: kind.into(), + name: name.into(), + namespace_path, + } + } +} + /// Correlation and classification metadata reported by a remote server. /// /// Fields are absent when the server did not supply them. Diagnostic text is @@ -41,6 +70,8 @@ pub struct RemoteErrorMetadata { pub job_id: Option, /// Request executed by the isolated SQL worker. pub job_request_id: Option, + /// Structured resource identity supplied by the server, when available. + pub resource: Option, } #[cfg(feature = "remote")] @@ -56,10 +87,37 @@ impl RemoteErrorMetadata { error_code: header("x-lancedb-error-code"), job_id: header("x-lancedb-sql-job-id"), job_request_id: header("x-job-request-id"), + resource: header("x-lancedb-resource") + .as_deref() + .and_then(parse_flight_resource_identity), } } } +#[cfg(feature = "remote")] +#[derive(Deserialize)] +struct FlightResourceIdentity { + kind: String, + name: String, + namespace_path: Vec, +} + +/// Decode the ASCII-safe structured resource metadata used by Flight SQL. +#[cfg(feature = "remote")] +pub(crate) fn parse_flight_resource_identity(value: &str) -> Option { + use base64::Engine as _; + + let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(value) + .ok()?; + let identity = serde_json::from_slice::(&bytes).ok()?; + Some(RemoteResourceIdentity { + kind: identity.kind, + name: identity.name, + namespace_path: identity.namespace_path, + }) +} + #[cfg(feature = "remote")] #[derive(Debug)] pub(crate) struct HttpErrorSource { diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index b9a3f8ae1..9af99dced 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -52,6 +52,34 @@ use super::util::parse_server_version; use super::{ARROW_STREAM_CONTENT_TYPE, extract_job_id}; use futures::TryStreamExt; +#[cfg(feature = "remote")] +fn direct_catalog_resource_error(error: Error) -> Error { + let Error::JobFailed { failure, .. } = &error else { + return error; + }; + let Some(Error::Sql { + message, + request_id, + grpc_code, + metadata, + }) = failure.source.as_deref() + else { + return error; + }; + let named_catalog_resource = metadata.resource.as_ref().is_some_and(|resource| { + matches!(resource.kind.as_str(), "function" | "namespace" | "secret") + }); + if !named_catalog_resource { + return error; + } + Error::Sql { + message: message.clone(), + request_id: request_id.clone(), + grpc_code: *grpc_code, + metadata: metadata.clone(), + } +} + // Request structure for the remote clone table API #[derive(serde::Serialize)] struct RemoteCloneTableRequest { @@ -861,6 +889,7 @@ impl Database for RemoteDatabase { self.sql_client()? .describe_function(name, version, namespace_path) .await + .map_err(direct_catalog_resource_error) } async fn list_functions( @@ -991,6 +1020,7 @@ impl Database for RemoteDatabase { self.sql_client()? .write_secret(name, value, namespace_path, true) .await + .map_err(direct_catalog_resource_error) } async fn alter_secret(&self, name: &str, value: &str, namespace_path: &[String]) -> Result<()> { crate::utils::validate_secret_name(name)?; @@ -998,6 +1028,7 @@ impl Database for RemoteDatabase { self.sql_client()? .write_secret(name, value, namespace_path, false) .await + .map_err(direct_catalog_resource_error) } async fn list_secrets( &self, @@ -1008,11 +1039,15 @@ impl Database for RemoteDatabase { self.sql_client()? .list_secrets(namespace_path, options) .await + .map_err(direct_catalog_resource_error) } async fn drop_secret(&self, name: &str, namespace_path: &[String]) -> Result<()> { crate::utils::validate_secret_name(name)?; validate_namespace(namespace_path)?; - self.sql_client()?.drop_secret(name, namespace_path).await + self.sql_client()? + .drop_secret(name, namespace_path) + .await + .map_err(direct_catalog_resource_error) } async fn describe_secret(&self, name: &str, namespace_path: &[String]) -> Result { crate::utils::validate_secret_name(name)?; @@ -1020,6 +1055,7 @@ impl Database for RemoteDatabase { self.sql_client()? .describe_secret(name, namespace_path) .await + .map_err(direct_catalog_resource_error) } async fn create_view( @@ -1910,6 +1946,7 @@ mod tests { use crate::connection::ConnectBuilder; use crate::database::Database; + use crate::error::{JobFailure, RemoteErrorMetadata, RemoteResourceIdentity}; use crate::graph::{EdgeTable, Endpoint, NodeTable, PropertyGraphDefinition}; use crate::materialized_view::CreateMaterializedViewRequest; use crate::{ @@ -1919,6 +1956,53 @@ mod tests { remote::{ARROW_STREAM_CONTENT_TYPE, ClientConfig, HeaderProvider, JSON_CONTENT_TYPE}, }; + #[test] + fn catalog_helpers_unwrap_only_a_named_resource_job_failure() { + let job_failed = |source: Error| Error::JobFailed { + job_id: Some("job-1".into()), + failure: Box::new(JobFailure { + sql_job_id: Some("job-1".into()), + source: Some(Arc::new(source)), + ..Default::default() + }), + }; + + for (kind, name) in [ + ("function", "missing"), + ("namespace", "missing_ns"), + ("secret", "missing_secret"), + ] { + let resource = RemoteResourceIdentity::new(kind, name, vec!["ns".into()]); + let sql = Error::Sql { + message: format!("{kind} '{name}' was not found"), + request_id: "worker-request".into(), + grpc_code: 5, + metadata: Box::new(RemoteErrorMetadata { + error_code: Some("not_found".into()), + job_id: Some("job-1".into()), + resource: Some(resource.clone()), + ..Default::default() + }), + }; + let Error::Sql { metadata, .. } = super::direct_catalog_resource_error(job_failed(sql)) + else { + panic!("catalog helper must surface the {kind} resource failure"); + }; + assert_eq!(metadata.resource, Some(resource)); + } + + let generic = Error::Sql { + message: "query failed".into(), + request_id: "worker-request".into(), + grpc_code: 3, + metadata: Box::new(RemoteErrorMetadata::default()), + }; + assert!(matches!( + super::direct_catalog_resource_error(job_failed(generic)), + Error::JobFailed { .. } + )); + } + #[tokio::test] async fn test_catalog_namespace_identifier_validates_all_namespace_components() { let mut db = super::RemoteDatabase::new_mock(|_| -> http::Response { diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index 83dedcb6d..724c813b3 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -1655,6 +1655,9 @@ pub(super) fn sql_status_error(request_id: &str, status: &tonic::Status) -> Erro error_code: header("x-lancedb-error-code"), job_id: header(SQL_JOB_ID_METADATA_KEY), job_request_id: header("x-job-request-id"), + resource: header("x-lancedb-resource") + .as_deref() + .and_then(crate::error::parse_flight_resource_identity), }; let sql_job_id = metadata.job_id.clone(); let error = Error::Sql { diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index cbca4e9f6..37125c1fd 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -3545,3 +3545,31 @@ fn sql_status_keeps_input_category_and_request_ids() { let error = super::sql_status_error("request", &tonic::Status::not_found("column_not_found")); assert!(error.remote_error_metadata().unwrap().error_code.is_none()); } + +#[test] +fn sql_status_reports_structured_resource_identity() { + let mut status = tonic::Status::not_found("Function 'missing' was not found"); + status + .metadata_mut() + .insert("x-lancedb-error-code", "not_found".parse().unwrap()); + status.metadata_mut().insert( + "x-lancedb-resource", + "eyJraW5kIjoiZnVuY3Rpb24iLCJuYW1lIjoibWlzc2luZyIsIm5hbWVzcGFjZV9wYXRoIjpbIm5zIl19" + .parse() + .unwrap(), + ); + status + .metadata_mut() + .insert("x-lancedb-sql-job-id", "sql-job-1".parse().unwrap()); + + let error = super::sql_status_error("current-request", &status); + let resource = error + .remote_error_metadata() + .unwrap() + .resource + .as_ref() + .unwrap(); + assert_eq!(resource.kind, "function"); + assert_eq!(resource.name, "missing"); + assert_eq!(resource.namespace_path, ["ns"]); +}