From d80f21d7d8e8de96c99f42564442bed6ad6fd0b5 Mon Sep 17 00:00:00 2001
From: David Li <li.davidm96@gmail.com>
Date: Thu, 22 Jan 2026 16:56:27 +0900
Subject: [PATCH] feat(rust/core)!: define cancellation in a sensible way

Closes #3454.
---
 rust/core/src/sync.rs                         | 39 +++++++--
 rust/driver/datafusion/src/lib.rs             |  8 --
 rust/driver/dummy/src/lib.rs                  | 32 +++----
 .../dummy/tests/driver_exporter_dummy.rs      | 22 +++--
 rust/driver/snowflake/src/connection.rs       |  4 +-
 rust/driver/snowflake/src/statement.rs        |  4 +-
 rust/driver_manager/src/lib.rs                | 84 +++++++++++++------
 .../tests/driver_manager_sqlite.rs            | 10 ++-
 rust/ffi/src/driver_exporter.rs               | 52 +++++++++---
 9 files changed, 173 insertions(+), 82 deletions(-)

diff --git a/rust/core/src/sync.rs b/rust/core/src/sync.rs
index de26918b95..aaa4ab4888 100644
--- a/rust/core/src/sync.rs
+++ b/rust/core/src/sync.rs
@@ -44,6 +44,24 @@ pub trait Optionable {
     fn get_option_double(&self, key: Self::Option) -> Result<f64>;
 }
 
+/// A handle to cancel an in-progress operation on a connection.
+///
+/// This is a separated handle because otherwise it would be impossible to
+/// call a `cancel` method on a connection or statement itself.
+pub trait CancelHandle: Send {
+    /// Cancel the in-progress operation on a connection.
+    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<()> {
+        Ok(())
+    }
+}
+
 /// A handle to an ADBC driver.
 pub trait Driver {
     type DatabaseType: Database;
@@ -76,6 +94,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.
@@ -94,8 +117,10 @@ pub trait Connection: Optionable<Option = OptionConnection> {
     /// Allocate and initialize a new statement.
     fn new_statement(&mut self) -> Result<Self::StatementType>;
 
-    /// Cancel the in-progress operation on a connection.
-    fn cancel(&mut self) -> Result<()>;
+    /// 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.
     ///
@@ -455,13 +480,15 @@ pub trait Statement: Optionable<Option = OptionStatement> {
     /// expected to be executed repeatedly, call [Statement::prepare] first.
     fn set_substrait_plan(&mut self, plan: impl AsRef<[u8]>) -> Result<()>;
 
-    /// Cancel execution of an in-progress query.
+    /// 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/datafusion/src/lib.rs b/rust/driver/datafusion/src/lib.rs
index 0db317a728..0f4af7331f 100644
--- a/rust/driver/datafusion/src/lib.rs
+++ b/rust/driver/datafusion/src/lib.rs
@@ -735,10 +735,6 @@ impl Connection for DataFusionConnection {
         })
     }
 
-    fn cancel(&mut self) -> adbc_core::error::Result<()> {
-        todo!()
-    }
-
     fn get_info(
         &self,
         codes: Option<std::collections::HashSet<adbc_core::options::InfoCode>>,
@@ -984,10 +980,6 @@ impl Statement for DataFusionStatement {
         self.substrait_plan = Some(Plan::decode(plan.as_ref()).unwrap());
         Ok(())
     }
-
-    fn cancel(&mut self) -> adbc_core::error::Result<()> {
-        todo!()
-    }
 }
 
 #[cfg(feature = "ffi")]
diff --git a/rust/driver/dummy/src/lib.rs b/rust/driver/dummy/src/lib.rs
index 72c81f1785..5660b57180 100644
--- a/rust/driver/dummy/src/lib.rs
+++ b/rust/driver/dummy/src/lib.rs
@@ -294,16 +294,24 @@ impl Connection for DummyConnection {
         Ok(Self::StatementType::default())
     }
 
-    // This method is used to test that errors round-trip correctly.
-    fn cancel(&mut 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)
+    /// 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<()> {
@@ -854,10 +862,6 @@ impl Statement for DummyStatement {
         Ok(())
     }
 
-    fn cancel(&mut self) -> Result<()> {
-        Ok(())
-    }
-
     fn execute(&mut self) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
         maybe_panic("StatementExecuteQuery");
         let batch = get_table_data();
diff --git a/rust/driver/dummy/tests/driver_exporter_dummy.rs b/rust/driver/dummy/tests/driver_exporter_dummy.rs
index 6d1445a45b..beb1f56a22 100644
--- a/rust/driver/dummy/tests/driver_exporter_dummy.rs
+++ b/rust/driver/dummy/tests/driver_exporter_dummy.rs
@@ -366,11 +366,14 @@ fn test_connection_get_info() {
 
 #[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_handle = exported_connection.get_cancel_handle();
+    let native_handle = native_connection.get_cancel_handle();
 
-    let exported_error = exported_connection.cancel().unwrap_err();
-    let native_error = native_connection.cancel().unwrap_err();
+    let exported_error = exported_handle.try_cancel().unwrap_err();
+    let native_error = native_handle.try_cancel().unwrap_err();
 
     assert_eq!(exported_error, native_error);
 }
@@ -569,11 +572,14 @@ 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();
+    exported_handle.try_cancel().unwrap();
+    native_handle.try_cancel().unwrap();
 }
 
 #[test]
