This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new d162dd8a5 feat(rust): define cancellation in a sensible way (#3905)
d162dd8a5 is described below
commit d162dd8a5d90623b593097c6c037de6a176b405c
Author: David Li <[email protected]>
AuthorDate: Fri Aug 7 08:39:50 2026 +0900
feat(rust): define cancellation in a sensible way (#3905)
Closes #3454.
---
.github/workflows/rust.yml | 6 +
rust/core/src/sync.rs | 53 ++-
rust/driver/dummy/src/lib.rs | 20 ++
rust/driver/dummy/tests/driver_exporter_dummy.rs | 35 +-
rust/driver_manager/src/lib.rs | 392 +++++++++++++++++----
rust/driver_manager/tests/driver_manager_sqlite.rs | 10 +-
rust/ffi/src/driver_exporter.rs | 52 ++-
7 files changed, 461 insertions(+), 107 deletions(-)
diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index 74d685853..024bfd833 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -181,14 +181,20 @@ jobs:
-Zcrate-attr='feature(non_exhaustive_omitted_patterns_lint)' \
-Zcrate-attr='allow(unused_features)'
- name: Test (Default Features)
+ env:
+ RUST_BACKTRACE: "1"
working-directory: rust
run: >
cargo test --all-targets --workspace
- name: Test (All Features)
+ env:
+ RUST_BACKTRACE: "1"
working-directory: rust
run: >
cargo test --all-targets --all-features --workspace
- name: Doctests
+ env:
+ RUST_BACKTRACE: "1"
working-directory: rust
run: >
cargo test --doc --all-features --workspace
diff --git a/rust/core/src/sync.rs b/rust/core/src/sync.rs
index e14c6b224..55c5f19bd 100644
--- a/rust/core/src/sync.rs
+++ b/rust/core/src/sync.rs
@@ -21,7 +21,7 @@ use arrow_array::{RecordBatch, RecordBatchReader};
use arrow_schema::Schema;
use crate::PartitionedResult;
-use crate::error::Result;
+use crate::error::{Error, Result, Status};
use crate::options::{self, OptionConnection, OptionDatabase, OptionStatement,
OptionValue};
/// Ability to configure an object by setting/getting options.
@@ -44,6 +44,28 @@ pub trait Optionable {
fn get_option_double(&self, key: Self::Option) -> Result<f64>;
}
+/// A handle to cancel an in-progress operation.
+///
+/// This is a separated handle because otherwise it would be impossible to
+/// safely call a `cancel` method on a database/connection/statement itself
+/// due to the borrow checker.
+pub trait CancelHandle: Send + Sync {
+ /// Attempt to cancel the in-progress operation (best-effort).
+ fn try_cancel(&self) -> Result<()>;
+}
+
+/// A cancellation handle that does nothing (because cancellation is
unsupported).
+pub struct NoOpCancellationHandle;
+
+impl CancelHandle for NoOpCancellationHandle {
+ fn try_cancel(&self) -> Result<()> {
+ Err(Error::with_message_and_status(
+ "cancellation not implemented",
+ Status::Unknown,
+ ))
+ }
+}
+
/// A handle to an ADBC driver.
pub trait Driver {
type DatabaseType: Database;
@@ -76,6 +98,11 @@ pub trait Database: Optionable<Option = OptionDatabase> {
&self,
opts: impl IntoIterator<Item = (options::OptionConnection,
OptionValue)>,
) -> Result<Self::ConnectionType>;
+
+ /// Get a handle to cancel operations on this database.
+ fn get_cancel_handle(&self) -> Box<dyn CancelHandle> {
+ Box::new(NoOpCancellationHandle {})
+ }
}
/// A handle to an ADBC connection.
@@ -95,7 +122,15 @@ pub trait Connection: Optionable<Option = OptionConnection>
{
fn new_statement(&mut self) -> Result<Self::StatementType>;
/// Cancel the in-progress operation on a connection.
- fn cancel(&mut self) -> Result<()>;
+ #[deprecated(since = "0.25.0", note = "Use get_cancel_handle() instead")]
+ fn cancel(&mut self) -> Result<()> {
+ self.get_cancel_handle().try_cancel()
+ }
+
+ /// Get a handle to cancel operations on this connection.
+ fn get_cancel_handle(&self) -> Box<dyn CancelHandle> {
+ Box::new(NoOpCancellationHandle {})
+ }
/// Get metadata about the database/driver.
///
@@ -456,12 +491,20 @@ pub trait Statement: Optionable<Option = OptionStatement>
{
fn set_substrait_plan(&mut self, plan: impl AsRef<[u8]>) -> Result<()>;
/// Cancel execution of an in-progress query.
+ #[deprecated(since = "0.25.0", note = "Use get_cancel_handle() instead")]
+ fn cancel(&mut self) -> Result<()> {
+ self.get_cancel_handle().try_cancel()
+ }
+
+ /// Get a handle to cancel operations on this statement.
///
- /// This can be called during [Statement::execute] (or similar), or while
- /// consuming a result set returned from such.
+ /// The resulting handle can be called during [Statement::execute] (or
+ /// similar), or while consuming a result set returned from such.
///
/// # Since
///
/// ADBC API revision 1.1.0
- fn cancel(&mut self) -> Result<()>;
+ fn get_cancel_handle(&self) -> Box<dyn CancelHandle> {
+ Box::new(NoOpCancellationHandle {})
+ }
}
diff --git a/rust/driver/dummy/src/lib.rs b/rust/driver/dummy/src/lib.rs
index 8c3b92e76..ad8fedd34 100644
--- a/rust/driver/dummy/src/lib.rs
+++ b/rust/driver/dummy/src/lib.rs
@@ -308,6 +308,26 @@ impl Connection for DummyConnection {
Err(error)
}
+ /// This method is used to test that errors round-trip correctly.
+ fn get_cancel_handle(&self) -> Box<dyn adbc_core::CancelHandle> {
+ struct CancelHandle;
+
+ impl adbc_core::CancelHandle for CancelHandle {
+ fn try_cancel(&self) -> Result<()> {
+ let mut error = Error::with_message_and_status("message",
Status::Cancelled);
+ error.vendor_code =
constants::ADBC_ERROR_VENDOR_CODE_PRIVATE_DATA;
+ error.sqlstate = [1, 2, 3, 4, 5];
+ error.details = Some(vec![
+ ("key1".into(), b"AAA".into()),
+ ("key2".into(), b"ZZZZZ".into()),
+ ]);
+ Err(error)
+ }
+ }
+
+ Box::new(CancelHandle)
+ }
+
fn commit(&mut self) -> Result<()> {
Ok(())
}
diff --git a/rust/driver/dummy/tests/driver_exporter_dummy.rs
b/rust/driver/dummy/tests/driver_exporter_dummy.rs
index 36e7657f9..ecb7ff893 100644
--- a/rust/driver/dummy/tests/driver_exporter_dummy.rs
+++ b/rust/driver/dummy/tests/driver_exporter_dummy.rs
@@ -465,11 +465,14 @@ fn test_connection_get_info_ignores_unrecognized_codes() {
#[test]
fn test_connection_cancel() {
- let (_, _, mut exported_connection, _) = get_exported();
- let (_, _, mut native_connection, _) = get_native();
+ let (_, _, exported_connection, _) = get_exported();
+ let (_, _, native_connection, _) = get_native();
- let exported_error = exported_connection.cancel().unwrap_err();
- let native_error = native_connection.cancel().unwrap_err();
+ let exported_handle = exported_connection.get_cancel_handle();
+ let native_handle = native_connection.get_cancel_handle();
+
+ let exported_error = exported_handle.try_cancel().unwrap_err();
+ let native_error = native_handle.try_cancel().unwrap_err();
assert_eq!(exported_error, native_error);
}
@@ -668,11 +671,27 @@ fn test_statement_bind_stream() {
#[test]
fn test_statement_cancel() {
- let (_, _, _, mut exported_statement) = get_exported();
- let (_, _, _, mut native_statement) = get_native();
+ let (_, _, _, exported_statement) = get_exported();
+ let (_, _, _, native_statement) = get_native();
+
+ let exported_handle = exported_statement.get_cancel_handle();
+ let native_handle = native_statement.get_cancel_handle();
- exported_statement.cancel().unwrap();
- native_statement.cancel().unwrap();
+ let res = exported_handle.try_cancel();
+ assert!(res.is_err());
+ assert!(
+ res.unwrap_err()
+ .to_string()
+ .contains("cancellation not implemented")
+ );
+
+ let res = native_handle.try_cancel();
+ assert!(res.is_err());
+ assert!(
+ res.unwrap_err()
+ .to_string()
+ .contains("cancellation not implemented")
+ );
}
#[test]
diff --git a/rust/driver_manager/src/lib.rs b/rust/driver_manager/src/lib.rs
index bb2cc700d..811886828 100644
--- a/rust/driver_manager/src/lib.rs
+++ b/rust/driver_manager/src/lib.rs
@@ -346,11 +346,16 @@ struct ManagedDatabaseInner {
impl Drop for ManagedDatabaseInner {
fn drop(&mut self) {
let driver = &self.driver.driver;
- let mut database = self.database.lock().unwrap();
- let method = driver_method!(driver, DatabaseRelease);
- // TODO(alexandreyc): how should we handle `DatabaseRelease` failing?
- // See:
https://github.com/apache/arrow-adbc/pull/1742#discussion_r1574388409
- unsafe { method(database.deref_mut(), null_mut()) };
+ if let Ok(mut database) = self.database.lock() {
+ let method = driver_method!(driver, DatabaseRelease);
+ // TODO(alexandreyc): how should we handle `DatabaseRelease`
failing?
+ // See:
https://github.com/apache/arrow-adbc/pull/1742#discussion_r1574388409
+ unsafe { method(database.deref_mut(), null_mut()) };
+ }
+ // We could still drop here but if the lock is poisoned, we have no
+ // clue what the status is. Since a panic comes from Rust code,
+ // _probably_ the FFI handle is unharmed. But I think it's safer to
+ // leak than to try to release.
}
}
@@ -637,8 +642,12 @@ impl ManagedDatabase {
mut connection: adbc_ffi::FFI_AdbcConnection,
) -> Result<adbc_ffi::FFI_AdbcConnection> {
let driver = self.ffi_driver();
- let mut database = self.inner.database.lock().unwrap();
-
+ let mut database = self.inner.database.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] database is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
// ConnectionInit
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionInit);
@@ -654,7 +663,12 @@ impl Optionable for ManagedDatabase {
fn get_option_bytes(&self, key: Self::Option) -> Result<Vec<u8>> {
let driver = self.ffi_driver();
- let database = &mut self.inner.database.lock().unwrap();
+ let mut database = self.inner.database.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] database is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let method = driver_method!(driver, DatabaseGetOptionBytes);
let populate = |key: *const c_char,
value: *mut u8,
@@ -667,7 +681,12 @@ impl Optionable for ManagedDatabase {
fn get_option_double(&self, key: Self::Option) -> Result<f64> {
let driver = self.ffi_driver();
- let mut database = self.inner.database.lock().unwrap();
+ let mut database = self.inner.database.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] database is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let key = CString::new(key.as_ref())?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let mut value: f64 = f64::default();
@@ -679,7 +698,12 @@ impl Optionable for ManagedDatabase {
fn get_option_int(&self, key: Self::Option) -> Result<i64> {
let driver = self.ffi_driver();
- let mut database = self.inner.database.lock().unwrap();
+ let mut database = self.inner.database.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] database is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let key = CString::new(key.as_ref())?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let mut value: i64 = 0;
@@ -691,7 +715,12 @@ impl Optionable for ManagedDatabase {
fn get_option_string(&self, key: Self::Option) -> Result<String> {
let driver = self.ffi_driver();
- let mut database = self.inner.database.lock().unwrap();
+ let mut database = self.inner.database.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] database is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let method = driver_method!(driver, DatabaseGetOption);
let populate = |key: *const c_char,
value: *mut c_char,
@@ -704,7 +733,12 @@ impl Optionable for ManagedDatabase {
fn set_option(&mut self, key: Self::Option, value: OptionValue) ->
Result<()> {
let driver = self.ffi_driver();
- let mut database = self.inner.database.lock().unwrap();
+ let mut database = self.inner.database.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] database is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
set_option_database(
driver,
database.deref_mut(),
@@ -765,11 +799,16 @@ struct ManagedConnectionInner {
impl Drop for ManagedConnectionInner {
fn drop(&mut self) {
let driver = &self.database.driver.driver;
- let mut connection = self.connection.lock().unwrap();
- let method = driver_method!(driver, ConnectionRelease);
- // TODO(alexandreyc): how should we handle `ConnectionRelease` failing?
- // See:
https://github.com/apache/arrow-adbc/pull/1742#discussion_r1574388409
- unsafe { method(connection.deref_mut(), null_mut()) };
+ if let Ok(mut connection) = self.connection.lock() {
+ let method = driver_method!(driver, ConnectionRelease);
+ // TODO(alexandreyc): how should we handle `ConnectionRelease`
failing?
+ // See:
https://github.com/apache/arrow-adbc/pull/1742#discussion_r1574388409
+ unsafe { method(connection.deref_mut(), null_mut()) };
+ }
+ // We could still drop here but if the lock is poisoned, we have no
+ // clue what the status is. Since a panic comes from Rust code,
+ // _probably_ the FFI handle is unharmed. But I think it's safer to
+ // leak than to try to release.
}
}
@@ -779,6 +818,36 @@ pub struct ManagedConnection {
inner: Arc<ManagedConnectionInner>,
}
+struct ConnectionCancelHandle {
+ inner: std::sync::Weak<ManagedConnectionInner>,
+}
+
+impl adbc_core::CancelHandle for ConnectionCancelHandle {
+ fn try_cancel(&self) -> Result<()> {
+ if let Some(inner) = self.inner.upgrade() {
+ if let AdbcVersion::V100 = inner.database.driver.version {
+ return Err(Error::with_message_and_status(
+ ERR_CANCEL_UNSUPPORTED,
+ Status::NotImplemented,
+ ));
+ }
+ let driver = &inner.database.driver.driver;
+ let mut connection = inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
+ let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
+ let method = driver_method!(driver, ConnectionCancel);
+ let status = unsafe { method(connection.deref_mut(), &mut error) };
+ check_status(status, error)
+ } else {
+ Ok(())
+ }
+ }
+}
+
impl ManagedConnection {
fn ffi_driver(&self) -> &adbc_ffi::FFI_AdbcDriver {
&self.inner.database.driver.driver
@@ -794,7 +863,12 @@ impl Optionable for ManagedConnection {
fn get_option_bytes(&self, key: Self::Option) -> Result<Vec<u8>> {
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let method = driver_method!(driver, ConnectionGetOptionBytes);
let populate = |key: *const c_char,
value: *mut u8,
@@ -809,7 +883,12 @@ impl Optionable for ManagedConnection {
let key = CString::new(key.as_ref())?;
let mut value: f64 = f64::default();
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetOptionDouble);
let status =
@@ -822,7 +901,12 @@ impl Optionable for ManagedConnection {
let key = CString::new(key.as_ref())?;
let mut value: i64 = 0;
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetOptionInt);
let status =
@@ -833,7 +917,12 @@ impl Optionable for ManagedConnection {
fn get_option_string(&self, key: Self::Option) -> Result<String> {
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let method = driver_method!(driver, ConnectionGetOption);
let populate = |key: *const c_char,
value: *mut c_char,
@@ -846,7 +935,12 @@ impl Optionable for ManagedConnection {
fn set_option(&mut self, key: Self::Option, value: OptionValue) ->
Result<()> {
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
set_option_connection(
driver,
connection.deref_mut(),
@@ -862,7 +956,12 @@ impl Connection for ManagedConnection {
fn new_statement(&mut self) -> Result<Self::StatementType> {
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut statement = adbc_ffi::FFI_AdbcStatement::default();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementNew);
@@ -877,24 +976,20 @@ impl Connection for ManagedConnection {
Ok(Self::StatementType { inner })
}
- fn cancel(&mut self) -> Result<()> {
- if let AdbcVersion::V100 = self.driver_version() {
- return Err(Error::with_message_and_status(
- ERR_CANCEL_UNSUPPORTED,
- Status::NotImplemented,
- ));
- }
- let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
- let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
- let method = driver_method!(driver, ConnectionCancel);
- let status = unsafe { method(connection.deref_mut(), &mut error) };
- check_status(status, error)
+ fn get_cancel_handle(&self) -> Box<dyn adbc_core::CancelHandle> {
+ Box::new(ConnectionCancelHandle {
+ inner: Arc::downgrade(&self.inner),
+ })
}
fn commit(&mut self) -> Result<()> {
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionCommit);
let status = unsafe { method(connection.deref_mut(), &mut error) };
@@ -903,7 +998,12 @@ impl Connection for ManagedConnection {
fn rollback(&mut self) -> Result<()> {
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionRollback);
let status = unsafe { method(connection.deref_mut(), &mut error) };
@@ -922,7 +1022,12 @@ impl Connection for ManagedConnection {
.map(|c| (c.as_ptr(), c.len()))
.unwrap_or((null(), 0));
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetInfo);
let status = unsafe {
@@ -978,7 +1083,12 @@ impl Connection for ManagedConnection {
};
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetObjects);
let mut stream = FFI_ArrowArrayStream::empty();
@@ -1026,7 +1136,12 @@ impl Connection for ManagedConnection {
let mut stream = FFI_ArrowArrayStream::empty();
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetStatistics);
let status = unsafe {
@@ -1054,7 +1169,12 @@ impl Connection for ManagedConnection {
}
let mut stream = FFI_ArrowArrayStream::empty();
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetStatisticNames);
let status = unsafe { method(connection.deref_mut(), &mut stream, &mut
error) };
@@ -1079,7 +1199,12 @@ impl Connection for ManagedConnection {
let mut schema = FFI_ArrowSchema::empty();
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetTableSchema);
let status = unsafe {
@@ -1099,7 +1224,12 @@ impl Connection for ManagedConnection {
fn get_table_types(&self) -> Result<Box<dyn RecordBatchReader + Send +
'static>> {
let mut stream = FFI_ArrowArrayStream::empty();
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetTableTypes);
let status = unsafe { method(connection.deref_mut(), &mut stream, &mut
error) };
@@ -1114,7 +1244,12 @@ impl Connection for ManagedConnection {
) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
let mut stream = FFI_ArrowArrayStream::empty();
let driver = self.ffi_driver();
- let mut connection = self.inner.connection.lock().unwrap();
+ let mut connection = self.inner.connection.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] connection is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionReadPartition);
let partition = partition.as_ref();
@@ -1153,10 +1288,45 @@ impl ManagedStatement {
}
}
+struct StatementCancelHandle {
+ inner: std::sync::Weak<ManagedStatementInner>,
+}
+
+impl adbc_core::CancelHandle for StatementCancelHandle {
+ fn try_cancel(&self) -> Result<()> {
+ if let Some(inner) = self.inner.upgrade() {
+ if let AdbcVersion::V100 =
inner.connection.database.driver.version {
+ return Err(Error::with_message_and_status(
+ ERR_CANCEL_UNSUPPORTED,
+ Status::NotImplemented,
+ ));
+ }
+ let driver = &inner.connection.database.driver.driver;
+ let mut statement = inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
+ let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
+ let method = driver_method!(driver, StatementCancel);
+ let status = unsafe { method(statement.deref_mut(), &mut error) };
+ check_status(status, error)
+ } else {
+ Ok(())
+ }
+ }
+}
+
impl Statement for ManagedStatement {
fn bind(&mut self, batch: RecordBatch) -> Result<()> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementBind);
let batch: StructArray = batch.into();
@@ -1168,7 +1338,12 @@ impl Statement for ManagedStatement {
fn bind_stream(&mut self, reader: Box<dyn RecordBatchReader + Send>) ->
Result<()> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementBindStream);
let mut stream = FFI_ArrowArrayStream::new(reader);
@@ -1177,24 +1352,20 @@ impl Statement for ManagedStatement {
Ok(())
}
- fn cancel(&mut self) -> Result<()> {
- if let AdbcVersion::V100 = self.driver_version() {
- return Err(Error::with_message_and_status(
- ERR_CANCEL_UNSUPPORTED,
- Status::NotImplemented,
- ));
- }
- let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
- let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
- let method = driver_method!(driver, StatementCancel);
- let status = unsafe { method(statement.deref_mut(), &mut error) };
- check_status(status, error)
+ fn get_cancel_handle(&self) -> Box<dyn adbc_core::CancelHandle> {
+ Box::new(StatementCancelHandle {
+ inner: Arc::downgrade(&self.inner),
+ })
}
fn execute(&mut self) -> Result<Box<dyn RecordBatchReader + Send +
'static>> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementExecuteQuery);
let mut stream = FFI_ArrowArrayStream::empty();
@@ -1206,7 +1377,12 @@ impl Statement for ManagedStatement {
fn execute_schema(&mut self) -> Result<arrow_schema::Schema> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementExecuteSchema);
let mut schema = FFI_ArrowSchema::empty();
@@ -1217,7 +1393,12 @@ impl Statement for ManagedStatement {
fn execute_update(&mut self) -> Result<Option<i64>> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementExecuteQuery);
let mut rows_affected: i64 = -1;
@@ -1235,7 +1416,12 @@ impl Statement for ManagedStatement {
fn execute_partitions(&mut self) -> Result<PartitionedResult> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementExecutePartitions);
let mut schema = FFI_ArrowSchema::empty();
@@ -1263,7 +1449,12 @@ impl Statement for ManagedStatement {
fn get_parameter_schema(&self) -> Result<arrow_schema::Schema> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementGetParameterSchema);
let mut schema = FFI_ArrowSchema::empty();
@@ -1274,7 +1465,12 @@ impl Statement for ManagedStatement {
fn prepare(&mut self) -> Result<()> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementPrepare);
let status = unsafe { method(statement.deref_mut(), &mut error) };
@@ -1285,7 +1481,12 @@ impl Statement for ManagedStatement {
fn set_sql_query(&mut self, query: impl AsRef<str>) -> Result<()> {
let query = CString::new(query.as_ref())?;
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementSetSqlQuery);
let status = unsafe { method(statement.deref_mut(), query.as_ptr(),
&mut error) };
@@ -1295,7 +1496,12 @@ impl Statement for ManagedStatement {
fn set_substrait_plan(&mut self, plan: impl AsRef<[u8]>) -> Result<()> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementSetSubstraitPlan);
let plan = plan.as_ref();
@@ -1311,7 +1517,12 @@ impl Optionable for ManagedStatement {
fn get_option_bytes(&self, key: Self::Option) -> Result<Vec<u8>> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let method = driver_method!(driver, StatementGetOptionBytes);
let populate = |key: *const c_char,
value: *mut u8,
@@ -1326,7 +1537,12 @@ impl Optionable for ManagedStatement {
let key = CString::new(key.as_ref())?;
let mut value: f64 = f64::default();
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementGetOptionDouble);
let status = unsafe { method(statement.deref_mut(), key.as_ptr(), &mut
value, &mut error) };
@@ -1338,7 +1554,12 @@ impl Optionable for ManagedStatement {
let key = CString::new(key.as_ref())?;
let mut value: i64 = 0;
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementGetOptionInt);
let status = unsafe { method(statement.deref_mut(), key.as_ptr(), &mut
value, &mut error) };
@@ -1348,7 +1569,12 @@ impl Optionable for ManagedStatement {
fn get_option_string(&self, key: Self::Option) -> Result<String> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
let method = driver_method!(driver, StatementGetOption);
let populate = |key: *const c_char,
value: *mut c_char,
@@ -1361,7 +1587,12 @@ impl Optionable for ManagedStatement {
fn set_option(&mut self, key: Self::Option, value: OptionValue) ->
Result<()> {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
+ let mut statement = self.inner.statement.lock().map_err(|e| {
+ Error::with_message_and_status(
+ format!("[Driver Manager] statement is poisoned: {e:?}"),
+ Status::Internal,
+ )
+ })?;
set_option_statement(
driver,
statement.deref_mut(),
@@ -1375,10 +1606,15 @@ impl Optionable for ManagedStatement {
impl Drop for ManagedStatement {
fn drop(&mut self) {
let driver = self.ffi_driver();
- let mut statement = self.inner.statement.lock().unwrap();
- let method = driver_method!(driver, StatementRelease);
- // TODO(alexandreyc): how should we handle `StatementRelease` failing?
- // See:
https://github.com/apache/arrow-adbc/pull/1742#discussion_r1574388409
- unsafe { method(statement.deref_mut(), null_mut()) };
+ if let Ok(mut statement) = self.inner.statement.lock() {
+ let method = driver_method!(driver, StatementRelease);
+ // TODO(alexandreyc): how should we handle `StatementRelease`
failing?
+ // See:
https://github.com/apache/arrow-adbc/pull/1742#discussion_r1574388409
+ unsafe { method(statement.deref_mut(), null_mut()) };
+ }
+ // We could still drop here but if the lock is poisoned, we have no
+ // clue what the status is. Since a panic comes from Rust code,
+ // _probably_ the FFI handle is unharmed. But I think it's safer to
+ // leak than to try to release.
}
}
diff --git a/rust/driver_manager/tests/driver_manager_sqlite.rs
b/rust/driver_manager/tests/driver_manager_sqlite.rs
index 0f6d89fa2..a84d8cb84 100644
--- a/rust/driver_manager/tests/driver_manager_sqlite.rs
+++ b/rust/driver_manager/tests/driver_manager_sqlite.rs
@@ -128,9 +128,10 @@ fn test_connection_get_option() {
fn test_connection_cancel() {
let mut driver = get_driver();
let database = get_database(&mut driver);
- let mut connection = database.new_connection().unwrap();
+ let connection = database.new_connection().unwrap();
- let error = connection.cancel().unwrap_err();
+ let handle = connection.get_cancel_handle();
+ let error = handle.try_cancel().unwrap_err();
assert_eq!(error.status, Status::NotImplemented);
}
@@ -285,9 +286,10 @@ fn test_statement_cancel() {
let mut driver = get_driver();
let database = get_database(&mut driver);
let mut connection = database.new_connection().unwrap();
- let mut statement = connection.new_statement().unwrap();
+ let statement = connection.new_statement().unwrap();
- let error = statement.cancel().unwrap_err();
+ let handle = statement.get_cancel_handle();
+ let error = handle.try_cancel().unwrap_err();
assert_eq!(error.status, Status::NotImplemented);
}
diff --git a/rust/ffi/src/driver_exporter.rs b/rust/ffi/src/driver_exporter.rs
index 429f0caf3..00a29ef1b 100644
--- a/rust/ffi/src/driver_exporter.rs
+++ b/rust/ffi/src/driver_exporter.rs
@@ -61,11 +61,16 @@ impl<DriverType: Driver> ExportedDatabase<DriverType> {
}
}
+struct InitializedConnection<DriverType: Driver> {
+ connection: ConnectionType<DriverType>,
+ cancel_handle: Box<dyn adbc_core::CancelHandle>,
+}
+
enum ExportedConnection<DriverType: Driver> {
/// Pre-init options
Options(HashMap<OptionConnection, OptionValue>),
/// Initialized connection
- Connection(ConnectionType<DriverType>),
+ Connection(InitializedConnection<DriverType>),
}
impl<DriverType: Driver> ExportedConnection<DriverType> {
@@ -77,13 +82,23 @@ impl<DriverType: Driver> ExportedConnection<DriverType> {
) {
match self {
Self::Options(options) => (Some(options), None),
- Self::Connection(connection) => (None, Some(connection)),
+ Self::Connection(connection) => (None, Some(&mut
connection.connection)),
}
}
fn try_connection(&mut self) -> Result<&mut ConnectionType<DriverType>> {
match self {
- Self::Connection(connection) => Ok(connection),
+ Self::Connection(connection) => Ok(&mut connection.connection),
+ _ => Err(Error::with_message_and_status(
+ "Connection not initialized",
+ Status::InvalidState,
+ )),
+ }
+ }
+
+ fn try_cancel(&mut self) -> Result<&mut dyn adbc_core::CancelHandle> {
+ match self {
+ Self::Connection(connection) =>
Ok(connection.cancel_handle.as_mut()),
_ => Err(Error::with_message_and_status(
"Connection not initialized",
Status::InvalidState,
@@ -92,7 +107,10 @@ impl<DriverType: Driver> ExportedConnection<DriverType> {
}
}
-struct ExportedStatement<DriverType: Driver>(StatementType<DriverType>);
+struct ExportedStatement<DriverType: Driver>(
+ StatementType<DriverType>,
+ Box<dyn adbc_core::CancelHandle>,
+);
pub trait FFIDriver {
fn ffi_driver() -> FFI_AdbcDriver;
@@ -828,7 +846,10 @@ unsafe fn connection_set_option_impl<DriverType: Driver,
Value: Into<OptionValue
options.insert(key.into(), value.into());
}
ExportedConnection::Connection(connection) => {
- check_err!(connection.set_option(key.into(), value.into()), error);
+ check_err!(
+ connection.connection.set_option(key.into(), value.into()),
+ error
+ );
}
}
@@ -878,7 +899,11 @@ extern "C" fn connection_init<DriverType: Driver>(
)),
};
let connection = check_err!(connection, error);
- *exported_connection = ExportedConnection::Connection(connection);
+ let cancel_handle = connection.get_cancel_handle();
+ *exported_connection =
ExportedConnection::Connection(InitializedConnection {
+ connection,
+ cancel_handle,
+ });
} else {
check_err!(
Err(Error::with_message_and_status(
@@ -1224,8 +1249,8 @@ extern "C" fn connection_cancel<DriverType: Driver>(
unsafe { connection_private_data::<DriverType>(connection) },
error
);
- let connection = check_err!(exported.try_connection(), error);
- check_err!(connection.cancel(), error);
+ let handle = check_err!(exported.try_cancel(), error);
+ check_err!(handle.try_cancel(), error);
ADBC_STATUS_OK
})
@@ -1421,8 +1446,12 @@ extern "C" fn statement_new<DriverType: Driver>(
let inner_connection =
check_err!(exported_connection.try_connection(), error);
let inner_statement = check_err!(inner_connection.new_statement(),
error);
+ let cancel_handle = inner_statement.get_cancel_handle();
- let exported =
Box::new(ExportedStatement::<DriverType>(inner_statement));
+ let exported = Box::new(ExportedStatement::<DriverType>(
+ inner_statement,
+ cancel_handle,
+ ));
statement.private_data = Box::into_raw(exported) as *mut c_void;
ADBC_STATUS_OK
@@ -1685,9 +1714,8 @@ extern "C" fn statement_cancel<DriverType: Driver>(
unsafe { statement_private_data::<DriverType>(statement) },
error
);
- let statement = &mut exported.0;
-
- check_err!(statement.cancel(), error);
+ let cancel_handle = &mut exported.1;
+ check_err!(cancel_handle.try_cancel(), error);
ADBC_STATUS_OK
})
