From 3ec70f5d795001ef4991f22698dae4cf14637469 Mon Sep 17 00:00:00 2001 From: Fredrik Fornwall Date: Tue, 7 Jul 2026 14:32:13 +0200 Subject: [PATCH] feat(rust/ffi): support ADBC 1.1.0 rich errors from exported streams The driver exporter exported result readers via arrow-rs's FFI_ArrowArrayStream::new, which reduces every mid-stream failure to an errno plus a message string: rich adbc_core errors surfaced through ArrowError::ExternalError landed in the catch-all EINVAL, and the ADBC 1.1.0 ErrorFromArrayStream slot in the driver table was left as a TODO. Replace the export machinery (adapted from arrow-rs's ffi_stream.rs) with one that stashes the full adbc_core::error::Error found in the error's source chain (or synthesized from the ArrowError variant): - get_next/get_schema return the errno for the error's ADBC status via a port of the canonical InternalAdbcStatusCodeToErrno table, so e.g. cancelled operations report ECANCELED as the validation suite expects. - ErrorFromArrayStream is now implemented, handing consumers the stream-owned FFI_AdbcError (message, SQLSTATE, vendor code, details) together with its status code, and NULL for foreign streams. Own streams are recognized via a registry of live private-data addresses, since Rust does not guarantee the function-pointer identity the C/C++ drivers rely on (Miri fails such comparisons). Tested by unit tests, a new end-to-end test driving the exported C driver table against a failing dummy-driver stream, and a clean run of the adbc_ffi suite under Miri. Signed-off-by: Fredrik Fornwall --- rust/Cargo.lock | 1 + rust/Cargo.toml | 5 + rust/driver/dummy/src/lib.rs | 64 ++- .../dummy/tests/driver_exporter_dummy.rs | 110 ++++ rust/ffi/Cargo.toml | 1 + rust/ffi/src/driver_exporter.rs | 16 +- rust/ffi/src/exported_stream.rs | 494 ++++++++++++++++++ rust/ffi/src/lib.rs | 1 + 8 files changed, 683 insertions(+), 9 deletions(-) create mode 100644 rust/ffi/src/exported_stream.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 91c94481d4..b32cc4b81d 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -66,6 +66,7 @@ dependencies = [ "adbc_core", "arrow-array", "arrow-schema", + "libc", ] [[package]] diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 53ba331488..592a10562d 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -43,3 +43,8 @@ arrow-array = { version = ">=58, <60", default-features = false, features = [ arrow-buffer = { version = ">=58, <60", default-features = false } arrow-schema = { version = ">=58, <60", default-features = false } arrow-select = { version = ">=58, <60", default-features = false } +# The floor matters for the minimal-versions CI job: libc 0.2.0 (the "0.2" +# minimum) predates the errno constants used by adbc_ffi and conflicts with +# rustix's own floor. 0.2.182 matches the lowest version the rest of the +# workspace already requires. +libc = { version = "0.2.182", default-features = false } diff --git a/rust/driver/dummy/src/lib.rs b/rust/driver/dummy/src/lib.rs index 7879873ca9..7ef52ba7dc 100644 --- a/rust/driver/dummy/src/lib.rs +++ b/rust/driver/dummy/src/lib.rs @@ -67,6 +67,63 @@ impl RecordBatchReader for SingleBatchReader { } } +/// SQL query that makes [`DummyStatement::execute`] return a stream which fails on +/// the first call to `next`, carrying the rich [`Error`] from [`stream_error`]. +/// +/// Used to exercise the driver exporter's `ErrorFromArrayStream` implementation. +pub const STREAM_ERROR_QUERY: &str = "TRIGGER_STREAM_ERROR"; + +/// The [`Error`] emitted by the stream produced for [`STREAM_ERROR_QUERY`]. +pub fn stream_error() -> Error { + Error { + message: "the stream failed midway".into(), + status: Status::IO, + vendor_code: 42, + sqlstate: [ + b'0' as std::os::raw::c_char, + b'8' as std::os::raw::c_char, + b'0' as std::os::raw::c_char, + b'0' as std::os::raw::c_char, + b'6' as std::os::raw::c_char, + ], + details: Some(vec![("detail-key".into(), b"detail-value".to_vec())]), + } +} + +/// A [`RecordBatchReader`] whose first `next` yields the [`Error`] from +/// [`stream_error`], wrapped as [`ArrowError::ExternalError`] exactly as a real +/// driver would surface a mid-stream failure. +#[derive(Debug)] +pub struct FailingReader { + schema: SchemaRef, + error: Option, +} + +impl FailingReader { + fn new() -> Self { + Self { + schema: Arc::new(get_table_schema()), + error: Some(stream_error()), + } + } +} + +impl Iterator for FailingReader { + type Item = std::result::Result; + + fn next(&mut self) -> Option { + self.error + .take() + .map(|error| Err(ArrowError::ExternalError(Box::new(error)))) + } +} + +impl RecordBatchReader for FailingReader { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + fn get_table_schema() -> Schema { Schema::new(vec![ Field::new("a", DataType::UInt32, true), @@ -817,6 +874,7 @@ impl Connection for DummyConnection { #[derive(Default)] pub struct DummyStatement { options: HashMap, + query: Option, } impl Optionable for DummyStatement { @@ -858,6 +916,9 @@ impl Statement for DummyStatement { fn execute(&mut self) -> Result> { maybe_panic("StatementExecuteQuery"); + if self.query.as_deref() == Some(STREAM_ERROR_QUERY) { + return Ok(Box::new(FailingReader::new())); + } let batch = get_table_data(); let reader = SingleBatchReader::new(batch); Ok(Box::new(reader)) @@ -887,7 +948,8 @@ impl Statement for DummyStatement { Ok(()) } - fn set_sql_query(&mut self, _query: impl AsRef) -> Result<()> { + fn set_sql_query(&mut self, query: impl AsRef) -> Result<()> { + self.query = Some(query.as_ref().to_string()); Ok(()) } diff --git a/rust/driver/dummy/tests/driver_exporter_dummy.rs b/rust/driver/dummy/tests/driver_exporter_dummy.rs index b64d2016c2..31c2364ebc 100644 --- a/rust/driver/dummy/tests/driver_exporter_dummy.rs +++ b/rust/driver/dummy/tests/driver_exporter_dummy.rs @@ -650,3 +650,113 @@ fn test_statement_get_parameter_schema() { let native_schema = native_statement.get_parameter_schema().unwrap(); assert_eq!(exported_schema, native_schema); } + +// ErrorFromArrayStream (ADBC 1.1.0): when a result stream fails mid-iteration, a +// consumer must be able to recover the full `AdbcError` (status, vendor code and +// structured details), not just an errno plus a message string. +// +// We drive the exported C driver directly here because the Rust driver manager +// imports streams via arrow's `ArrowArrayStreamReader`, which has no notion of +// this ADBC-specific entry point and so cannot exercise it. +#[test] +fn test_error_from_array_stream() { + use std::ffi::{CStr, CString}; + use std::ptr::null_mut; + + use adbc_core::constants; + use adbc_core::error::{AdbcStatusCode, Status}; + use adbc_ffi::{ + FFIDriver, FFI_AdbcConnection, FFI_AdbcDatabase, FFI_AdbcError, FFI_AdbcStatement, + }; + use arrow_array::ffi::FFI_ArrowArray; + use arrow_array::ffi_stream::FFI_ArrowArrayStream; + + use adbc_dummy::{stream_error, STREAM_ERROR_QUERY}; + + let driver = DummyDriver::ffi_driver(); + + unsafe { + let mut error = FFI_AdbcError::default(); + + let mut database = FFI_AdbcDatabase::default(); + assert_eq!( + (driver.DatabaseNew.unwrap())(&mut database, &mut error), + constants::ADBC_STATUS_OK + ); + assert_eq!( + (driver.DatabaseInit.unwrap())(&mut database, &mut error), + constants::ADBC_STATUS_OK + ); + + let mut connection = FFI_AdbcConnection::default(); + assert_eq!( + (driver.ConnectionNew.unwrap())(&mut connection, &mut error), + constants::ADBC_STATUS_OK + ); + assert_eq!( + (driver.ConnectionInit.unwrap())(&mut connection, &mut database, &mut error), + constants::ADBC_STATUS_OK + ); + + let mut statement = FFI_AdbcStatement::default(); + assert_eq!( + (driver.StatementNew.unwrap())(&mut connection, &mut statement, &mut error), + constants::ADBC_STATUS_OK + ); + + let query = CString::new(STREAM_ERROR_QUERY).unwrap(); + assert_eq!( + (driver.StatementSetSqlQuery.unwrap())(&mut statement, query.as_ptr(), &mut error), + constants::ADBC_STATUS_OK + ); + + // Execute to obtain the result stream. + let mut stream = FFI_ArrowArrayStream::empty(); + assert_eq!( + (driver.StatementExecuteQuery.unwrap())( + &mut statement, + &mut stream, + null_mut(), + &mut error, + ), + constants::ADBC_STATUS_OK + ); + + // Iterating the stream fails on the first batch. + let mut array = FFI_ArrowArray::empty(); + let rc = (stream.get_next.unwrap())(&mut stream, &mut array); + assert_ne!(rc, 0, "the stream was expected to fail"); + + // Recover the full error via ErrorFromArrayStream. + let mut status: AdbcStatusCode = constants::ADBC_STATUS_OK; + let error_ptr = (driver.ErrorFromArrayStream.unwrap())(&mut stream, &mut status); + assert!(!error_ptr.is_null(), "expected a recovered error"); + + let expected = stream_error(); + assert_eq!(status, AdbcStatusCode::from(Status::IO)); + assert_eq!( + CStr::from_ptr((*error_ptr).message).to_str().unwrap(), + expected.message + ); + assert_eq!((*error_ptr).vendor_code, expected.vendor_code); + assert_eq!((*error_ptr).sqlstate, expected.sqlstate); + + // Structured details survive too, retrievable via ErrorGetDetail*. + assert_eq!((driver.ErrorGetDetailCount.unwrap())(error_ptr), 1); + let detail = (driver.ErrorGetDetail.unwrap())(error_ptr, 0); + assert_eq!(CStr::from_ptr(detail.key).to_str().unwrap(), "detail-key"); + assert_eq!( + std::slice::from_raw_parts(detail.value, detail.value_length), + b"detail-value" + ); + + // Cleanup. Releasing the stream frees the recovered error; the other + // handles are released via the driver. + if let Some(release) = stream.release { + release(&mut stream); + } + (driver.StatementRelease.unwrap())(&mut statement, &mut error); + (driver.ConnectionRelease.unwrap())(&mut connection, &mut error); + (driver.DatabaseRelease.unwrap())(&mut database, &mut error); + } +} diff --git a/rust/ffi/Cargo.toml b/rust/ffi/Cargo.toml index 553a51cf6e..01a71b93a4 100644 --- a/rust/ffi/Cargo.toml +++ b/rust/ffi/Cargo.toml @@ -37,6 +37,7 @@ default = [] adbc_core.workspace = true arrow-array.workspace = true arrow-schema.workspace = true +libc.workspace = true [package.metadata.docs.rs] all-features = true diff --git a/rust/ffi/src/driver_exporter.rs b/rust/ffi/src/driver_exporter.rs index 1a41d100b4..a5541125f8 100644 --- a/rust/ffi/src/driver_exporter.rs +++ b/rust/ffi/src/driver_exporter.rs @@ -132,7 +132,7 @@ impl FFIDriver for DriverType { StatementSetSubstraitPlan: Some(statement_set_substrait_plan::), ErrorGetDetailCount: Some(error_get_detail_count), ErrorGetDetail: Some(error_get_detail), - ErrorFromArrayStream: None, // TODO(alexandreyc): what to do with this? + ErrorFromArrayStream: Some(crate::exported_stream::error_from_array_stream), DatabaseGetOption: Some(database_get_option::), DatabaseGetOptionBytes: Some(database_get_option_bytes::), DatabaseGetOptionDouble: Some(database_get_option_double::), @@ -1117,7 +1117,7 @@ extern "C" fn connection_get_table_types( let reader = check_err!(connection.get_table_types(), error); let reader = Box::new(reader); - let reader = FFI_ArrowArrayStream::new(reader); + let reader = crate::exported_stream::export_reader(reader); unsafe { std::ptr::write_unaligned(out, reader) }; ADBC_STATUS_OK @@ -1185,7 +1185,7 @@ extern "C" fn connection_get_info( let reader = check_err!(connection.get_info(info_codes), error); let reader = Box::new(reader); - let reader = FFI_ArrowArrayStream::new(reader); + let reader = crate::exported_stream::export_reader(reader); unsafe { std::ptr::write_unaligned(out, reader) }; ADBC_STATUS_OK @@ -1263,7 +1263,7 @@ extern "C" fn connection_get_statistic_names( let reader = check_err!(connection.get_statistic_names(), error); let reader = Box::new(reader); - let reader = FFI_ArrowArrayStream::new(reader); + let reader = crate::exported_stream::export_reader(reader); unsafe { std::ptr::write_unaligned(out, reader) }; ADBC_STATUS_OK @@ -1292,7 +1292,7 @@ extern "C" fn connection_read_partition( unsafe { std::slice::from_raw_parts(serialized_partition, serialized_length) }; let reader = check_err!(connection.read_partition(partition), error); let reader = Box::new(reader); - let reader = FFI_ArrowArrayStream::new(reader); + let reader = crate::exported_stream::export_reader(reader); unsafe { std::ptr::write_unaligned(out, reader) }; ADBC_STATUS_OK @@ -1326,7 +1326,7 @@ extern "C" fn connection_get_statistics( let reader = connection.get_statistics(catalog, db_schema, table_name, approximate); let reader = check_err!(reader, error); let reader = Box::new(reader); - let reader = FFI_ArrowArrayStream::new(reader); + let reader = crate::exported_stream::export_reader(reader); unsafe { std::ptr::write_unaligned(out, reader) }; ADBC_STATUS_OK @@ -1383,7 +1383,7 @@ extern "C" fn connection_get_objects( ); let reader = check_err!(reader, error); let reader = Box::new(reader); - let reader = FFI_ArrowArrayStream::new(reader); + let reader = crate::exported_stream::export_reader(reader); unsafe { std::ptr::write_unaligned(out, reader) }; ADBC_STATUS_OK @@ -1726,7 +1726,7 @@ extern "C" fn statement_execute_query( if !out.is_null() { let reader = check_err!(statement.execute(), error); let reader = Box::new(reader); - let reader = FFI_ArrowArrayStream::new(reader); + let reader = crate::exported_stream::export_reader(reader); unsafe { std::ptr::write_unaligned(out, reader) }; } else { let rows_affected_value = check_err!(statement.execute_update(), error).unwrap_or(-1); diff --git a/rust/ffi/src/exported_stream.rs b/rust/ffi/src/exported_stream.rs new file mode 100644 index 0000000000..401c8a8cdb --- /dev/null +++ b/rust/ffi/src/exported_stream.rs @@ -0,0 +1,494 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Export a [`RecordBatchReader`] as an [`FFI_ArrowArrayStream`] with full ADBC error +//! reporting. +//! +//! The Arrow C stream interface can only report an errno-compatible error code and a message +//! from its callbacks, while ADBC 1.1.0 defines a richer error model (status code, SQLSTATE, +//! vendor code and key/value details) and the `AdbcErrorFromArrayStream` API to recover it +//! from a stream returned by a driver. +//! +//! `arrow-rs`'s own exporter ([`FFI_ArrowArrayStream::new`]) cannot support this: it stashes +//! only the error message, and it maps every error to one of `ENOSYS`/`ENOMEM`/`EIO`/`EINVAL` +//! based solely on the [`ArrowError`] variant. Rust ADBC drivers surface their rich +//! [`adbc_core::error::Error`] through [`ArrowError::ExternalError`], which lands in the +//! catch-all `EINVAL` — so not even a cancelled query could report `ECANCELED`. +//! +//! This module is the same export machinery, adapted from `arrow-rs`'s `ffi_stream.rs` +//! (Apache-2.0), except that on failure the [`adbc_core::error::Error`] carried by the error's +//! [source chain](std::error::Error::source) (or one synthesized from the [`ArrowError`] +//! variant when the chain carries none) is stashed in the stream's private data, and: +//! +//! * `get_schema`/`get_next` return the errno for the error's [`Status`], through a direct +//! port of the C driver framework's canonical `InternalAdbcStatusCodeToErrno` table +//! (`c/driver/common/utils.c`) — the same table the in-tree C/C++ drivers (which implement +//! their stream callbacks by hand, e.g. the PostgreSQL driver's `TupleReader`) return from +//! `get_next`. In particular a cancelled operation reports `ECANCELED`, as the C++ +//! validation suite's `SqlQueryCancel` expects after `AdbcStatementCancel`. +//! * [`error_from_array_stream`] — wired into the exported driver table as the ADBC 1.1.0 +//! `ErrorFromArrayStream` function — hands the consumer the stashed error as an +//! [`FFI_AdbcError`] (message, SQLSTATE, vendor code and details) together with its +//! [`AdbcStatusCode`], with none of the lossy errno projection. + +use std::collections::BTreeSet; +use std::ffi::CString; +use std::os::raw::{c_char, c_int, c_void}; +use std::ptr::addr_of; +use std::sync::Mutex; + +use adbc_core::constants::ADBC_STATUS_OK; +use adbc_core::error::{AdbcStatusCode, Error, Status}; +use arrow_array::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; +use arrow_array::ffi_stream::FFI_ArrowArrayStream; +use arrow_array::{Array, RecordBatchReader, StructArray}; +use arrow_schema::ArrowError; +use libc::{EACCES, ECANCELED, EEXIST, EINVAL, EIO, ENOENT, ENOTSUP, ETIMEDOUT}; + +use crate::FFI_AdbcError; + +/// Export `reader` as an [`FFI_ArrowArrayStream`] that retains the rich ADBC error behind any +/// failure for retrieval through [`error_from_array_stream`], and whose callbacks report the +/// errno for the error's ADBC status. +pub(crate) fn export_reader(reader: Box) -> FFI_ArrowArrayStream { + let private_data = Box::new(StreamPrivateData { + reader, + last_message: None, + last_status: ADBC_STATUS_OK, + last_error: FFI_AdbcError::default(), + }); + let private_data = Box::into_raw(private_data); + exported_streams() + .expect("Poisoned exported-streams registry") + .insert(private_data as usize); + FFI_ArrowArrayStream { + get_schema: Some(get_schema), + get_next: Some(get_next), + get_last_error: Some(get_last_error), + release: Some(release_stream), + private_data: private_data as *mut c_void, + } +} + +/// The private data of every live stream created by [`export_reader`], so that +/// [`error_from_array_stream`] can recognize this exporter's streams. +/// +/// The C/C++ drivers recognize their own streams by comparing the stream's release callback +/// against their own by address, but that technique is not sound in Rust: pointers to the same +/// function are not guaranteed to compare equal (the compiler may duplicate or merge +/// functions), and Miri in fact fails such comparisons. A registry keyed by the private-data +/// address is exact: entries are removed in `release`, so an address cannot be occupied by a +/// foreign allocation while it is in the set. +fn exported_streams() -> std::sync::LockResult>> { + static EXPORTED_STREAMS: Mutex> = Mutex::new(BTreeSet::new()); + EXPORTED_STREAMS.lock() +} + +struct StreamPrivateData { + reader: Box, + /// The last error's message, for `get_last_error`. + last_message: Option, + /// The last error's ADBC status, for [`error_from_array_stream`]. + last_status: AdbcStatusCode, + /// The last error in full, for [`error_from_array_stream`]. An empty [`Default`] error + /// until a callback fails. + last_error: FFI_AdbcError, +} + +/// The driver's ADBC 1.1.0 `ErrorFromArrayStream` implementation. +/// +/// For a stream exported through [`export_reader`], returns the rich error stashed by the +/// last failed callback (an empty error with `ADBC_STATUS_OK` while no call has failed) and +/// writes its status code to `status` (if not NULL). Returns NULL for foreign or already +/// released streams. +/// +/// Per the ADBC specification the returned error remains owned by the stream: the caller +/// must not release it, and it is invalidated when the stream itself is released. +pub(crate) unsafe extern "C" fn error_from_array_stream( + stream: *mut FFI_ArrowArrayStream, + status: *mut AdbcStatusCode, +) -> *const FFI_AdbcError { + let Some(stream) = (unsafe { stream.as_mut() }) else { + return std::ptr::null(); + }; + let is_exported_here = exported_streams() + .expect("Poisoned exported-streams registry") + .contains(&(stream.private_data as usize)); + if !is_exported_here { + return std::ptr::null(); + } + let data = stream.private_data as *mut StreamPrivateData; + if !status.is_null() { + unsafe { std::ptr::write_unaligned(status, (*data).last_status) }; + } + // Hand out a raw-derived pointer, not `&(*data).last_error`: consumers (e.g. the Python + // bindings' `convert_error`) call the error's own release callback through this pointer, + // which mutates the error in place — undefined behavior through a shared reference. + unsafe { &raw mut (*data).last_error as *const FFI_AdbcError } +} + +/// The rich ADBC error behind `err`: the first [`adbc_core::error::Error`] in its +/// [source chain](std::error::Error::source), or an error synthesized from the [`ArrowError`] +/// variant when the chain carries none. +fn adbc_error(err: &ArrowError) -> Error { + let mut source: Option<&(dyn std::error::Error + 'static)> = Some(err); + while let Some(error) = source { + if let Some(adbc) = error.downcast_ref::() { + return adbc.clone(); + } + source = error.source(); + } + let status = match err { + ArrowError::NotYetImplemented(_) => Status::NotImplemented, + ArrowError::IoError(_, _) => Status::IO, + _ => Status::Internal, + }; + Error::with_message_and_status(err.to_string(), status) +} + +/// A direct port of the C driver framework's `InternalAdbcStatusCodeToErrno` +/// (`c/driver/common/utils.c`) — the canonical AdbcStatusCode → errno table that the in-tree +/// C/C++ drivers return from their stream callbacks. +fn errno_for_status(status: Status) -> i32 { + match status { + Status::Ok => 0, + Status::Unknown => EIO, + Status::NotImplemented => ENOTSUP, + Status::NotFound => ENOENT, + Status::AlreadyExists => EEXIST, + Status::InvalidArguments | Status::InvalidState => EINVAL, + Status::InvalidData | Status::Integrity | Status::Internal | Status::IO => EIO, + Status::Cancelled => ECANCELED, + Status::Timeout => ETIMEDOUT, + Status::Unauthenticated | Status::Unauthorized => EACCES, + } +} + +fn private_data(stream: *mut FFI_ArrowArrayStream) -> &'static mut StreamPrivateData { + unsafe { &mut *((*stream).private_data as *mut StreamPrivateData) } +} + +/// Stash the rich error behind `err` for [`error_from_array_stream`] and `get_last_error`, +/// and return the errno to report through the stream callback. +fn set_error(data: &mut StreamPrivateData, err: &ArrowError) -> c_int { + let error = adbc_error(err); + let status = error.status; + // A NUL byte inside the message would truncate it, never fail the conversion. + let message = err.to_string().replace('\0', " "); + data.last_message = Some(CString::new(message).expect("NUL bytes were replaced")); + data.last_status = status.into(); + // Dropping the previous value releases any error stashed by an earlier failure. + data.last_error = FFI_AdbcError::try_from(error).unwrap_or_else(Into::into); + errno_for_status(status) +} + +unsafe extern "C" fn get_schema( + stream: *mut FFI_ArrowArrayStream, + out: *mut FFI_ArrowSchema, +) -> c_int { + let data = private_data(stream); + match FFI_ArrowSchema::try_from(data.reader.schema().as_ref()) { + Ok(schema) => { + unsafe { std::ptr::copy(addr_of!(schema), out, 1) }; + std::mem::forget(schema); + 0 + } + Err(ref err) => set_error(data, err), + } +} + +unsafe extern "C" fn get_next( + stream: *mut FFI_ArrowArrayStream, + out: *mut FFI_ArrowArray, +) -> c_int { + let data = private_data(stream); + match data.reader.next() { + // End of stream: a released ArrowArray marks completion. + None => { + unsafe { std::ptr::write(out, FFI_ArrowArray::empty()) }; + 0 + } + Some(Ok(batch)) => { + let array = FFI_ArrowArray::new(&StructArray::from(batch).to_data()); + unsafe { std::ptr::write_unaligned(out, array) }; + 0 + } + Some(Err(ref err)) => set_error(data, err), + } +} + +unsafe extern "C" fn get_last_error(stream: *mut FFI_ArrowArrayStream) -> *const c_char { + match &private_data(stream).last_message { + Some(err) => err.as_ptr(), + None => std::ptr::null(), + } +} + +unsafe extern "C" fn release_stream(stream: *mut FFI_ArrowArrayStream) { + if stream.is_null() { + return; + } + let stream = unsafe { &mut *stream }; + stream.get_schema = None; + stream.get_next = None; + stream.get_last_error = None; + exported_streams() + .expect("Poisoned exported-streams registry") + .remove(&(stream.private_data as usize)); + // Also releases the stashed error, invalidating pointers previously handed out by + // `error_from_array_stream`, as the ADBC specification prescribes. + drop(unsafe { Box::from_raw(stream.private_data as *mut StreamPrivateData) }); + stream.private_data = std::ptr::null_mut(); + stream.release = None; +} + +#[cfg(test)] +mod tests { + use std::ffi::CStr; + use std::sync::Arc; + + use adbc_core::constants::{ + ADBC_ERROR_VENDOR_CODE_PRIVATE_DATA, ADBC_STATUS_CANCELLED, ADBC_STATUS_IO, + }; + use arrow_array::{Int64Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema, SchemaRef}; + + use crate::types::ErrorPrivateData; + + use super::*; + + /// A reader yielding the given batches and errors in order. + struct FailingReader { + schema: SchemaRef, + items: std::vec::IntoIter>, + } + + impl Iterator for FailingReader { + type Item = Result; + fn next(&mut self) -> Option { + self.items.next() + } + } + + impl RecordBatchReader for FailingReader { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + } + + fn reader_with(items: Vec>) -> Box { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + Box::new(FailingReader { + schema, + items: items.into_iter(), + }) + } + + fn batch() -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1]))]).unwrap() + } + + fn cancelled_error() -> ArrowError { + ArrowError::ExternalError(Box::new(Error::with_message_and_status( + "operation cancelled", + Status::Cancelled, + ))) + } + + /// Drive the exported stream exactly as a C consumer would, collecting each get_next errno. + fn drive(stream: &mut FFI_ArrowArrayStream, calls: usize) -> Vec { + let get_next = stream.get_next.unwrap(); + (0..calls) + .map(|_| { + let mut out = FFI_ArrowArray::empty(); + unsafe { get_next(stream as *mut _, &mut out as *mut _) } + }) + .collect() + } + + /// Call the driver's `ErrorFromArrayStream` exactly as the driver manager would. + fn rich_error( + stream: &mut FFI_ArrowArrayStream, + ) -> (*const FFI_AdbcError, Option) { + // A sentinel that no ADBC status code uses, to detect whether `status` was written. + let mut status: AdbcStatusCode = u8::MAX; + let error = unsafe { error_from_array_stream(stream as *mut _, &mut status as *mut _) }; + (error, (status != u8::MAX).then_some(status)) + } + + #[test] + fn stream_ends_cleanly() { + let mut stream = export_reader(reader_with(vec![Ok(batch())])); + assert_eq!(drive(&mut stream, 2), vec![0, 0]); + } + + #[test] + fn cancelled_adbc_error_maps_to_ecanceled() { + let mut stream = export_reader(reader_with(vec![Ok(batch()), Err(cancelled_error())])); + assert_eq!(drive(&mut stream, 2), vec![0, ECANCELED]); + // The message is preserved for the consumer's get_last_error. + let last_error = stream.get_last_error.unwrap(); + let message = unsafe { CStr::from_ptr(last_error(&mut stream as *mut _)) }; + assert!(message.to_string_lossy().contains("operation cancelled")); + } + + #[test] + fn rich_error_is_available_through_error_from_array_stream() { + let failure = ArrowError::ExternalError(Box::new(Error { + message: "operation cancelled".into(), + status: Status::Cancelled, + vendor_code: ADBC_ERROR_VENDOR_CODE_PRIVATE_DATA, + sqlstate: [b'5' as _, b'7' as _, b'0' as _, b'1' as _, b'4' as _], + details: Some(vec![("reason".into(), b"user requested".to_vec())]), + })); + let mut stream = export_reader(reader_with(vec![Err(failure)])); + assert_eq!(drive(&mut stream, 1), vec![ECANCELED]); + + let (error, status) = rich_error(&mut stream); + assert_eq!(status, Some(ADBC_STATUS_CANCELLED)); + let error = unsafe { error.as_ref() }.unwrap(); + let message = unsafe { CStr::from_ptr(error.message) }; + assert!(message.to_string_lossy().contains("operation cancelled")); + assert_eq!(error.vendor_code, ADBC_ERROR_VENDOR_CODE_PRIVATE_DATA); + assert_eq!( + error.sqlstate, + [b'5' as _, b'7' as _, b'0' as _, b'1' as _, b'4' as _] + ); + // The 1.1.0 key/value details survive the conversion. + let details = unsafe { &*(error.private_data as *const ErrorPrivateData) }; + assert_eq!(details.keys, vec![CString::new("reason").unwrap()]); + assert_eq!(details.values, vec![b"user requested".to_vec()]); + } + + #[test] + fn consumer_may_release_the_returned_error_early() { + // The ADBC spec says the stream owns the error returned by ErrorFromArrayStream, but + // the Python bindings' `convert_error` releases it anyway after copying its contents. + // That must not double-free when the stream is released afterwards. + let mut stream = export_reader(reader_with(vec![Err(cancelled_error())])); + assert_eq!(drive(&mut stream, 1), vec![ECANCELED]); + + let (error, _) = rich_error(&mut stream); + let error = error as *mut FFI_AdbcError; + unsafe { ((*error).release.unwrap())(error) }; + // The stream is dropped (released) here, freeing the stashed error a second time. + } + + #[test] + fn errors_without_an_adbc_status_are_synthesized() { + let mut stream = export_reader(reader_with(vec![Err(ArrowError::IoError( + "connection reset".into(), + std::io::Error::other("connection reset"), + ))])); + assert_eq!(drive(&mut stream, 1), vec![EIO]); + + let (error, status) = rich_error(&mut stream); + assert_eq!(status, Some(ADBC_STATUS_IO)); + let error = unsafe { error.as_ref() }.unwrap(); + let message = unsafe { CStr::from_ptr(error.message) }; + assert!(message.to_string_lossy().contains("connection reset")); + } + + #[test] + fn stream_without_failures_reports_ok() { + let mut stream = export_reader(reader_with(vec![Ok(batch())])); + let (error, status) = rich_error(&mut stream); + assert_eq!(status, Some(ADBC_STATUS_OK)); + let error = unsafe { error.as_ref() }.unwrap(); + assert!(error.message.is_null()); + } + + #[test] + fn foreign_and_released_streams_are_not_recognized() { + // A stream exported by arrow-rs, not by this module. + let mut foreign = FFI_ArrowArrayStream::new(reader_with(vec![])); + let (error, status) = rich_error(&mut foreign); + assert!(error.is_null()); + assert_eq!(status, None, "status must not be written"); + + // A stream exported by this module, after release. + let mut released = export_reader(reader_with(vec![])); + unsafe { released.release.unwrap()(&mut released as *mut _) }; + let (error, status) = rich_error(&mut released); + assert!(error.is_null()); + assert_eq!(status, None, "status must not be written"); + } + + #[test] + fn wrapped_cancelled_error_is_found_through_the_source_chain() { + // A Cancelled ADBC error nested one level deeper still maps to ECANCELED. + #[derive(Debug)] + struct Wrapper(Error); + impl std::fmt::Display for Wrapper { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "wrapped: {}", self.0) + } + } + impl std::error::Error for Wrapper { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(&self.0) + } + } + let err = ArrowError::ExternalError(Box::new(Wrapper(Error::with_message_and_status( + "cancelled", + Status::Cancelled, + )))); + assert_eq!(adbc_error(&err).status, Status::Cancelled); + } + + #[test] + fn arrow_errors_synthesize_a_status() { + let synthesized = |err: &ArrowError| adbc_error(err).status; + assert_eq!( + synthesized(&ArrowError::NotYetImplemented("x".into())), + Status::NotImplemented + ); + assert_eq!( + synthesized(&ArrowError::IoError("x".into(), std::io::Error::other("x"))), + Status::IO + ); + assert_eq!( + synthesized(&ArrowError::ComputeError("x".into())), + Status::Internal + ); + } + + #[test] + fn adbc_statuses_map_through_the_canonical_table() { + // Spot-check the InternalAdbcStatusCodeToErrno port (c/driver/common/utils.c). + assert_eq!(errno_for_status(Status::Cancelled), ECANCELED); + assert_eq!(errno_for_status(Status::NotFound), ENOENT); + assert_eq!(errno_for_status(Status::NotImplemented), ENOTSUP); + assert_eq!(errno_for_status(Status::AlreadyExists), EEXIST); + assert_eq!(errno_for_status(Status::Timeout), ETIMEDOUT); + assert_eq!(errno_for_status(Status::Unauthorized), EACCES); + assert_eq!(errno_for_status(Status::InvalidArguments), EINVAL); + assert_eq!(errno_for_status(Status::Internal), EIO); + } + + #[test] + fn schema_round_trips() { + let mut stream = export_reader(reader_with(vec![])); + let get_schema = stream.get_schema.unwrap(); + let mut out = FFI_ArrowSchema::empty(); + assert_eq!( + unsafe { get_schema(&mut stream as *mut _, &mut out as *mut _) }, + 0 + ); + } +} diff --git a/rust/ffi/src/lib.rs b/rust/ffi/src/lib.rs index 0805d2f0e5..f80fd871cf 100644 --- a/rust/ffi/src/lib.rs +++ b/rust/ffi/src/lib.rs @@ -44,6 +44,7 @@ pub mod driver_exporter; pub mod options; #[doc(hidden)] pub use driver_exporter::FFIDriver; +pub(crate) mod exported_stream; pub mod methods; pub(crate) mod types; pub use types::{