diff --git a/rust/driver/snowflake/src/connection.rs b/rust/driver/snowflake/src/connection.rs
index 2e916ac4b2..ba3116c3ef 100644
--- a/rust/driver/snowflake/src/connection.rs
+++ b/rust/driver/snowflake/src/connection.rs
@@ -70,8 +70,8 @@ impl adbc_core::Connection for Connection {
         self.0.new_statement().map(Statement)
     }
 
-    fn cancel(&mut self) -> Result<()> {
-        self.0.cancel()
+    fn get_cancel_handle(&self) -> Box<dyn adbc_core::CancelHandle> {
+        self.0.get_cancel_handle()
     }
 
     fn get_info(
diff --git a/rust/driver/snowflake/src/statement.rs b/rust/driver/snowflake/src/statement.rs
index f956852b26..00109bd7cb 100644
--- a/rust/driver/snowflake/src/statement.rs
+++ b/rust/driver/snowflake/src/statement.rs
@@ -96,7 +96,7 @@ impl adbc_core::Statement for Statement {
         self.0.set_substrait_plan(plan)
     }
 
-    fn cancel(&mut self) -> Result<()> {
-        self.0.cancel()
+    fn get_cancel_handle(&self) -> Box<dyn adbc_core::CancelHandle> {
+        self.0.get_cancel_handle()
     }
 }
diff --git a/rust/driver_manager/src/lib.rs b/rust/driver_manager/src/lib.rs
index b2cc03f3ad..85be7a4faf 100644
--- a/rust/driver_manager/src/lib.rs
+++ b/rust/driver_manager/src/lib.rs
@@ -1027,6 +1027,31 @@ 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().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)
+        } else {
+            Ok(())
+        }
+    }
+}
+
 impl ManagedConnection {
     fn ffi_driver(&self) -> &adbc_ffi::FFI_AdbcDriver {
         &self.inner.database.driver.driver
@@ -1125,19 +1150,10 @@ 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<()> {
@@ -1401,6 +1417,31 @@ 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().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)
+        } else {
+            Ok(())
+        }
+    }
+}
+
 impl Statement for ManagedStatement {
     fn bind(&mut self, batch: RecordBatch) -> Result<()> {
         let driver = self.ffi_driver();
@@ -1425,19 +1466,10 @@ 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>> {
diff --git a/rust/driver_manager/tests/driver_manager_sqlite.rs b/rust/driver_manager/tests/driver_manager_sqlite.rs
index ed462d3852..465ea1ecf1 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 e4573ac556..3212bcb420 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;
@@ -844,7 +862,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
+            );
         }
     }
 
@@ -894,7 +915,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(
@@ -1241,8 +1266,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
     })
@@ -1439,8 +1464,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
@@ -1703,9 +1732,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
     })