pub mod error;
pub mod profile;
pub mod search;
use std::cell::UnsafeCell;
use std::collections::HashSet;
use std::ffi::{CString, OsStr};
use std::ops::DerefMut;
use std::os::raw::c_char;
use std::path::PathBuf;
use std::pin::Pin;
use std::ptr::{null, null_mut};
use std::sync::{Arc, LockResult, Mutex, MutexGuard};
use adbc_ffi::options::{
check_status, get_option_bytes, get_option_string, set_option_connection,
set_option_connection_raw, set_option_database, set_option_statement_raw,
};
use arrow_array::ffi::{FFI_ArrowSchema, to_ffi};
use arrow_array::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream};
use arrow_array::{Array, RecordBatch, RecordBatchReader, StructArray};
use adbc_core::{
Connection, Database, Driver, LoadFlags, Optionable, PartitionedResult, Statement,
error::{Error, Result, Status},
options::{self, AdbcVersion, InfoCode, OptionDatabase, OptionValue},
};
use adbc_ffi::driver_method;
use self::search::{DriverLibrary, DriverLocator, parse_driver_uri};
use crate::profile::{
ConnectionProfile, ConnectionProfileProvider, FilesystemProfileProvider, process_profile_value,
};
const ERR_CANCEL_UNSUPPORTED: &str =
"Canceling connection or statement is not supported with ADBC 1.0.0";
const ERR_STATISTICS_UNSUPPORTED: &str = "Statistics are not supported with ADBC 1.0.0";
#[derive(Debug)]
struct ManagedDriverInner {
driver: adbc_ffi::FFI_AdbcDriver,
version: AdbcVersion, _library: Option<libloading::Library>,
}
#[derive(Clone, Debug)]
pub struct ManagedDriver {
inner: Pin<Arc<ManagedDriverInner>>,
}
impl ManagedDriver {
pub fn version(&self) -> AdbcVersion {
self.inner.version
}
pub fn load_static(
init: &adbc_ffi::FFI_AdbcDriverInitFunc,
version: AdbcVersion,
) -> Result<Self> {
let driver = DriverLibrary::from_static_init(init).init_driver(version)?;
let inner = Arc::pin(ManagedDriverInner {
driver,
version,
_library: None,
});
Ok(ManagedDriver { inner })
}
pub fn load_from_name(
name: impl AsRef<OsStr>,
entrypoint: Option<&[u8]>,
version: AdbcVersion,
load_flags: LoadFlags,
additional_search_paths: Option<Vec<PathBuf>>,
) -> Result<Self> {
let search_hit = DriverLibrary::search(name, load_flags, additional_search_paths)?;
let entrypoint = search_hit.resolve_entrypoint(entrypoint).to_vec();
Self::load_from_library(search_hit.library, entrypoint.as_ref(), version)
}
pub fn load_dynamic_from_filename(
filename: impl AsRef<OsStr>,
entrypoint: Option<&[u8]>,
version: AdbcVersion,
) -> Result<Self> {
let entrypoint = DriverLibrary::derive_entrypoint(entrypoint, filename.as_ref());
let library = DriverLibrary::load_library(filename)?;
Self::load_from_library(library, entrypoint.as_ref(), version)
}
fn load_from_library(
library: libloading::Library,
entrypoint: &[u8],
version: AdbcVersion,
) -> Result<ManagedDriver> {
let driver =
DriverLibrary::try_from_dynamic_library(&library, entrypoint)?.init_driver(version)?;
let inner = Arc::pin(ManagedDriverInner {
driver,
version,
_library: Some(library),
});
Ok(ManagedDriver { inner })
}
pub fn load_dynamic_from_name(
name: impl AsRef<str>,
entrypoint: Option<&[u8]>,
version: AdbcVersion,
) -> Result<Self> {
let entrypoint = DriverLibrary::derive_entrypoint_from_name(entrypoint, name.as_ref());
let library = DriverLibrary::load_library_from_name(name.as_ref())?;
Self::load_from_library(library, entrypoint.as_ref(), version)
}
fn inner_ffi_driver(&self) -> &adbc_ffi::FFI_AdbcDriver {
&self.inner.driver
}
fn database_new(&self) -> Result<adbc_ffi::FFI_AdbcDatabase> {
let driver = self.inner_ffi_driver();
let mut database = adbc_ffi::FFI_AdbcDatabase {
private_driver: driver,
..Default::default()
};
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(*driver, DatabaseNew);
let status = unsafe { method(&mut database, &mut error) };
check_status(status, error)?;
Ok(database)
}
fn database_init(
&self,
mut database: adbc_ffi::FFI_AdbcDatabase,
) -> Result<adbc_ffi::FFI_AdbcDatabase> {
let driver = self.inner_ffi_driver();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, DatabaseInit);
let status = unsafe { method(&mut database, &mut error) };
check_status(status, error)?;
Ok(database)
}
}
impl Driver for ManagedDriver {
type DatabaseType = ManagedDatabase;
fn new_database(&mut self) -> Result<Self::DatabaseType> {
let database = self.database_new()?;
let database = self.database_init(database)?;
let inner = Arc::new(ManagedDatabaseInner {
database: Mutex::new(database),
driver: self.inner.clone(),
});
Ok(Self::DatabaseType { inner })
}
fn new_database_with_opts(
&mut self,
opts: impl IntoIterator<Item = (<Self::DatabaseType as Optionable>::Option, OptionValue)>,
) -> Result<Self::DatabaseType> {
let mut database = self.database_new()?;
{
let driver = self.inner_ffi_driver();
for (key, value) in opts {
set_option_database(driver, &mut database, self.inner.version, key, value)?;
}
}
let database = self.database_init(database)?;
let inner = Arc::new(ManagedDatabaseInner {
database: Mutex::new(database),
driver: self.inner.clone(),
});
Ok(Self::DatabaseType { inner })
}
}
struct ManagedDatabaseInner {
database: Mutex<adbc_ffi::FFI_AdbcDatabase>,
driver: Pin<Arc<ManagedDriverInner>>,
}
impl Drop for ManagedDatabaseInner {
fn drop(&mut self) {
let driver = &self.driver.driver;
if let Ok(mut database) = self.database.lock() {
let method = driver_method!(driver, DatabaseRelease);
unsafe { method(database.deref_mut(), null_mut()) };
}
}
}
#[derive(Clone)]
pub struct ManagedDatabase {
inner: Arc<ManagedDatabaseInner>,
}
impl ManagedDatabase {
pub fn from_uri(
uri: &str,
entrypoint: Option<&[u8]>,
version: AdbcVersion,
load_flags: LoadFlags,
additional_search_paths: Option<Vec<PathBuf>>,
) -> Result<Self> {
Self::from_uri_with_opts(
uri,
entrypoint,
version,
load_flags,
additional_search_paths,
std::iter::empty(),
)
}
pub fn from_uri_with_opts(
uri: &str,
entrypoint: Option<&[u8]>,
version: AdbcVersion,
load_flags: LoadFlags,
additional_search_paths: Option<Vec<PathBuf>>,
opts: impl IntoIterator<Item = (<Self as Optionable>::Option, OptionValue)>,
) -> Result<Self> {
Self::from_uri_with_profile_provider(
uri,
entrypoint,
version,
load_flags,
additional_search_paths,
FilesystemProfileProvider::default(),
opts,
)
}
pub fn from_uri_with_profile_provider(
uri: &str,
entrypoint: Option<&[u8]>,
version: AdbcVersion,
load_flags: LoadFlags,
additional_search_paths: Option<Vec<PathBuf>>,
profile_provider: impl ConnectionProfileProvider,
opts: impl IntoIterator<Item = (<Self as Optionable>::Option, OptionValue)>,
) -> Result<Self> {
let result = parse_driver_uri(uri)?;
let (mut drv, default_opts) = match result {
DriverLocator::Uri(driver, final_uri) => {
let drv = ManagedDriver::load_from_name(
driver,
entrypoint,
version,
load_flags,
additional_search_paths,
)?;
let final_opts = vec![(
OptionDatabase::Uri,
OptionValue::String(final_uri.to_string()),
)];
(drv, final_opts)
}
DriverLocator::Profile(profile) => {
let profile = profile_provider.get_profile(profile)?;
let (driver_name, init_func) = profile.get_driver_name()?;
let drv: ManagedDriver;
if let Some(init_fn) = init_func {
drv = ManagedDriver::load_static(init_fn, version)?;
} else {
drv = ManagedDriver::load_from_name(
driver_name,
entrypoint,
version,
load_flags,
additional_search_paths,
)?;
}
let profile_opts: Vec<(OptionDatabase, OptionValue)> = profile
.get_options()?
.into_iter()
.map(|(k, v)| -> Result<(OptionDatabase, OptionValue)> {
if let OptionValue::String(s) = v {
let result = process_profile_value(&s)?;
Ok((k, result))
} else {
Ok((k, v))
}
})
.collect::<Result<Vec<_>>>()?;
(drv, profile_opts)
}
};
drv.new_database_with_opts(default_opts.into_iter().chain(opts))
}
fn ffi_driver(&self) -> &adbc_ffi::FFI_AdbcDriver {
&self.inner.driver.driver
}
fn driver_version(&self) -> AdbcVersion {
self.inner.driver.version
}
fn connection_new(&self) -> Result<adbc_ffi::FFI_AdbcConnection> {
let driver = self.ffi_driver();
let mut connection = adbc_ffi::FFI_AdbcConnection::default();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(*driver, ConnectionNew);
let status = unsafe { method(&mut connection, &mut error) };
check_status(status, error)?;
Ok(connection)
}
fn connection_init(
&self,
mut connection: adbc_ffi::FFI_AdbcConnection,
) -> Result<adbc_ffi::FFI_AdbcConnection> {
let driver = self.ffi_driver();
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 mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionInit);
let status = unsafe { method(&mut connection, &mut *database, &mut error) };
check_status(status, error)?;
Ok(connection)
}
}
impl Optionable for ManagedDatabase {
type Option = options::OptionDatabase;
fn get_option_bytes(&self, key: Self::Option) -> Result<Vec<u8>> {
let driver = self.ffi_driver();
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,
length: *mut usize,
error: *mut adbc_ffi::FFI_AdbcError| unsafe {
method(database.deref_mut(), key, value, length, error)
};
get_option_bytes(key, populate, driver)
}
fn get_option_double(&self, key: Self::Option) -> Result<f64> {
let driver = self.ffi_driver();
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();
let method = driver_method!(driver, DatabaseGetOptionDouble);
let status = unsafe { method(database.deref_mut(), key.as_ptr(), &mut value, &mut error) };
check_status(status, error)?;
Ok(value)
}
fn get_option_int(&self, key: Self::Option) -> Result<i64> {
let driver = self.ffi_driver();
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;
let method = driver_method!(driver, DatabaseGetOptionInt);
let status = unsafe { method(database.deref_mut(), key.as_ptr(), &mut value, &mut error) };
check_status(status, error)?;
Ok(value)
}
fn get_option_string(&self, key: Self::Option) -> Result<String> {
let driver = self.ffi_driver();
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,
length: *mut usize,
error: *mut adbc_ffi::FFI_AdbcError| unsafe {
method(database.deref_mut(), key, value, length, error)
};
get_option_string(key, populate, driver)
}
fn set_option(&mut self, key: Self::Option, value: OptionValue) -> Result<()> {
let driver = self.ffi_driver();
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(),
self.driver_version(),
key,
value,
)
}
}
impl Database for ManagedDatabase {
type ConnectionType = ManagedConnection;
fn new_connection(&self) -> Result<Self::ConnectionType> {
let connection = self.connection_new()?;
let connection = self.connection_init(connection)?;
let inner = ManagedConnectionInner {
connection: NativeHandle::new(connection),
database: self.inner.clone(),
};
Ok(Self::ConnectionType {
inner: Arc::new(inner),
})
}
fn new_connection_with_opts(
&self,
opts: impl IntoIterator<Item = (<Self::ConnectionType as Optionable>::Option, OptionValue)>,
) -> Result<Self::ConnectionType> {
let mut connection = self.connection_new()?;
{
let driver = self.ffi_driver();
for (key, value) in opts {
set_option_connection(driver, &mut connection, self.driver_version(), key, value)?;
}
}
let connection = self.connection_init(connection)?;
let inner = ManagedConnectionInner {
connection: NativeHandle::new(connection),
database: self.inner.clone(),
};
Ok(Self::ConnectionType {
inner: Arc::new(inner),
})
}
}
struct NativeHandle<T> {
value: UnsafeCell<T>,
mu: Mutex<()>,
}
unsafe impl<T: Send> Sync for NativeHandle<T> {}
impl<T> NativeHandle<T> {
fn new(value: T) -> Self {
Self {
value: UnsafeCell::new(value),
mu: Mutex::new(()),
}
}
fn lock(&self) -> LockResult<MutexGuard<'_, ()>> {
self.mu.lock()
}
fn get(&self) -> *mut T {
self.value.get()
}
}
struct ManagedConnectionInner {
connection: NativeHandle<adbc_ffi::FFI_AdbcConnection>,
database: Arc<ManagedDatabaseInner>,
}
impl Drop for ManagedConnectionInner {
fn drop(&mut self) {
let driver = &self.database.driver.driver;
if !self.connection.mu.is_poisoned() {
let connection = self.connection.get();
let method = driver_method!(driver, ConnectionRelease);
unsafe { method(connection, null_mut()) };
}
}
}
#[derive(Clone)]
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 Some(method) = driver.ConnectionCancel else {
return Err(Error::with_message_and_status(
"[Driver Manager] ConnectionCancel not supported by driver",
Status::NotImplemented,
));
};
let connection = inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let status = unsafe { method(connection, &mut error) };
check_status(status, error)
} else {
Ok(())
}
}
}
impl ManagedConnection {
fn ffi_driver(&self) -> &adbc_ffi::FFI_AdbcDriver {
&self.inner.database.driver.driver
}
fn driver_version(&self) -> AdbcVersion {
self.inner.database.driver.version
}
}
impl Optionable for ManagedConnection {
type Option = options::OptionConnection;
fn get_option_bytes(&self, key: Self::Option) -> Result<Vec<u8>> {
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let method = driver_method!(driver, ConnectionGetOptionBytes);
let populate = |key: *const c_char,
value: *mut u8,
length: *mut usize,
error: *mut adbc_ffi::FFI_AdbcError| unsafe {
method(connection, key, value, length, error)
};
get_option_bytes(key, populate, driver)
}
fn get_option_double(&self, key: Self::Option) -> Result<f64> {
let key = CString::new(key.as_ref())?;
let mut value: f64 = f64::default();
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetOptionDouble);
let status = unsafe { method(connection, key.as_ptr(), &mut value, &mut error) };
check_status(status, error)?;
Ok(value)
}
fn get_option_int(&self, key: Self::Option) -> Result<i64> {
let key = CString::new(key.as_ref())?;
let mut value: i64 = 0;
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetOptionInt);
let status = unsafe { method(connection, key.as_ptr(), &mut value, &mut error) };
check_status(status, error)?;
Ok(value)
}
fn get_option_string(&self, key: Self::Option) -> Result<String> {
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let method = driver_method!(driver, ConnectionGetOption);
let populate = |key: *const c_char,
value: *mut c_char,
length: *mut usize,
error: *mut adbc_ffi::FFI_AdbcError| unsafe {
method(connection, key, value, length, error)
};
get_option_string(key, populate, driver)
}
fn set_option(&mut self, key: Self::Option, value: OptionValue) -> Result<()> {
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
unsafe { set_option_connection_raw(driver, connection, self.driver_version(), key, value) }
}
}
impl Connection for ManagedConnection {
type StatementType = ManagedStatement;
fn new_statement(&mut self) -> Result<Self::StatementType> {
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let mut statement = adbc_ffi::FFI_AdbcStatement::default();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementNew);
let status = unsafe { method(connection, &mut statement, &mut error) };
check_status(status, error)?;
let inner = Arc::new(ManagedStatementInner {
statement: NativeHandle::new(statement),
connection: self.inner.clone(),
});
Ok(Self::StatementType { inner })
}
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 _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionCommit);
let status = unsafe { method(connection, &mut error) };
check_status(status, error)
}
fn rollback(&mut self) -> Result<()> {
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionRollback);
let status = unsafe { method(connection, &mut error) };
check_status(status, error)
}
fn get_info(
&self,
codes: Option<HashSet<InfoCode>>,
) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
let mut stream = FFI_ArrowArrayStream::empty();
let codes: Option<Vec<u32>> =
codes.map(|codes| codes.iter().map(|code| code.into()).collect());
let (codes_ptr, codes_len) = codes
.as_ref()
.map(|c| (c.as_ptr(), c.len()))
.unwrap_or((null(), 0));
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetInfo);
let status = unsafe { method(connection, codes_ptr, codes_len, &mut stream, &mut error) };
check_status(status, error)?;
let reader = ArrowArrayStreamReader::try_new(stream)?;
Ok(Box::new(reader))
}
fn get_objects(
&self,
depth: crate::options::ObjectDepth,
catalog: Option<&str>,
db_schema: Option<&str>,
table_name: Option<&str>,
table_type: Option<Vec<&str>>,
column_name: Option<&str>,
) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
let catalog = catalog.map(CString::new).transpose()?;
let db_schema = db_schema.map(CString::new).transpose()?;
let table_name = table_name.map(CString::new).transpose()?;
let column_name = column_name.map(CString::new).transpose()?;
let table_type = table_type
.map(|t| {
t.iter()
.map(|x| CString::new(*x))
.collect::<std::result::Result<Vec<CString>, _>>()
})
.transpose()?;
let catalog_ptr = catalog.as_ref().map(|c| c.as_ptr()).unwrap_or(null());
let db_schema_ptr = db_schema.as_ref().map(|c| c.as_ptr()).unwrap_or(null());
let table_name_ptr = table_name.as_ref().map(|c| c.as_ptr()).unwrap_or(null());
let column_name_ptr = column_name.as_ref().map(|c| c.as_ptr()).unwrap_or(null());
let mut table_type_ptrs = table_type
.as_ref()
.map(|v| v.iter().map(|c| c.as_ptr()))
.map(|c| c.collect::<Vec<_>>());
let table_type_ptr = match table_type_ptrs.as_mut() {
None => null(),
Some(t) => {
t.push(null());
t.as_ptr()
}
};
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetObjects);
let mut stream = FFI_ArrowArrayStream::empty();
let status = unsafe {
method(
connection,
depth.into(),
catalog_ptr,
db_schema_ptr,
table_name_ptr,
table_type_ptr,
column_name_ptr,
&mut stream,
&mut error,
)
};
check_status(status, error)?;
let reader = ArrowArrayStreamReader::try_new(stream)?;
Ok(Box::new(reader))
}
fn get_statistics(
&self,
catalog: Option<&str>,
db_schema: Option<&str>,
table_name: Option<&str>,
approximate: bool,
) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
if let AdbcVersion::V100 = self.driver_version() {
return Err(Error::with_message_and_status(
ERR_STATISTICS_UNSUPPORTED,
Status::NotImplemented,
));
}
let catalog = catalog.map(CString::new).transpose()?;
let db_schema = db_schema.map(CString::new).transpose()?;
let table_name = table_name.map(CString::new).transpose()?;
let catalog_ptr = catalog.as_ref().map(|c| c.as_ptr()).unwrap_or(null());
let db_schema_ptr = db_schema.as_ref().map(|c| c.as_ptr()).unwrap_or(null());
let table_name_ptr = table_name.as_ref().map(|c| c.as_ptr()).unwrap_or(null());
let mut stream = FFI_ArrowArrayStream::empty();
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetStatistics);
let status = unsafe {
method(
connection,
catalog_ptr,
db_schema_ptr,
table_name_ptr,
approximate as std::os::raw::c_char,
&mut stream,
&mut error,
)
};
check_status(status, error)?;
let reader = ArrowArrayStreamReader::try_new(stream)?;
Ok(Box::new(reader))
}
fn get_statistic_names(&self) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
if let AdbcVersion::V100 = self.driver_version() {
return Err(Error::with_message_and_status(
ERR_STATISTICS_UNSUPPORTED,
Status::NotImplemented,
));
}
let mut stream = FFI_ArrowArrayStream::empty();
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetStatisticNames);
let status = unsafe { method(connection, &mut stream, &mut error) };
check_status(status, error)?;
let reader = ArrowArrayStreamReader::try_new(stream)?;
Ok(Box::new(reader))
}
fn get_table_schema(
&self,
catalog: Option<&str>,
db_schema: Option<&str>,
table_name: &str,
) -> Result<arrow_schema::Schema> {
let catalog = catalog.map(CString::new).transpose()?;
let db_schema = db_schema.map(CString::new).transpose()?;
let table_name = CString::new(table_name)?;
let catalog_ptr = catalog.as_ref().map(|c| c.as_ptr()).unwrap_or(null());
let db_schema_ptr = db_schema.as_ref().map(|c| c.as_ptr()).unwrap_or(null());
let table_name_ptr = table_name.as_ptr();
let mut schema = FFI_ArrowSchema::empty();
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetTableSchema);
let status = unsafe {
method(
connection,
catalog_ptr,
db_schema_ptr,
table_name_ptr,
&mut schema,
&mut error,
)
};
check_status(status, error)?;
Ok((&schema).try_into()?)
}
fn get_table_types(&self) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
let mut stream = FFI_ArrowArrayStream::empty();
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionGetTableTypes);
let status = unsafe { method(connection, &mut stream, &mut error) };
check_status(status, error)?;
let reader = ArrowArrayStreamReader::try_new(stream)?;
Ok(Box::new(reader))
}
fn read_partition(
&self,
partition: impl AsRef<[u8]>,
) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
let mut stream = FFI_ArrowArrayStream::empty();
let driver = self.ffi_driver();
let _guard = self.inner.connection.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] connection is poisoned: {e:?}"),
Status::Internal,
)
})?;
let connection = self.inner.connection.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, ConnectionReadPartition);
let partition = partition.as_ref();
let status = unsafe {
method(
connection,
partition.as_ptr(),
partition.len(),
&mut stream,
&mut error,
)
};
check_status(status, error)?;
let reader = ArrowArrayStreamReader::try_new(stream)?;
Ok(Box::new(reader))
}
}
struct ManagedStatementInner {
statement: NativeHandle<adbc_ffi::FFI_AdbcStatement>,
connection: Arc<ManagedConnectionInner>,
}
#[derive(Clone)]
pub struct ManagedStatement {
inner: Arc<ManagedStatementInner>,
}
impl ManagedStatement {
fn driver_version(&self) -> AdbcVersion {
self.inner.connection.database.driver.version
}
fn ffi_driver(&self) -> &adbc_ffi::FFI_AdbcDriver {
&self.inner.connection.database.driver.driver
}
}
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 Some(method) = driver.StatementCancel else {
return Err(Error::with_message_and_status(
"[Driver Manager] StatementCancel not supported by driver",
Status::NotImplemented,
));
};
let statement = inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let status = unsafe { method(statement, &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 _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementBind);
let batch: StructArray = batch.into();
let (mut array, mut schema) = to_ffi(&batch.to_data())?;
let status = unsafe { method(statement, &mut array, &mut schema, &mut error) };
check_status(status, error)?;
Ok(())
}
fn bind_stream(&mut self, reader: Box<dyn RecordBatchReader + Send>) -> Result<()> {
let driver = self.ffi_driver();
let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementBindStream);
let mut stream = FFI_ArrowArrayStream::new(reader);
let status = unsafe { method(statement, &mut stream, &mut error) };
check_status(status, error)?;
Ok(())
}
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 _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementExecuteQuery);
let mut stream = FFI_ArrowArrayStream::empty();
let status = unsafe { method(statement, &mut stream, null_mut(), &mut error) };
check_status(status, error)?;
let reader = ArrowArrayStreamReader::try_new(stream)?;
Ok(Box::new(reader))
}
fn execute_schema(&mut self) -> Result<arrow_schema::Schema> {
let driver = self.ffi_driver();
let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementExecuteSchema);
let mut schema = FFI_ArrowSchema::empty();
let status = unsafe { method(statement, &mut schema, &mut error) };
check_status(status, error)?;
Ok((&schema).try_into()?)
}
fn execute_update(&mut self) -> Result<Option<i64>> {
let driver = self.ffi_driver();
let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementExecuteQuery);
let mut rows_affected: i64 = -1;
let status = unsafe { method(statement, null_mut(), &mut rows_affected, &mut error) };
check_status(status, error)?;
Ok((rows_affected != -1).then_some(rows_affected))
}
fn execute_partitions(&mut self) -> Result<PartitionedResult> {
let driver = self.ffi_driver();
let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementExecutePartitions);
let mut schema = FFI_ArrowSchema::empty();
let mut partitions = adbc_ffi::FFI_AdbcPartitions::default();
let mut rows_affected: i64 = -1;
let status = unsafe {
method(
statement,
&mut schema,
&mut partitions,
&mut rows_affected,
&mut error,
)
};
check_status(status, error)?;
let result = PartitionedResult {
partitions: unsafe { adbc_ffi::import_partitions(partitions) }?,
schema: (&schema).try_into()?,
rows_affected,
};
Ok(result)
}
fn get_parameter_schema(&self) -> Result<arrow_schema::Schema> {
let driver = self.ffi_driver();
let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementGetParameterSchema);
let mut schema = FFI_ArrowSchema::empty();
let status = unsafe { method(statement, &mut schema, &mut error) };
check_status(status, error)?;
Ok((&schema).try_into()?)
}
fn prepare(&mut self) -> Result<()> {
let driver = self.ffi_driver();
let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementPrepare);
let status = unsafe { method(statement, &mut error) };
check_status(status, error)?;
Ok(())
}
fn set_sql_query(&mut self, query: impl AsRef<str>) -> Result<()> {
let query = CString::new(query.as_ref())?;
let driver = self.ffi_driver();
let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementSetSqlQuery);
let status = unsafe { method(statement, query.as_ptr(), &mut error) };
check_status(status, error)?;
Ok(())
}
fn set_substrait_plan(&mut self, plan: impl AsRef<[u8]>) -> Result<()> {
let driver = self.ffi_driver();
let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementSetSubstraitPlan);
let plan = plan.as_ref();
let status = unsafe { method(statement, plan.as_ptr(), plan.len(), &mut error) };
check_status(status, error)?;
Ok(())
}
}
impl Optionable for ManagedStatement {
type Option = options::OptionStatement;
fn get_option_bytes(&self, key: Self::Option) -> Result<Vec<u8>> {
let driver = self.ffi_driver();
let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let method = driver_method!(driver, StatementGetOptionBytes);
let populate = |key: *const c_char,
value: *mut u8,
length: *mut usize,
error: *mut adbc_ffi::FFI_AdbcError| unsafe {
method(statement, key, value, length, error)
};
get_option_bytes(key, populate, driver)
}
fn get_option_double(&self, key: Self::Option) -> Result<f64> {
let key = CString::new(key.as_ref())?;
let mut value: f64 = f64::default();
let driver = self.ffi_driver();
let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementGetOptionDouble);
let status = unsafe { method(statement, key.as_ptr(), &mut value, &mut error) };
check_status(status, error)?;
Ok(value)
}
fn get_option_int(&self, key: Self::Option) -> Result<i64> {
let key = CString::new(key.as_ref())?;
let mut value: i64 = 0;
let driver = self.ffi_driver();
let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let mut error = adbc_ffi::FFI_AdbcError::with_driver(driver);
let method = driver_method!(driver, StatementGetOptionInt);
let status = unsafe { method(statement, key.as_ptr(), &mut value, &mut error) };
check_status(status, error)?;
Ok(value)
}
fn get_option_string(&self, key: Self::Option) -> Result<String> {
let driver = self.ffi_driver();
let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
let method = driver_method!(driver, StatementGetOption);
let populate = |key: *const c_char,
value: *mut c_char,
length: *mut usize,
error: *mut adbc_ffi::FFI_AdbcError| unsafe {
method(statement, key, value, length, error)
};
get_option_string(key, populate, driver)
}
fn set_option(&mut self, key: Self::Option, value: OptionValue) -> Result<()> {
let driver = self.ffi_driver();
let _guard = self.inner.statement.lock().map_err(|e| {
Error::with_message_and_status(
format!("[Driver Manager] statement is poisoned: {e:?}"),
Status::Internal,
)
})?;
let statement = self.inner.statement.get();
unsafe { set_option_statement_raw(driver, statement, self.driver_version(), key, value) }
}
}
impl Drop for ManagedStatementInner {
fn drop(&mut self) {
let driver = &self.connection.database.driver.driver;
if !self.statement.mu.is_poisoned() {
let statement = self.statement.get();
let method = driver_method!(driver, StatementRelease);
unsafe { method(statement, null_mut()) };
}
}
}