From 889f78899ce4b386b345cebcf5b004cd832c673a Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Wed, 14 Oct 2026 13:31:52 +0300 Subject: [PATCH 01/27] feat: Use a known Schema instead of inferring it from the DB data --- crates/clickhouse/src/conn.rs | 25 ++++- crates/clickhouse/src/lib.rs | 3 +- crates/clickhouse/src/pool.rs | 31 +++++- .../sql/db_connection_pool/dbconnection.rs | 8 ++ crates/mysql/src/conn.rs | 13 +++ crates/mysql/src/pool.rs | 21 ++++- crates/odbc/src/conn.rs | 14 +++ crates/odbc/src/pool.rs | 12 +++ crates/oracle/src/conn.rs | 16 +++- crates/oracle/src/pool.rs | 31 +++++- crates/postgres/src/conn.rs | 14 +++ crates/postgres/src/pool.rs | 23 ++++- crates/sqlite/src/conn.rs | 16 +++- crates/sqlite/src/pool.rs | 94 +++++++++++++++++-- 14 files changed, 292 insertions(+), 29 deletions(-) diff --git a/crates/clickhouse/src/conn.rs b/crates/clickhouse/src/conn.rs index 352e87cf2..c07d4684d 100644 --- a/crates/clickhouse/src/conn.rs +++ b/crates/clickhouse/src/conn.rs @@ -20,6 +20,7 @@ use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::{ #[derive(Clone)] pub struct ClickHouseConnection { pub client: Client, + pub schema: Option, } impl std::fmt::Debug for ClickHouseConnection { @@ -32,7 +33,10 @@ impl std::fmt::Debug for ClickHouseConnection { impl ClickHouseConnection { #[must_use] pub fn new(client: Client) -> Self { - Self { client } + Self { + client, + schema: None, + } } } @@ -60,7 +64,20 @@ impl AsyncDbConnection for ClickHouseConnection { where Self: Sized, { - Self { client: conn } + Self { + client: conn, + schema: None, + } + } + + fn new_with_schema(conn: Client, schema: SchemaRef) -> Self + where + Self: Sized, + { + Self { + client: conn, + schema: Some(schema), + } } async fn tables(&self, schema: &str) -> Result, Error> { @@ -103,6 +120,10 @@ impl AsyncDbConnection for ClickHouseConnection { /// /// * `table_reference` - The table reference. async fn get_schema(&self, table_reference: &TableReference) -> Result { + if let Some(schema) = &self.schema { + return Ok(Arc::clone(schema)); + } + #[derive(Row, Deserialize)] struct CatalogRow { db: String, diff --git a/crates/clickhouse/src/lib.rs b/crates/clickhouse/src/lib.rs index 4f4e09ba6..9d0a29d64 100644 --- a/crates/clickhouse/src/lib.rs +++ b/crates/clickhouse/src/lib.rs @@ -27,7 +27,6 @@ use datafusion::sql::unparser; use datafusion::{common::Constraints, common::TableReference}; use std::sync::Arc; -use crate::conn::ClickHouseConnection; use crate::pool::ClickHouseConnectionPool; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::AsyncDbConnection; @@ -50,7 +49,7 @@ impl ClickHouseTableFactory { args: Option>, ) -> Result, Box> { - let connection = ClickHouseConnection::new(self.pool.client()); + let connection = self.pool.connect_direct(); let client: &dyn AsyncDbConnection = &connection; let schema = client.get_schema(&table_reference).await?; let table_provider = Arc::new(ClickHouseTable::new( diff --git a/crates/clickhouse/src/pool.rs b/crates/clickhouse/src/pool.rs index 481a7bc44..eff0b03fb 100644 --- a/crates/clickhouse/src/pool.rs +++ b/crates/clickhouse/src/pool.rs @@ -1,12 +1,14 @@ -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; use clickhouse::{Client, Compression}; +use datafusion::arrow::datatypes::SchemaRef; use secrecy::{ExposeSecret, SecretString}; use snafu::{ResultExt, Snafu}; use crate::conn::ClickHouseConnection; use datafusion_table_providers_common::sql::db_connection_pool::{ - dbconnection::DbConnection, DbConnectionPool, JoinPushDown, + dbconnection::{AsyncDbConnection, DbConnection}, + DbConnectionPool, JoinPushDown, }; type PoolResult> = std::result::Result; @@ -26,6 +28,7 @@ pub enum Error { pub struct ClickHouseConnectionPool { pub client: Client, pub join_push_down: JoinPushDown, + pub schema: Option, } impl std::fmt::Debug for ClickHouseConnectionPool { @@ -108,18 +111,40 @@ impl ClickHouseConnectionPool { Ok(Self { client, join_push_down, + schema: None, }) } + /// Use the provided `schema` for all connections instead of inferring it from the database. + #[must_use] + pub fn with_schema(mut self, schema: SchemaRef) -> Self { + self.schema = Some(schema); + self + } + pub fn client(&self) -> Client { self.client.clone() } + + /// Returns a direct connection to the underlying database. + #[must_use] + pub fn connect_direct(&self) -> ClickHouseConnection { + match &self.schema { + Some(schema) => { + >::new_with_schema( + self.client(), + Arc::clone(schema), + ) + } + None => ClickHouseConnection::new(self.client()), + } + } } #[async_trait::async_trait] impl DbConnectionPool for ClickHouseConnectionPool { async fn connect(&self) -> PoolResult>> { - Ok(Box::new(ClickHouseConnection::new(self.client()))) + Ok(Box::new(self.connect_direct())) } fn join_push_down(&self) -> JoinPushDown { diff --git a/crates/common/src/sql/db_connection_pool/dbconnection.rs b/crates/common/src/sql/db_connection_pool/dbconnection.rs index 45f1f9d38..babc42e4d 100644 --- a/crates/common/src/sql/db_connection_pool/dbconnection.rs +++ b/crates/common/src/sql/db_connection_pool/dbconnection.rs @@ -100,6 +100,14 @@ pub trait AsyncDbConnection: DbConnection + Sync { where Self: Sized; + /// Creates a new connection that uses the provided `schema` instead of inferring it. + /// + /// When a schema is provided, [`AsyncDbConnection::get_schema`] returns it as-is + /// instead of querying the database for it. + fn new_with_schema(conn: T, schema: SchemaRef) -> Self + where + Self: Sized; + async fn tables(&self, schema: &str) -> Result, Error>; async fn schemas(&self) -> Result, Error>; diff --git a/crates/mysql/src/conn.rs b/crates/mysql/src/conn.rs index cb2bd65e8..628f6b3af 100644 --- a/crates/mysql/src/conn.rs +++ b/crates/mysql/src/conn.rs @@ -46,6 +46,7 @@ pub enum Error { pub struct MySQLConnection { pub conn: Arc>, + pub schema: Option, } impl MySQLConnection { @@ -91,6 +92,14 @@ impl<'a> AsyncDbConnection for MySQLConnection { fn new(conn: Conn) -> Self { MySQLConnection { conn: Arc::new(Mutex::new(conn)), + schema: None, + } + } + + fn new_with_schema(conn: Conn, schema: SchemaRef) -> Self { + MySQLConnection { + conn: Arc::new(Mutex::new(conn)), + schema: Some(schema), } } @@ -153,6 +162,10 @@ impl<'a> AsyncDbConnection for MySQLConnection { SchemaRef, datafusion_table_providers_common::sql::db_connection_pool::dbconnection::Error, > { + if let Some(schema) = &self.schema { + return Ok(Arc::clone(schema)); + } + let mut conn = self.conn.lock().await; let conn = &mut *conn; diff --git a/crates/mysql/src/pool.rs b/crates/mysql/src/pool.rs index 26de708ea..c97ea6aed 100644 --- a/crates/mysql/src/pool.rs +++ b/crates/mysql/src/pool.rs @@ -1,4 +1,5 @@ use crate::conn::MySQLConnection; +use datafusion::arrow::datatypes::SchemaRef; use std::{collections::HashMap, path::PathBuf, sync::Arc}; use async_trait::async_trait; @@ -56,6 +57,7 @@ pub enum Error { pub struct MySQLConnectionPool { pool: Arc, join_push_down: JoinPushDown, + schema: Option, } const SETUP_QUERIES: [&str; 4] = [ @@ -209,9 +211,24 @@ impl MySQLConnectionPool { Ok(Self { pool: Arc::new(pool), join_push_down, + schema: None, }) } + /// Use the provided `schema` for all connections instead of inferring it from the database. + #[must_use] + pub fn with_schema(mut self, schema: SchemaRef) -> Self { + self.schema = Some(schema); + self + } + + fn new_connection(&self, conn: mysql_async::Conn) -> MySQLConnection { + match &self.schema { + Some(schema) => MySQLConnection::new_with_schema(conn, Arc::clone(schema)), + None => MySQLConnection::new(conn), + } + } + /// Returns a direct connection to the underlying database. /// /// # Errors @@ -221,7 +238,7 @@ impl MySQLConnectionPool { let pool = Arc::clone(&self.pool); let conn = pool.get_conn().await.context(MySQLConnectionSnafu)?; - Ok(MySQLConnection::new(conn)) + Ok(self.new_connection(conn)) } pub fn metrics(&self) -> Arc { @@ -291,7 +308,7 @@ impl DbConnectionPool for MySQ .context(MySQLConnectionSnafu) .map_err(|e| Box::new(e) as Box)?; - Ok(Box::new(MySQLConnection::new(conn))) + Ok(Box::new(self.new_connection(conn))) } fn join_push_down(&self) -> JoinPushDown { diff --git a/crates/odbc/src/conn.rs b/crates/odbc/src/conn.rs index be4fdcaf0..2e3825030 100644 --- a/crates/odbc/src/conn.rs +++ b/crates/odbc/src/conn.rs @@ -93,6 +93,7 @@ pub enum Error { pub struct ODBCConnection<'a> { pub conn: Arc>>, pub params: Arc>, + pub schema: Option, } impl<'a> DbConnection, ODBCParameter> for ODBCConnection<'a> @@ -131,6 +132,15 @@ where ODBCConnection { conn: Arc::new(conn.into()), params: Arc::new(HashMap::new()), + schema: None, + } + } + + fn new_with_schema(conn: Connection<'a>, schema: SchemaRef) -> Self { + ODBCConnection { + conn: Arc::new(conn.into()), + params: Arc::new(HashMap::new()), + schema: Some(schema), } } @@ -157,6 +167,10 @@ where &self, table_reference: &TableReference, ) -> Result { + if let Some(schema) = &self.schema { + return Ok(Arc::clone(schema)); + } + let conn = Arc::clone(&self.conn); let sql = format!( "SELECT * FROM {} LIMIT 1", diff --git a/crates/odbc/src/pool.rs b/crates/odbc/src/pool.rs index aab7db6ff..2d4fd3921 100644 --- a/crates/odbc/src/pool.rs +++ b/crates/odbc/src/pool.rs @@ -17,6 +17,7 @@ limitations under the License. use crate::conn::ODBCConnection; use crate::conn::{ODBCDbConnection, ODBCParameter}; use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; use datafusion_table_providers_common::sql::db_connection_pool::{ runtime::run_async_with_tokio, DbConnectionPool, JoinPushDown, }; @@ -59,6 +60,7 @@ pub struct ODBCPool { params: Arc>, connection_string: String, connection_id: String, + schema: Option, } fn hash_string(val: &str) -> String { @@ -92,9 +94,17 @@ impl ODBCPool { connection_string, connection_id, pool: &ENV, + schema: None, }) } + /// Use the provided `schema` for all connections instead of inferring it from the database. + #[must_use] + pub fn with_schema(mut self, schema: SchemaRef) -> Self { + self.schema = Some(schema); + self + } + #[must_use] pub fn odbc_environment(&self) -> &'static Environment { self.pool @@ -112,6 +122,7 @@ where let pool: &'static Environment = self.pool; let connection_string = self.connection_string.clone(); let params = Arc::clone(&self.params); + let schema = self.schema.clone(); let connect = async move || -> Result< Box>, @@ -129,6 +140,7 @@ where Ok(Box::new(ODBCConnection { conn: Arc::new(cxn.into()), params, + schema, }) as Box>) }; run_async_with_tokio(connect).await diff --git a/crates/oracle/src/conn.rs b/crates/oracle/src/conn.rs index 2e40758cf..1b700d557 100644 --- a/crates/oracle/src/conn.rs +++ b/crates/oracle/src/conn.rs @@ -21,11 +21,12 @@ pub type OraclePooledConnection = bb8::PooledConnection<'static, OracleConnectio pub struct OracleConnection { pub conn: OraclePooledConnection, + pub schema: Option, } impl OracleConnection { pub fn new(conn: OraclePooledConnection) -> Self { - Self { conn } + Self { conn, schema: None } } } @@ -48,13 +49,24 @@ impl DbConnection for Orac #[async_trait] impl AsyncDbConnection for OracleConnection { fn new(conn: OraclePooledConnection) -> Self { - Self { conn } + Self { conn, schema: None } + } + + fn new_with_schema(conn: OraclePooledConnection, schema: SchemaRef) -> Self { + Self { + conn, + schema: Some(schema), + } } async fn get_schema( &self, table_reference: &TableReference, ) -> std::result::Result { + if let Some(schema) = &self.schema { + return Ok(Arc::clone(schema)); + } + let table_name = table_reference.table().to_uppercase(); let schema_name = table_reference.schema().map(|s| s.to_uppercase()); diff --git a/crates/oracle/src/pool.rs b/crates/oracle/src/pool.rs index 65122e085..d268ef3b5 100644 --- a/crates/oracle/src/pool.rs +++ b/crates/oracle/src/pool.rs @@ -6,14 +6,18 @@ use std::sync::Arc; use async_trait::async_trait; use bb8::CustomizeConnection; use bb8_oracle::OracleConnectionManager; +use datafusion::arrow::datatypes::SchemaRef; use oracle::{Connection, Connector}; use secrecy::{ExposeSecret, SecretString}; use snafu::prelude::*; use datafusion_table_providers_common::sql::db_connection_pool::{ - dbconnection::DbConnection, DbConnectionPool, JoinPushDown, + dbconnection::{AsyncDbConnection, DbConnection}, + DbConnectionPool, JoinPushDown, }; +use crate::conn::{OracleConnection, OraclePooledConnection}; + /// Default TCP port for Oracle Database const DEFAULT_ORACLE_PORT: u16 = 1521; @@ -80,6 +84,7 @@ impl CustomizeConnection, bb8_oracle::Error> for SetTimezoneCust pub struct OracleConnectionPool { pool: Arc>, + schema: Option, } impl std::fmt::Debug for OracleConnectionPool { @@ -154,16 +159,34 @@ impl OracleConnectionPool { Ok(Self { pool: Arc::new(pool), + schema: None, }) } - pub async fn connect_direct(&self) -> Result { + /// Use the provided `schema` for all connections instead of inferring it from the database. + #[must_use] + pub fn with_schema(mut self, schema: SchemaRef) -> Self { + self.schema = Some(schema); + self + } + + fn new_connection(&self, conn: OraclePooledConnection) -> OracleConnection { + match &self.schema { + Some(schema) => >::new_with_schema( + conn, + Arc::clone(schema), + ), + None => OracleConnection::new(conn), + } + } + + pub async fn connect_direct(&self) -> Result { let conn = Arc::clone(&self.pool) .get_owned() .await .map_err(|e| Error::PoolRunError { source: e })?; - Ok(crate::conn::OracleConnection::new(conn)) + Ok(self.new_connection(conn)) } } @@ -189,7 +212,7 @@ impl Box::new(Error::PoolRunError { source: e }) as Box })?; - Ok(Box::new(crate::conn::OracleConnection::new(conn))) + Ok(Box::new(self.new_connection(conn))) } fn join_push_down(&self) -> JoinPushDown { diff --git a/crates/postgres/src/conn.rs b/crates/postgres/src/conn.rs index ac3dc3179..d03af07f0 100644 --- a/crates/postgres/src/conn.rs +++ b/crates/postgres/src/conn.rs @@ -298,6 +298,7 @@ fn format_postgres_query_error(source: &bb8_postgres::tokio_postgres::Error) -> pub struct PostgresConnection { pub conn: PostgresPooledConnection, unsupported_type_action: UnsupportedTypeAction, + schema: Option, } impl SchemaValidator for PostgresConnection { @@ -339,6 +340,15 @@ impl<'a> AsyncDbConnection PostgresConnection { conn, unsupported_type_action: UnsupportedTypeAction::default(), + schema: None, + } + } + + fn new_with_schema(conn: PostgresPooledConnection, schema: SchemaRef) -> Self { + PostgresConnection { + conn, + unsupported_type_action: UnsupportedTypeAction::default(), + schema: Some(schema), } } @@ -392,6 +402,10 @@ impl<'a> AsyncDbConnection SchemaRef, datafusion_table_providers_common::sql::db_connection_pool::dbconnection::Error, > { + if let Some(schema) = &self.schema { + return Ok(Arc::clone(schema)); + } + let (variant, columns) = self.query_variant_and_schema(table_reference).await?; // Native inference can return zero rows even though the table exists and is diff --git a/crates/postgres/src/pool.rs b/crates/postgres/src/pool.rs index b1b7e3ca7..4cfa0991e 100644 --- a/crates/postgres/src/pool.rs +++ b/crates/postgres/src/pool.rs @@ -1,4 +1,5 @@ -use crate::conn::PostgresConnection; +use crate::conn::{PostgresConnection, PostgresPooledConnection}; +use arrow::datatypes::SchemaRef; use std::{collections::HashMap, path::PathBuf, str::FromStr, sync::Arc}; use async_trait::async_trait; @@ -231,6 +232,7 @@ pub struct PostgresConnectionPool { join_push_down: JoinPushDown, unsupported_type_action: UnsupportedTypeAction, io_handle: Option, + schema: Option, } impl PostgresConnectionPool { @@ -422,6 +424,7 @@ impl PostgresConnectionPool { join_push_down, unsupported_type_action: UnsupportedTypeAction::default(), io_handle: None, + schema: None, }) } @@ -439,6 +442,20 @@ impl PostgresConnectionPool { self } + /// Use the provided `schema` for all connections instead of inferring it from the database. + #[must_use] + pub fn with_schema(mut self, schema: SchemaRef) -> Self { + self.schema = Some(schema); + self + } + + fn new_connection(&self, conn: PostgresPooledConnection) -> PostgresConnection { + match &self.schema { + Some(schema) => PostgresConnection::new_with_schema(conn, Arc::clone(schema)), + None => PostgresConnection::new(conn), + } + } + /// Returns a direct connection to the underlying database. /// /// # Errors @@ -454,7 +471,7 @@ impl PostgresConnectionPool { } else { pool.get_owned().await.map_err(map_pool_run_error)? }; - Ok(PostgresConnection::new(conn)) + Ok(self.new_connection(conn)) } } @@ -666,7 +683,7 @@ impl run_async_with_tokio(get_conn).await? }; Ok(Box::new( - PostgresConnection::new(conn) + self.new_connection(conn) .with_unsupported_type_action(self.unsupported_type_action), )) } diff --git a/crates/sqlite/src/conn.rs b/crates/sqlite/src/conn.rs index 52c65aa74..1198791ef 100644 --- a/crates/sqlite/src/conn.rs +++ b/crates/sqlite/src/conn.rs @@ -1,4 +1,4 @@ -use std::any::Any; +use std::{any::Any, sync::Arc}; use crate::arrow_sql_gen::rows_to_arrow; use arrow::datatypes::SchemaRef; @@ -33,6 +33,7 @@ pub enum Error { pub struct SqliteConnection { pub conn: Connection, + pub schema: Option, } impl SchemaValidator for SqliteConnection { @@ -86,7 +87,14 @@ impl DbConnection for SqliteConnection #[async_trait] impl AsyncDbConnection for SqliteConnection { fn new(conn: Connection) -> Self { - SqliteConnection { conn } + SqliteConnection { conn, schema: None } + } + + fn new_with_schema(conn: Connection, schema: SchemaRef) -> Self { + SqliteConnection { + conn, + schema: Some(schema), + } } async fn tables( @@ -129,6 +137,10 @@ impl AsyncDbConnection for SqliteConnec SchemaRef, datafusion_table_providers_common::sql::db_connection_pool::dbconnection::Error, > { + if let Some(schema) = &self.schema { + return Ok(Arc::clone(schema)); + } + let table_reference = table_reference.to_quoted_string(); let schema: SchemaRef = self .conn diff --git a/crates/sqlite/src/pool.rs b/crates/sqlite/src/pool.rs index 2d36168c7..be83c3966 100644 --- a/crates/sqlite/src/pool.rs +++ b/crates/sqlite/src/pool.rs @@ -2,6 +2,7 @@ use crate::conn::SqliteConnection; use std::{sync::Arc, time::Duration}; use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; use snafu::prelude::*; use tokio_rusqlite::{Connection, ToSql}; @@ -29,6 +30,7 @@ pub struct SqliteConnectionPoolFactory { mode: Mode, attach_databases: Option>>, busy_timeout: Duration, + schema: Option, } impl SqliteConnectionPoolFactory { @@ -38,9 +40,17 @@ impl SqliteConnectionPoolFactory { mode, attach_databases: None, busy_timeout, + schema: None, } } + /// Use the provided `schema` for all connections instead of inferring it from the table data. + #[must_use] + pub fn with_schema(mut self, schema: SchemaRef) -> Self { + self.schema = Some(schema); + self + } + #[must_use] pub fn with_databases(mut self, attach_databases: Option>>) -> Self { self.attach_databases = attach_databases; @@ -84,7 +94,7 @@ impl SqliteConnectionPoolFactory { vec![] }; - let pool = SqliteConnectionPool::new( + let mut pool = SqliteConnectionPool::new( &self.path, self.mode, join_push_down, @@ -92,6 +102,7 @@ impl SqliteConnectionPoolFactory { self.busy_timeout, ) .await?; + pool.schema = self.schema.clone(); pool.setup().await?; @@ -107,6 +118,7 @@ pub struct SqliteConnectionPool { path: Arc, attach_databases: Vec>, busy_timeout: Duration, + schema: Option, } impl SqliteConnectionPool { @@ -143,6 +155,7 @@ impl SqliteConnectionPool { attach_databases, path: path.into(), busy_timeout, + schema: None, }) } @@ -233,7 +246,15 @@ impl SqliteConnectionPool { #[must_use] pub fn connect_sync(&self) -> Box> { - Box::new(SqliteConnection::new(self.conn.clone())) + Box::new(self.new_connection()) + } + + fn new_connection(&self) -> SqliteConnection { + let conn = self.conn.clone(); + match &self.schema { + Some(schema) => SqliteConnection::new_with_schema(conn, Arc::clone(schema)), + None => SqliteConnection::new(conn), + } } /// Will attempt to clone the connection pool. This will always succeed for in-memory mode. @@ -250,6 +271,7 @@ impl SqliteConnectionPool { path: Arc::clone(&self.path), attach_databases: self.attach_databases.clone(), busy_timeout: self.busy_timeout, + schema: self.schema.clone(), }), Mode::File => { let attach_databases = if self.attach_databases.is_empty() { @@ -258,10 +280,11 @@ impl SqliteConnectionPool { Some(self.attach_databases.clone()) }; - SqliteConnectionPoolFactory::new(&self.path, self.mode, self.busy_timeout) - .with_databases(attach_databases) - .build() - .await + let mut factory = + SqliteConnectionPoolFactory::new(&self.path, self.mode, self.busy_timeout) + .with_databases(attach_databases); + factory.schema = self.schema.clone(); + factory.build().await } } } @@ -272,9 +295,7 @@ impl DbConnectionPool for SqliteConnect async fn connect( &self, ) -> PoolResult>> { - let conn = self.conn.clone(); - - Ok(Box::new(SqliteConnection::new(conn))) + Ok(Box::new(self.new_connection())) } fn join_push_down(&self) -> JoinPushDown { @@ -416,4 +437,59 @@ mod tests { db_names[1] ))); } + + #[tokio::test] + async fn test_sqlite_connection_pool_factory_with_schema() { + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::TableReference; + + let factory = SqliteConnectionPoolFactory::new( + "./test.sqlite", + Mode::Memory, + Duration::from_millis(5000), + ); + let pool = factory.build().await.unwrap(); + let conn = pool.connect().await.unwrap(); + let conn = conn.as_async().unwrap(); + conn.execute("CREATE TABLE t (id INTEGER, name TEXT)", &[]) + .await + .unwrap(); + conn.execute("INSERT INTO t VALUES (1, 'a')", &[]) + .await + .unwrap(); + + // the schema is inferred from the data when none is provided + let inferred = conn.get_schema(&TableReference::bare("t")).await.unwrap(); + assert_eq!(inferred.field(0).data_type(), &DataType::Int64); + + // the provided schema is used as-is + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::LargeUtf8, true), + ])); + let factory = SqliteConnectionPoolFactory::new( + "./test.sqlite", + Mode::Memory, + Duration::from_millis(5000), + ) + .with_schema(Arc::clone(&schema)); + let pool = factory.build().await.unwrap(); + let conn = pool.connect().await.unwrap(); + let actual = conn + .as_async() + .unwrap() + .get_schema(&TableReference::bare("t")) + .await + .unwrap(); + assert_eq!(actual, schema); + + let actual = pool + .connect_sync() + .as_async() + .unwrap() + .get_schema(&TableReference::bare("t")) + .await + .unwrap(); + assert_eq!(actual, schema); + } } From de8d00437f43e74ec8b47838409d4d294dcc019b Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Wed, 14 Oct 2026 16:03:11 +0300 Subject: [PATCH 02/27] Add IT tests for the engines requiring Docker --- core/tests/clickhouse/mod.rs | 93 ++++++++++++++++++++++++++++++++++++ core/tests/integration.rs | 2 + core/tests/mysql/mod.rs | 77 +++++++++++++++++++++++++++++ core/tests/odbc/mod.rs | 88 ++++++++++++++++++++++++++++++++++ core/tests/oracle/mod.rs | 80 +++++++++++++++++++++++++++++++ core/tests/postgres/mod.rs | 75 +++++++++++++++++++++++++++++ crates/sqlite/src/pool.rs | 26 +++++++--- 7 files changed, 435 insertions(+), 6 deletions(-) create mode 100644 core/tests/odbc/mod.rs diff --git a/core/tests/clickhouse/mod.rs b/core/tests/clickhouse/mod.rs index c8d75bfc0..c3323900c 100644 --- a/core/tests/clickhouse/mod.rs +++ b/core/tests/clickhouse/mod.rs @@ -150,3 +150,96 @@ async fn clickhouse_insert_and_read() { assert_eq!(record_batch[0].num_rows(), 2); assert_eq!(record_batch[0].num_columns(), 9); } + +/// A pool created with `with_schema` must return the provided schema from `get_schema` +/// instead of inferring it from the database. +#[tokio::test] +async fn clickhouse_connection_pool_with_schema() { + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion_table_providers::sql::db_connection_pool::{ + dbconnection::AsyncDbConnection, DbConnectionPool, + }; + use std::sync::Arc; + + let port = { + let mut guard = CONTAINER_MANAGER_INSTANCE.lock().unwrap(); + let container_manager = guard.as_mut().unwrap(); + container_manager + .start_container(common::start_clickhouse_docker_container) + .await; + container_manager.port + }; + + let table_name = "with_schema_table"; + let table_reference = TableReference::bare(table_name); + let pool = ClickHouseConnectionPool::new(get_clickhouse_params(port)) + .await + .unwrap(); + pool.client() + .query(&format!( + "CREATE TABLE IF NOT EXISTS {table_name} (id Int64, name String) ENGINE = MergeTree() ORDER BY id" + )) + .execute() + .await + .unwrap(); + pool.client() + .query(&format!( + "INSERT INTO {table_name} VALUES (1, NULL), (2, 'a')" + )) + .execute() + .await + .unwrap(); + + // the schema is inferred from the database when none is provided + let inferred = pool + .connect_direct() + .get_schema(&table_reference) + .await + .expect("schema should be inferred"); + assert_eq!( + inferred.fields.len(), + 2, + "There should be 2 fields in the inferred table" + ); + let field1 = inferred.field(0); + assert_eq!(field1.name(), "id"); + assert_eq!(field1.data_type(), &DataType::Int64); + assert_eq!(field1.is_nullable(), false); + let field2 = inferred.field(1); + assert_eq!(field2.name(), "name"); + assert_eq!(field2.data_type(), &DataType::Utf8); // ClickHouseConnection::get_schema() infers it correctly + assert_eq!(field2.is_nullable(), false); + + // the provided schema is used as-is + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + ])); + let pool = ClickHouseConnectionPool::new(get_clickhouse_params(port)) + .await + .unwrap() + .with_schema(Arc::clone(&schema)); + + let actual = pool + .connect_direct() + .get_schema(&table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, schema); + + let conn = pool.connect().await.expect("connect should work"); + let actual = conn + .as_async() + .expect("should be async connection") + .get_schema(&table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, schema); + + // the table factory uses the pool's schema too + let table_provider = ClickHouseTableFactory::new(pool) + .table_provider(table_reference, None) + .await + .expect("table provider should be created"); + assert_eq!(table_provider.schema(), schema); +} diff --git a/core/tests/integration.rs b/core/tests/integration.rs index 889c6bae1..067067a24 100644 --- a/core/tests/integration.rs +++ b/core/tests/integration.rs @@ -20,6 +20,8 @@ mod flight; mod mongodb; #[cfg(feature = "mysql")] mod mysql; +#[cfg(feature = "odbc")] +mod odbc; #[cfg(feature = "oracle")] mod oracle; #[cfg(feature = "postgres")] diff --git a/core/tests/mysql/mod.rs b/core/tests/mysql/mod.rs index 9992b5d20..79ff224f2 100644 --- a/core/tests/mysql/mod.rs +++ b/core/tests/mysql/mod.rs @@ -1000,4 +1000,81 @@ async fn test_mysql_arrow_oneway() { test_mysql_zero_date_type(port).await; test_mysql_projected_schema_column_reorder(port).await; test_mysql_sort_limit(port).await; + test_mysql_connection_pool_with_schema(port).await; +} + +/// A pool created with `with_schema` must return the provided schema from `get_schema` +/// instead of inferring it from the database. +async fn test_mysql_connection_pool_with_schema(port: u16) { + use datafusion::common::TableReference; + + let table_reference = TableReference::bare("with_schema_table"); + + let pool = common::get_mysql_connection_pool(port) + .await + .expect("MySQL connection pool should be created"); + let conn = pool + .connect_direct() + .await + .expect("Connection should be established"); + conn.execute( + "CREATE TABLE with_schema_table (id INT PRIMARY KEY, name VARCHAR(50))", + &[], + ) + .await + .expect("MySQL table should be created"); + conn.execute("INSERT INTO with_schema_table VALUES (1, NULL), (2, 'a')", &[]) + .await + .expect("MySQL table data should be inserted"); + + // the schema is inferred from the database when none is provided + let inferred = conn + .get_schema(&table_reference) + .await + .expect("schema should be inferred"); + assert_eq!( + inferred.fields.len(), + 2, + "There should be 2 fields in the inferred table" + ); + let field1 = inferred.field(0); + assert_eq!(field1.name(), "id"); + assert_eq!(field1.data_type(), &DataType::Int32); + assert_eq!(field1.is_nullable(), true); + let field2 = inferred.field(1); + assert_eq!(field2.name(), "name"); + assert_eq!(field2.data_type(), &DataType::Utf8); // MySQLConnection::get_schema() infers it correctly + assert_eq!(field2.is_nullable(), true); + + // the provided schema is used as-is + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, true), + ])); + let pool = common::get_mysql_connection_pool(port) + .await + .expect("MySQL connection pool should be created") + .with_schema(Arc::clone(&schema)); + + let direct_conn = pool + .connect_direct() + .await + .expect("Connection should be established"); + let actual = direct_conn + .get_schema(&table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, schema); + + let conn = pool + .connect() + .await + .expect("Connection should be established"); + let actual = conn + .as_async() + .expect("should be async connection") + .get_schema(&table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, schema); } diff --git a/core/tests/odbc/mod.rs b/core/tests/odbc/mod.rs new file mode 100644 index 000000000..775bd6c1b --- /dev/null +++ b/core/tests/odbc/mod.rs @@ -0,0 +1,88 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::common::TableReference; +use datafusion_table_providers::sql::db_connection_pool::{odbcpool::ODBCPool, DbConnectionPool}; +use datafusion_table_providers::util::secrets::to_secret_map; + +/// The SQLite ODBC driver to use: a driver name registered in `odbcinst.ini` +/// or a path to the driver library. +fn sqlite_odbc_driver() -> String { + std::env::var("ODBC_SQLITE_DRIVER").unwrap_or_else(|_| "SQLite3".to_string()) +} + +fn get_odbc_pool(db_path: &str) -> ODBCPool { + let params = to_secret_map(HashMap::from([( + "connection_string".to_string(), + format!("driver={};database={db_path};", sqlite_odbc_driver()), + )])); + ODBCPool::new(params).expect("ODBC pool should be created") +} + +/// A pool created with `with_schema` must return the provided schema from `get_schema` +/// instead of inferring it from the database. +#[tokio::test] +async fn test_odbc_connection_pool_with_schema() { + let dir = tempfile::tempdir().expect("temp dir should be created"); + let db_path = dir.path().join("with_schema.db"); + let db_path = db_path.to_str().expect("path should be valid UTF-8"); + let table_reference = TableReference::bare("with_schema_table"); + + let pool = get_odbc_pool(db_path); + let conn = match pool.connect().await { + Ok(conn) => conn, + Err(e) => { + // The suite skips rather than fails when no SQLite ODBC driver is installed. + eprintln!( + "Skipping: unable to connect with ODBC driver '{}': {e}", + sqlite_odbc_driver() + ); + return; + } + }; + let conn = conn.as_async().expect("should be async connection"); + conn.execute( + "CREATE TABLE with_schema_table (id INTEGER, name TEXT)", + &[], + ) + .await + .expect("table should be created"); + conn.execute("INSERT INTO with_schema_table VALUES (1, NULL), (2, 'a')", &[]) + .await + .expect("row should be inserted"); + + // the schema is inferred from the database when none is provided + let inferred = conn + .get_schema(&table_reference) + .await + .expect("schema should be inferred"); + assert_eq!( + inferred.fields.len(), + 2, + "There should be 2 fields in the inferred table" + ); + let field1 = inferred.field(0); + assert_eq!(field1.name(), "id"); + assert_eq!(field1.data_type(), &DataType::Int32); + assert_eq!(field1.is_nullable(), true); + let field2 = inferred.field(1); + assert_eq!(field2.name(), "name"); + assert_eq!(field2.data_type(), &DataType::Utf8); // ODBCConnection::get_schema() infers it correctly + assert_eq!(field2.is_nullable(), true); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int8, false), + Field::new("name", DataType::LargeUtf8, true), + ])); + // the provided schema is used as-is + let pool = get_odbc_pool(db_path).with_schema(Arc::clone(&schema)); + let conn = pool.connect().await.expect("connect should work"); + let actual = conn + .as_async() + .expect("should be async connection") + .get_schema(&table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, schema); +} diff --git a/core/tests/oracle/mod.rs b/core/tests/oracle/mod.rs index 53c317bea..9f38ac9a6 100644 --- a/core/tests/oracle/mod.rs +++ b/core/tests/oracle/mod.rs @@ -951,3 +951,83 @@ async fn test_oracle_large_query_streaming() { batch_count ); } + +/// A pool created with `with_schema` must return the provided schema from `get_schema` +/// instead of inferring it from the database. +#[tokio::test] +async fn test_oracle_connection_pool_with_schema() { + use datafusion_table_providers::sql::db_connection_pool::dbconnection::AsyncDbConnection; + + let Some(pool) = common::get_oracle_connection_pool().await else { + return; + }; + let conn = pool + .connect_direct() + .await + .expect("Failed to get connection"); + let oracle_conn = conn.conn.clone(); + + // Create table with test data + let table_name = "TEST_TABLE_WITH_SCHEMA"; + let _ = oracle_conn.execute(&format!("DROP TABLE {table_name}"), &[]); + oracle_conn + .execute( + &format!("CREATE TABLE {table_name} (id NUMBER PRIMARY KEY, name VARCHAR2(50))"), + &[], + ) + .expect("Failed to create table"); + + let sql = format!("INSERT INTO {table_name} (id, name) VALUES (1, NULL), (2, 'a')"); + oracle_conn.execute(&sql, &[]).expect("Execute failed"); + + let table_reference = TableReference::bare(table_name); + // the schema is inferred from the database when none is provided + let inferred = conn + .get_schema(&table_reference) + .await + .expect("schema should be inferred"); + assert_eq!( + inferred.fields.len(), + 2, + "There should be 2 fields in the inferred table" + ); + let field1 = inferred.field(0); + assert_eq!(field1.name(), "id"); + assert_eq!(field1.data_type(), &DataType::Int32); + assert_eq!(field1.is_nullable(), true); + let field2 = inferred.field(1); + assert_eq!(field2.name(), "name"); + assert_eq!(field2.data_type(), &DataType::Utf8); // OracleConnection::get_schema() infers it correctly + assert_eq!(field2.is_nullable(), true); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int8, false), + Field::new("name", DataType::LargeUtf8, true), + ])); + + // the provided schema is used as-is + let pool = OracleConnectionPool::new(common::get_oracle_params()) + .await + .expect("Failed to create Oracle connection pool") + .with_schema(Arc::clone(&schema)); + + let direct_conn = pool + .connect_direct() + .await + .expect("Failed to get connection"); + let actual = direct_conn + .get_schema(&table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, schema); + + let conn: Box> = + pool.connect().await.expect("Failed to get connection"); + let actual = conn + .as_async() + .expect("should be async connection") + .get_schema(&table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, schema); +} diff --git a/core/tests/postgres/mod.rs b/core/tests/postgres/mod.rs index 02cca9dc7..65fb7e2d9 100644 --- a/core/tests/postgres/mod.rs +++ b/core/tests/postgres/mod.rs @@ -179,6 +179,81 @@ async fn test_arrow_postgres_one_way() { test_postgres_composite_array_list_struct(port).await; test_postgres_sort_limit(port).await; test_postgres_unconstrained_numeric_precision(port).await; + test_postgres_connection_pool_with_schema(port).await; +} + +/// A pool created with `with_schema` must return the provided schema from `get_schema` +/// instead of inferring it from the database. +async fn test_postgres_connection_pool_with_schema(port: u16) { + use datafusion::common::TableReference; + use datafusion_table_providers::sql::db_connection_pool::dbconnection::AsyncDbConnection; + + let table_reference = TableReference::bare("with_schema_table"); + + let pool = common::get_postgres_connection_pool(port) + .await + .expect("Postgres connection pool should be created"); + let conn = pool.connect_direct().await.expect("connect should work"); + conn.conn + .execute( + "CREATE TABLE with_schema_table (id INTEGER, name TEXT)", + &[], + ) + .await + .expect("table should be created"); + conn.conn + .execute("INSERT INTO with_schema_table VALUES (1, NULL), (2, 'a')", &[]) + .await + .expect("row should be inserted"); + + // the schema is inferred from the database when none is provided + let sqltable_pool: Arc = Arc::new(pool); + let conn = sqltable_pool.connect().await.expect("connect should work"); + let inferred = conn + .as_async() + .expect("should be async connection") + .get_schema(&table_reference) + .await + .expect("schema should be inferred"); + assert_eq!( + inferred.fields.len(), + 2, + "There should be 2 fields in the inferred table" + ); + let field1 = inferred.field(0); + assert_eq!(field1.name(), "id"); + assert_eq!(field1.data_type(), &DataType::Int32); + assert_eq!(field1.is_nullable(), true); + let field2 = inferred.field(1); + assert_eq!(field2.name(), "name"); + assert_eq!(field2.data_type(), &DataType::Utf8); // PostgresConnection::get_schema() infers it correctly + assert_eq!(field2.is_nullable(), true); + + // the provided schema is used as-is + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::LargeUtf8, true), + ])); + let pool = common::get_postgres_connection_pool(port) + .await + .expect("Postgres connection pool should be created") + .with_schema(Arc::clone(&schema)); + + let direct_conn = pool.connect_direct().await.expect("connect should work"); + let actual = AsyncDbConnection::get_schema(&direct_conn, &table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, schema); + + let sqltable_pool: Arc = Arc::new(pool); + let conn = sqltable_pool.connect().await.expect("connect should work"); + let actual = conn + .as_async() + .expect("should be async connection") + .get_schema(&table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, schema); } /// An unconstrained `numeric` column (what `max()`, `avg()` and arithmetic over `numeric` diff --git a/crates/sqlite/src/pool.rs b/crates/sqlite/src/pool.rs index be83c3966..44f3f6984 100644 --- a/crates/sqlite/src/pool.rs +++ b/crates/sqlite/src/pool.rs @@ -451,21 +451,33 @@ mod tests { let pool = factory.build().await.unwrap(); let conn = pool.connect().await.unwrap(); let conn = conn.as_async().unwrap(); - conn.execute("CREATE TABLE t (id INTEGER, name TEXT)", &[]) + conn.execute("CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT)", &[]) .await .unwrap(); - conn.execute("INSERT INTO t VALUES (1, 'a')", &[]) + conn.execute("INSERT INTO t VALUES (1, NULL), (2, 'a')", &[]) .await .unwrap(); // the schema is inferred from the data when none is provided let inferred = conn.get_schema(&TableReference::bare("t")).await.unwrap(); - assert_eq!(inferred.field(0).data_type(), &DataType::Int64); + assert_eq!( + inferred.fields.len(), + 2, + "There should be 2 fields in the inferred table" + ); + let field1 = inferred.field(0); + assert_eq!(field1.name(), "id"); + assert_eq!(field1.data_type(), &DataType::Int64); + assert_eq!(field1.is_nullable(), true); + let field2 = inferred.field(1); + assert_eq!(field2.name(), "name"); + assert_eq!(field2.data_type(), &DataType::Null); + assert_eq!(field2.is_nullable(), true); // the provided schema is used as-is let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, false), - Field::new("name", DataType::LargeUtf8, true), + Field::new("name", DataType::Utf8, true), ])); let factory = SqliteConnectionPoolFactory::new( "./test.sqlite", @@ -473,12 +485,14 @@ mod tests { Duration::from_millis(5000), ) .with_schema(Arc::clone(&schema)); + + let table_reference = TableReference::bare("t"); let pool = factory.build().await.unwrap(); let conn = pool.connect().await.unwrap(); let actual = conn .as_async() .unwrap() - .get_schema(&TableReference::bare("t")) + .get_schema(&table_reference) .await .unwrap(); assert_eq!(actual, schema); @@ -487,7 +501,7 @@ mod tests { .connect_sync() .as_async() .unwrap() - .get_schema(&TableReference::bare("t")) + .get_schema(&table_reference) .await .unwrap(); assert_eq!(actual, schema); From e5ed94e6007491b2cdcafca4cc887309ff6b0ac0 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Wed, 14 Oct 2026 17:13:39 +0300 Subject: [PATCH 03/27] fmt --- core/tests/mysql/mod.rs | 9 ++++++--- core/tests/odbc/mod.rs | 9 ++++++--- core/tests/postgres/mod.rs | 5 ++++- 3 files changed, 16 insertions(+), 7 deletions(-) diff --git a/core/tests/mysql/mod.rs b/core/tests/mysql/mod.rs index 79ff224f2..9f2325e9f 100644 --- a/core/tests/mysql/mod.rs +++ b/core/tests/mysql/mod.rs @@ -1023,9 +1023,12 @@ async fn test_mysql_connection_pool_with_schema(port: u16) { ) .await .expect("MySQL table should be created"); - conn.execute("INSERT INTO with_schema_table VALUES (1, NULL), (2, 'a')", &[]) - .await - .expect("MySQL table data should be inserted"); + conn.execute( + "INSERT INTO with_schema_table VALUES (1, NULL), (2, 'a')", + &[], + ) + .await + .expect("MySQL table data should be inserted"); // the schema is inferred from the database when none is provided let inferred = conn diff --git a/core/tests/odbc/mod.rs b/core/tests/odbc/mod.rs index 775bd6c1b..c02e14a58 100644 --- a/core/tests/odbc/mod.rs +++ b/core/tests/odbc/mod.rs @@ -48,9 +48,12 @@ async fn test_odbc_connection_pool_with_schema() { ) .await .expect("table should be created"); - conn.execute("INSERT INTO with_schema_table VALUES (1, NULL), (2, 'a')", &[]) - .await - .expect("row should be inserted"); + conn.execute( + "INSERT INTO with_schema_table VALUES (1, NULL), (2, 'a')", + &[], + ) + .await + .expect("row should be inserted"); // the schema is inferred from the database when none is provided let inferred = conn diff --git a/core/tests/postgres/mod.rs b/core/tests/postgres/mod.rs index 65fb7e2d9..9b983a785 100644 --- a/core/tests/postgres/mod.rs +++ b/core/tests/postgres/mod.rs @@ -202,7 +202,10 @@ async fn test_postgres_connection_pool_with_schema(port: u16) { .await .expect("table should be created"); conn.conn - .execute("INSERT INTO with_schema_table VALUES (1, NULL), (2, 'a')", &[]) + .execute( + "INSERT INTO with_schema_table VALUES (1, NULL), (2, 'a')", + &[], + ) .await .expect("row should be inserted"); From 670d963573ba807bd7be13e4fb8b5701a4ffa994 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Wed, 14 Oct 2026 17:21:11 +0300 Subject: [PATCH 04/27] clippy --- crates/common/src/sql/sql_provider_datafusion/mod.rs | 1 - crates/sqlite/src/pool.rs | 4 ++-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/crates/common/src/sql/sql_provider_datafusion/mod.rs b/crates/common/src/sql/sql_provider_datafusion/mod.rs index 839e52c9e..6644a655d 100644 --- a/crates/common/src/sql/sql_provider_datafusion/mod.rs +++ b/crates/common/src/sql/sql_provider_datafusion/mod.rs @@ -738,7 +738,6 @@ mod tests { use std::{error::Error, sync::Arc}; use datafusion::common::TableReference; - use datafusion::execution::context::SessionContext; use tracing::{level_filters::LevelFilter, subscriber::DefaultGuard, Dispatch}; use crate::sql::sql_provider_datafusion::SqlTable; diff --git a/crates/sqlite/src/pool.rs b/crates/sqlite/src/pool.rs index 44f3f6984..0acdd941e 100644 --- a/crates/sqlite/src/pool.rs +++ b/crates/sqlite/src/pool.rs @@ -468,11 +468,11 @@ mod tests { let field1 = inferred.field(0); assert_eq!(field1.name(), "id"); assert_eq!(field1.data_type(), &DataType::Int64); - assert_eq!(field1.is_nullable(), true); + assert!(field1.is_nullable()); let field2 = inferred.field(1); assert_eq!(field2.name(), "name"); assert_eq!(field2.data_type(), &DataType::Null); - assert_eq!(field2.is_nullable(), true); + assert!(field2.is_nullable()); // the provided schema is used as-is let schema = Arc::new(Schema::new(vec![ From dd33ffca189d1e277eaf485d17ec1dcd7bbc0f73 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Thu, 15 Oct 2026 09:34:45 +0300 Subject: [PATCH 05/27] uppercase the inferred column names for Oracle --- core/tests/oracle/mod.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/core/tests/oracle/mod.rs b/core/tests/oracle/mod.rs index 9f38ac9a6..7eb25f68e 100644 --- a/core/tests/oracle/mod.rs +++ b/core/tests/oracle/mod.rs @@ -969,7 +969,7 @@ async fn test_oracle_connection_pool_with_schema() { // Create table with test data let table_name = "TEST_TABLE_WITH_SCHEMA"; - let _ = oracle_conn.execute(&format!("DROP TABLE {table_name}"), &[]); + let _ = oracle_conn.execute(&format!("DROP TABLE IF EXISTS {table_name}"), &[]); oracle_conn .execute( &format!("CREATE TABLE {table_name} (id NUMBER PRIMARY KEY, name VARCHAR2(50))"), @@ -992,11 +992,11 @@ async fn test_oracle_connection_pool_with_schema() { "There should be 2 fields in the inferred table" ); let field1 = inferred.field(0); - assert_eq!(field1.name(), "id"); + assert_eq!(field1.name(), "ID"); assert_eq!(field1.data_type(), &DataType::Int32); assert_eq!(field1.is_nullable(), true); let field2 = inferred.field(1); - assert_eq!(field2.name(), "name"); + assert_eq!(field2.name(), "NAME"); assert_eq!(field2.data_type(), &DataType::Utf8); // OracleConnection::get_schema() infers it correctly assert_eq!(field2.is_nullable(), true); From af01f374a1e4bde48d633ef9e230a449203a61dc Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Thu, 15 Oct 2026 09:52:59 +0300 Subject: [PATCH 06/27] Use LONG for the `id` column in Oracle --- core/tests/oracle/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core/tests/oracle/mod.rs b/core/tests/oracle/mod.rs index 7eb25f68e..802b2cd21 100644 --- a/core/tests/oracle/mod.rs +++ b/core/tests/oracle/mod.rs @@ -972,7 +972,7 @@ async fn test_oracle_connection_pool_with_schema() { let _ = oracle_conn.execute(&format!("DROP TABLE IF EXISTS {table_name}"), &[]); oracle_conn .execute( - &format!("CREATE TABLE {table_name} (id NUMBER PRIMARY KEY, name VARCHAR2(50))"), + &format!("CREATE TABLE {table_name} (id LONG PRIMARY KEY, name VARCHAR2(50))"), &[], ) .expect("Failed to create table"); @@ -993,7 +993,7 @@ async fn test_oracle_connection_pool_with_schema() { ); let field1 = inferred.field(0); assert_eq!(field1.name(), "ID"); - assert_eq!(field1.data_type(), &DataType::Int32); + assert_eq!(field1.data_type(), &DataType::Int64); assert_eq!(field1.is_nullable(), true); let field2 = inferred.field(1); assert_eq!(field2.name(), "NAME"); From 11b721d9719f3fcda430ab5867b72dc5af5fc98d Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Thu, 15 Oct 2026 11:14:20 +0300 Subject: [PATCH 07/27] Start an embedded Docker container for the Oracle IT tests --- .github/workflows/pr.yaml | 35 +------- core/tests/integration.rs | 1 + core/tests/oracle/common.rs | 165 ++++++++++++++++++------------------ core/tests/oracle/mod.rs | 93 ++++++++++++-------- 4 files changed, 143 insertions(+), 151 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index bb4e58816..cb703ed83 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -201,7 +201,9 @@ jobs: timeout-minutes: 45 env: - ORACLE_DOCKER_IMAGE: docker.io/gvenzl/oracle-free:latest + # Must match the tag the integration tests start (they pull it only if + # an image with exactly this tag is missing). + ORACLE_DOCKER_IMAGE: gvenzl/oracle-free:latest RUST_LOG: info steps: @@ -244,37 +246,6 @@ jobs: docker save -o /tmp/oracle-image/oracle.tar "${ORACLE_DOCKER_IMAGE}" fi - # Build the test binary first so the Rust compile does not eat into the - # Oracle container's boot/health window below or the job timeout. - - name: Build integration test binary - run: cargo build --tests -p datafusion-table-providers --no-default-features --features oracle - - - name: Start Oracle container - run: | - docker run -d --name oracle-it \ - -e ORACLE_PASSWORD=password \ - -p 1521:1521 \ - "${ORACLE_DOCKER_IMAGE}" - echo "Waiting for Oracle to become ready (cold boot can take ~10 minutes on CI runners)..." - # gvenzl/oracle-free ships no Docker HEALTHCHECK, so readiness is - # detected via the image's own "DATABASE IS READY TO USE!" marker - # plus a successful TCP probe of the published listener port. - for i in $(seq 1 360); do - if docker logs oracle-it 2>&1 | grep -q "DATABASE IS READY TO USE!" \ - && (exec 3<>/dev/tcp/127.0.0.1/1521) 2>/dev/null; then - echo "Oracle listener is up on 1521" - exit 0 - fi - if docker inspect --format '{{.State.Status}}' oracle-it | grep -q exited; then - docker logs oracle-it - exit 1 - fi - sleep 5 - done - echo "Oracle never became ready" >&2 - docker logs oracle-it >&2 || true - exit 1 - - name: Run Oracle integration tests run: cargo test -p datafusion-table-providers --test integration --no-default-features --features oracle -- --test-threads=1 --nocapture diff --git a/core/tests/integration.rs b/core/tests/integration.rs index 067067a24..e59622079 100644 --- a/core/tests/integration.rs +++ b/core/tests/integration.rs @@ -9,6 +9,7 @@ mod clickhouse; feature = "clickhouse", feature = "mongodb", feature = "mysql", + feature = "oracle", feature = "postgres" ))] mod docker; diff --git a/core/tests/oracle/common.rs b/core/tests/oracle/common.rs index 19efd2284..98ace4d40 100644 --- a/core/tests/oracle/common.rs +++ b/core/tests/oracle/common.rs @@ -1,112 +1,113 @@ +use bollard::models::HealthConfig; use datafusion_table_providers::sql::db_connection_pool::oraclepool::OracleConnectionPool; use secrecy::SecretString; +use tracing::instrument; use std::collections::HashMap; -use std::env; -use std::net::ToSocketAddrs; use std::sync::Arc; use std::time::Duration; -const ORACLE_PASSWORD: &str = "password"; +use crate::docker::{ContainerRunnerBuilder, RunningContainer}; + +const ORACLE_PASSWORD: &str = "integration-test-pw"; const ORACLE_USER: &str = "system"; const ORACLE_SERVICE: &str = "FREEPDB1"; -const DEFAULT_ORACLE_PORT: u16 = 1521; +const ORACLE_DOCKER_CONTAINER: &str = "runtime-integration-test-oracle"; -pub fn get_oracle_params() -> HashMap { +pub(super) fn get_oracle_params(port: u16) -> HashMap { let mut params = HashMap::new(); - - // Default to strict env vars or defaults - let host = env::var("ORACLE_HOST").unwrap_or_else(|_| "localhost".to_string()); - let port = env::var("ORACLE_PORT").unwrap_or_else(|_| DEFAULT_ORACLE_PORT.to_string()); - let user = env::var("ORACLE_USER").unwrap_or_else(|_| ORACLE_USER.to_string()); - let pass = env::var("ORACLE_PASSWORD").unwrap_or_else(|_| ORACLE_PASSWORD.to_string()); - let service = env::var("ORACLE_SERVICE").unwrap_or_else(|_| ORACLE_SERVICE.to_string()); - - params.insert("host".to_string(), SecretString::from(host)); - params.insert("port".to_string(), SecretString::from(port)); - params.insert("user".to_string(), SecretString::from(user)); - params.insert("password".to_string(), SecretString::from(pass)); - params.insert("service_name".to_string(), SecretString::from(service)); - - // Optional wallet params - if let Ok(wallet) = env::var("ORACLE_WALLET_PATH") { - params.insert("wallet_path".to_string(), SecretString::from(wallet)); - } - if let Ok(wpass) = env::var("ORACLE_WALLET_PASSWORD") { - params.insert("wallet_password".to_string(), SecretString::from(wpass)); - } - + params.insert( + "host".to_string(), + SecretString::from("localhost".to_string()), + ); + params.insert("port".to_string(), SecretString::from(port.to_string())); + params.insert( + "user".to_string(), + SecretString::from(ORACLE_USER.to_string()), + ); + params.insert( + "password".to_string(), + SecretString::from(ORACLE_PASSWORD.to_string()), + ); + params.insert( + "service_name".to_string(), + SecretString::from(ORACLE_SERVICE.to_string()), + ); params } -/// Probes the configured Oracle listener with a short TCP connect. -/// -/// Returns `false` when no Oracle server is listening, so tests can skip -/// instead of failing in environments (like most CI jobs) that have no -/// Oracle instance available. -async fn listener_reachable() -> bool { - let host = env::var("ORACLE_HOST").unwrap_or_else(|_| "localhost".to_string()); - let port: u16 = env::var("ORACLE_PORT") - .ok() - .and_then(|p| p.parse().ok()) - .unwrap_or(DEFAULT_ORACLE_PORT); +#[instrument] +pub(super) async fn start_oracle_docker_container( + port: u16, +) -> Result { + let container_name = format!("{ORACLE_DOCKER_CONTAINER}-{port}"); + + // gvenzl/oracle-free is not mirrored in the default (ECR public library) + // registry, so it is pulled from Docker Hub unless overridden. + let oracle_docker_image = std::env::var("ORACLE_DOCKER_IMAGE") + .unwrap_or_else(|_| "gvenzl/oracle-free:latest".to_string()); - let addrs = match (host.as_str(), port).to_socket_addrs() { - Ok(addrs) => addrs.collect::>(), - Err(_) => return false, - }; + let running_container = ContainerRunnerBuilder::new(container_name) + .image(oracle_docker_image) + .add_port_binding(1521, port) + .add_env_var("ORACLE_PASSWORD", ORACLE_PASSWORD) + .healthcheck(HealthConfig { + // The image ships `healthcheck.sh` (on its PATH) but no Docker + // HEALTHCHECK; it succeeds once FREEPDB1 is open for READ WRITE. + test: Some(vec!["CMD".to_string(), "healthcheck.sh".to_string()]), + interval: Some(5_000_000_000), // 5s + timeout: Some(5_000_000_000), // 5s + retries: Some(10), + start_period: Some(30_000_000_000), // 30s + start_interval: Some(1_000_000_000), // 1s + }) + // A cold boot can take several minutes on CI runners. + .health_timeout(Duration::from_secs(900)) + .build()? + .run() + .await?; - for addr in addrs { - if tokio::time::timeout(Duration::from_secs(3), tokio::net::TcpStream::connect(addr)) - .await - .map(|r| r.is_ok()) - .unwrap_or(false) - { - return true; - } - } - false + Ok(running_container) } -/// Waits for the configured Oracle listener to accept TCP connections. +/// Starts the Oracle container shared by all Oracle tests on first use and +/// returns its host port. /// -/// The CI job starts gvenzl/oracle-free as its own step (with its own -/// health-wait loop) before invoking the tests, so here we only wait. Keeping -/// container startup out of the test process means the slow Oracle boot runs -/// against the step's own budget instead of the first test's. -async fn wait_for_listener(timeout: Duration) -> bool { - let deadline = tokio::time::Instant::now() + timeout; - loop { - if listener_reachable().await { - return true; - } - if tokio::time::Instant::now() >= deadline { - eprintln!("Oracle listener not reachable after {timeout:?}; skipping Oracle integration tests"); - return false; - } - tokio::time::sleep(Duration::from_secs(5)).await; - } +/// Memoized separately from the pool so that a failed pool creation does not +/// start (and leak) another container on the next test's retry. +async fn ensure_oracle_container() -> u16 { + static PORT: tokio::sync::OnceCell = tokio::sync::OnceCell::const_new(); + + *PORT + .get_or_init(|| async { + let port = super::container_port(); + let running_container = start_oracle_docker_container(port) + .await + .expect("Oracle Docker container to start"); + super::register_container(running_container); + port + }) + .await } -/// Creates the shared test connection pool, or `None` when no Oracle is -/// available (the Oracle suite then skips rather than fails). +/// Returns the connection pool shared by all Oracle tests, starting the +/// Oracle container on first use. /// -/// The whole attempt is memoized: with `--test-threads=1` every test calls -/// this, and only the first may pay the availability wait. -pub async fn get_oracle_connection_pool() -> Option> { - static POOL: tokio::sync::OnceCell>> = +/// Each `#[tokio::test]` runs in its own runtime, so the pool and the +/// container are memoized across tests: only the first caller pays for the +/// (slow) Oracle boot, the others wait for it and reuse the same container. +pub async fn get_oracle_connection_pool() -> Arc { + static POOL: tokio::sync::OnceCell> = tokio::sync::OnceCell::const_new(); POOL.get_or_init(|| async { - if !wait_for_listener(Duration::from_secs(900)).await { - return None; - } - let params = get_oracle_params(); - Some(Arc::new( - OracleConnectionPool::new(params) + let port = ensure_oracle_container().await; + + Arc::new( + OracleConnectionPool::new(get_oracle_params(port)) .await .expect("Failed to create Oracle connection pool"), - )) + ) }) .await .clone() diff --git a/core/tests/oracle/mod.rs b/core/tests/oracle/mod.rs index 802b2cd21..5b5ef44b1 100644 --- a/core/tests/oracle/mod.rs +++ b/core/tests/oracle/mod.rs @@ -12,10 +12,51 @@ use datafusion_table_providers::sql::db_connection_pool::oraclepool::OracleConne use datafusion_table_providers::sql::db_connection_pool::DbConnectionPool; use datafusion_table_providers::sql::sql_provider_datafusion::SqlTable; use futures::StreamExt; -use std::sync::Arc; +use linktime::{ctor, dtor}; +use std::sync::{Arc, Mutex}; + +use crate::docker::{ContainerManager, RunningContainer}; mod common; +/// Owns the single Oracle container shared by every test in this module; the +/// container is removed when the test binary exits. +static CONTAINER_MANAGER_INSTANCE: Mutex> = Mutex::new(None); + +#[ctor(unsafe)] +fn global_setup() { + let mut guard = CONTAINER_MANAGER_INSTANCE.lock().unwrap(); + *guard = Some(ContainerManager::default()); +} + +#[dtor(unsafe)] +fn global_teardown() { + let mut guard = CONTAINER_MANAGER_INSTANCE + .lock() + .unwrap_or_else(|e| e.into_inner()); + if let Some(container_manager) = guard.take() { + drop(container_manager); + } +} + +/// The host port the shared Oracle container is (or will be) published on. +fn container_port() -> u16 { + CONTAINER_MANAGER_INSTANCE + .lock() + .unwrap() + .as_ref() + .expect("ContainerManager to be initialized") + .port +} + +/// Hands the started container to the manager so it is cleaned up at exit. +fn register_container(running_container: RunningContainer) { + let mut guard = CONTAINER_MANAGER_INSTANCE.lock().unwrap(); + let container_manager = guard.as_mut().expect("ContainerManager to be initialized"); + container_manager.claimed = true; + container_manager.running_container = Some(running_container); +} + /// A user-owned fixture table for schema-scoped tests. /// /// Schema resolution is scoped to the connecting user's schema @@ -79,9 +120,7 @@ impl Drop for FixtureTable { #[tokio::test] async fn test_oracle_connection_pool() { - let Some(pool) = common::get_oracle_connection_pool().await else { - return; - }; + let pool = common::get_oracle_connection_pool().await; let conn = pool .connect_direct() .await @@ -108,9 +147,7 @@ async fn test_oracle_connection_pool() { /// owning schema (`CURRENT_SCHEMA`) and would not see the synonym. #[tokio::test] async fn test_oracle_table_provider_registration() { - let Some(pool) = common::get_oracle_connection_pool().await else { - return; - }; + let pool = common::get_oracle_connection_pool().await; let factory = OracleTableFactory::new(Arc::clone(&pool)); let provider = factory @@ -135,9 +172,7 @@ async fn test_oracle_table_provider_registration() { /// connecting user's schema. Validates the full Provider -> DataFusion query path. #[tokio::test] async fn test_oracle_query_with_data() { - let Some(pool) = common::get_oracle_connection_pool().await else { - return; - }; + let pool = common::get_oracle_connection_pool().await; let fixture = FixtureTable::create(pool, "QUERY_DATA_TEST").await; let factory = OracleTableFactory::new(Arc::clone(&fixture.pool)); let provider = factory @@ -169,9 +204,7 @@ async fn test_oracle_query_with_data() { /// Test schema inference for a user table (mixed types), covering nullable flags. #[tokio::test] async fn test_oracle_explain_plan() { - let Some(pool) = common::get_oracle_connection_pool().await else { - return; - }; + let pool = common::get_oracle_connection_pool().await; let fixture = FixtureTable::create(pool, "SCHEMA_TEST").await; let factory = OracleTableFactory::new(Arc::clone(&fixture.pool)); @@ -301,9 +334,7 @@ async fn insert_test_rows( async fn test_oracle_insert_and_read() { let table_name = "TEST_EMPLOYEES"; - let Some(pool) = common::get_oracle_connection_pool().await else { - return; - }; + let pool = common::get_oracle_connection_pool().await; let conn = pool .connect_direct() .await @@ -319,9 +350,7 @@ async fn test_oracle_insert_and_read() { drop(conn); drop(pool); - let Some(pool_query) = common::get_oracle_connection_pool().await else { - return; - }; + let pool_query = common::get_oracle_connection_pool().await; let factory = OracleTableFactory::new(pool_query); let ctx = SessionContext::new(); @@ -425,9 +454,7 @@ async fn arrow_oracle_one_way( insert_table_stmt: &str, expected_record: RecordBatch, ) { - let Some(pool) = common::get_oracle_connection_pool().await else { - return; - }; + let pool = common::get_oracle_connection_pool().await; let conn = pool .connect_direct() .await @@ -811,9 +838,7 @@ async fn test_oracle_null_handling() { } #[tokio::test] async fn test_oracle_sql_generation_limit() { - let Some(pool) = common::get_oracle_connection_pool().await else { - return; - }; + let pool = common::get_oracle_connection_pool().await; let fixture = FixtureTable::create(pool, "LIMIT_GEN_TEST").await; let factory = OracleTableFactory::new(Arc::clone(&fixture.pool)); @@ -846,9 +871,7 @@ async fn test_oracle_sql_generation_limit() { #[tokio::test] async fn test_oracle_sql_generation_filter() { - let Some(pool) = common::get_oracle_connection_pool().await else { - return; - }; + let pool = common::get_oracle_connection_pool().await; let fixture = FixtureTable::create(pool, "FILTER_GEN_TEST").await; let factory = OracleTableFactory::new(Arc::clone(&fixture.pool)); @@ -884,9 +907,7 @@ async fn test_oracle_sql_generation_filter() { #[tokio::test] async fn test_oracle_large_query_streaming() { - let Some(pool) = common::get_oracle_connection_pool().await else { - return; - }; + let pool = common::get_oracle_connection_pool().await; let conn = pool .connect_direct() .await @@ -958,9 +979,7 @@ async fn test_oracle_large_query_streaming() { async fn test_oracle_connection_pool_with_schema() { use datafusion_table_providers::sql::db_connection_pool::dbconnection::AsyncDbConnection; - let Some(pool) = common::get_oracle_connection_pool().await else { - return; - }; + let pool = common::get_oracle_connection_pool().await; let conn = pool .connect_direct() .await @@ -972,7 +991,7 @@ async fn test_oracle_connection_pool_with_schema() { let _ = oracle_conn.execute(&format!("DROP TABLE IF EXISTS {table_name}"), &[]); oracle_conn .execute( - &format!("CREATE TABLE {table_name} (id LONG PRIMARY KEY, name VARCHAR2(50))"), + &format!("CREATE TABLE {table_name} (id LONG, name VARCHAR2(50))"), &[], ) .expect("Failed to create table"); @@ -993,7 +1012,7 @@ async fn test_oracle_connection_pool_with_schema() { ); let field1 = inferred.field(0); assert_eq!(field1.name(), "ID"); - assert_eq!(field1.data_type(), &DataType::Int64); + assert_eq!(field1.data_type(), &DataType::Utf8); // the inferring is wrong here ! assert_eq!(field1.is_nullable(), true); let field2 = inferred.field(1); assert_eq!(field2.name(), "NAME"); @@ -1006,7 +1025,7 @@ async fn test_oracle_connection_pool_with_schema() { ])); // the provided schema is used as-is - let pool = OracleConnectionPool::new(common::get_oracle_params()) + let pool = OracleConnectionPool::new(common::get_oracle_params(container_port())) .await .expect("Failed to create Oracle connection pool") .with_schema(Arc::clone(&schema)); From 73a319455f7d0a65c2dec09ba40b6a949ce1c3a8 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Thu, 15 Oct 2026 12:09:01 +0300 Subject: [PATCH 08/27] Add a check for successful connect to Oracle and log any errors --- core/tests/oracle/common.rs | 68 +++++++++++++++++++++++++------------ 1 file changed, 46 insertions(+), 22 deletions(-) diff --git a/core/tests/oracle/common.rs b/core/tests/oracle/common.rs index 98ace4d40..dc388d749 100644 --- a/core/tests/oracle/common.rs +++ b/core/tests/oracle/common.rs @@ -70,11 +70,39 @@ pub(super) async fn start_oracle_docker_container( Ok(running_container) } -/// Starts the Oracle container shared by all Oracle tests on first use and -/// returns its host port. +/// Polls until a real login to the container succeeds, so tests never see a +/// database that is "healthy" but not yet accepting sessions (e.g. before the +/// listener has registered the PDB service). /// -/// Memoized separately from the pool so that a failed pool creation does not -/// start (and leak) another container on the next test's retry. +/// Connects directly instead of through the bb8 pool, which only reports +/// `TimedOut` and hides the Oracle error; every failed attempt is printed so CI +/// logs show the real cause. +async fn wait_until_connectable(port: u16, timeout: Duration) { + let connect_string = format!("//localhost:{port}/{ORACLE_SERVICE}"); + let deadline = tokio::time::Instant::now() + timeout; + loop { + let cs = connect_string.clone(); + let result = tokio::task::spawn_blocking(move || { + oracle::Connection::connect(ORACLE_USER, ORACLE_PASSWORD, cs) + }) + .await + .expect("Oracle connect task panicked"); + match result { + Ok(conn) => { + let _ = conn.close(); + return; + } + Err(e) if tokio::time::Instant::now() < deadline => { + eprintln!("Oracle at {connect_string} not accepting sessions yet: {e}"); + tokio::time::sleep(Duration::from_secs(2)).await; + } + Err(e) => panic!("Oracle at {connect_string} never accepted a session: {e}"), + } + } +} + +/// Starts the Oracle container shared by all Oracle tests on first use and +/// returns its host port once the database accepts logins. async fn ensure_oracle_container() -> u16 { static PORT: tokio::sync::OnceCell = tokio::sync::OnceCell::const_new(); @@ -85,30 +113,26 @@ async fn ensure_oracle_container() -> u16 { .await .expect("Oracle Docker container to start"); super::register_container(running_container); + wait_until_connectable(port, Duration::from_secs(300)).await; port }) .await } -/// Returns the connection pool shared by all Oracle tests, starting the -/// Oracle container on first use. +/// Returns a new connection pool to the Oracle container shared by all Oracle +/// tests, starting the container on first use. /// -/// Each `#[tokio::test]` runs in its own runtime, so the pool and the -/// container are memoized across tests: only the first caller pays for the -/// (slow) Oracle boot, the others wait for it and reuse the same container. +/// Only the container is shared, never the pool: each `#[tokio::test]` runs on +/// its own runtime, and bb8 spawns its connection tasks on the runtime of the +/// test that triggered them. Those tasks die with that runtime, and the +/// connections they were opening stay counted as pending, so a pool shared +/// across tests can end up permanently timing out. pub async fn get_oracle_connection_pool() -> Arc { - static POOL: tokio::sync::OnceCell> = - tokio::sync::OnceCell::const_new(); + let port = ensure_oracle_container().await; - POOL.get_or_init(|| async { - let port = ensure_oracle_container().await; - - Arc::new( - OracleConnectionPool::new(get_oracle_params(port)) - .await - .expect("Failed to create Oracle connection pool"), - ) - }) - .await - .clone() + Arc::new( + OracleConnectionPool::new(get_oracle_params(port)) + .await + .expect("Failed to create Oracle connection pool"), + ) } From 6d5e31a018465d275e2b87a59bff39792f966335 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Thu, 15 Oct 2026 12:50:03 +0300 Subject: [PATCH 09/27] expose the ports, for Docker images like oracle-lite that do not EXPOSE --- core/tests/docker/mod.rs | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/core/tests/docker/mod.rs b/core/tests/docker/mod.rs index 27d5520c7..301ad04ce 100644 --- a/core/tests/docker/mod.rs +++ b/core/tests/docker/mod.rs @@ -153,10 +153,16 @@ impl ContainerRunner<'_> { } tracing::debug!("Port bindings: {port_bindings_map:?}"); - let port_bindings = if port_bindings_map.is_empty() { - None + // A binding only takes effect for an exposed port. Images like + // gvenzl/oracle-free declare no EXPOSE, and unlike `docker run -p`, the + // API does not expose bound ports implicitly on every Docker version. + let (exposed_ports, port_bindings) = if port_bindings_map.is_empty() { + (None, None) } else { - Some(port_bindings_map) + ( + Some(port_bindings_map.keys().cloned().collect()), + Some(port_bindings_map), + ) }; let host_config = Some(HostConfig { @@ -173,6 +179,7 @@ impl ContainerRunner<'_> { let config = ContainerCreateBody { image: Some(self.image.clone()), env: Some(env_vars), + exposed_ports, host_config, healthcheck: self.healthcheck, ..Default::default() From a22ff9555cb99eb8d71544f6a98e5e78414aabd5 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Thu, 15 Oct 2026 16:37:04 +0300 Subject: [PATCH 10/27] Minor improvements --- core/tests/clickhouse/mod.rs | 4 ++-- core/tests/odbc/mod.rs | 4 +++- crates/clickhouse/src/lib.rs | 2 +- crates/clickhouse/src/pool.rs | 9 +++++---- crates/oracle/src/pool.rs | 2 +- crates/sqlite/src/pool.rs | 25 +++++++++++++++++++------ 6 files changed, 31 insertions(+), 15 deletions(-) diff --git a/core/tests/clickhouse/mod.rs b/core/tests/clickhouse/mod.rs index c3323900c..43c395f3d 100644 --- a/core/tests/clickhouse/mod.rs +++ b/core/tests/clickhouse/mod.rs @@ -192,7 +192,7 @@ async fn clickhouse_connection_pool_with_schema() { // the schema is inferred from the database when none is provided let inferred = pool - .connect_direct() + .connect_direct().unwrap() .get_schema(&table_reference) .await .expect("schema should be inferred"); @@ -221,7 +221,7 @@ async fn clickhouse_connection_pool_with_schema() { .with_schema(Arc::clone(&schema)); let actual = pool - .connect_direct() + .connect_direct().unwrap() .get_schema(&table_reference) .await .expect("schema should be returned"); diff --git a/core/tests/odbc/mod.rs b/core/tests/odbc/mod.rs index c02e14a58..3680c7a08 100644 --- a/core/tests/odbc/mod.rs +++ b/core/tests/odbc/mod.rs @@ -33,7 +33,9 @@ async fn test_odbc_connection_pool_with_schema() { let conn = match pool.connect().await { Ok(conn) => conn, Err(e) => { - // The suite skips rather than fails when no SQLite ODBC driver is installed. + if std::env::var("CI").is_ok() { + panic!("SQLite ODBC driver unavailable in CI: {e}"); + } eprintln!( "Skipping: unable to connect with ODBC driver '{}': {e}", sqlite_odbc_driver() diff --git a/crates/clickhouse/src/lib.rs b/crates/clickhouse/src/lib.rs index 9d0a29d64..47aa95ddc 100644 --- a/crates/clickhouse/src/lib.rs +++ b/crates/clickhouse/src/lib.rs @@ -49,7 +49,7 @@ impl ClickHouseTableFactory { args: Option>, ) -> Result, Box> { - let connection = self.pool.connect_direct(); + let connection = self.pool.connect_direct()?; let client: &dyn AsyncDbConnection = &connection; let schema = client.get_schema(&table_reference).await?; let table_provider = Arc::new(ClickHouseTable::new( diff --git a/crates/clickhouse/src/pool.rs b/crates/clickhouse/src/pool.rs index eff0b03fb..3fe7f7ca8 100644 --- a/crates/clickhouse/src/pool.rs +++ b/crates/clickhouse/src/pool.rs @@ -128,8 +128,8 @@ impl ClickHouseConnectionPool { /// Returns a direct connection to the underlying database. #[must_use] - pub fn connect_direct(&self) -> ClickHouseConnection { - match &self.schema { + pub fn connect_direct(&self) -> PoolResult { + let connection = match &self.schema { Some(schema) => { >::new_with_schema( self.client(), @@ -137,14 +137,15 @@ impl ClickHouseConnectionPool { ) } None => ClickHouseConnection::new(self.client()), - } + }; + Ok(connection) } } #[async_trait::async_trait] impl DbConnectionPool for ClickHouseConnectionPool { async fn connect(&self) -> PoolResult>> { - Ok(Box::new(self.connect_direct())) + self.connect_direct().map(|c| Box::new(c) as Box>) } fn join_push_down(&self) -> JoinPushDown { diff --git a/crates/oracle/src/pool.rs b/crates/oracle/src/pool.rs index d268ef3b5..cea8633ec 100644 --- a/crates/oracle/src/pool.rs +++ b/crates/oracle/src/pool.rs @@ -172,7 +172,7 @@ impl OracleConnectionPool { fn new_connection(&self, conn: OraclePooledConnection) -> OracleConnection { match &self.schema { - Some(schema) => >::new_with_schema( + Some(schema) => OracleConnection::new_with_schema( conn, Arc::clone(schema), ), diff --git a/crates/sqlite/src/pool.rs b/crates/sqlite/src/pool.rs index 0acdd941e..269e9b0e4 100644 --- a/crates/sqlite/src/pool.rs +++ b/crates/sqlite/src/pool.rs @@ -94,15 +94,15 @@ impl SqliteConnectionPoolFactory { vec![] }; - let mut pool = SqliteConnectionPool::new( + let mut pool = SqliteConnectionPool::new_with_schema( &self.path, self.mode, join_push_down, attach_databases, self.busy_timeout, + self.schema.clone() ) .await?; - pool.schema = self.schema.clone(); pool.setup().await?; @@ -138,7 +138,18 @@ impl SqliteConnectionPool { attach_databases: Vec>, busy_timeout: Duration, ) -> PoolResult { - let conn = match mode { + Self::new_with_schema(path, mode, join_push_down, attach_databases, busy_timeout, None).await + } + + pub async fn new_with_schema( + path: &str, + mode: Mode, + join_push_down: JoinPushDown, + attach_databases: Vec>, + busy_timeout: Duration, + schema: Option, + ) -> PoolResult { + let conn = match mode { Mode::Memory => Connection::open_in_memory() .await .map_err(|e| Error::ConnectionPoolError { source: e })?, @@ -155,7 +166,7 @@ impl SqliteConnectionPool { attach_databases, path: path.into(), busy_timeout, - schema: None, + schema, }) } @@ -283,7 +294,9 @@ impl SqliteConnectionPool { let mut factory = SqliteConnectionPoolFactory::new(&self.path, self.mode, self.busy_timeout) .with_databases(attach_databases); - factory.schema = self.schema.clone(); + if let Some(schema) = &self.schema { + factory = factory.with_schema(schema.clone()); + } factory.build().await } } @@ -471,7 +484,7 @@ mod tests { assert!(field1.is_nullable()); let field2 = inferred.field(1); assert_eq!(field2.name(), "name"); - assert_eq!(field2.data_type(), &DataType::Null); + assert_eq!(field2.data_type(), &DataType::Null); // inferring uses only the first row assert!(field2.is_nullable()); // the provided schema is used as-is From 32ce197c5da1952ba45fe1621f3eb66fd10945a6 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Thu, 15 Oct 2026 19:36:50 +0300 Subject: [PATCH 11/27] Pass a map to the pool --- core/tests/clickhouse/mod.rs | 61 ++++++-- core/tests/mysql/mod.rs | 44 +++++- core/tests/odbc/mod.rs | 36 ++++- core/tests/oracle/mod.rs | 45 +++++- core/tests/postgres/mod.rs | 39 +++++- crates/clickhouse/src/conn.rs | 14 +- crates/clickhouse/src/pool.rs | 37 ++--- .../sql/db_connection_pool/dbconnection.rs | 17 ++- crates/mysql/src/conn.rs | 12 +- crates/mysql/src/pool.rs | 25 ++-- crates/odbc/src/conn.rs | 12 +- crates/odbc/src/pool.rs | 24 ++-- crates/oracle/src/conn.rs | 24 ++-- crates/oracle/src/pool.rs | 28 ++-- crates/postgres/src/conn.rs | 11 +- crates/postgres/src/pool.rs | 25 ++-- crates/sqlite/src/conn.rs | 15 +- crates/sqlite/src/pool.rs | 132 ++++++++++++------ 18 files changed, 421 insertions(+), 180 deletions(-) diff --git a/core/tests/clickhouse/mod.rs b/core/tests/clickhouse/mod.rs index 43c395f3d..3e89c47e3 100644 --- a/core/tests/clickhouse/mod.rs +++ b/core/tests/clickhouse/mod.rs @@ -151,8 +151,8 @@ async fn clickhouse_insert_and_read() { assert_eq!(record_batch[0].num_columns(), 9); } -/// A pool created with `with_schema` must return the provided schema from `get_schema` -/// instead of inferring it from the database. +/// A pool created with `with_schema` must return the schema provided for a table from +/// `get_schema` instead of inferring it from the database. #[tokio::test] async fn clickhouse_connection_pool_with_schema() { use datafusion::arrow::datatypes::{DataType, Field, Schema}; @@ -192,7 +192,8 @@ async fn clickhouse_connection_pool_with_schema() { // the schema is inferred from the database when none is provided let inferred = pool - .connect_direct().unwrap() + .connect_direct() + .unwrap() .get_schema(&table_reference) .await .expect("schema should be inferred"); @@ -210,36 +211,72 @@ async fn clickhouse_connection_pool_with_schema() { assert_eq!(field2.data_type(), &DataType::Utf8); // ClickHouseConnection::get_schema() infers it correctly assert_eq!(field2.is_nullable(), false); - // the provided schema is used as-is + // the provided schemas are used as-is, each one for its own table let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, false), Field::new("name", DataType::Utf8, true), ])); + // not created in the database, so its schema can only come from the pool + let other_table_reference = TableReference::bare("other_table"); + let other_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Float64, + true, + )])); let pool = ClickHouseConnectionPool::new(get_clickhouse_params(port)) .await .unwrap() - .with_schema(Arc::clone(&schema)); + .with_schema(table_reference.clone(), Arc::clone(&schema)) + .with_schema(other_table_reference.clone(), Arc::clone(&other_schema)); - let actual = pool - .connect_direct().unwrap() + let direct_conn = pool.connect_direct().unwrap(); + let actual = direct_conn .get_schema(&table_reference) .await .expect("schema should be returned"); assert_eq!(actual, schema); + let actual = direct_conn + .get_schema(&other_table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, other_schema); let conn = pool.connect().await.expect("connect should work"); + let conn = conn.as_async().expect("should be async connection"); let actual = conn - .as_async() - .expect("should be async connection") .get_schema(&table_reference) .await .expect("schema should be returned"); assert_eq!(actual, schema); + let actual = conn + .get_schema(&other_table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, other_schema); - // the table factory uses the pool's schema too - let table_provider = ClickHouseTableFactory::new(pool) - .table_provider(table_reference, None) + // the table factory uses the pool's schemas too + let table_factory = ClickHouseTableFactory::new(pool); + let table_provider = table_factory + .table_provider(table_reference.clone(), None) .await .expect("table provider should be created"); assert_eq!(table_provider.schema(), schema); + let table_provider = table_factory + .table_provider(other_table_reference.clone(), None) + .await + .expect("table provider should be created"); + assert_eq!(table_provider.schema(), other_schema); + + // the schema of a table without a provided schema is still inferred + let pool = ClickHouseConnectionPool::new(get_clickhouse_params(port)) + .await + .unwrap() + .with_schema(other_table_reference, other_schema); + let actual = pool + .connect_direct() + .unwrap() + .get_schema(&table_reference) + .await + .expect("schema should be inferred"); + assert_eq!(actual, inferred); } diff --git a/core/tests/mysql/mod.rs b/core/tests/mysql/mod.rs index 9f2325e9f..54e423834 100644 --- a/core/tests/mysql/mod.rs +++ b/core/tests/mysql/mod.rs @@ -1003,8 +1003,8 @@ async fn test_mysql_arrow_oneway() { test_mysql_connection_pool_with_schema(port).await; } -/// A pool created with `with_schema` must return the provided schema from `get_schema` -/// instead of inferring it from the database. +/// A pool created with `with_schema` must return the schema provided for a table from +/// `get_schema` instead of inferring it from the database. async fn test_mysql_connection_pool_with_schema(port: u16) { use datafusion::common::TableReference; @@ -1049,15 +1049,23 @@ async fn test_mysql_connection_pool_with_schema(port: u16) { assert_eq!(field2.data_type(), &DataType::Utf8); // MySQLConnection::get_schema() infers it correctly assert_eq!(field2.is_nullable(), true); - // the provided schema is used as-is + // the provided schemas are used as-is, each one for its own table let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int64, false), Field::new("name", DataType::Utf8, true), ])); + // not created in the database, so its schema can only come from the pool + let other_table_reference = TableReference::bare("other_table"); + let other_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Float64, + true, + )])); let pool = common::get_mysql_connection_pool(port) .await .expect("MySQL connection pool should be created") - .with_schema(Arc::clone(&schema)); + .with_schema(table_reference.clone(), Arc::clone(&schema)) + .with_schema(other_table_reference.clone(), Arc::clone(&other_schema)); let direct_conn = pool .connect_direct() @@ -1068,16 +1076,40 @@ async fn test_mysql_connection_pool_with_schema(port: u16) { .await .expect("schema should be returned"); assert_eq!(actual, schema); + let actual = direct_conn + .get_schema(&other_table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, other_schema); let conn = pool .connect() .await .expect("Connection should be established"); + let conn = conn.as_async().expect("should be async connection"); let actual = conn - .as_async() - .expect("should be async connection") .get_schema(&table_reference) .await .expect("schema should be returned"); assert_eq!(actual, schema); + let actual = conn + .get_schema(&other_table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, other_schema); + + // the schema of a table without a provided schema is still inferred + let pool = common::get_mysql_connection_pool(port) + .await + .expect("MySQL connection pool should be created") + .with_schema(other_table_reference, other_schema); + let direct_conn = pool + .connect_direct() + .await + .expect("Connection should be established"); + let actual = direct_conn + .get_schema(&table_reference) + .await + .expect("schema should be inferred"); + assert_eq!(actual, inferred); } diff --git a/core/tests/odbc/mod.rs b/core/tests/odbc/mod.rs index 3680c7a08..56e1e0d7a 100644 --- a/core/tests/odbc/mod.rs +++ b/core/tests/odbc/mod.rs @@ -20,8 +20,8 @@ fn get_odbc_pool(db_path: &str) -> ODBCPool { ODBCPool::new(params).expect("ODBC pool should be created") } -/// A pool created with `with_schema` must return the provided schema from `get_schema` -/// instead of inferring it from the database. +/// A pool created with `with_schema` must return the schema provided for a table from +/// `get_schema` instead of inferring it from the database. #[tokio::test] async fn test_odbc_connection_pool_with_schema() { let dir = tempfile::tempdir().expect("temp dir should be created"); @@ -76,18 +76,42 @@ async fn test_odbc_connection_pool_with_schema() { assert_eq!(field2.data_type(), &DataType::Utf8); // ODBCConnection::get_schema() infers it correctly assert_eq!(field2.is_nullable(), true); + // the provided schemas are used as-is, each one for its own table let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int8, false), Field::new("name", DataType::LargeUtf8, true), ])); - // the provided schema is used as-is - let pool = get_odbc_pool(db_path).with_schema(Arc::clone(&schema)); + // not created in the database, so its schema can only come from the pool + let other_table_reference = TableReference::bare("other_table"); + let other_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Float64, + true, + )])); + let pool = get_odbc_pool(db_path) + .with_schema(table_reference.clone(), Arc::clone(&schema)) + .with_schema(other_table_reference.clone(), Arc::clone(&other_schema)); let conn = pool.connect().await.expect("connect should work"); + let conn = conn.as_async().expect("should be async connection"); let actual = conn - .as_async() - .expect("should be async connection") .get_schema(&table_reference) .await .expect("schema should be returned"); assert_eq!(actual, schema); + let actual = conn + .get_schema(&other_table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, other_schema); + + // the schema of a table without a provided schema is still inferred + let pool = get_odbc_pool(db_path).with_schema(other_table_reference, other_schema); + let conn = pool.connect().await.expect("connect should work"); + let actual = conn + .as_async() + .expect("should be async connection") + .get_schema(&table_reference) + .await + .expect("schema should be inferred"); + assert_eq!(actual, inferred); } diff --git a/core/tests/oracle/mod.rs b/core/tests/oracle/mod.rs index 5b5ef44b1..5241c0e20 100644 --- a/core/tests/oracle/mod.rs +++ b/core/tests/oracle/mod.rs @@ -973,8 +973,8 @@ async fn test_oracle_large_query_streaming() { ); } -/// A pool created with `with_schema` must return the provided schema from `get_schema` -/// instead of inferring it from the database. +/// A pool created with `with_schema` must return the schema provided for a table from +/// `get_schema` instead of inferring it from the database. #[tokio::test] async fn test_oracle_connection_pool_with_schema() { use datafusion_table_providers::sql::db_connection_pool::dbconnection::AsyncDbConnection; @@ -1019,16 +1019,23 @@ async fn test_oracle_connection_pool_with_schema() { assert_eq!(field2.data_type(), &DataType::Utf8); // OracleConnection::get_schema() infers it correctly assert_eq!(field2.is_nullable(), true); + // the provided schemas are used as-is, each one for its own table let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int8, false), Field::new("name", DataType::LargeUtf8, true), ])); - - // the provided schema is used as-is + // not created in the database, so its schema can only come from the pool + let other_table_reference = TableReference::bare("OTHER_TABLE"); + let other_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Float64, + true, + )])); let pool = OracleConnectionPool::new(common::get_oracle_params(container_port())) .await .expect("Failed to create Oracle connection pool") - .with_schema(Arc::clone(&schema)); + .with_schema(table_reference.clone(), Arc::clone(&schema)) + .with_schema(other_table_reference.clone(), Arc::clone(&other_schema)); let direct_conn = pool .connect_direct() @@ -1039,14 +1046,38 @@ async fn test_oracle_connection_pool_with_schema() { .await .expect("schema should be returned"); assert_eq!(actual, schema); + let actual = direct_conn + .get_schema(&other_table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, other_schema); let conn: Box> = pool.connect().await.expect("Failed to get connection"); + let conn = conn.as_async().expect("should be async connection"); let actual = conn - .as_async() - .expect("should be async connection") .get_schema(&table_reference) .await .expect("schema should be returned"); assert_eq!(actual, schema); + let actual = conn + .get_schema(&other_table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, other_schema); + + // the schema of a table without a provided schema is still inferred + let pool = OracleConnectionPool::new(common::get_oracle_params(container_port())) + .await + .expect("Failed to create Oracle connection pool") + .with_schema(other_table_reference, other_schema); + let direct_conn = pool + .connect_direct() + .await + .expect("Failed to get connection"); + let actual = direct_conn + .get_schema(&table_reference) + .await + .expect("schema should be inferred"); + assert_eq!(actual, inferred); } diff --git a/core/tests/postgres/mod.rs b/core/tests/postgres/mod.rs index 9b983a785..eb21ca494 100644 --- a/core/tests/postgres/mod.rs +++ b/core/tests/postgres/mod.rs @@ -182,8 +182,8 @@ async fn test_arrow_postgres_one_way() { test_postgres_connection_pool_with_schema(port).await; } -/// A pool created with `with_schema` must return the provided schema from `get_schema` -/// instead of inferring it from the database. +/// A pool created with `with_schema` must return the schema provided for a table from +/// `get_schema` instead of inferring it from the database. async fn test_postgres_connection_pool_with_schema(port: u16) { use datafusion::common::TableReference; use datafusion_table_providers::sql::db_connection_pool::dbconnection::AsyncDbConnection; @@ -232,31 +232,58 @@ async fn test_postgres_connection_pool_with_schema(port: u16) { assert_eq!(field2.data_type(), &DataType::Utf8); // PostgresConnection::get_schema() infers it correctly assert_eq!(field2.is_nullable(), true); - // the provided schema is used as-is + // the provided schemas are used as-is, each one for its own table let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int64, false), Field::new("name", DataType::LargeUtf8, true), ])); + // not created in the database, so its schema can only come from the pool + let other_table_reference = TableReference::bare("other_table"); + let other_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Float64, + true, + )])); let pool = common::get_postgres_connection_pool(port) .await .expect("Postgres connection pool should be created") - .with_schema(Arc::clone(&schema)); + .with_schema(table_reference.clone(), Arc::clone(&schema)) + .with_schema(other_table_reference.clone(), Arc::clone(&other_schema)); let direct_conn = pool.connect_direct().await.expect("connect should work"); let actual = AsyncDbConnection::get_schema(&direct_conn, &table_reference) .await .expect("schema should be returned"); assert_eq!(actual, schema); + let actual = AsyncDbConnection::get_schema(&direct_conn, &other_table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, other_schema); let sqltable_pool: Arc = Arc::new(pool); let conn = sqltable_pool.connect().await.expect("connect should work"); + let conn = conn.as_async().expect("should be async connection"); let actual = conn - .as_async() - .expect("should be async connection") .get_schema(&table_reference) .await .expect("schema should be returned"); assert_eq!(actual, schema); + let actual = conn + .get_schema(&other_table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, other_schema); + + // the schema of a table without a provided schema is still inferred + let pool = common::get_postgres_connection_pool(port) + .await + .expect("Postgres connection pool should be created") + .with_schema(other_table_reference, other_schema); + let direct_conn = pool.connect_direct().await.expect("connect should work"); + let actual = AsyncDbConnection::get_schema(&direct_conn, &table_reference) + .await + .expect("schema should be inferred"); + assert_eq!(actual, inferred); } /// An unconstrained `numeric` column (what `max()`, `avg()` and arithmetic over `numeric` diff --git a/crates/clickhouse/src/conn.rs b/crates/clickhouse/src/conn.rs index c07d4684d..1e22467ad 100644 --- a/crates/clickhouse/src/conn.rs +++ b/crates/clickhouse/src/conn.rs @@ -14,13 +14,13 @@ use serde::Deserialize; use snafu::ResultExt; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::{ - AsyncDbConnection, DbConnection, Error, SyncDbConnection, + AsyncDbConnection, DbConnection, Error, SyncDbConnection, TableSchemas, }; #[derive(Clone)] pub struct ClickHouseConnection { pub client: Client, - pub schema: Option, + pub schemas: Arc, } impl std::fmt::Debug for ClickHouseConnection { @@ -35,7 +35,7 @@ impl ClickHouseConnection { pub fn new(client: Client) -> Self { Self { client, - schema: None, + schemas: Arc::default(), } } } @@ -66,17 +66,17 @@ impl AsyncDbConnection for ClickHouseConnection { { Self { client: conn, - schema: None, + schemas: Arc::default(), } } - fn new_with_schema(conn: Client, schema: SchemaRef) -> Self + fn new_with_schemas(conn: Client, schemas: Arc) -> Self where Self: Sized, { Self { client: conn, - schema: Some(schema), + schemas, } } @@ -120,7 +120,7 @@ impl AsyncDbConnection for ClickHouseConnection { /// /// * `table_reference` - The table reference. async fn get_schema(&self, table_reference: &TableReference) -> Result { - if let Some(schema) = &self.schema { + if let Some(schema) = self.schemas.get(table_reference) { return Ok(Arc::clone(schema)); } diff --git a/crates/clickhouse/src/pool.rs b/crates/clickhouse/src/pool.rs index 3fe7f7ca8..8ee3ffec1 100644 --- a/crates/clickhouse/src/pool.rs +++ b/crates/clickhouse/src/pool.rs @@ -2,12 +2,13 @@ use std::{collections::HashMap, sync::Arc}; use clickhouse::{Client, Compression}; use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::TableReference; use secrecy::{ExposeSecret, SecretString}; use snafu::{ResultExt, Snafu}; use crate::conn::ClickHouseConnection; use datafusion_table_providers_common::sql::db_connection_pool::{ - dbconnection::{AsyncDbConnection, DbConnection}, + dbconnection::{AsyncDbConnection, DbConnection, TableSchemas}, DbConnectionPool, JoinPushDown, }; type PoolResult> = std::result::Result; @@ -28,7 +29,7 @@ pub enum Error { pub struct ClickHouseConnectionPool { pub client: Client, pub join_push_down: JoinPushDown, - pub schema: Option, + pub schemas: Arc, } impl std::fmt::Debug for ClickHouseConnectionPool { @@ -111,14 +112,21 @@ impl ClickHouseConnectionPool { Ok(Self { client, join_push_down, - schema: None, + schemas: Arc::default(), }) } - /// Use the provided `schema` for all connections instead of inferring it from the database. + /// Use the provided `schema` for `table_reference` instead of inferring it from the database. + /// + /// Call it once per table to provide the schemas of several tables. Table references are + /// matched exactly, e.g. `users` and `public.users` are different tables. #[must_use] - pub fn with_schema(mut self, schema: SchemaRef) -> Self { - self.schema = Some(schema); + pub fn with_schema( + mut self, + table_reference: impl Into, + schema: SchemaRef, + ) -> Self { + Arc::make_mut(&mut self.schemas).insert(table_reference.into(), schema); self } @@ -127,17 +135,11 @@ impl ClickHouseConnectionPool { } /// Returns a direct connection to the underlying database. - #[must_use] pub fn connect_direct(&self) -> PoolResult { - let connection = match &self.schema { - Some(schema) => { - >::new_with_schema( - self.client(), - Arc::clone(schema), - ) - } - None => ClickHouseConnection::new(self.client()), - }; + let connection = >::new_with_schemas( + self.client(), + Arc::clone(&self.schemas), + ); Ok(connection) } } @@ -145,7 +147,8 @@ impl ClickHouseConnectionPool { #[async_trait::async_trait] impl DbConnectionPool for ClickHouseConnectionPool { async fn connect(&self) -> PoolResult>> { - self.connect_direct().map(|c| Box::new(c) as Box>) + self.connect_direct() + .map(|c| Box::new(c) as Box>) } fn join_push_down(&self) -> JoinPushDown { diff --git a/crates/common/src/sql/db_connection_pool/dbconnection.rs b/crates/common/src/sql/db_connection_pool/dbconnection.rs index babc42e4d..9da6235ce 100644 --- a/crates/common/src/sql/db_connection_pool/dbconnection.rs +++ b/crates/common/src/sql/db_connection_pool/dbconnection.rs @@ -1,4 +1,4 @@ -use std::{any::Any, sync::Arc}; +use std::{any::Any, collections::HashMap, sync::Arc}; use datafusion::{ arrow::datatypes::SchemaRef, common::TableReference, execution::SendableRecordBatchStream, @@ -10,6 +10,12 @@ use crate::sql::db_connection_pool::runtime::run_async_with_tokio; pub type GenericError = Box; pub type Result = std::result::Result; +/// Known Arrow schemas, keyed by the table they belong to. +/// +/// Connections created with [`AsyncDbConnection::new_with_schemas`] return these schemas +/// from [`AsyncDbConnection::get_schema`] instead of inferring them from the database. +pub type TableSchemas = HashMap; + #[derive(Debug, Snafu)] #[snafu(visibility(pub))] pub enum Error { @@ -100,11 +106,12 @@ pub trait AsyncDbConnection: DbConnection + Sync { where Self: Sized; - /// Creates a new connection that uses the provided `schema` instead of inferring it. + /// Creates a new connection that uses the provided `schemas` instead of inferring them. /// - /// When a schema is provided, [`AsyncDbConnection::get_schema`] returns it as-is - /// instead of querying the database for it. - fn new_with_schema(conn: T, schema: SchemaRef) -> Self + /// When `schemas` contains the requested table reference, [`AsyncDbConnection::get_schema`] + /// returns its schema as-is instead of querying the database for it. Table references are + /// matched exactly, e.g. `users` and `public.users` are different keys. + fn new_with_schemas(conn: T, schemas: Arc) -> Self where Self: Sized; diff --git a/crates/mysql/src/conn.rs b/crates/mysql/src/conn.rs index 628f6b3af..60a3a2e36 100644 --- a/crates/mysql/src/conn.rs +++ b/crates/mysql/src/conn.rs @@ -17,7 +17,7 @@ use snafu::prelude::*; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::Result; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::{ - AsyncDbConnection, DbConnection, + AsyncDbConnection, DbConnection, TableSchemas, }; #[derive(Debug, Snafu)] @@ -46,7 +46,7 @@ pub enum Error { pub struct MySQLConnection { pub conn: Arc>, - pub schema: Option, + pub schemas: Arc, } impl MySQLConnection { @@ -92,14 +92,14 @@ impl<'a> AsyncDbConnection for MySQLConnection { fn new(conn: Conn) -> Self { MySQLConnection { conn: Arc::new(Mutex::new(conn)), - schema: None, + schemas: Arc::default(), } } - fn new_with_schema(conn: Conn, schema: SchemaRef) -> Self { + fn new_with_schemas(conn: Conn, schemas: Arc) -> Self { MySQLConnection { conn: Arc::new(Mutex::new(conn)), - schema: Some(schema), + schemas, } } @@ -162,7 +162,7 @@ impl<'a> AsyncDbConnection for MySQLConnection { SchemaRef, datafusion_table_providers_common::sql::db_connection_pool::dbconnection::Error, > { - if let Some(schema) = &self.schema { + if let Some(schema) = self.schemas.get(table_reference) { return Ok(Arc::clone(schema)); } diff --git a/crates/mysql/src/pool.rs b/crates/mysql/src/pool.rs index c97ea6aed..435d082d0 100644 --- a/crates/mysql/src/pool.rs +++ b/crates/mysql/src/pool.rs @@ -1,5 +1,6 @@ use crate::conn::MySQLConnection; use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::TableReference; use std::{collections::HashMap, path::PathBuf, sync::Arc}; use async_trait::async_trait; @@ -13,7 +14,7 @@ use snafu::{ResultExt, Snafu}; use datafusion_table_providers_common::{ sql::db_connection_pool::{ - dbconnection::{AsyncDbConnection, DbConnection}, + dbconnection::{AsyncDbConnection, DbConnection, TableSchemas}, JoinPushDown, }, util::{self, ns_lookup::verify_ns_lookup_and_tcp_connect}, @@ -57,7 +58,7 @@ pub enum Error { pub struct MySQLConnectionPool { pool: Arc, join_push_down: JoinPushDown, - schema: Option, + schemas: Arc, } const SETUP_QUERIES: [&str; 4] = [ @@ -211,22 +212,26 @@ impl MySQLConnectionPool { Ok(Self { pool: Arc::new(pool), join_push_down, - schema: None, + schemas: Arc::default(), }) } - /// Use the provided `schema` for all connections instead of inferring it from the database. + /// Use the provided `schema` for `table_reference` instead of inferring it from the database. + /// + /// Call it once per table to provide the schemas of several tables. Table references are + /// matched exactly, e.g. `users` and `public.users` are different tables. #[must_use] - pub fn with_schema(mut self, schema: SchemaRef) -> Self { - self.schema = Some(schema); + pub fn with_schema( + mut self, + table_reference: impl Into, + schema: SchemaRef, + ) -> Self { + Arc::make_mut(&mut self.schemas).insert(table_reference.into(), schema); self } fn new_connection(&self, conn: mysql_async::Conn) -> MySQLConnection { - match &self.schema { - Some(schema) => MySQLConnection::new_with_schema(conn, Arc::clone(schema)), - None => MySQLConnection::new(conn), - } + MySQLConnection::new_with_schemas(conn, Arc::clone(&self.schemas)) } /// Returns a direct connection to the underlying database. diff --git a/crates/odbc/src/conn.rs b/crates/odbc/src/conn.rs index 2e3825030..eb7c96863 100644 --- a/crates/odbc/src/conn.rs +++ b/crates/odbc/src/conn.rs @@ -31,7 +31,7 @@ use datafusion::error::DataFusionError; use datafusion::execution::SendableRecordBatchStream; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion_table_providers_common::sql::db_connection_pool::{ - dbconnection::{self, AsyncDbConnection, DbConnection, GenericError}, + dbconnection::{self, AsyncDbConnection, DbConnection, GenericError, TableSchemas}, runtime::run_async_with_tokio, DbConnectionPool, }; @@ -93,7 +93,7 @@ pub enum Error { pub struct ODBCConnection<'a> { pub conn: Arc>>, pub params: Arc>, - pub schema: Option, + pub schemas: Arc, } impl<'a> DbConnection, ODBCParameter> for ODBCConnection<'a> @@ -132,15 +132,15 @@ where ODBCConnection { conn: Arc::new(conn.into()), params: Arc::new(HashMap::new()), - schema: None, + schemas: Arc::default(), } } - fn new_with_schema(conn: Connection<'a>, schema: SchemaRef) -> Self { + fn new_with_schemas(conn: Connection<'a>, schemas: Arc) -> Self { ODBCConnection { conn: Arc::new(conn.into()), params: Arc::new(HashMap::new()), - schema: Some(schema), + schemas, } } @@ -167,7 +167,7 @@ where &self, table_reference: &TableReference, ) -> Result { - if let Some(schema) = &self.schema { + if let Some(schema) = self.schemas.get(table_reference) { return Ok(Arc::clone(schema)); } diff --git a/crates/odbc/src/pool.rs b/crates/odbc/src/pool.rs index 2d4fd3921..975f0cc2c 100644 --- a/crates/odbc/src/pool.rs +++ b/crates/odbc/src/pool.rs @@ -18,8 +18,9 @@ use crate::conn::ODBCConnection; use crate::conn::{ODBCDbConnection, ODBCParameter}; use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::TableReference; use datafusion_table_providers_common::sql::db_connection_pool::{ - runtime::run_async_with_tokio, DbConnectionPool, JoinPushDown, + dbconnection::TableSchemas, runtime::run_async_with_tokio, DbConnectionPool, JoinPushDown, }; use odbc_api::{sys::AttrConnectionPooling, Connection, ConnectionOptions, Environment}; use secrecy::{ExposeSecret, SecretBox, SecretString}; @@ -60,7 +61,7 @@ pub struct ODBCPool { params: Arc>, connection_string: String, connection_id: String, - schema: Option, + schemas: Arc, } fn hash_string(val: &str) -> String { @@ -94,14 +95,21 @@ impl ODBCPool { connection_string, connection_id, pool: &ENV, - schema: None, + schemas: Arc::default(), }) } - /// Use the provided `schema` for all connections instead of inferring it from the database. + /// Use the provided `schema` for `table_reference` instead of inferring it from the database. + /// + /// Call it once per table to provide the schemas of several tables. Table references are + /// matched exactly, e.g. `users` and `public.users` are different tables. #[must_use] - pub fn with_schema(mut self, schema: SchemaRef) -> Self { - self.schema = Some(schema); + pub fn with_schema( + mut self, + table_reference: impl Into, + schema: SchemaRef, + ) -> Self { + Arc::make_mut(&mut self.schemas).insert(table_reference.into(), schema); self } @@ -122,7 +130,7 @@ where let pool: &'static Environment = self.pool; let connection_string = self.connection_string.clone(); let params = Arc::clone(&self.params); - let schema = self.schema.clone(); + let schemas = Arc::clone(&self.schemas); let connect = async move || -> Result< Box>, @@ -140,7 +148,7 @@ where Ok(Box::new(ODBCConnection { conn: Arc::new(cxn.into()), params, - schema, + schemas, }) as Box>) }; run_async_with_tokio(connect).await diff --git a/crates/oracle/src/conn.rs b/crates/oracle/src/conn.rs index 1b700d557..cc3a0f04c 100644 --- a/crates/oracle/src/conn.rs +++ b/crates/oracle/src/conn.rs @@ -13,20 +13,24 @@ use tokio::task; use crate::arrow_sql_gen::{map_oracle_type_to_arrow_type, rows_to_arrow}; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::{ - AsyncDbConnection, DbConnection, Error, GenericError, Result, UnableToGetSchemaSnafu, - UnableToGetSchemasSnafu, UnableToGetTablesSnafu, UnableToQueryArrowSnafu, + AsyncDbConnection, DbConnection, Error, GenericError, Result, TableSchemas, + UnableToGetSchemaSnafu, UnableToGetSchemasSnafu, UnableToGetTablesSnafu, + UnableToQueryArrowSnafu, }; pub type OraclePooledConnection = bb8::PooledConnection<'static, OracleConnectionManager>; pub struct OracleConnection { pub conn: OraclePooledConnection, - pub schema: Option, + pub schemas: Arc, } impl OracleConnection { pub fn new(conn: OraclePooledConnection) -> Self { - Self { conn, schema: None } + Self { + conn, + schemas: Arc::default(), + } } } @@ -49,21 +53,21 @@ impl DbConnection for Orac #[async_trait] impl AsyncDbConnection for OracleConnection { fn new(conn: OraclePooledConnection) -> Self { - Self { conn, schema: None } - } - - fn new_with_schema(conn: OraclePooledConnection, schema: SchemaRef) -> Self { Self { conn, - schema: Some(schema), + schemas: Arc::default(), } } + fn new_with_schemas(conn: OraclePooledConnection, schemas: Arc) -> Self { + Self { conn, schemas } + } + async fn get_schema( &self, table_reference: &TableReference, ) -> std::result::Result { - if let Some(schema) = &self.schema { + if let Some(schema) = self.schemas.get(table_reference) { return Ok(Arc::clone(schema)); } diff --git a/crates/oracle/src/pool.rs b/crates/oracle/src/pool.rs index cea8633ec..c14058484 100644 --- a/crates/oracle/src/pool.rs +++ b/crates/oracle/src/pool.rs @@ -7,12 +7,13 @@ use async_trait::async_trait; use bb8::CustomizeConnection; use bb8_oracle::OracleConnectionManager; use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::TableReference; use oracle::{Connection, Connector}; use secrecy::{ExposeSecret, SecretString}; use snafu::prelude::*; use datafusion_table_providers_common::sql::db_connection_pool::{ - dbconnection::{AsyncDbConnection, DbConnection}, + dbconnection::{AsyncDbConnection, DbConnection, TableSchemas}, DbConnectionPool, JoinPushDown, }; @@ -84,7 +85,7 @@ impl CustomizeConnection, bb8_oracle::Error> for SetTimezoneCust pub struct OracleConnectionPool { pool: Arc>, - schema: Option, + schemas: Arc, } impl std::fmt::Debug for OracleConnectionPool { @@ -159,25 +160,26 @@ impl OracleConnectionPool { Ok(Self { pool: Arc::new(pool), - schema: None, + schemas: Arc::default(), }) } - /// Use the provided `schema` for all connections instead of inferring it from the database. + /// Use the provided `schema` for `table_reference` instead of inferring it from the database. + /// + /// Call it once per table to provide the schemas of several tables. Table references are + /// matched exactly, e.g. `users` and `public.users` are different tables. #[must_use] - pub fn with_schema(mut self, schema: SchemaRef) -> Self { - self.schema = Some(schema); + pub fn with_schema( + mut self, + table_reference: impl Into, + schema: SchemaRef, + ) -> Self { + Arc::make_mut(&mut self.schemas).insert(table_reference.into(), schema); self } fn new_connection(&self, conn: OraclePooledConnection) -> OracleConnection { - match &self.schema { - Some(schema) => OracleConnection::new_with_schema( - conn, - Arc::clone(schema), - ), - None => OracleConnection::new(conn), - } + OracleConnection::new_with_schemas(conn, Arc::clone(&self.schemas)) } pub async fn connect_direct(&self) -> Result { diff --git a/crates/postgres/src/conn.rs b/crates/postgres/src/conn.rs index d03af07f0..98a590c58 100644 --- a/crates/postgres/src/conn.rs +++ b/crates/postgres/src/conn.rs @@ -46,6 +46,7 @@ use tokio_postgres::Row; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::AsyncDbConnection; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::DbConnection; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::Result; +use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::TableSchemas; const SCHEMA_QUERY: &str = r" WITH custom_type_details AS ( @@ -298,7 +299,7 @@ fn format_postgres_query_error(source: &bb8_postgres::tokio_postgres::Error) -> pub struct PostgresConnection { pub conn: PostgresPooledConnection, unsupported_type_action: UnsupportedTypeAction, - schema: Option, + schemas: Arc, } impl SchemaValidator for PostgresConnection { @@ -340,15 +341,15 @@ impl<'a> AsyncDbConnection PostgresConnection { conn, unsupported_type_action: UnsupportedTypeAction::default(), - schema: None, + schemas: Arc::default(), } } - fn new_with_schema(conn: PostgresPooledConnection, schema: SchemaRef) -> Self { + fn new_with_schemas(conn: PostgresPooledConnection, schemas: Arc) -> Self { PostgresConnection { conn, unsupported_type_action: UnsupportedTypeAction::default(), - schema: Some(schema), + schemas, } } @@ -402,7 +403,7 @@ impl<'a> AsyncDbConnection SchemaRef, datafusion_table_providers_common::sql::db_connection_pool::dbconnection::Error, > { - if let Some(schema) = &self.schema { + if let Some(schema) = self.schemas.get(table_reference) { return Ok(Arc::clone(schema)); } diff --git a/crates/postgres/src/pool.rs b/crates/postgres/src/pool.rs index 4cfa0991e..14702d9b2 100644 --- a/crates/postgres/src/pool.rs +++ b/crates/postgres/src/pool.rs @@ -1,5 +1,6 @@ use crate::conn::{PostgresConnection, PostgresPooledConnection}; use arrow::datatypes::SchemaRef; +use datafusion::common::TableReference; use std::{collections::HashMap, path::PathBuf, str::FromStr, sync::Arc}; use async_trait::async_trait; @@ -17,7 +18,7 @@ use tokio::runtime::Handle; use tokio_postgres; use datafusion_table_providers_common::sql::db_connection_pool::{ - dbconnection::{AsyncDbConnection, DbConnection}, + dbconnection::{AsyncDbConnection, DbConnection, TableSchemas}, runtime::run_async_with_tokio, DbConnectionPool, JoinPushDown, PasswordProvider, StaticPasswordProvider, }; @@ -232,7 +233,7 @@ pub struct PostgresConnectionPool { join_push_down: JoinPushDown, unsupported_type_action: UnsupportedTypeAction, io_handle: Option, - schema: Option, + schemas: Arc, } impl PostgresConnectionPool { @@ -424,7 +425,7 @@ impl PostgresConnectionPool { join_push_down, unsupported_type_action: UnsupportedTypeAction::default(), io_handle: None, - schema: None, + schemas: Arc::default(), }) } @@ -442,18 +443,22 @@ impl PostgresConnectionPool { self } - /// Use the provided `schema` for all connections instead of inferring it from the database. + /// Use the provided `schema` for `table_reference` instead of inferring it from the database. + /// + /// Call it once per table to provide the schemas of several tables. Table references are + /// matched exactly, e.g. `users` and `public.users` are different tables. #[must_use] - pub fn with_schema(mut self, schema: SchemaRef) -> Self { - self.schema = Some(schema); + pub fn with_schema( + mut self, + table_reference: impl Into, + schema: SchemaRef, + ) -> Self { + Arc::make_mut(&mut self.schemas).insert(table_reference.into(), schema); self } fn new_connection(&self, conn: PostgresPooledConnection) -> PostgresConnection { - match &self.schema { - Some(schema) => PostgresConnection::new_with_schema(conn, Arc::clone(schema)), - None => PostgresConnection::new(conn), - } + PostgresConnection::new_with_schemas(conn, Arc::clone(&self.schemas)) } /// Returns a direct connection to the underlying database. diff --git a/crates/sqlite/src/conn.rs b/crates/sqlite/src/conn.rs index 1198791ef..88236333a 100644 --- a/crates/sqlite/src/conn.rs +++ b/crates/sqlite/src/conn.rs @@ -16,6 +16,7 @@ use tokio_rusqlite::Connection; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::AsyncDbConnection; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::DbConnection; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::Result; +use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::TableSchemas; #[derive(Debug, Snafu)] pub enum Error { @@ -33,7 +34,7 @@ pub enum Error { pub struct SqliteConnection { pub conn: Connection, - pub schema: Option, + pub schemas: Arc, } impl SchemaValidator for SqliteConnection { @@ -87,16 +88,16 @@ impl DbConnection for SqliteConnection #[async_trait] impl AsyncDbConnection for SqliteConnection { fn new(conn: Connection) -> Self { - SqliteConnection { conn, schema: None } - } - - fn new_with_schema(conn: Connection, schema: SchemaRef) -> Self { SqliteConnection { conn, - schema: Some(schema), + schemas: Arc::default(), } } + fn new_with_schemas(conn: Connection, schemas: Arc) -> Self { + SqliteConnection { conn, schemas } + } + async fn tables( &self, _schema: &str, @@ -137,7 +138,7 @@ impl AsyncDbConnection for SqliteConnec SchemaRef, datafusion_table_providers_common::sql::db_connection_pool::dbconnection::Error, > { - if let Some(schema) = &self.schema { + if let Some(schema) = self.schemas.get(table_reference) { return Ok(Arc::clone(schema)); } diff --git a/crates/sqlite/src/pool.rs b/crates/sqlite/src/pool.rs index 269e9b0e4..fb799bc57 100644 --- a/crates/sqlite/src/pool.rs +++ b/crates/sqlite/src/pool.rs @@ -3,13 +3,14 @@ use std::{sync::Arc, time::Duration}; use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::TableReference; use snafu::prelude::*; use tokio_rusqlite::{Connection, ToSql}; use datafusion_table_providers_common::sql::db_connection_pool::DbConnectionPool; type PoolResult = std::result::Result>; use datafusion_table_providers_common::sql::db_connection_pool::{ - dbconnection::{AsyncDbConnection, DbConnection}, + dbconnection::{AsyncDbConnection, DbConnection, TableSchemas}, JoinPushDown, Mode, }; @@ -30,7 +31,7 @@ pub struct SqliteConnectionPoolFactory { mode: Mode, attach_databases: Option>>, busy_timeout: Duration, - schema: Option, + schemas: Arc, } impl SqliteConnectionPoolFactory { @@ -40,14 +41,21 @@ impl SqliteConnectionPoolFactory { mode, attach_databases: None, busy_timeout, - schema: None, + schemas: Arc::default(), } } - /// Use the provided `schema` for all connections instead of inferring it from the table data. + /// Use the provided `schema` for `table_reference` instead of inferring it from the table data. + /// + /// Call it once per table to provide the schemas of several tables. Table references are + /// matched exactly, e.g. `users` and `main.users` are different tables. #[must_use] - pub fn with_schema(mut self, schema: SchemaRef) -> Self { - self.schema = Some(schema); + pub fn with_schema( + mut self, + table_reference: impl Into, + schema: SchemaRef, + ) -> Self { + Arc::make_mut(&mut self.schemas).insert(table_reference.into(), schema); self } @@ -94,13 +102,13 @@ impl SqliteConnectionPoolFactory { vec![] }; - let mut pool = SqliteConnectionPool::new_with_schema( + let pool = SqliteConnectionPool::new_with_schemas( &self.path, self.mode, join_push_down, attach_databases, self.busy_timeout, - self.schema.clone() + Arc::clone(&self.schemas), ) .await?; @@ -118,7 +126,7 @@ pub struct SqliteConnectionPool { path: Arc, attach_databases: Vec>, busy_timeout: Duration, - schema: Option, + schemas: Arc, } impl SqliteConnectionPool { @@ -138,18 +146,32 @@ impl SqliteConnectionPool { attach_databases: Vec>, busy_timeout: Duration, ) -> PoolResult { - Self::new_with_schema(path, mode, join_push_down, attach_databases, busy_timeout, None).await + Self::new_with_schemas( + path, + mode, + join_push_down, + attach_databases, + busy_timeout, + Arc::default(), + ) + .await } - pub async fn new_with_schema( + /// Creates a new instance of `SqliteConnectionPool` whose connections return the provided + /// `schemas` from `get_schema` instead of inferring them from the table data. + /// + /// # Errors + /// + /// Returns an error if there is a problem creating the connection pool. + pub async fn new_with_schemas( path: &str, mode: Mode, join_push_down: JoinPushDown, attach_databases: Vec>, busy_timeout: Duration, - schema: Option, + schemas: Arc, ) -> PoolResult { - let conn = match mode { + let conn = match mode { Mode::Memory => Connection::open_in_memory() .await .map_err(|e| Error::ConnectionPoolError { source: e })?, @@ -166,7 +188,7 @@ impl SqliteConnectionPool { attach_databases, path: path.into(), busy_timeout, - schema, + schemas, }) } @@ -262,10 +284,7 @@ impl SqliteConnectionPool { fn new_connection(&self) -> SqliteConnection { let conn = self.conn.clone(); - match &self.schema { - Some(schema) => SqliteConnection::new_with_schema(conn, Arc::clone(schema)), - None => SqliteConnection::new(conn), - } + SqliteConnection::new_with_schemas(conn, Arc::clone(&self.schemas)) } /// Will attempt to clone the connection pool. This will always succeed for in-memory mode. @@ -282,7 +301,7 @@ impl SqliteConnectionPool { path: Arc::clone(&self.path), attach_databases: self.attach_databases.clone(), busy_timeout: self.busy_timeout, - schema: self.schema.clone(), + schemas: Arc::clone(&self.schemas), }), Mode::File => { let attach_databases = if self.attach_databases.is_empty() { @@ -294,9 +313,7 @@ impl SqliteConnectionPool { let mut factory = SqliteConnectionPoolFactory::new(&self.path, self.mode, self.busy_timeout) .with_databases(attach_databases); - if let Some(schema) = &self.schema { - factory = factory.with_schema(schema.clone()); - } + factory.schemas = Arc::clone(&self.schemas); factory.build().await } } @@ -487,36 +504,73 @@ mod tests { assert_eq!(field2.data_type(), &DataType::Null); // inferring uses only the first row assert!(field2.is_nullable()); - // the provided schema is used as-is - let schema = Arc::new(Schema::new(vec![ + // the provided schemas are used as-is, each one for its own table + let schema_t = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, false), Field::new("name", DataType::Utf8, true), ])); + let schema_u = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Float64, + true, + )])); let factory = SqliteConnectionPoolFactory::new( "./test.sqlite", Mode::Memory, Duration::from_millis(5000), ) - .with_schema(Arc::clone(&schema)); + .with_schema("t", Arc::clone(&schema_t)) + .with_schema(TableReference::bare("u"), Arc::clone(&schema_u)); - let table_reference = TableReference::bare("t"); let pool = factory.build().await.unwrap(); let conn = pool.connect().await.unwrap(); - let actual = conn - .as_async() - .unwrap() - .get_schema(&table_reference) - .await - .unwrap(); - assert_eq!(actual, schema); + let conn = conn.as_async().unwrap(); + assert_eq!( + conn.get_schema(&TableReference::bare("t")).await.unwrap(), + schema_t + ); + assert_eq!( + conn.get_schema(&TableReference::bare("u")).await.unwrap(), + schema_u + ); + + let sync_conn = pool.connect_sync(); + let sync_conn = sync_conn.as_async().unwrap(); + assert_eq!( + sync_conn + .get_schema(&TableReference::bare("t")) + .await + .unwrap(), + schema_t + ); - let actual = pool - .connect_sync() - .as_async() - .unwrap() - .get_schema(&table_reference) + // the schema of a table without a provided schema is still inferred + conn.execute("CREATE TABLE v (flag INTEGER)", &[]) .await .unwrap(); - assert_eq!(actual, schema); + conn.execute("INSERT INTO v VALUES (1)", &[]).await.unwrap(); + let inferred = conn.get_schema(&TableReference::bare("v")).await.unwrap(); + assert_eq!(inferred.fields.len(), 1); + assert_eq!(inferred.field(0).name(), "flag"); + assert_eq!(inferred.field(0).data_type(), &DataType::Int64); + + // table references are matched exactly + assert!(conn + .get_schema(&TableReference::partial("main", "t")) + .await + .is_err()); + + // a cloned pool keeps the provided schemas + let cloned = pool.try_clone().await.unwrap(); + let cloned_conn = cloned.connect().await.unwrap(); + assert_eq!( + cloned_conn + .as_async() + .unwrap() + .get_schema(&TableReference::bare("u")) + .await + .unwrap(), + schema_u + ); } } From d658276b87e0d47d21b19f8b80ccb07c6f4fb8ee Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Fri, 16 Oct 2026 10:06:16 +0300 Subject: [PATCH 12/27] optimize imports --- .github/workflows/pr.yaml | 4 +--- core/tests/clickhouse/mod.rs | 13 +++++-------- core/tests/mysql/mod.rs | 5 +---- core/tests/postgres/mod.rs | 7 ++----- crates/sqlite/src/pool.rs | 5 ++--- 5 files changed, 11 insertions(+), 23 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index cb703ed83..3d4793be7 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -201,9 +201,7 @@ jobs: timeout-minutes: 45 env: - # Must match the tag the integration tests start (they pull it only if - # an image with exactly this tag is missing). - ORACLE_DOCKER_IMAGE: gvenzl/oracle-free:latest + ORACLE_DOCKER_IMAGE: docker.io/gvenzl/oracle-free:latest RUST_LOG: info steps: diff --git a/core/tests/clickhouse/mod.rs b/core/tests/clickhouse/mod.rs index 3e89c47e3..fe896bfe3 100644 --- a/core/tests/clickhouse/mod.rs +++ b/core/tests/clickhouse/mod.rs @@ -1,12 +1,15 @@ use clickhouse::Client; use common::get_clickhouse_params; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::{common::TableReference, prelude::SessionContext}; use datafusion_table_providers::{ clickhouse::ClickHouseTableFactory, - sql::db_connection_pool::clickhousepool::ClickHouseConnectionPool, + sql::db_connection_pool::{ + clickhousepool::ClickHouseConnectionPool, dbconnection::AsyncDbConnection, DbConnectionPool, + }, }; use linktime::{ctor, dtor}; -use std::sync::Mutex; +use std::sync::{Arc, Mutex}; mod common; @@ -155,12 +158,6 @@ async fn clickhouse_insert_and_read() { /// `get_schema` instead of inferring it from the database. #[tokio::test] async fn clickhouse_connection_pool_with_schema() { - use datafusion::arrow::datatypes::{DataType, Field, Schema}; - use datafusion_table_providers::sql::db_connection_pool::{ - dbconnection::AsyncDbConnection, DbConnectionPool, - }; - use std::sync::Arc; - let port = { let mut guard = CONTAINER_MANAGER_INSTANCE.lock().unwrap(); let container_manager = guard.as_mut().unwrap(); diff --git a/core/tests/mysql/mod.rs b/core/tests/mysql/mod.rs index 54e423834..51f7f2b10 100644 --- a/core/tests/mysql/mod.rs +++ b/core/tests/mysql/mod.rs @@ -11,13 +11,12 @@ use arrow::{ array::*, datatypes::{i256, DataType, Field, Schema, TimeUnit, UInt16Type}, }; - use datafusion_table_providers::sql::db_connection_pool::dbconnection::AsyncDbConnection; use crate::arrow_record_batch_gen::*; use datafusion::arrow::datatypes::SchemaRef; use datafusion::catalog::TableProviderFactory; -use datafusion::common::{Constraints, ToDFSchema}; +use datafusion::common::{Constraints, TableReference, ToDFSchema}; use datafusion::logical_expr::dml::InsertOp; use datafusion::logical_expr::CreateExternalTable; use datafusion::physical_plan::collect; @@ -1006,8 +1005,6 @@ async fn test_mysql_arrow_oneway() { /// A pool created with `with_schema` must return the schema provided for a table from /// `get_schema` instead of inferring it from the database. async fn test_mysql_connection_pool_with_schema(port: u16) { - use datafusion::common::TableReference; - let table_reference = TableReference::bare("with_schema_table"); let pool = common::get_mysql_connection_pool(port) diff --git a/core/tests/postgres/mod.rs b/core/tests/postgres/mod.rs index eb21ca494..3c30655fa 100644 --- a/core/tests/postgres/mod.rs +++ b/core/tests/postgres/mod.rs @@ -11,7 +11,7 @@ use datafusion::physical_plan::collect; use datafusion::{catalog::TableProvider, execution::context::SessionContext}; use datafusion::{catalog::TableProviderFactory, logical_expr::dml::InsertOp}; use datafusion::{ - common::{Constraints, ToDFSchema}, + common::{Constraints, TableReference, ToDFSchema}, datasource::memory::MemorySourceConfig, }; @@ -19,7 +19,7 @@ use datafusion_federation::schema_cast::record_convert::try_cast_to; use datafusion_table_providers::{ postgres::{DynPostgresConnectionPool, PostgresTableProviderFactory}, - sql::sql_provider_datafusion::SqlTable, + sql::{db_connection_pool::dbconnection::AsyncDbConnection, sql_provider_datafusion::SqlTable}, UnsupportedTypeAction, }; @@ -185,9 +185,6 @@ async fn test_arrow_postgres_one_way() { /// A pool created with `with_schema` must return the schema provided for a table from /// `get_schema` instead of inferring it from the database. async fn test_postgres_connection_pool_with_schema(port: u16) { - use datafusion::common::TableReference; - use datafusion_table_providers::sql::db_connection_pool::dbconnection::AsyncDbConnection; - let table_reference = TableReference::bare("with_schema_table"); let pool = common::get_postgres_connection_pool(port) diff --git a/crates/sqlite/src/pool.rs b/crates/sqlite/src/pool.rs index fb799bc57..189c30cb3 100644 --- a/crates/sqlite/src/pool.rs +++ b/crates/sqlite/src/pool.rs @@ -336,6 +336,8 @@ impl DbConnectionPool for SqliteConnect #[cfg(test)] mod tests { use super::*; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::TableReference; use datafusion_table_providers_common::sql::db_connection_pool::Mode; use rand::RngExt; use rstest::rstest; @@ -470,9 +472,6 @@ mod tests { #[tokio::test] async fn test_sqlite_connection_pool_factory_with_schema() { - use datafusion::arrow::datatypes::{DataType, Field, Schema}; - use datafusion::common::TableReference; - let factory = SqliteConnectionPoolFactory::new( "./test.sqlite", Mode::Memory, From 8101d9db256d781719429c6d7a43c1e4a17c2e4a Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Wed, 14 Oct 2026 17:13:03 +0300 Subject: [PATCH 13/27] Add a provider for Cassandra --- .github/workflows/pr.yaml | 5 +- Cargo.lock | 109 +++- Cargo.toml | 3 +- Makefile | 3 +- README.md | 2 + core/Cargo.toml | 3 + core/src/lib.rs | 6 + core/tests/cassandra/common.rs | 59 ++ core/tests/cassandra/mod.rs | 454 +++++++++++++++ core/tests/integration.rs | 3 + crates/cassandra/Cargo.toml | 27 + crates/cassandra/src/conn.rs | 262 +++++++++ crates/cassandra/src/lib.rs | 57 ++ crates/cassandra/src/pool.rs | 330 +++++++++++ crates/cassandra/src/table.rs | 328 +++++++++++ crates/cassandra/src/types.rs | 981 +++++++++++++++++++++++++++++++++ 16 files changed, 2628 insertions(+), 4 deletions(-) create mode 100644 core/tests/cassandra/common.rs create mode 100644 core/tests/cassandra/mod.rs create mode 100644 crates/cassandra/Cargo.toml create mode 100644 crates/cassandra/src/conn.rs create mode 100644 crates/cassandra/src/lib.rs create mode 100644 crates/cassandra/src/pool.rs create mode 100644 crates/cassandra/src/table.rs create mode 100644 crates/cassandra/src/types.rs diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 3d4793be7..cf3aa8227 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -65,6 +65,7 @@ jobs: crate: - datafusion-table-providers-common - datafusion-table-providers-adbc + - datafusion-table-providers-cassandra - datafusion-table-providers-clickhouse - datafusion-table-providers-duckdb - datafusion-table-providers-flightsql @@ -129,6 +130,7 @@ jobs: PG_DOCKER_IMAGE: postgres:16 MYSQL_DOCKER_IMAGE: mysql:8.0 CLICKHOUSE_DOCKER_IMAGE: clickhouse/clickhouse-server:24.8 + CASSANDRA_DOCKER_IMAGE: cassandra:5.0 # Keep integration test logs concise for faster CI runs. RUST_LOG: info # Serialize Docker-backed tests a bit to avoid stampeding many Postgres containers. @@ -156,7 +158,7 @@ jobs: uses: actions/cache@v6 with: path: /tmp/docker-images - key: docker-images-${{ env.PG_DOCKER_IMAGE }}-${{ env.MYSQL_DOCKER_IMAGE }}-${{ env.CLICKHOUSE_DOCKER_IMAGE }} + key: docker-images-${{ env.PG_DOCKER_IMAGE }}-${{ env.MYSQL_DOCKER_IMAGE }}-${{ env.CLICKHOUSE_DOCKER_IMAGE }}-${{ env.CASSANDRA_DOCKER_IMAGE }} - name: Load or pull Docker images run: | @@ -177,6 +179,7 @@ jobs: load_or_pull postgres "${{ env.PG_DOCKER_IMAGE }}" load_or_pull mysql "${{ env.MYSQL_DOCKER_IMAGE }}" load_or_pull clickhouse "${{ env.CLICKHOUSE_DOCKER_IMAGE }}" + load_or_pull cassandra "${{ env.CASSANDRA_DOCKER_IMAGE }}" - name: Install ODBC & Sqlite run: | diff --git a/Cargo.lock b/Cargo.lock index 9ff3d6a6f..a4cbcd230 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2439,6 +2439,7 @@ dependencies = [ "datafusion", "datafusion-federation", "datafusion-table-providers-adbc", + "datafusion-table-providers-cassandra", "datafusion-table-providers-clickhouse", "datafusion-table-providers-common", "datafusion-table-providers-duckdb", @@ -2460,6 +2461,7 @@ dependencies = [ "rand 0.10.3", "reqwest", "rstest", + "scylla", "secrecy", "serde", "serde_json", @@ -2495,6 +2497,23 @@ dependencies = [ "tracing", ] +[[package]] +name = "datafusion-table-providers-cassandra" +version = "0.14.0" +dependencies = [ + "arrow 59.2.0", + "async-trait", + "bigdecimal", + "datafusion", + "datafusion-table-providers-common", + "futures", + "num-bigint 0.4.8", + "scylla", + "secrecy", + "snafu", + "tokio", +] + [[package]] name = "datafusion-table-providers-clickhouse" version = "0.14.0" @@ -5545,7 +5564,7 @@ dependencies = [ "getrandom 0.4.3", "lru-slab", "rand 0.10.3", - "rand_pcg", + "rand_pcg 0.10.2", "ring", "rustc-hash", "rustls", @@ -5695,6 +5714,15 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" +[[package]] +name = "rand_pcg" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b48ac3f7ffaab7fac4d2376632268aa5f89abdb55f7ebf8f4d11fffccb2320f7" +dependencies = [ + "rand_core 0.9.5", +] + [[package]] name = "rand_pcg" version = "0.10.2" @@ -6108,6 +6136,85 @@ version = "0.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c1257cd4248b4132760d6524d6dda4e053bc648c9070b960929bf50cfb1e7add" +[[package]] +name = "scylla" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29eebcb7e34257f8ce01aaa1469644825bd4b4995d401c362be65d994e3feaa3" +dependencies = [ + "arc-swap", + "async-trait", + "bytes", + "chrono", + "dashmap", + "futures", + "hashbrown 0.17.1", + "itertools 0.15.0", + "rand 0.9.5", + "rand_pcg 0.9.0", + "scylla-cql", + "scylla-cql-core", + "serde", + "serde_json", + "smallvec", + "socket2", + "thiserror 2.0.21", + "tokio", + "tracing", + "uuid", +] + +[[package]] +name = "scylla-cql" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a398b0e78fb3872c4d5afc74e461d6b80dd69d030e13ddc7442e5f32c482e0da" +dependencies = [ + "bigdecimal", + "byteorder", + "bytes", + "chrono", + "itertools 0.15.0", + "lz4_flex 0.14.0", + "num-bigint 0.4.8", + "scylla-cql-core", + "snap", + "stable_deref_trait", + "thiserror 2.0.21", + "tokio", + "uuid", + "yoke", +] + +[[package]] +name = "scylla-cql-core" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72ca9de2eb08d04a9c85002ac179c51353c91152cd2cd58fc3d7cbc1eee8ab16" +dependencies = [ + "bigdecimal", + "byteorder", + "bytes", + "chrono", + "itertools 0.15.0", + "num-bigint 0.4.8", + "scylla-macros", + "thiserror 2.0.21", + "uuid", +] + +[[package]] +name = "scylla-macros" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c82c9c67cb4912cefc8cb2a68ebdcbe0273b66510c2bb06581ae6b68d98498aa" +dependencies = [ + "darling 0.24.1", + "proc-macro2", + "quote", + "syn 3.0.6", +] + [[package]] name = "sea-query" version = "1.0.2" diff --git a/Cargo.toml b/Cargo.toml index 9a135a2bf..9b6a6ec52 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,7 @@ # under the License. [workspace] -members = ["core", "crates/common", "python", "crates/odbc", "crates/clickhouse", "crates/flightsql", "crates/mongodb", "crates/mysql", "crates/sqlite", "crates/postgres", "crates/duckdb", "crates/adbc", "crates/oracle"] +members = ["core", "crates/common", "python", "crates/odbc", "crates/clickhouse", "crates/flightsql", "crates/mongodb", "crates/mysql", "crates/sqlite", "crates/postgres", "crates/duckdb", "crates/adbc", "crates/oracle", "crates/cassandra"] resolver = "2" [workspace.package] @@ -57,6 +57,7 @@ datafusion-table-providers-postgres = { version = "0.14.0", path = "crates/postg datafusion-table-providers-duckdb = { version = "0.14.0", path = "crates/duckdb" } datafusion-table-providers-adbc = { version = "0.14.0", path = "crates/adbc" } datafusion-table-providers-oracle = { version = "0.14.0", path = "crates/oracle" } +datafusion-table-providers-cassandra = { version = "0.14.0", path = "crates/cassandra" } duckdb = { version = "1.10505.0" } adbc_core = { version = "0.24" } adbc_driver_manager = { version = "0.24" } diff --git a/Makefile b/Makefile index a1cabd4d1..75d4d58ca 100644 --- a/Makefile +++ b/Makefile @@ -5,6 +5,7 @@ all: test: cargo test --features clickhouse,duckdb,flight,mysql,postgres,sqlite,adbc -p datafusion-table-providers --lib cargo test -p datafusion-table-providers-oracle + cargo test -p datafusion-table-providers-cassandra .PHONY: lint lint: @@ -12,4 +13,4 @@ lint: .PHONY: test-integration test-integration: - RUST_LOG=$${RUST_LOG:-info} cargo test -p datafusion-table-providers --test integration --no-default-features --features postgres,sqlite,mysql,flight,clickhouse,duckdb,adbc -- --nocapture + RUST_LOG=$${RUST_LOG:-info} cargo test -p datafusion-table-providers --test integration --no-default-features --features postgres,sqlite,mysql,flight,clickhouse,duckdb,adbc,cassandra -- --nocapture diff --git a/README.md b/README.md index ade35297c..18ac453a9 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,7 @@ Prefer a single-provider leaf crate when you only need one backend — for examp - Leaf crates enable `federation` by default (opt out with `default-features = false`). SQLite also defaults to the `bundled` rusqlite feature. - The existing `datafusion-table-providers` facade keeps the same module paths and provider feature names (`duckdb`, `postgres`, `sqlite`, etc.). Enabling a provider feature pulls in that leaf crate with its defaults (including federation). There is no separate `federation` / `*-federation` feature on the facade. - ClickHouse pool connections are exposed as `ClickHouseConnection` (wrapping the ClickHouse `Client`) under `sql::db_connection_pool::dbconnection::clickhouseconn`. +- Cassandra pool connections are exposed as `CassandraConnection` (wrapping a shared `scylla` driver `Session`) under `sql::db_connection_pool::dbconnection::cassandraconn`. The Cassandra crate has no `federation` feature: CQL cannot run the queries federation pushes down, so only projections and `LIMIT` are pushed down to Cassandra. Existing examples continue to use the facade crate and its feature flags. @@ -40,6 +41,7 @@ Existing examples continue to use the facade crate and its feature flags. - ADBC (`datafusion-table-providers-adbc`) - ODBC (`datafusion-table-providers-odbc`) - Oracle (`datafusion-table-providers-oracle`) +- Apache Cassandra (`datafusion-table-providers-cassandra`) ## Development diff --git a/core/Cargo.toml b/core/Cargo.toml index a1a43c3cb..4df6dc954 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -14,6 +14,7 @@ datafusion-table-providers-common = { workspace = true } tokio = { version = "1.52", features = ["macros", "fs", "rt-multi-thread"] } datafusion-table-providers-adbc = { workspace = true, optional = true } +datafusion-table-providers-cassandra = { workspace = true, optional = true } datafusion-table-providers-clickhouse = { workspace = true, optional = true } datafusion-table-providers-duckdb = { workspace = true, optional = true } datafusion-table-providers-flightsql = { workspace = true, optional = true } @@ -58,6 +59,7 @@ oracle = { version = "0.6", features = ["chrono"] } prost = { version = "=0.14.4" } rand = "0.10" reqwest = "0.13" +scylla = "1.9" rstest = "0.26.1" secrecy = "0.10" serde = { version = "1.0", features = ["derive"] } @@ -74,6 +76,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] } default = [] adbc = ["dep:datafusion-table-providers-adbc"] +cassandra = ["dep:datafusion-table-providers-cassandra"] clickhouse = ["dep:datafusion-table-providers-clickhouse"] duckdb = ["dep:datafusion-table-providers-duckdb"] flight = ["dep:datafusion-table-providers-flightsql"] diff --git a/core/src/lib.rs b/core/src/lib.rs index d36704fb8..37382d177 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -39,6 +39,8 @@ pub mod sql { #[cfg(feature = "adbc")] pub use datafusion_table_providers_adbc::pool as adbcpool; + #[cfg(feature = "cassandra")] + pub use datafusion_table_providers_cassandra::pool as cassandrapool; #[cfg(feature = "clickhouse")] pub use datafusion_table_providers_clickhouse::pool as clickhousepool; #[cfg(feature = "duckdb")] @@ -59,6 +61,8 @@ pub mod sql { #[cfg(feature = "adbc")] pub use datafusion_table_providers_adbc::conn as adbcconn; + #[cfg(feature = "cassandra")] + pub use datafusion_table_providers_cassandra::conn as cassandraconn; #[cfg(feature = "clickhouse")] pub use datafusion_table_providers_clickhouse::conn as clickhouseconn; #[cfg(feature = "duckdb")] @@ -79,6 +83,8 @@ pub mod sql { #[cfg(feature = "adbc")] pub use datafusion_table_providers_adbc as adbc; +#[cfg(feature = "cassandra")] +pub use datafusion_table_providers_cassandra as cassandra; #[cfg(feature = "clickhouse")] pub use datafusion_table_providers_clickhouse as clickhouse; #[cfg(feature = "duckdb")] diff --git a/core/tests/cassandra/common.rs b/core/tests/cassandra/common.rs new file mode 100644 index 000000000..eb0954903 --- /dev/null +++ b/core/tests/cassandra/common.rs @@ -0,0 +1,59 @@ +use bollard::models::HealthConfig; +use datafusion_table_providers::sql::db_connection_pool::cassandrapool::CassandraConnectionPool; +use datafusion_table_providers::util::secrets::to_secret_map; +use std::collections::HashMap; +use tracing::instrument; + +use crate::docker::{container_registry, ContainerRunnerBuilder, RunningContainer}; + +const CASSANDRA_DOCKER_CONTAINER: &str = "runtime-integration-test-cassandra"; + +pub(super) fn get_cassandra_params(port: u16) -> HashMap { + HashMap::from([ + ("hosts".to_string(), "127.0.0.1".to_string()), + ("port".to_string(), port.to_string()), + ]) +} + +#[instrument] +pub(super) async fn start_cassandra_docker_container( + port: u16, +) -> Result { + let cassandra_docker_image = std::env::var("CASSANDRA_DOCKER_IMAGE") + .unwrap_or_else(|_| format!("{}cassandra:5.0", container_registry())); + + let running_container = + ContainerRunnerBuilder::new(format!("{CASSANDRA_DOCKER_CONTAINER}-{port}")) + .image(cassandra_docker_image) + .add_port_binding(9042, port) + // Keep the JVM small so that the node starts quickly in CI + .add_env_var("MAX_HEAP_SIZE", "512M") + .add_env_var("HEAP_NEWSIZE", "128M") + .healthcheck(HealthConfig { + test: Some(vec![ + "CMD-SHELL".to_string(), + "cqlsh -e 'DESCRIBE KEYSPACES' || exit 1".to_string(), + ]), + interval: Some(2_000_000_000), // 2s + timeout: Some(10_000_000_000), // 10s + retries: Some(90), + start_period: Some(20_000_000_000), // 20s + start_interval: None, + }) + // Cassandra takes much longer to boot than the other databases + .health_timeout(std::time::Duration::from_secs(300)) + .build()? + .run() + .await?; + + Ok(running_container) +} + +#[instrument] +pub(super) async fn get_cassandra_connection_pool( + port: u16, +) -> Result { + CassandraConnectionPool::new(to_secret_map(get_cassandra_params(port))) + .await + .map_err(|e| anyhow::anyhow!(e)) +} diff --git a/core/tests/cassandra/mod.rs b/core/tests/cassandra/mod.rs new file mode 100644 index 000000000..818c61255 --- /dev/null +++ b/core/tests/cassandra/mod.rs @@ -0,0 +1,454 @@ +use std::sync::{Arc, Mutex}; + +use arrow::array::RecordBatch; +use arrow::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit}; +use arrow::util::display::array_value_to_string; +use datafusion::common::TableReference; +use datafusion::execution::context::SessionContext; +use datafusion_table_providers::cassandra::{CassandraConnectionPool, CassandraTableFactory}; +use datafusion_table_providers::sql::db_connection_pool::dbconnection::AsyncDbConnection; +use datafusion_table_providers::sql::db_connection_pool::DbConnectionPool; +use linktime::{ctor, dtor}; + +use crate::docker::ContainerManager; + +mod common; + +const KEYSPACE: &str = "it_ks"; + +static CONTAINER_MANAGER_INSTANCE: Mutex> = Mutex::new(None); + +#[ctor(unsafe)] +fn global_setup() { + let mut guard = CONTAINER_MANAGER_INSTANCE.lock().unwrap(); + *guard = Some(ContainerManager::default()); +} + +#[dtor(unsafe)] +fn global_teardown() { + let mut guard = CONTAINER_MANAGER_INSTANCE + .lock() + .unwrap_or_else(|e| e.into_inner()); + if let Some(container_manager) = guard.take() { + drop(container_manager); + } +} + +/// All Cassandra tests share one container: Cassandra takes too long to start for one per test. +#[test_log::test(tokio::test)] +async fn test_cassandra() { + let port = { + let mut guard = CONTAINER_MANAGER_INSTANCE.lock().unwrap(); + let container_manager = guard.as_mut().unwrap(); + container_manager + .start_container(common::start_cassandra_docker_container) + .await; + container_manager.port + }; + + let pool = Arc::new( + common::get_cassandra_connection_pool(port) + .await + .expect("Cassandra connection pool should be created"), + ); + create_test_data(&pool).await; + + test_cassandra_types(&pool).await; + test_cassandra_projection_filter_and_limit(&pool).await; + test_cassandra_count(&pool).await; + test_cassandra_explain_shows_pushed_down_cql(&pool).await; + test_cassandra_tables_and_schemas(&pool).await; + test_cassandra_missing_table(&pool).await; + test_cassandra_connection_pool_with_schema(port).await; +} + +async fn execute(pool: &CassandraConnectionPool, cql: &str) { + pool.connect_direct() + .execute(cql, &[]) + .await + .unwrap_or_else(|e| panic!("CQL should be executed: {cql}: {e}")); +} + +async fn create_test_data(pool: &CassandraConnectionPool) { + execute( + pool, + &format!( + "CREATE KEYSPACE IF NOT EXISTS {KEYSPACE} WITH replication = {{'class': 'SimpleStrategy', 'replication_factor': 1}}" + ), + ) + .await; + execute( + pool, + &format!("CREATE TYPE IF NOT EXISTS {KEYSPACE}.address (street text, zip int)"), + ) + .await; + execute( + pool, + &format!( + "CREATE TABLE IF NOT EXISTS {KEYSPACE}.all_types ( + id int PRIMARY KEY, + a_ascii ascii, + a_bigint bigint, + a_blob blob, + a_boolean boolean, + a_date date, + a_decimal decimal, + a_double double, + a_duration duration, + a_float float, + a_inet inet, + a_list list, + a_map map, + a_set set, + a_smallint smallint, + a_text text, + a_time time, + a_timestamp timestamp, + a_timeuuid timeuuid, + a_tinyint tinyint, + a_tuple tuple, + a_udt frozen
, + a_uuid uuid, + a_varint varint, + a_vector vector + )" + ), + ) + .await; + execute( + pool, + &format!( + "INSERT INTO {KEYSPACE}.all_types ( + id, a_ascii, a_bigint, a_blob, a_boolean, a_date, a_decimal, a_double, a_duration, + a_float, a_inet, a_list, a_map, a_set, a_smallint, a_text, a_time, a_timestamp, + a_timeuuid, a_tinyint, a_tuple, a_udt, a_uuid, a_varint, a_vector + ) VALUES ( + 1, 'ascii', 9000000000, 0xcafe, true, '2024-01-15', 123.45, 1.5, 1mo2d3ns, + 2.5, '127.0.0.1', [1, 2, 3], {{'k': 1}}, {{'a', 'b'}}, 32000, 'text', + '13:14:15.123456789', '2024-01-15T10:20:30.123Z', + 50554d6e-29bb-11e5-b345-feff819cdc9f, 100, (7, 'seven'), + {{street: 'Main', zip: 12345}}, 123e4567-e89b-12d3-a456-426614174000, + 123456789012345678901234567890, [1.0, 2.0, 3.0] + )" + ), + ) + .await; + // all columns but the key are null + execute( + pool, + &format!("INSERT INTO {KEYSPACE}.all_types (id) VALUES (2)"), + ) + .await; + + execute( + pool, + &format!( + "CREATE TABLE IF NOT EXISTS {KEYSPACE}.users (id int PRIMARY KEY, name text, age int)" + ), + ) + .await; + for (id, name, age) in [ + (1, "alice", 31), + (2, "bob", 25), + (3, "carol", 42), + (4, "dave", 35), + (5, "eve", 19), + ] { + execute( + pool, + &format!("INSERT INTO {KEYSPACE}.users (id, name, age) VALUES ({id}, '{name}', {age})"), + ) + .await; + } +} + +async fn register_table(ctx: &SessionContext, pool: &Arc, table: &str) { + let table_provider = CassandraTableFactory::new(Arc::clone(pool)) + .table_provider(TableReference::partial(KEYSPACE, table)) + .await + .expect("table provider should be created"); + ctx.register_table(table, table_provider) + .expect("table should be registered"); +} + +async fn query(ctx: &SessionContext, sql: &str) -> Vec { + ctx.sql(sql) + .await + .expect("query should be planned") + .collect() + .await + .expect("query should be executed") +} + +async fn test_cassandra_types(pool: &Arc) { + let ctx = SessionContext::new(); + register_table(&ctx, pool, "all_types").await; + + let batches = query(&ctx, "SELECT * FROM all_types ORDER BY id").await; + let batch = arrow::compute::concat_batches(&batches[0].schema(), &batches) + .expect("batches should be concatenated"); + assert_eq!(batch.num_rows(), 2); + + let item = |data_type| Arc::new(Field::new("item", data_type, true)); + let expected: Vec<(&str, DataType, &str)> = vec![ + ("id", DataType::Int32, "1"), + ("a_ascii", DataType::Utf8, "ascii"), + ("a_bigint", DataType::Int64, "9000000000"), + ("a_blob", DataType::Binary, "cafe"), + ("a_boolean", DataType::Boolean, "true"), + ("a_date", DataType::Date32, "2024-01-15"), + ("a_decimal", DataType::Utf8, "123.45"), + ("a_double", DataType::Float64, "1.5"), + ( + "a_duration", + DataType::Interval(IntervalUnit::MonthDayNano), + "1 mons 2 days 0.000000003 secs", + ), + ("a_float", DataType::Float32, "2.5"), + ("a_inet", DataType::Utf8, "127.0.0.1"), + ("a_list", DataType::List(item(DataType::Int32)), "[1, 2, 3]"), + ( + "a_map", + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int32, true), + ] + .into(), + ), + false, + )), + false, + ), + "{k: 1}", + ), + ("a_set", DataType::List(item(DataType::Utf8)), "[a, b]"), + ("a_smallint", DataType::Int16, "32000"), + ("a_text", DataType::Utf8, "text"), + ( + "a_time", + DataType::Time64(TimeUnit::Nanosecond), + "13:14:15.123456789", + ), + ( + "a_timestamp", + DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())), + "2024-01-15T10:20:30.123Z", + ), + ( + "a_timeuuid", + DataType::Utf8, + "50554d6e-29bb-11e5-b345-feff819cdc9f", + ), + ("a_tinyint", DataType::Int8, "100"), + ( + "a_tuple", + DataType::Struct( + vec![ + Field::new("0", DataType::Int32, true), + Field::new("1", DataType::Utf8, true), + ] + .into(), + ), + "{0: 7, 1: seven}", + ), + ( + "a_udt", + DataType::Struct( + vec![ + Field::new("street", DataType::Utf8, true), + Field::new("zip", DataType::Int32, true), + ] + .into(), + ), + "{street: Main, zip: 12345}", + ), + ( + "a_uuid", + DataType::Utf8, + "123e4567-e89b-12d3-a456-426614174000", + ), + ("a_varint", DataType::Utf8, "123456789012345678901234567890"), + ( + "a_vector", + DataType::FixedSizeList(item(DataType::Float32), 3), + "[1.0, 2.0, 3.0]", + ), + ]; + + let schema = batch.schema(); + assert_eq!(schema.fields().len(), expected.len()); + for (i, (name, data_type, value)) in expected.into_iter().enumerate() { + let field = schema.field(i); + assert_eq!(field.name(), name); + assert_eq!(field.data_type(), &data_type, "type of {name}"); + assert!(field.is_nullable(), "{name} should be nullable"); + + let column = batch.column(i); + assert_eq!( + array_value_to_string(column, 0).expect("value should be formatted"), + value, + "value of {name}" + ); + if name != "id" { + assert!(column.is_null(1), "{name} should be null in the 2nd row"); + } + } +} + +async fn test_cassandra_projection_filter_and_limit(pool: &Arc) { + let ctx = SessionContext::new(); + register_table(&ctx, pool, "users").await; + + // the filter isn't valid CQL (age is not a key column), so DataFusion must apply it + let batches = query( + &ctx, + "SELECT name, age FROM users WHERE age > 30 ORDER BY age DESC", + ) + .await; + let rows: Vec<(String, String)> = batches + .iter() + .flat_map(|batch| { + (0..batch.num_rows()).map(move |row| { + ( + array_value_to_string(batch.column(0), row).unwrap(), + array_value_to_string(batch.column(1), row).unwrap(), + ) + }) + }) + .collect(); + assert_eq!( + rows, + vec![ + ("carol".to_string(), "42".to_string()), + ("dave".to_string(), "35".to_string()), + ("alice".to_string(), "31".to_string()), + ] + ); + + let batches = query(&ctx, "SELECT id FROM users LIMIT 2").await; + let num_rows: usize = batches.iter().map(RecordBatch::num_rows).sum(); + assert_eq!(num_rows, 2); + assert_eq!(batches[0].schema().fields().len(), 1); + + let batches = query(&ctx, "SELECT id FROM users LIMIT 0").await; + let num_rows: usize = batches.iter().map(RecordBatch::num_rows).sum(); + assert_eq!(num_rows, 0); +} + +async fn test_cassandra_count(pool: &Arc) { + let ctx = SessionContext::new(); + register_table(&ctx, pool, "users").await; + + let batches = query(&ctx, "SELECT count(*) FROM users").await; + assert_eq!(array_value_to_string(batches[0].column(0), 0).unwrap(), "5"); + + let batches = query(&ctx, "SELECT count(*) FROM users WHERE age < 30").await; + assert_eq!(array_value_to_string(batches[0].column(0), 0).unwrap(), "2"); +} + +async fn test_cassandra_explain_shows_pushed_down_cql(pool: &Arc) { + let ctx = SessionContext::new(); + register_table(&ctx, pool, "users").await; + + let batches = query(&ctx, "EXPLAIN SELECT name FROM users LIMIT 3").await; + let plan = arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string(); + assert!( + plan.contains(r#"CassandraExec cql=SELECT "name" FROM "it_ks"."users" LIMIT 3"#), + "the projection and limit should be pushed down:\n{plan}" + ); +} + +async fn test_cassandra_tables_and_schemas(pool: &Arc) { + let conn = pool.connect().await.expect("connection should be created"); + let conn = conn.as_async().expect("connection should be async"); + + let keyspaces = conn.schemas().await.expect("keyspaces should be listed"); + assert!(keyspaces.contains(&KEYSPACE.to_string()), "{keyspaces:?}"); + assert!(keyspaces.contains(&"system".to_string()), "{keyspaces:?}"); + + let mut tables = conn + .tables(KEYSPACE) + .await + .expect("tables should be listed"); + tables.sort(); + assert_eq!(tables, vec!["all_types".to_string(), "users".to_string()]); +} + +async fn test_cassandra_missing_table(pool: &Arc) { + let result = CassandraTableFactory::new(Arc::clone(pool)) + .table_provider(TableReference::partial(KEYSPACE, "does_not_exist")) + .await; + let error = result.expect_err("a missing table should be an error"); + assert!( + error.to_string().contains("does_not_exist"), + "the error should name the table: {error}" + ); +} + +/// A pool created with `with_schema` must return the provided schema from `get_schema` +/// instead of inferring it from the database. +async fn test_cassandra_connection_pool_with_schema(port: u16) { + let table_reference = TableReference::partial(KEYSPACE, "users"); + + // the schema is inferred from the database when none is provided + let pool = common::get_cassandra_connection_pool(port) + .await + .expect("Cassandra connection pool should be created"); + let inferred = pool + .connect_direct() + .get_schema(&table_reference) + .await + .expect("schema should be inferred"); + assert_eq!( + inferred.as_ref(), + &Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("age", DataType::Int32, true), + Field::new("name", DataType::Utf8, true), + ]) + ); + + // the provided schema is used as-is, also to convert the values + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::LargeUtf8, true), + ])); + let pool = Arc::new( + common::get_cassandra_connection_pool(port) + .await + .expect("Cassandra connection pool should be created") + .with_schema(Arc::clone(&schema)), + ); + + let actual = pool + .connect_direct() + .get_schema(&table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, schema); + + let conn = pool.connect().await.expect("connection should be created"); + let actual = conn + .as_async() + .expect("connection should be async") + .get_schema(&table_reference) + .await + .expect("schema should be returned"); + assert_eq!(actual, schema); + + let ctx = SessionContext::new(); + register_table(&ctx, &pool, "users").await; + let batches = query(&ctx, "SELECT id, name FROM users WHERE id = 3").await; + assert_eq!(batches[0].schema(), schema); + assert_eq!(array_value_to_string(batches[0].column(0), 0).unwrap(), "3"); + assert_eq!( + array_value_to_string(batches[0].column(1), 0).unwrap(), + "carol" + ); +} diff --git a/core/tests/integration.rs b/core/tests/integration.rs index e59622079..62fbfa9b7 100644 --- a/core/tests/integration.rs +++ b/core/tests/integration.rs @@ -3,9 +3,12 @@ use rand::RngExt; #[cfg(feature = "adbc")] mod adbc; mod arrow_record_batch_gen; +#[cfg(feature = "cassandra")] +mod cassandra; #[cfg(feature = "clickhouse")] mod clickhouse; #[cfg(any( + feature = "cassandra", feature = "clickhouse", feature = "mongodb", feature = "mysql", diff --git a/crates/cassandra/Cargo.toml b/crates/cassandra/Cargo.toml new file mode 100644 index 000000000..e4053a983 --- /dev/null +++ b/crates/cassandra/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "datafusion-table-providers-cassandra" +version = { workspace = true } +readme = { workspace = true } +edition = { workspace = true } +repository = { workspace = true } +license = { workspace = true } +description = "DataFusion Apache Cassandra table provider" + +[dependencies] +arrow = { workspace = true } +async-trait = "0.1" +bigdecimal = "0.4" +datafusion = { workspace = true, default-features = false, features = ["sql"] } +datafusion-table-providers-common = { workspace = true } +futures = "0.3" +num-bigint = "0.4" +scylla = { version = "1.9", features = ["bigdecimal-04", "num-bigint-04"] } +secrecy = "0.10" +snafu = "0.9" + +[dev-dependencies] +tokio = { version = "1.52", features = ["macros", "rt-multi-thread"] } + +[package.metadata.docs.rs] +all-features = true +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/cassandra/src/conn.rs b/crates/cassandra/src/conn.rs new file mode 100644 index 000000000..168f6e76f --- /dev/null +++ b/crates/cassandra/src/conn.rs @@ -0,0 +1,262 @@ +use std::{any::Any, sync::Arc}; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::TableReference; +use datafusion::execution::SendableRecordBatchStream; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use futures::{StreamExt, TryStreamExt}; +use scylla::client::session::Session; +use scylla::value::{CqlValue, Row}; +use snafu::prelude::*; + +use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::{ + self, AsyncDbConnection, DbConnection, GenericError, Result, +}; + +use crate::types::{column_specs_to_schema, rows_to_record_batch}; + +/// Number of rows converted to Arrow at a time; Cassandra pages are 5000 rows by default. +const BATCH_SIZE: usize = 8192; + +#[derive(Debug, Snafu)] +pub enum Error { + #[snafu(display("Failed to query Cassandra: {source}"))] + Query { + source: scylla::errors::PagerExecutionError, + }, + + #[snafu(display("Failed to execute the Cassandra statement: {source}"))] + Execute { + source: scylla::errors::ExecutionError, + }, + + #[snafu(display("Failed to read rows from Cassandra: {source}"))] + ReadRows { + source: scylla::errors::NextRowError, + }, + + #[snafu(display("Unexpected Cassandra result: {source}"))] + TypeCheck { + source: scylla::deserialize::TypeCheckError, + }, + + #[snafu(display("Failed to convert the Cassandra result to Arrow: {source}"))] + Conversion { source: crate::types::Error }, +} + +/// A connection to Cassandra. +/// +/// The driver's [`Session`] manages its own connection pool, so a connection is a cheap handle +/// to a shared session. +pub struct CassandraConnection { + pub session: Arc, + schema: Option, +} + +impl std::fmt::Debug for CassandraConnection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CassandraConnection") + .field("schema", &self.schema) + .finish_non_exhaustive() + } +} + +impl CassandraConnection { + /// Runs a CQL query that returns a single text column and collects its values. + async fn query_strings(&self, cql: &str, params: Vec) -> Result, Error> { + self.session + .query_iter(cql, params) + .await + .context(QuerySnafu)? + .rows_stream::<(String,)>() + .context(TypeCheckSnafu)? + .map_ok(|(value,)| value) + .try_collect() + .await + .context(ReadRowsSnafu) + } +} + +impl DbConnection, CqlValue> for CassandraConnection { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn as_async(&self) -> Option<&dyn AsyncDbConnection, CqlValue>> { + Some(self) + } +} + +#[async_trait] +impl AsyncDbConnection, CqlValue> for CassandraConnection { + fn new(session: Arc) -> Self { + Self { + session, + schema: None, + } + } + + fn new_with_schema(session: Arc, schema: SchemaRef) -> Self { + Self { + session, + schema: Some(schema), + } + } + + /// Returns the tables of the keyspace `schema`. + async fn tables(&self, schema: &str) -> Result, dbconnection::Error> { + self.query_strings( + "SELECT table_name FROM system_schema.tables WHERE keyspace_name = ?", + vec![CqlValue::Text(schema.to_string())], + ) + .await + .boxed() + .context(dbconnection::UnableToGetTablesSnafu) + } + + /// Returns the keyspaces. + async fn schemas(&self) -> Result, dbconnection::Error> { + self.query_strings("SELECT keyspace_name FROM system_schema.keyspaces", vec![]) + .await + .boxed() + .context(dbconnection::UnableToGetSchemasSnafu) + } + + /// Returns the schema of a table, read from the result metadata of a prepared + /// `SELECT *` statement so that no table data is read. + async fn get_schema( + &self, + table_reference: &TableReference, + ) -> Result { + if let Some(schema) = &self.schema { + return Ok(Arc::clone(schema)); + } + + let cql = format!("SELECT * FROM {}", quote_table_reference(table_reference)); + let prepared = self.session.prepare(cql).await.map_err(|e| { + let message = e.to_string(); + let source: GenericError = Box::new(e); + if message.contains("unconfigured table") || message.contains("does not exist") { + dbconnection::Error::UndefinedTable { + table_name: table_reference.to_string(), + source, + } + } else { + dbconnection::Error::UnableToGetSchema { source } + } + })?; + + let col_specs = prepared.get_current_result_set_col_specs(); + column_specs_to_schema(col_specs.get().iter()) + .boxed() + .context(dbconnection::UnableToGetSchemaSnafu) + } + + /// Runs a CQL query and streams its result as Arrow record batches. + /// + /// The batches have `projected_schema` when it is given, otherwise the schema of the result + /// set. The fields of `projected_schema` are matched to the result set columns by name. + async fn query_arrow( + &self, + sql: &str, + params: &[CqlValue], + projected_schema: Option, + ) -> Result { + let pager = self + .session + .query_iter(sql, params.to_vec()) + .await + .context(QuerySnafu)?; + + let column_names: Vec = pager + .column_specs() + .iter() + .map(|spec| spec.name().to_string()) + .collect(); + let schema = match projected_schema { + Some(schema) => schema, + None => column_specs_to_schema(pager.column_specs().iter()).context(ConversionSnafu)?, + }; + + let rows = pager.rows_stream::().context(TypeCheckSnafu)?; + let stream_schema = Arc::clone(&schema); + let batches = rows.chunks(BATCH_SIZE).map(move |chunk| { + let rows = chunk + .into_iter() + .collect::, _>>() + .context(ReadRowsSnafu) + .map_err(to_datafusion_error)?; + rows_to_record_batch(&stream_schema, &column_names, &rows) + .context(ConversionSnafu) + .map_err(to_datafusion_error) + }); + + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, batches))) + } + + /// Executes a CQL statement. + /// + /// Cassandra doesn't report the number of affected rows, so this always returns 0. + async fn execute(&self, sql: &str, params: &[CqlValue]) -> Result { + self.session + .query_unpaged(sql, params.to_vec()) + .await + .context(ExecuteSnafu)?; + Ok(0) + } +} + +fn to_datafusion_error(e: Error) -> datafusion::error::DataFusionError { + datafusion::error::DataFusionError::External(Box::new(e)) +} + +/// Quotes a CQL identifier, so that it is matched case-sensitively. +pub(crate) fn quote_identifier(identifier: &str) -> String { + format!("\"{}\"", identifier.replace('"', "\"\"")) +} + +/// Quotes all parts of a table reference. The catalog part is ignored: Cassandra only has +/// keyspaces and tables. +pub(crate) fn quote_table_reference(table_reference: &TableReference) -> String { + match table_reference.schema() { + Some(keyspace) => format!( + "{}.{}", + quote_identifier(keyspace), + quote_identifier(table_reference.table()) + ), + None => quote_identifier(table_reference.table()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn quotes_identifiers() { + assert_eq!(quote_identifier("users"), "\"users\""); + assert_eq!(quote_identifier("MixedCase"), "\"MixedCase\""); + assert_eq!(quote_identifier("we\"ird"), "\"we\"\"ird\""); + } + + #[test] + fn quotes_table_references() { + assert_eq!( + quote_table_reference(&TableReference::bare("users")), + "\"users\"" + ); + assert_eq!( + quote_table_reference(&TableReference::partial("ks", "Users")), + "\"ks\".\"Users\"" + ); + assert_eq!( + quote_table_reference(&TableReference::full("catalog", "ks", "users")), + "\"ks\".\"users\"" + ); + } +} diff --git a/crates/cassandra/src/lib.rs b/crates/cassandra/src/lib.rs new file mode 100644 index 000000000..b515830ce --- /dev/null +++ b/crates/cassandra/src/lib.rs @@ -0,0 +1,57 @@ +//! A DataFusion [`TableProvider`] for Apache Cassandra, built on the +//! [ScyllaDB Rust driver](https://docs.rs/scylla), so it works with ScyllaDB too. +//! +//! CQL is far more restrictive than SQL: it has no joins or aggregations, and `WHERE` clauses +//! are limited to key columns. Scans therefore only push the projection and `LIMIT` down to +//! Cassandra. Filters, aggregations and joins are evaluated by DataFusion. + +pub mod conn; +pub mod pool; +mod table; +pub mod types; + +use std::sync::Arc; + +use datafusion::catalog::TableProvider; +use datafusion::common::TableReference; + +use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::AsyncDbConnection; + +pub use crate::pool::CassandraConnectionPool; +pub use crate::table::{CassandraExec, CassandraTable}; + +/// Creates [`TableProvider`]s for Cassandra tables. +#[derive(Debug)] +pub struct CassandraTableFactory { + pool: Arc, +} + +impl CassandraTableFactory { + pub fn new(pool: impl Into>) -> Self { + Self { pool: pool.into() } + } + + /// Creates a table provider for `table_reference`, given as `table` (in the pool's + /// keyspace) or `keyspace.table`. Names are case-sensitive. + /// + /// # Errors + /// + /// Returns an error if the schema of the table cannot be read. + pub async fn table_provider( + &self, + table_reference: TableReference, + ) -> Result, Box> + { + let schema = self + .pool + .connect_direct() + .get_schema(&table_reference) + .await?; + + Ok(Arc::new(CassandraTable::new( + table_reference, + Arc::clone(&self.pool), + schema, + ))) + } +} diff --git a/crates/cassandra/src/pool.rs b/crates/cassandra/src/pool.rs new file mode 100644 index 000000000..a7a40edd9 --- /dev/null +++ b/crates/cassandra/src/pool.rs @@ -0,0 +1,330 @@ +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; +use scylla::client::session::Session; +use scylla::client::session_builder::SessionBuilder; +use scylla::value::CqlValue; +use secrecy::{ExposeSecret, SecretString}; +use snafu::prelude::*; + +use datafusion_table_providers_common::sql::db_connection_pool::{ + dbconnection::{AsyncDbConnection, DbConnection}, + DbConnectionPool, JoinPushDown, +}; + +use crate::conn::CassandraConnection; + +type PoolResult> = std::result::Result; + +const DEFAULT_HOST: &str = "localhost"; +const DEFAULT_PORT: u16 = 9042; + +#[derive(Debug, Snafu)] +pub enum Error { + #[snafu(display("Failed to connect to Cassandra: {source}"))] + Connection { + source: Box, + }, + + #[snafu(display("Invalid value '{value}' for parameter '{parameter_name}'"))] + InvalidParameter { + parameter_name: String, + value: String, + }, + + #[snafu(display( + "Parameter '{parameter_name}' requires parameter '{required}' to be set too" + ))] + MissingParameter { + parameter_name: String, + required: String, + }, +} + +/// The Cassandra connection settings, parsed from the pool parameters. +#[derive(Debug, PartialEq)] +struct CassandraConfig { + /// Contact points as `host:port`. + known_nodes: Vec, + keyspace: Option, + credentials: Option<(String, String)>, + local_datacenter: Option, + connection_timeout: Option, +} + +impl CassandraConfig { + fn from_params(params: &HashMap) -> Result { + let param = |name: &str| { + params + .get(name) + .map(|v| v.expose_secret().trim().to_string()) + .filter(|v| !v.is_empty()) + }; + + let port = match param("port") { + Some(port) => port.parse::().map_err(|_| Error::InvalidParameter { + parameter_name: "port".to_string(), + value: port, + })?, + None => DEFAULT_PORT, + }; + + let hosts = param("hosts") + .or_else(|| param("host")) + .unwrap_or_else(|| DEFAULT_HOST.to_string()); + let known_nodes = hosts + .split(',') + .map(str::trim) + .filter(|host| !host.is_empty()) + .map(|host| with_default_port(host, port)) + .collect(); + + let credentials = match (param("user"), param("password")) { + (Some(user), Some(password)) => Some((user, password)), + (None, None) => None, + (Some(_), None) => { + return MissingParameterSnafu { + parameter_name: "user", + required: "password", + } + .fail() + } + (None, Some(_)) => { + return MissingParameterSnafu { + parameter_name: "password", + required: "user", + } + .fail() + } + }; + + let connection_timeout = match param("connection_timeout_ms") { + Some(timeout) => Some(Duration::from_millis(timeout.parse::().map_err( + |_| Error::InvalidParameter { + parameter_name: "connection_timeout_ms".to_string(), + value: timeout, + }, + )?)), + None => None, + }; + + Ok(Self { + known_nodes, + keyspace: param("keyspace"), + credentials, + local_datacenter: param("local_datacenter"), + connection_timeout, + }) + } + + fn session_builder(&self) -> SessionBuilder { + let mut builder = SessionBuilder::new().known_nodes(&self.known_nodes); + if let Some(keyspace) = &self.keyspace { + builder = builder.use_keyspace(keyspace, true); + } + if let Some((user, password)) = &self.credentials { + builder = builder.user(user, password); + } + if let Some(datacenter) = &self.local_datacenter { + builder = builder.prefer_datacenter(datacenter.clone()); + } + if let Some(timeout) = self.connection_timeout { + builder = builder.connection_timeout(timeout); + } + builder + } +} + +/// Appends `port` to `host` unless it already has one. IPv6 addresses must be given in +/// brackets, e.g. `[::1]:9042`, to carry a port. +fn with_default_port(host: &str, port: u16) -> String { + let has_port = match host.rsplit_once(':') { + Some((address, port)) => { + port.parse::().is_ok() && (!address.contains(':') || address.ends_with(']')) + } + None => false, + }; + if has_port { + host.to_string() + } else if host.contains(':') && !host.starts_with('[') { + format!("[{host}]:{port}") + } else { + format!("{host}:{port}") + } +} + +/// A connection pool for Apache Cassandra (and other CQL-compatible databases). +/// +/// The underlying driver [`Session`] maintains connections to every node of the cluster, so the +/// pool shares one session among all its connections. +#[derive(Clone)] +pub struct CassandraConnectionPool { + session: Arc, + join_push_down: JoinPushDown, + schema: Option, +} + +impl std::fmt::Debug for CassandraConnectionPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CassandraConnectionPool") + .field("join_push_down", &self.join_push_down) + .field("schema", &self.schema) + .finish_non_exhaustive() + } +} + +impl CassandraConnectionPool { + /// Creates a new `CassandraConnectionPool`. + /// + /// # Arguments + /// + /// * `params` - The connection parameters: + /// * `hosts` - Comma-separated contact points, as `host` or `host:port`. Defaults to `localhost`. + /// `host` is accepted as an alias. + /// * `port` - The port of contact points that don't specify one. Defaults to `9042`. + /// * `keyspace` - The keyspace that unqualified table names are resolved in. + /// * `user` and `password` - Credentials for password authentication. + /// * `local_datacenter` - The datacenter to prefer when routing requests. + /// * `connection_timeout_ms` - The timeout for establishing connections. + /// + /// # Errors + /// + /// Returns an error if the parameters are invalid or the cluster cannot be reached. + pub async fn new(params: HashMap) -> PoolResult { + let config = CassandraConfig::from_params(¶ms)?; + let session = config + .session_builder() + .build() + .await + .map_err(|e| Error::Connection { + source: Box::new(e), + })?; + + Ok(Self::from_session(Arc::new(session))) + } + + /// Creates a pool from an existing driver session. + #[must_use] + pub fn from_session(session: Arc) -> Self { + // Joins can't be pushed down: CQL has no joins. + Self { + session, + join_push_down: JoinPushDown::Disallow, + schema: None, + } + } + + /// Use the provided `schema` for all connections instead of inferring it from the database. + #[must_use] + pub fn with_schema(mut self, schema: SchemaRef) -> Self { + self.schema = Some(schema); + self + } + + /// Returns the driver session. + #[must_use] + pub fn session(&self) -> Arc { + Arc::clone(&self.session) + } + + /// Returns a direct connection to the underlying database. + #[must_use] + pub fn connect_direct(&self) -> CassandraConnection { + match &self.schema { + Some(schema) => { + CassandraConnection::new_with_schema(self.session(), Arc::clone(schema)) + } + None => CassandraConnection::new(self.session()), + } + } +} + +#[async_trait] +impl DbConnectionPool, CqlValue> for CassandraConnectionPool { + async fn connect(&self) -> PoolResult, CqlValue>>> { + Ok(Box::new(self.connect_direct())) + } + + fn join_push_down(&self) -> JoinPushDown { + self.join_push_down.clone() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn params(pairs: &[(&str, &str)]) -> HashMap { + pairs + .iter() + .map(|(k, v)| (k.to_string(), SecretString::from(v.to_string()))) + .collect() + } + + #[test] + fn uses_defaults() { + let config = CassandraConfig::from_params(&HashMap::new()).unwrap(); + assert_eq!( + config, + CassandraConfig { + known_nodes: vec!["localhost:9042".to_string()], + keyspace: None, + credentials: None, + local_datacenter: None, + connection_timeout: None, + } + ); + } + + #[test] + fn parses_all_params() { + let config = CassandraConfig::from_params(¶ms(&[ + ("hosts", "node1, node2:9043,10.0.0.1"), + ("port", "19042"), + ("keyspace", "ks"), + ("user", "cassandra"), + ("password", "secret"), + ("local_datacenter", "dc1"), + ("connection_timeout_ms", "1500"), + ])) + .unwrap(); + assert_eq!( + config, + CassandraConfig { + known_nodes: vec![ + "node1:19042".to_string(), + "node2:9043".to_string(), + "10.0.0.1:19042".to_string(), + ], + keyspace: Some("ks".to_string()), + credentials: Some(("cassandra".to_string(), "secret".to_string())), + local_datacenter: Some("dc1".to_string()), + connection_timeout: Some(Duration::from_millis(1500)), + } + ); + } + + #[test] + fn accepts_host_alias() { + let config = CassandraConfig::from_params(¶ms(&[("host", "db")])).unwrap(); + assert_eq!(config.known_nodes, vec!["db:9042".to_string()]); + } + + #[test] + fn rejects_invalid_params() { + assert!(CassandraConfig::from_params(¶ms(&[("port", "abc")])).is_err()); + assert!(CassandraConfig::from_params(¶ms(&[("connection_timeout_ms", "-1")])).is_err()); + assert!(CassandraConfig::from_params(¶ms(&[("user", "cassandra")])).is_err()); + assert!(CassandraConfig::from_params(¶ms(&[("password", "secret")])).is_err()); + } + + #[test] + fn adds_default_port() { + assert_eq!(with_default_port("db", 9042), "db:9042"); + assert_eq!(with_default_port("db:9043", 9042), "db:9043"); + assert_eq!(with_default_port("::1", 9042), "[::1]:9042"); + assert_eq!(with_default_port("[::1]", 9042), "[::1]:9042"); + assert_eq!(with_default_port("[::1]:9043", 9042), "[::1]:9043"); + } +} diff --git a/crates/cassandra/src/table.rs b/crates/cassandra/src/table.rs new file mode 100644 index 000000000..eca07abd0 --- /dev/null +++ b/crates/cassandra/src/table.rs @@ -0,0 +1,328 @@ +use std::fmt; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::catalog::Session; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{project_schema, TableReference}; +use datafusion::datasource::TableType; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::logical_expr::Expr; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, +}; +use futures::TryStreamExt; + +use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::AsyncDbConnection; + +use crate::conn::{quote_identifier, quote_table_reference}; +use crate::pool::CassandraConnectionPool; + +/// A Cassandra table. +/// +/// Filters are not pushed down (see the crate documentation), so DataFusion applies them. +pub struct CassandraTable { + table_reference: TableReference, + pool: Arc, + schema: SchemaRef, +} + +impl fmt::Debug for CassandraTable { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("CassandraTable") + .field("table_reference", &self.table_reference) + .field("schema", &self.schema) + .finish_non_exhaustive() + } +} + +impl CassandraTable { + #[must_use] + pub fn new( + table_reference: TableReference, + pool: Arc, + schema: SchemaRef, + ) -> Self { + Self { + table_reference, + pool, + schema, + } + } +} + +#[async_trait] +impl datafusion::catalog::TableProvider for CassandraTable { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + limit: Option, + ) -> DataFusionResult> { + Ok(Arc::new(CassandraExec::try_new( + self.table_reference.clone(), + Arc::clone(&self.pool), + &self.schema, + projection, + limit, + )?)) + } +} + +/// Scans a Cassandra table with `SELECT FROM [LIMIT n]`. +pub struct CassandraExec { + table_reference: TableReference, + pool: Arc, + /// The columns to select from Cassandra. + columns: Vec, + projected_schema: SchemaRef, + limit: Option, + properties: Arc, +} + +impl CassandraExec { + /// Creates a scan of the `projection` columns of `schema`. + /// + /// # Errors + /// + /// Returns an error if the projection doesn't match the schema. + pub fn try_new( + table_reference: TableReference, + pool: Arc, + schema: &SchemaRef, + projection: Option<&Vec>, + limit: Option, + ) -> DataFusionResult { + let projected_schema = project_schema(schema, projection)?; + let columns = scan_columns(&table_reference, schema, &projected_schema)?; + + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&projected_schema)), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + + Ok(Self { + table_reference, + pool, + columns, + projected_schema, + limit, + properties, + }) + } + + /// The CQL statement this scan runs. + #[must_use] + pub fn cql(&self) -> String { + build_cql(&self.table_reference, &self.columns, self.limit) + } +} + +/// Returns the columns to select from Cassandra to produce `projected_schema`. +fn scan_columns( + table_reference: &TableReference, + schema: &SchemaRef, + projected_schema: &SchemaRef, +) -> DataFusionResult> { + if projected_schema.fields().is_empty() { + // CQL can't select zero columns. Queries like `SELECT count(*)` only need the row + // count, so select the first column (a partition key column for `SELECT *`) and + // drop it when converting the rows. + let first = schema.fields().first().ok_or_else(|| { + DataFusionError::Plan(format!("Cassandra table {table_reference} has no columns")) + })?; + Ok(vec![first.name().clone()]) + } else { + Ok(projected_schema + .fields() + .iter() + .map(|f| f.name().clone()) + .collect()) + } +} + +fn build_cql(table_reference: &TableReference, columns: &[String], limit: Option) -> String { + let columns = columns + .iter() + .map(|c| quote_identifier(c)) + .collect::>() + .join(", "); + let mut cql = format!( + "SELECT {columns} FROM {}", + quote_table_reference(table_reference) + ); + // CQL requires a positive limit; `LIMIT 0` doesn't query Cassandra at all + if let Some(limit) = limit.filter(|l| *l > 0) { + cql.push_str(&format!(" LIMIT {limit}")); + } + cql +} + +impl fmt::Debug for CassandraExec { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "CassandraExec cql={}", self.cql()) + } +} + +impl DisplayAs for CassandraExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "CassandraExec cql={}", self.cql()) + } +} + +impl ExecutionPlan for CassandraExec { + fn name(&self) -> &'static str { + "CassandraExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.projected_schema) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DataFusionResult> { + Ok(self) + } + + fn supports_limit_pushdown(&self) -> bool { + true + } + + fn fetch(&self) -> Option { + self.limit + } + + fn with_fetch(&self, limit: Option) -> Option> { + Some(Arc::new(Self { + table_reference: self.table_reference.clone(), + pool: Arc::clone(&self.pool), + columns: self.columns.clone(), + projected_schema: Arc::clone(&self.projected_schema), + limit, + properties: Arc::clone(&self.properties), + })) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> DataFusionResult { + let schema = Arc::clone(&self.projected_schema); + if self.limit == Some(0) { + return Ok(Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::empty(), + ))); + } + + let cql = self.cql(); + let conn = self.pool.connect_direct(); + + let stream_schema = Arc::clone(&schema); + let stream = futures::stream::once(async move { + conn.query_arrow(&cql, &[], Some(stream_schema)) + .await + .map_err(DataFusionError::External) + }) + .try_flatten(); + + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + + fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("Name", DataType::Utf8, true), + Field::new("age", DataType::Int32, true), + ])) + } + + fn cql( + table_reference: TableReference, + projection: Option<&Vec>, + limit: Option, + ) -> String { + let schema = schema(); + let projected_schema = project_schema(&schema, projection).unwrap(); + let columns = scan_columns(&table_reference, &schema, &projected_schema).unwrap(); + build_cql(&table_reference, &columns, limit) + } + + #[test] + fn selects_all_columns() { + assert_eq!( + cql(TableReference::partial("ks", "users"), None, None), + "SELECT \"id\", \"Name\", \"age\" FROM \"ks\".\"users\"" + ); + } + + #[test] + fn selects_projected_columns_with_limit() { + assert_eq!( + cql(TableReference::bare("users"), Some(&vec![2, 1]), Some(10)), + "SELECT \"age\", \"Name\" FROM \"users\" LIMIT 10" + ); + } + + #[test] + fn selects_first_column_for_empty_projection() { + assert_eq!( + cql(TableReference::bare("users"), Some(&vec![]), None), + "SELECT \"id\" FROM \"users\"" + ); + } + + #[test] + fn omits_zero_limit() { + assert_eq!( + cql(TableReference::bare("users"), Some(&vec![0]), Some(0)), + "SELECT \"id\" FROM \"users\"" + ); + } + + #[test] + fn fails_for_table_without_columns() { + let schema = Arc::new(Schema::empty()); + assert!(scan_columns(&TableReference::bare("t"), &schema, &schema).is_err()); + } +} diff --git a/crates/cassandra/src/types.rs b/crates/cassandra/src/types.rs new file mode 100644 index 000000000..ecf3f1420 --- /dev/null +++ b/crates/cassandra/src/types.rs @@ -0,0 +1,981 @@ +//! Conversion between CQL and Arrow types and values. + +use std::sync::Arc; + +use arrow::array::{ + new_null_array, ArrayRef, BinaryArray, BooleanArray, Date32Array, Decimal128Array, + FixedSizeListArray, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, + IntervalMonthDayNanoArray, LargeBinaryArray, LargeStringArray, ListArray, MapArray, + RecordBatch, RecordBatchOptions, StringArray, StructArray, Time64MicrosecondArray, + Time64NanosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampNanosecondArray, TimestampSecondArray, +}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::datatypes::{ + DataType, Field, Fields, IntervalMonthDayNano, Schema, SchemaRef, TimeUnit, +}; +use bigdecimal::BigDecimal; +use num_bigint::BigInt; +use scylla::frame::response::result::{CollectionType, ColumnSpec, ColumnType, NativeType}; +use scylla::value::{CqlValue, Row}; +use snafu::prelude::*; + +#[derive(Debug, Snafu)] +pub enum Error { + #[snafu(display("Unsupported CQL type for column '{column_name}': {cql_type}"))] + UnsupportedCqlType { + column_name: String, + cql_type: String, + }, + + #[snafu(display("Unsupported Arrow type {data_type} for a Cassandra column"))] + UnsupportedArrowType { data_type: DataType }, + + #[snafu(display("Cannot convert CQL value {value:?} to Arrow type {data_type}"))] + IncompatibleValue { value: String, data_type: DataType }, + + #[snafu(display("Column '{column_name}' is not part of the query result"))] + MissingColumn { column_name: String }, + + #[snafu(display("Failed to build Arrow data: {source}"))] + Arrow { source: arrow::error::ArrowError }, +} + +pub type Result = std::result::Result; + +/// Number of days between -5877641-06-23 (the CQL `date` epoch) and the Unix epoch. +const CQL_DATE_UNIX_EPOCH_OFFSET: i64 = 1 << 31; + +/// Builds the Arrow schema for the columns of a CQL result set. +/// +/// Every field is nullable: Cassandra returns `null` for any regular column that has no value. +pub fn column_specs_to_schema<'a>( + specs: impl IntoIterator>, +) -> Result { + let fields = specs + .into_iter() + .map(|spec| { + Ok(Field::new( + spec.name(), + map_column_type(spec.name(), spec.typ())?, + true, + )) + }) + .collect::>>()?; + Ok(Arc::new(Schema::new(fields))) +} + +/// Maps a CQL column type to its Arrow equivalent. +/// +/// Types without an exact Arrow counterpart are mapped to their string representation: +/// `decimal` and `varint` (arbitrary precision), `uuid`, `timeuuid` and `inet`. +pub fn map_column_type(column_name: &str, typ: &ColumnType<'_>) -> Result { + let data_type = match typ { + ColumnType::Native(native) => match native { + NativeType::Ascii | NativeType::Text => DataType::Utf8, + NativeType::Boolean => DataType::Boolean, + NativeType::Blob => DataType::Binary, + NativeType::Counter | NativeType::BigInt => DataType::Int64, + NativeType::Date => DataType::Date32, + NativeType::Decimal | NativeType::Varint => DataType::Utf8, + NativeType::Double => DataType::Float64, + NativeType::Duration => { + DataType::Interval(arrow::datatypes::IntervalUnit::MonthDayNano) + } + NativeType::Float => DataType::Float32, + NativeType::Int => DataType::Int32, + NativeType::SmallInt => DataType::Int16, + NativeType::TinyInt => DataType::Int8, + NativeType::Timestamp => DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())), + NativeType::Time => DataType::Time64(TimeUnit::Nanosecond), + NativeType::Uuid | NativeType::Timeuuid | NativeType::Inet => DataType::Utf8, + _ => return unsupported(column_name, typ), + }, + ColumnType::Collection { + typ: collection, .. + } => match collection { + CollectionType::List(element) | CollectionType::Set(element) => { + DataType::List(list_item_field(map_column_type(column_name, element)?)) + } + CollectionType::Map(key, value) => map_data_type( + map_column_type(column_name, key)?, + map_column_type(column_name, value)?, + ), + _ => return unsupported(column_name, typ), + }, + ColumnType::Vector { + typ: element, + dimensions, + } => DataType::FixedSizeList( + list_item_field(map_column_type(column_name, element)?), + i32::from(*dimensions), + ), + ColumnType::UserDefinedType { definition, .. } => DataType::Struct( + definition + .field_types + .iter() + .map(|(name, typ)| { + Ok(Field::new( + name.as_ref(), + map_column_type(column_name, typ)?, + true, + )) + }) + .collect::>()?, + ), + ColumnType::Tuple(types) => DataType::Struct( + types + .iter() + .enumerate() + .map(|(i, typ)| { + Ok(Field::new( + i.to_string(), + map_column_type(column_name, typ)?, + true, + )) + }) + .collect::>()?, + ), + _ => return unsupported(column_name, typ), + }; + Ok(data_type) +} + +fn unsupported(column_name: &str, typ: &ColumnType<'_>) -> Result { + UnsupportedCqlTypeSnafu { + column_name, + cql_type: format!("{typ:?}"), + } + .fail() +} + +fn list_item_field(data_type: DataType) -> Arc { + Arc::new(Field::new("item", data_type, true)) +} + +fn map_data_type(key: DataType, value: DataType) -> DataType { + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("key", key, false), + Field::new("value", value, true), + ])), + false, + )), + false, + ) +} + +/// Converts CQL rows to a [`RecordBatch`] with the given `schema`. +/// +/// `column_names` are the names of the result set columns, in the order of the values in each +/// row. Every field of `schema` is looked up by name among them, so `schema` may be a projection +/// of the result set. A `schema` without fields produces a batch that only carries the row count. +pub fn rows_to_record_batch( + schema: &SchemaRef, + column_names: &[String], + rows: &[Row], +) -> Result { + let columns = schema + .fields() + .iter() + .map(|field| { + let index = column_names + .iter() + .position(|name| name == field.name()) + .context(MissingColumnSnafu { + column_name: field.name(), + })?; + let values: Vec> = rows + .iter() + .map(|row| row.columns.get(index).and_then(Option::as_ref)) + .collect(); + values_to_array(field.data_type(), &values) + }) + .collect::>>()?; + + RecordBatch::try_new_with_options( + Arc::clone(schema), + columns, + &RecordBatchOptions::new().with_row_count(Some(rows.len())), + ) + .context(ArrowSnafu) +} + +/// Converts a column of CQL values to an Arrow array of the given type. +pub fn values_to_array(data_type: &DataType, values: &[Option<&CqlValue>]) -> Result { + // CQL has a legacy "empty" value that is distinct from null; it carries no data. + let values: Vec> = values + .iter() + .map(|v| v.filter(|v| !matches!(v, CqlValue::Empty))) + .collect(); + let values = values.as_slice(); + + let array: ArrayRef = match data_type { + DataType::Null => new_null_array(data_type, values.len()), + DataType::Boolean => Arc::new( + values + .iter() + .map(|v| convert(*v, data_type, CqlValue::as_boolean)) + .collect::>()?, + ), + DataType::Int8 => Arc::new(collect_ints::(data_type, values)?), + DataType::Int16 => Arc::new(collect_ints::(data_type, values)?), + DataType::Int32 => Arc::new(collect_ints::(data_type, values)?), + DataType::Int64 => Arc::new(collect_ints::(data_type, values)?), + DataType::Float32 => Arc::new( + values + .iter() + .map(|v| convert(*v, data_type, |v| v.as_float())) + .collect::>()?, + ), + DataType::Float64 => Arc::new( + values + .iter() + .map(|v| { + convert(*v, data_type, |v| { + v.as_double().or_else(|| v.as_float().map(f64::from)) + }) + }) + .collect::>()?, + ), + DataType::Utf8 => Arc::new( + values + .iter() + .map(|v| convert(*v, data_type, value_to_string)) + .collect::>()?, + ), + DataType::LargeUtf8 => Arc::new( + values + .iter() + .map(|v| convert(*v, data_type, value_to_string)) + .collect::>()?, + ), + DataType::Binary => Arc::new( + values + .iter() + .map(|v| convert(*v, data_type, |v| v.as_blob().cloned())) + .collect::>>()? + .into_iter() + .collect::(), + ), + DataType::LargeBinary => Arc::new( + values + .iter() + .map(|v| convert(*v, data_type, |v| v.as_blob().cloned())) + .collect::>>()? + .into_iter() + .collect::(), + ), + DataType::Date32 => Arc::new( + values + .iter() + .map(|v| { + convert(*v, data_type, |v| { + v.as_cql_date().and_then(|d| { + i32::try_from(i64::from(d.0) - CQL_DATE_UNIX_EPOCH_OFFSET).ok() + }) + }) + }) + .collect::>()?, + ), + DataType::Time64(unit) => { + let nanos = values + .iter() + .map(|v| convert(*v, data_type, |v| v.as_cql_time().map(|t| t.0))) + .collect::>>()?; + match unit { + TimeUnit::Nanosecond => Arc::new(Time64NanosecondArray::from(nanos)), + TimeUnit::Microsecond => Arc::new(Time64MicrosecondArray::from( + nanos + .into_iter() + .map(|n| n.map(|n| n / 1_000)) + .collect::>(), + )), + _ => { + return UnsupportedArrowTypeSnafu { + data_type: data_type.clone(), + } + .fail() + } + } + } + DataType::Timestamp(unit, tz) => { + let millis = values + .iter() + .map(|v| convert(*v, data_type, |v| v.as_cql_timestamp().map(|t| t.0))) + .collect::>>()?; + let scaled = |factor: i64, divisor: i64| -> Vec> { + millis + .iter() + .map(|m| m.map(|m| m.saturating_mul(factor) / divisor)) + .collect() + }; + let array: ArrayRef = match unit { + TimeUnit::Second => Arc::new(TimestampSecondArray::from(scaled(1, 1_000))), + TimeUnit::Millisecond => Arc::new(TimestampMillisecondArray::from(millis)), + TimeUnit::Microsecond => { + Arc::new(TimestampMicrosecondArray::from(scaled(1_000, 1))) + } + TimeUnit::Nanosecond => { + Arc::new(TimestampNanosecondArray::from(scaled(1_000_000, 1))) + } + }; + match tz { + Some(tz) => { + arrow::compute::cast(&array, &DataType::Timestamp(*unit, Some(Arc::clone(tz)))) + .context(ArrowSnafu)? + } + None => array, + } + } + DataType::Interval(arrow::datatypes::IntervalUnit::MonthDayNano) => Arc::new( + values + .iter() + .map(|v| { + convert(*v, data_type, |v| { + v.as_cql_duration() + .map(|d| IntervalMonthDayNano::new(d.months, d.days, d.nanoseconds)) + }) + }) + .collect::>()?, + ), + DataType::Decimal128(precision, scale) => Arc::new( + values + .iter() + .map(|v| convert(*v, data_type, |v| value_to_i128_with_scale(v, *scale))) + .collect::>()? + .with_precision_and_scale(*precision, *scale) + .context(ArrowSnafu)?, + ), + DataType::List(item) => { + let mut offsets = Vec::with_capacity(values.len() + 1); + offsets.push(0i32); + let mut children = Vec::new(); + for value in values { + if let Some(value) = value { + let elements = collection_elements(value).context(IncompatibleValueSnafu { + value: format!("{value:?}"), + data_type: data_type.clone(), + })?; + children.extend(elements.iter().map(Some)); + } + offsets.push(i32::try_from(children.len()).map_err(|_| { + Error::IncompatibleValue { + value: "list with more than i32::MAX elements".to_string(), + data_type: data_type.clone(), + } + })?); + } + let child = values_to_array(item.data_type(), &children)?; + Arc::new( + ListArray::try_new( + Arc::clone(item), + OffsetBuffer::new(offsets.into()), + child, + nulls(values), + ) + .context(ArrowSnafu)?, + ) + } + DataType::FixedSizeList(item, size) => { + let size_usize = usize::try_from(*size).unwrap_or_default(); + let mut children = Vec::with_capacity(values.len() * size_usize); + for value in values { + match value { + Some(value) => { + let elements = collection_elements(value) + .filter(|elements| elements.len() == size_usize) + .context(IncompatibleValueSnafu { + value: format!("{value:?}"), + data_type: data_type.clone(), + })?; + children.extend(elements.iter().map(Some)); + } + None => children.extend(std::iter::repeat_n(None, size_usize)), + } + } + let child = values_to_array(item.data_type(), &children)?; + Arc::new( + FixedSizeListArray::try_new(Arc::clone(item), *size, child, nulls(values)) + .context(ArrowSnafu)?, + ) + } + DataType::Struct(fields) => { + let mut columns: Vec>> = + vec![Vec::with_capacity(values.len()); fields.len()]; + for value in values { + for (i, field) in fields.iter().enumerate() { + let field_value = match value { + None => None, + Some(CqlValue::UserDefinedType { + fields: udt_fields, .. + }) => udt_fields + .iter() + .find(|(name, _)| name == field.name()) + .and_then(|(_, v)| v.as_ref()), + Some(CqlValue::Tuple(elements)) => elements.get(i).and_then(Option::as_ref), + Some(value) => { + return IncompatibleValueSnafu { + value: format!("{value:?}"), + data_type: data_type.clone(), + } + .fail() + } + }; + columns[i].push(field_value); + } + } + let children = fields + .iter() + .zip(columns) + .map(|(field, column)| values_to_array(field.data_type(), &column)) + .collect::>>()?; + Arc::new( + StructArray::try_new(fields.clone(), children, nulls(values)) + .context(ArrowSnafu)?, + ) + } + DataType::Map(entries_field, sorted) => { + let DataType::Struct(entry_fields) = entries_field.data_type() else { + return UnsupportedArrowTypeSnafu { + data_type: data_type.clone(), + } + .fail(); + }; + let mut offsets = Vec::with_capacity(values.len() + 1); + offsets.push(0i32); + let mut keys = Vec::new(); + let mut items = Vec::new(); + for value in values { + match value { + Some(CqlValue::Map(entries)) => { + for (k, v) in entries { + keys.push(Some(k)); + items.push(Some(v)); + } + } + Some(value) => { + return IncompatibleValueSnafu { + value: format!("{value:?}"), + data_type: data_type.clone(), + } + .fail() + } + None => {} + } + offsets.push( + i32::try_from(keys.len()).map_err(|_| Error::IncompatibleValue { + value: "map with more than i32::MAX entries".to_string(), + data_type: data_type.clone(), + })?, + ); + } + let entries = StructArray::try_new( + entry_fields.clone(), + vec![ + values_to_array(entry_fields[0].data_type(), &keys)?, + values_to_array(entry_fields[1].data_type(), &items)?, + ], + None, + ) + .context(ArrowSnafu)?; + Arc::new( + MapArray::try_new( + Arc::clone(entries_field), + OffsetBuffer::new(offsets.into()), + entries, + nulls(values), + *sorted, + ) + .context(ArrowSnafu)?, + ) + } + _ => { + return UnsupportedArrowTypeSnafu { + data_type: data_type.clone(), + } + .fail() + } + }; + Ok(array) +} + +/// Converts a non-null CQL value with `f`, failing if `f` cannot handle it. +fn convert( + value: Option<&CqlValue>, + data_type: &DataType, + f: impl FnOnce(&CqlValue) -> Option, +) -> Result> { + match value { + None => Ok(None), + Some(value) => f(value).map(Some).context(IncompatibleValueSnafu { + value: format!("{value:?}"), + data_type: data_type.clone(), + }), + } +} + +/// Collects integer values into an array, widening or narrowing them to `N` when they fit. +fn collect_ints(data_type: &DataType, values: &[Option<&CqlValue>]) -> Result +where + N: TryFrom, + A: FromIterator>, +{ + values + .iter() + .map(|v| { + convert(*v, data_type, |v| { + value_to_i64(v).and_then(|i| N::try_from(i).ok()) + }) + }) + .collect() +} + +fn value_to_i64(value: &CqlValue) -> Option { + match value { + CqlValue::TinyInt(v) => Some(i64::from(*v)), + CqlValue::SmallInt(v) => Some(i64::from(*v)), + CqlValue::Int(v) => Some(i64::from(*v)), + CqlValue::BigInt(v) => Some(*v), + CqlValue::Counter(c) => Some(c.0), + CqlValue::Varint(v) => i64::try_from(BigInt::from(v.clone())).ok(), + _ => None, + } +} + +/// Converts a numeric CQL value to the unscaled `i128` of a decimal with the given `scale`. +fn value_to_i128_with_scale(value: &CqlValue, scale: i8) -> Option { + let decimal = match value { + CqlValue::Decimal(d) => BigDecimal::from(d.clone()), + CqlValue::Varint(v) => BigDecimal::from(BigInt::from(v.clone())), + other => BigDecimal::from(value_to_i64(other)?), + }; + let (unscaled, _) = decimal + .with_scale(i64::from(scale)) + .into_bigint_and_exponent(); + i128::try_from(unscaled).ok() +} + +/// Returns the string representation of a scalar CQL value. +fn value_to_string(value: &CqlValue) -> Option { + let s = match value { + CqlValue::Ascii(s) | CqlValue::Text(s) => s.clone(), + CqlValue::Uuid(u) => u.to_string(), + CqlValue::Timeuuid(u) => u.to_string(), + CqlValue::Inet(ip) => ip.to_string(), + CqlValue::Decimal(d) => BigDecimal::from(d.clone()).to_string(), + CqlValue::Varint(v) => BigInt::from(v.clone()).to_string(), + CqlValue::Boolean(b) => b.to_string(), + CqlValue::Double(d) => d.to_string(), + CqlValue::Float(f) => f.to_string(), + other => value_to_i64(other)?.to_string(), + }; + Some(s) +} + +fn collection_elements(value: &CqlValue) -> Option<&[CqlValue]> { + match value { + CqlValue::List(elements) | CqlValue::Set(elements) | CqlValue::Vector(elements) => { + Some(elements) + } + _ => None, + } +} + +fn nulls(values: &[Option<&CqlValue>]) -> Option { + if values.iter().all(Option::is_some) { + None + } else { + Some(NullBuffer::from_iter(values.iter().map(Option::is_some))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, AsArray}; + use arrow::datatypes::{Int32Type, IntervalUnit}; + use scylla::frame::response::result::{TableSpec, UserDefinedType}; + use scylla::value::{ + Counter, CqlDate, CqlDecimal, CqlDuration, CqlTime, CqlTimestamp, CqlVarint, + }; + use std::borrow::Cow; + + fn native(native: NativeType) -> ColumnType<'static> { + ColumnType::Native(native) + } + + fn list_of(typ: ColumnType<'static>) -> ColumnType<'static> { + ColumnType::Collection { + frozen: false, + typ: CollectionType::List(Box::new(typ)), + } + } + + #[test] + fn maps_native_types() { + let cases = [ + (NativeType::Ascii, DataType::Utf8), + (NativeType::Text, DataType::Utf8), + (NativeType::Boolean, DataType::Boolean), + (NativeType::Blob, DataType::Binary), + (NativeType::Counter, DataType::Int64), + (NativeType::BigInt, DataType::Int64), + (NativeType::Int, DataType::Int32), + (NativeType::SmallInt, DataType::Int16), + (NativeType::TinyInt, DataType::Int8), + (NativeType::Double, DataType::Float64), + (NativeType::Float, DataType::Float32), + (NativeType::Date, DataType::Date32), + (NativeType::Decimal, DataType::Utf8), + (NativeType::Varint, DataType::Utf8), + (NativeType::Uuid, DataType::Utf8), + (NativeType::Timeuuid, DataType::Utf8), + (NativeType::Inet, DataType::Utf8), + (NativeType::Time, DataType::Time64(TimeUnit::Nanosecond)), + ( + NativeType::Timestamp, + DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())), + ), + ( + NativeType::Duration, + DataType::Interval(IntervalUnit::MonthDayNano), + ), + ]; + for (cql, arrow) in cases { + assert_eq!( + map_column_type("c", &native(cql.clone())).unwrap(), + arrow, + "{cql:?}" + ); + } + } + + #[test] + fn maps_composite_types() { + assert_eq!( + map_column_type("c", &list_of(native(NativeType::Int))).unwrap(), + DataType::List(list_item_field(DataType::Int32)) + ); + + let set = ColumnType::Collection { + frozen: true, + typ: CollectionType::Set(Box::new(native(NativeType::Text))), + }; + assert_eq!( + map_column_type("c", &set).unwrap(), + DataType::List(list_item_field(DataType::Utf8)) + ); + + let map = ColumnType::Collection { + frozen: false, + typ: CollectionType::Map( + Box::new(native(NativeType::Text)), + Box::new(native(NativeType::BigInt)), + ), + }; + assert_eq!( + map_column_type("c", &map).unwrap(), + map_data_type(DataType::Utf8, DataType::Int64) + ); + + let vector = ColumnType::Vector { + typ: Box::new(native(NativeType::Float)), + dimensions: 3, + }; + assert_eq!( + map_column_type("c", &vector).unwrap(), + DataType::FixedSizeList(list_item_field(DataType::Float32), 3) + ); + + let tuple = ColumnType::Tuple(vec![native(NativeType::Int), native(NativeType::Text)]); + assert_eq!( + map_column_type("c", &tuple).unwrap(), + DataType::Struct(Fields::from(vec![ + Field::new("0", DataType::Int32, true), + Field::new("1", DataType::Utf8, true), + ])) + ); + + let udt = ColumnType::UserDefinedType { + frozen: true, + definition: Arc::new(UserDefinedType { + name: Cow::Borrowed("address"), + keyspace: Cow::Borrowed("ks"), + field_types: vec![ + (Cow::Borrowed("street"), native(NativeType::Text)), + (Cow::Borrowed("zip"), native(NativeType::Int)), + ], + }), + }; + assert_eq!( + map_column_type("c", &udt).unwrap(), + DataType::Struct(Fields::from(vec![ + Field::new("street", DataType::Utf8, true), + Field::new("zip", DataType::Int32, true), + ])) + ); + } + + #[test] + fn builds_schema_from_column_specs() { + let table = TableSpec::borrowed("ks", "t"); + let specs = [ + ColumnSpec::borrowed("id", native(NativeType::Int), table.clone()), + ColumnSpec::borrowed("tags", list_of(native(NativeType::Text)), table), + ]; + let schema = column_specs_to_schema(&specs).unwrap(); + assert_eq!( + schema.as_ref(), + &Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new( + "tags", + DataType::List(list_item_field(DataType::Utf8)), + true + ), + ]) + ); + } + + #[test] + fn converts_scalar_values() { + let text = CqlValue::Text("hello".to_string()); + let array = values_to_array(&DataType::Utf8, &[Some(&text), None]).unwrap(); + let array = array.as_string::(); + assert_eq!(array.value(0), "hello"); + assert!(array.is_null(1)); + + let values = [ + CqlValue::TinyInt(1), + CqlValue::SmallInt(2), + CqlValue::Int(3), + CqlValue::BigInt(4), + CqlValue::Counter(Counter(5)), + ]; + let refs: Vec<_> = values.iter().map(Some).collect(); + let array = values_to_array(&DataType::Int64, &refs).unwrap(); + assert_eq!( + array.as_primitive::().values(), + &[1, 2, 3, 4, 5] + ); + + // a value that doesn't fit into the target type is an error, not a silent truncation + let big = CqlValue::BigInt(i64::MAX); + assert!(values_to_array(&DataType::Int32, &[Some(&big)]).is_err()); + + // the CQL "empty" value is read as null + let array = values_to_array(&DataType::Int32, &[Some(&CqlValue::Empty)]).unwrap(); + assert!(array.is_null(0)); + } + + #[test] + fn converts_temporal_values() { + // 1970-01-02 is one day after the Unix epoch + let date = CqlValue::Date(CqlDate((1 << 31) + 1)); + let array = values_to_array(&DataType::Date32, &[Some(&date)]).unwrap(); + assert_eq!( + array + .as_primitive::() + .value(0), + 1 + ); + + let ts = CqlValue::Timestamp(CqlTimestamp(1_500)); + let utc = DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())); + let array = values_to_array(&utc, &[Some(&ts)]).unwrap(); + assert_eq!(array.data_type(), &utc); + assert_eq!( + array + .as_primitive::() + .value(0), + 1_500 + ); + let array = values_to_array( + &DataType::Timestamp(TimeUnit::Microsecond, None), + &[Some(&ts)], + ) + .unwrap(); + assert_eq!( + array + .as_primitive::() + .value(0), + 1_500_000 + ); + + let time = CqlValue::Time(CqlTime(42_000)); + let array = + values_to_array(&DataType::Time64(TimeUnit::Nanosecond), &[Some(&time)]).unwrap(); + assert_eq!( + array + .as_primitive::() + .value(0), + 42_000 + ); + + let duration = CqlValue::Duration(CqlDuration { + months: 1, + days: 2, + nanoseconds: 3, + }); + let array = values_to_array( + &DataType::Interval(IntervalUnit::MonthDayNano), + &[Some(&duration)], + ) + .unwrap(); + assert_eq!( + array + .as_primitive::() + .value(0), + IntervalMonthDayNano::new(1, 2, 3) + ); + } + + #[test] + fn converts_arbitrary_precision_numbers() { + // 12345 * 10^-2 + let decimal = CqlValue::Decimal(CqlDecimal::from_signed_be_bytes_slice_and_exponent( + &12345i32.to_be_bytes(), + 2, + )); + let varint = CqlValue::Varint(CqlVarint::from_signed_bytes_be_slice( + &(-42i64).to_be_bytes(), + )); + + let array = values_to_array(&DataType::Utf8, &[Some(&decimal), Some(&varint)]).unwrap(); + let array = array.as_string::(); + assert_eq!(array.value(0), "123.45"); + assert_eq!(array.value(1), "-42"); + + let array = values_to_array( + &DataType::Decimal128(10, 3), + &[Some(&decimal), Some(&varint)], + ) + .unwrap(); + let array = array.as_primitive::(); + assert_eq!(array.value(0), 123_450); + assert_eq!(array.value(1), -42_000); + } + + #[test] + fn converts_nested_values() { + let list = CqlValue::List(vec![CqlValue::Int(1), CqlValue::Int(2)]); + let empty = CqlValue::Set(vec![]); + let data_type = DataType::List(list_item_field(DataType::Int32)); + let array = values_to_array(&data_type, &[Some(&list), None, Some(&empty)]).unwrap(); + let array = array.as_list::(); + assert_eq!(array.len(), 3); + assert!(array.is_null(1)); + assert_eq!(array.value(0).as_primitive::().values(), &[1, 2]); + assert_eq!(array.value(2).len(), 0); + + let map = CqlValue::Map(vec![(CqlValue::Text("a".into()), CqlValue::BigInt(1))]); + let array = values_to_array( + &map_data_type(DataType::Utf8, DataType::Int64), + &[Some(&map)], + ) + .unwrap(); + let array = array.as_map(); + assert_eq!(array.keys().as_string::().value(0), "a"); + assert_eq!( + array + .values() + .as_primitive::() + .value(0), + 1 + ); + + let udt = CqlValue::UserDefinedType { + keyspace: "ks".into(), + name: "address".into(), + fields: vec![ + ("street".into(), Some(CqlValue::Text("Main".into()))), + ("zip".into(), None), + ], + }; + let fields = Fields::from(vec![ + Field::new("street", DataType::Utf8, true), + Field::new("zip", DataType::Int32, true), + ]); + let array = values_to_array(&DataType::Struct(fields), &[Some(&udt), None]).unwrap(); + let array = array.as_struct(); + assert!(array.is_null(1)); + assert_eq!(array.column(0).as_string::().value(0), "Main"); + assert!(array.column(1).is_null(0)); + + let tuple = CqlValue::Tuple(vec![Some(CqlValue::Int(7)), None]); + let fields = Fields::from(vec![ + Field::new("0", DataType::Int32, true), + Field::new("1", DataType::Utf8, true), + ]); + let array = values_to_array(&DataType::Struct(fields), &[Some(&tuple)]).unwrap(); + let array = array.as_struct(); + assert_eq!(array.column(0).as_primitive::().value(0), 7); + assert!(array.column(1).is_null(0)); + + let vector = CqlValue::Vector(vec![CqlValue::Float(1.0), CqlValue::Float(2.0)]); + let data_type = DataType::FixedSizeList(list_item_field(DataType::Float32), 2); + let array = values_to_array(&data_type, &[Some(&vector), None]).unwrap(); + let array = array.as_fixed_size_list(); + assert!(array.is_null(1)); + assert_eq!( + array + .value(0) + .as_primitive::() + .values(), + &[1.0, 2.0] + ); + } + + #[test] + fn builds_record_batch_by_column_name() { + let rows = vec![ + Row { + columns: vec![Some(CqlValue::Int(1)), Some(CqlValue::Text("a".into()))], + }, + Row { + columns: vec![Some(CqlValue::Int(2)), None], + }, + ]; + let column_names = vec!["id".to_string(), "name".to_string()]; + + // the schema may list the columns in a different order than the result set + let schema = Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, true), + Field::new("id", DataType::Int64, true), + ])); + let batch = rows_to_record_batch(&schema, &column_names, &rows).unwrap(); + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.column(0).as_string::().value(0), "a"); + assert!(batch.column(0).is_null(1)); + assert_eq!( + batch + .column(1) + .as_primitive::() + .values(), + &[1, 2] + ); + + // an empty schema keeps the row count, e.g. for `SELECT count(*)` + let batch = rows_to_record_batch(&Arc::new(Schema::empty()), &column_names, &rows).unwrap(); + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.num_columns(), 0); + + let schema = Arc::new(Schema::new(vec![Field::new( + "missing", + DataType::Int32, + true, + )])); + assert!(rows_to_record_batch(&schema, &column_names, &rows).is_err()); + + // a type that the values can't be converted to is an error + let schema = Arc::new(Schema::new(vec![Field::new( + "name", + DataType::Boolean, + true, + )])); + assert!(rows_to_record_batch(&schema, &column_names, &rows).is_err()); + } +} From a6eea01183bda78f69eb7ec1b03d15f371dfffda Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Fri, 16 Oct 2026 15:48:34 +0300 Subject: [PATCH 14/27] Fix the IT tests for Cassandra --- core/tests/cassandra/common.rs | 4 ++++ core/tests/cassandra/mod.rs | 2 +- crates/cassandra/src/conn.rs | 20 ++++++++++---------- crates/cassandra/src/pool.rs | 29 ++++++++++++++++------------- 4 files changed, 31 insertions(+), 24 deletions(-) diff --git a/core/tests/cassandra/common.rs b/core/tests/cassandra/common.rs index eb0954903..eb2c52920 100644 --- a/core/tests/cassandra/common.rs +++ b/core/tests/cassandra/common.rs @@ -26,6 +26,10 @@ pub(super) async fn start_cassandra_docker_container( ContainerRunnerBuilder::new(format!("{CASSANDRA_DOCKER_CONTAINER}-{port}")) .image(cassandra_docker_image) .add_port_binding(9042, port) + // The driver connects to the address each node advertises in `system.local`, not to + // the contact point. By default that is the container's internal IP, which, combined + // with the published host port, isn't reachable from the host. + .add_env_var("CASSANDRA_BROADCAST_RPC_ADDRESS", "127.0.0.1") // Keep the JVM small so that the node starts quickly in CI .add_env_var("MAX_HEAP_SIZE", "512M") .add_env_var("HEAP_NEWSIZE", "128M") diff --git a/core/tests/cassandra/mod.rs b/core/tests/cassandra/mod.rs index 818c61255..29f1f04c6 100644 --- a/core/tests/cassandra/mod.rs +++ b/core/tests/cassandra/mod.rs @@ -423,7 +423,7 @@ async fn test_cassandra_connection_pool_with_schema(port: u16) { common::get_cassandra_connection_pool(port) .await .expect("Cassandra connection pool should be created") - .with_schema(Arc::clone(&schema)), + .with_schema(table_reference.clone(), Arc::clone(&schema)), ); let actual = pool diff --git a/crates/cassandra/src/conn.rs b/crates/cassandra/src/conn.rs index 168f6e76f..79ef2ccaf 100644 --- a/crates/cassandra/src/conn.rs +++ b/crates/cassandra/src/conn.rs @@ -1,5 +1,4 @@ use std::{any::Any, sync::Arc}; - use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::TableReference; @@ -10,9 +9,7 @@ use scylla::client::session::Session; use scylla::value::{CqlValue, Row}; use snafu::prelude::*; -use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::{ - self, AsyncDbConnection, DbConnection, GenericError, Result, -}; +use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::{self, AsyncDbConnection, DbConnection, GenericError, Result, TableSchemas}; use crate::types::{column_specs_to_schema, rows_to_record_batch}; @@ -51,13 +48,13 @@ pub enum Error { /// to a shared session. pub struct CassandraConnection { pub session: Arc, - schema: Option, + schemas: Arc, } impl std::fmt::Debug for CassandraConnection { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("CassandraConnection") - .field("schema", &self.schema) + .field("schemas", &self.schemas) .finish_non_exhaustive() } } @@ -97,14 +94,17 @@ impl AsyncDbConnection, CqlValue> for CassandraConnection { fn new(session: Arc) -> Self { Self { session, - schema: None, + schemas: Default::default(), } } - fn new_with_schema(session: Arc, schema: SchemaRef) -> Self { + fn new_with_schemas(session: Arc, schemas: Arc) -> Self + where + Self: Sized + { Self { session, - schema: Some(schema), + schemas } } @@ -133,7 +133,7 @@ impl AsyncDbConnection, CqlValue> for CassandraConnection { &self, table_reference: &TableReference, ) -> Result { - if let Some(schema) = &self.schema { + if let Some(schema) = self.schemas.get(table_reference) { return Ok(Arc::clone(schema)); } diff --git a/crates/cassandra/src/pool.rs b/crates/cassandra/src/pool.rs index a7a40edd9..74c93223f 100644 --- a/crates/cassandra/src/pool.rs +++ b/crates/cassandra/src/pool.rs @@ -2,6 +2,7 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::TableReference; use scylla::client::session::Session; use scylla::client::session_builder::SessionBuilder; use scylla::value::CqlValue; @@ -12,7 +13,7 @@ use datafusion_table_providers_common::sql::db_connection_pool::{ dbconnection::{AsyncDbConnection, DbConnection}, DbConnectionPool, JoinPushDown, }; - +use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::TableSchemas; use crate::conn::CassandraConnection; type PoolResult> = std::result::Result; @@ -162,14 +163,14 @@ fn with_default_port(host: &str, port: u16) -> String { pub struct CassandraConnectionPool { session: Arc, join_push_down: JoinPushDown, - schema: Option, + schemas: Arc, } impl std::fmt::Debug for CassandraConnectionPool { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("CassandraConnectionPool") .field("join_push_down", &self.join_push_down) - .field("schema", &self.schema) + .field("schemas", &self.schemas) .finish_non_exhaustive() } } @@ -211,14 +212,21 @@ impl CassandraConnectionPool { Self { session, join_push_down: JoinPushDown::Disallow, - schema: None, + schemas: Default::default(), } } - /// Use the provided `schema` for all connections instead of inferring it from the database. + /// Use the provided `schema` for `table_reference` instead of inferring it from the database. + /// + /// Call it once per table to provide the schemas of several tables. Table references are + /// matched exactly, e.g. `users` and `public.users` are different tables. #[must_use] - pub fn with_schema(mut self, schema: SchemaRef) -> Self { - self.schema = Some(schema); + pub fn with_schema( + mut self, + table_reference: impl Into, + schema: SchemaRef, + ) -> Self { + Arc::make_mut(&mut self.schemas).insert(table_reference.into(), schema); self } @@ -231,12 +239,7 @@ impl CassandraConnectionPool { /// Returns a direct connection to the underlying database. #[must_use] pub fn connect_direct(&self) -> CassandraConnection { - match &self.schema { - Some(schema) => { - CassandraConnection::new_with_schema(self.session(), Arc::clone(schema)) - } - None => CassandraConnection::new(self.session()), - } + CassandraConnection::new_with_schemas(self.session(), Arc::clone(&self.schemas),) } } From 7a52f19e0ff60ab3de379bf959ab6986c46cb7e2 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Fri, 16 Oct 2026 16:11:46 +0300 Subject: [PATCH 15/27] fmt --- crates/cassandra/src/conn.rs | 13 ++++++------- crates/cassandra/src/pool.rs | 6 +++--- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/crates/cassandra/src/conn.rs b/crates/cassandra/src/conn.rs index 79ef2ccaf..3cad432bf 100644 --- a/crates/cassandra/src/conn.rs +++ b/crates/cassandra/src/conn.rs @@ -1,4 +1,3 @@ -use std::{any::Any, sync::Arc}; use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::TableReference; @@ -8,8 +7,11 @@ use futures::{StreamExt, TryStreamExt}; use scylla::client::session::Session; use scylla::value::{CqlValue, Row}; use snafu::prelude::*; +use std::{any::Any, sync::Arc}; -use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::{self, AsyncDbConnection, DbConnection, GenericError, Result, TableSchemas}; +use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::{ + self, AsyncDbConnection, DbConnection, GenericError, Result, TableSchemas, +}; use crate::types::{column_specs_to_schema, rows_to_record_batch}; @@ -100,12 +102,9 @@ impl AsyncDbConnection, CqlValue> for CassandraConnection { fn new_with_schemas(session: Arc, schemas: Arc) -> Self where - Self: Sized + Self: Sized, { - Self { - session, - schemas - } + Self { session, schemas } } /// Returns the tables of the keyspace `schema`. diff --git a/crates/cassandra/src/pool.rs b/crates/cassandra/src/pool.rs index 74c93223f..1439e8d17 100644 --- a/crates/cassandra/src/pool.rs +++ b/crates/cassandra/src/pool.rs @@ -9,12 +9,12 @@ use scylla::value::CqlValue; use secrecy::{ExposeSecret, SecretString}; use snafu::prelude::*; +use crate::conn::CassandraConnection; +use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::TableSchemas; use datafusion_table_providers_common::sql::db_connection_pool::{ dbconnection::{AsyncDbConnection, DbConnection}, DbConnectionPool, JoinPushDown, }; -use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::TableSchemas; -use crate::conn::CassandraConnection; type PoolResult> = std::result::Result; @@ -239,7 +239,7 @@ impl CassandraConnectionPool { /// Returns a direct connection to the underlying database. #[must_use] pub fn connect_direct(&self) -> CassandraConnection { - CassandraConnection::new_with_schemas(self.session(), Arc::clone(&self.schemas),) + CassandraConnection::new_with_schemas(self.session(), Arc::clone(&self.schemas)) } } From 618c5ef1ce06bdcb6b6056b7b8e54e634d9c4865 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Fri, 16 Oct 2026 17:36:04 +0300 Subject: [PATCH 16/27] Add info about the Cassandra table provider and an example --- README.md | 35 ++++++++++++++++ core/Cargo.toml | 5 +++ core/examples/cassandra.rs | 86 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 126 insertions(+) create mode 100644 core/examples/cassandra.rs diff --git a/README.md b/README.md index 18ac453a9..9cad8909d 100644 --- a/README.md +++ b/README.md @@ -192,6 +192,41 @@ EOF cargo run -p datafusion-table-providers --example mysql --features mysql ``` +### Cassandra + +The Cassandra provider is built on the [ScyllaDB Rust driver](https://crates.io/crates/scylla), so it works with ScyllaDB too. CQL has no joins or aggregations and can filter only on key columns, so only projections and `LIMIT` are pushed down to Cassandra; filters, aggregations and joins are evaluated by DataFusion. + +In order to run the Cassandra example, you need to have a Cassandra server running. You can use the following command to start a Cassandra server in a Docker container the example can use: + +```bash +# The driver connects to the address the node advertises, so make it advertise one reachable from the host +docker run --name cassandra \ + -e CASSANDRA_BROADCAST_RPC_ADDRESS=127.0.0.1 \ + -p 9042:9042 \ + -d cassandra:5.0 +# Wait for the Cassandra server to start (this may take a minute) +until docker exec cassandra cqlsh -e 'DESCRIBE KEYSPACES' >/dev/null 2>&1; do + sleep 5 +done + +# Create a table in the Cassandra server and insert some data +docker exec -i cassandra cqlsh </dev/null 2>&1; do sleep 5; done +/// +/// # Create a table and insert sample data +/// docker exec -i cassandra cqlsh < Date: Fri, 16 Oct 2026 18:41:47 +0300 Subject: [PATCH 17/27] Add write support to Cassandra TableProvider --- Cargo.lock | 1 + README.md | 2 +- core/examples/cassandra.rs | 16 +- core/tests/cassandra/mod.rs | 211 +++++++++-- crates/cassandra/Cargo.toml | 1 + crates/cassandra/src/lib.rs | 25 ++ crates/cassandra/src/types.rs | 663 +++++++++++++++++++++++++++++++++- crates/cassandra/src/write.rs | 236 ++++++++++++ 8 files changed, 1115 insertions(+), 40 deletions(-) create mode 100644 crates/cassandra/src/write.rs diff --git a/Cargo.lock b/Cargo.lock index a4cbcd230..13328f713 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2512,6 +2512,7 @@ dependencies = [ "secrecy", "snafu", "tokio", + "uuid", ] [[package]] diff --git a/README.md b/README.md index 9cad8909d..6d725bd20 100644 --- a/README.md +++ b/README.md @@ -194,7 +194,7 @@ cargo run -p datafusion-table-providers --example mysql --features mysql ### Cassandra -The Cassandra provider is built on the [ScyllaDB Rust driver](https://crates.io/crates/scylla), so it works with ScyllaDB too. CQL has no joins or aggregations and can filter only on key columns, so only projections and `LIMIT` are pushed down to Cassandra; filters, aggregations and joins are evaluated by DataFusion. +The Cassandra provider is built on the [ScyllaDB Rust driver](https://crates.io/crates/scylla), so it works with ScyllaDB too. CQL has no joins or aggregations and can filter only on key columns, so only projections and `LIMIT` are pushed down to Cassandra; filters, aggregations and joins are evaluated by DataFusion. Tables created with `CassandraTableFactory::read_write_table_provider` also support `INSERT INTO`, which writes each row with a CQL `INSERT` (an upsert), and `INSERT OVERWRITE`, which truncates the table first. In order to run the Cassandra example, you need to have a Cassandra server running. You can use the following command to start a Cassandra server in a Docker container the example can use: diff --git a/core/examples/cassandra.rs b/core/examples/cassandra.rs index 45ad72820..23f85d5c9 100644 --- a/core/examples/cassandra.rs +++ b/core/examples/cassandra.rs @@ -12,6 +12,7 @@ use datafusion_table_providers::{ /// 2. Create and use CassandraTableFactory to generate TableProvider /// 3. Register TableProvider with DataFusion /// 4. Use SQL queries to access Cassandra table data +/// 5. Write data to a Cassandra table with `INSERT INTO` /// /// Prerequisites: /// Start a Cassandra server using Docker: @@ -53,25 +54,34 @@ async fn main() { ); // Create Cassandra table provider factory - // Used to generate TableProvider instances that can read Cassandra table data + // Used to generate TableProvider instances that can read and write Cassandra table data let table_factory = CassandraTableFactory::new(cassandra_pool); // Create DataFusion session context let ctx = SessionContext::new(); // Register the Cassandra "companies" table as "companies" + // `read_write_table_provider` also supports INSERT INTO; use `table_provider` for read-only access ctx.register_table( "companies", table_factory - .table_provider(TableReference::bare("companies")) + .read_write_table_provider(TableReference::bare("companies")) .await .expect("failed to create table provider"), ) .expect("failed to register table"); + // Write Example: Insert rows into the table. + // As in CQL, inserting a row with an existing primary key replaces that row. + let df = ctx + .sql("INSERT INTO companies (id, name) VALUES (2, 'Initech'), (3, 'Umbrella Corporation')") + .await + .expect("insert failed"); + df.show().await.expect("show failed"); + // Query Example 1: Read the whole table let df = ctx - .sql("SELECT * FROM companies") + .sql("SELECT * FROM companies ORDER BY id") .await .expect("select failed"); df.show().await.expect("show failed"); diff --git a/core/tests/cassandra/mod.rs b/core/tests/cassandra/mod.rs index 29f1f04c6..435e6b1c0 100644 --- a/core/tests/cassandra/mod.rs +++ b/core/tests/cassandra/mod.rs @@ -16,6 +16,35 @@ mod common; const KEYSPACE: &str = "it_ks"; +/// The columns of the `all_types` table, which has a column of every supported CQL type. +const ALL_TYPES_COLUMNS: &str = " + id int PRIMARY KEY, + a_ascii ascii, + a_bigint bigint, + a_blob blob, + a_boolean boolean, + a_date date, + a_decimal decimal, + a_double double, + a_duration duration, + a_float float, + a_inet inet, + a_list list, + a_map map, + a_set set, + a_smallint smallint, + a_text text, + a_time time, + a_timestamp timestamp, + a_timeuuid timeuuid, + a_tinyint tinyint, + a_tuple tuple, + a_udt frozen
, + a_uuid uuid, + a_varint varint, + a_vector vector +"; + static CONTAINER_MANAGER_INSTANCE: Mutex> = Mutex::new(None); #[ctor(unsafe)] @@ -60,6 +89,8 @@ async fn test_cassandra() { test_cassandra_tables_and_schemas(&pool).await; test_cassandra_missing_table(&pool).await; test_cassandra_connection_pool_with_schema(port).await; + test_cassandra_insert_all_types(&pool).await; + test_cassandra_insert_values(&pool).await; } async fn execute(pool: &CassandraConnectionPool, cql: &str) { @@ -82,39 +113,13 @@ async fn create_test_data(pool: &CassandraConnectionPool) { &format!("CREATE TYPE IF NOT EXISTS {KEYSPACE}.address (street text, zip int)"), ) .await; - execute( - pool, - &format!( - "CREATE TABLE IF NOT EXISTS {KEYSPACE}.all_types ( - id int PRIMARY KEY, - a_ascii ascii, - a_bigint bigint, - a_blob blob, - a_boolean boolean, - a_date date, - a_decimal decimal, - a_double double, - a_duration duration, - a_float float, - a_inet inet, - a_list list, - a_map map, - a_set set, - a_smallint smallint, - a_text text, - a_time time, - a_timestamp timestamp, - a_timeuuid timeuuid, - a_tinyint tinyint, - a_tuple tuple, - a_udt frozen
, - a_uuid uuid, - a_varint varint, - a_vector vector - )" - ), - ) - .await; + for table in ["all_types", "all_types_written"] { + execute( + pool, + &format!("CREATE TABLE IF NOT EXISTS {KEYSPACE}.{table} ({ALL_TYPES_COLUMNS})"), + ) + .await; + } execute( pool, &format!( @@ -147,6 +152,13 @@ async fn create_test_data(pool: &CassandraConnectionPool) { ), ) .await; + execute( + pool, + &format!( + "CREATE TABLE IF NOT EXISTS {KEYSPACE}.users_written (id int PRIMARY KEY, name text, age int)" + ), + ) + .await; for (id, name, age) in [ (1, "alice", 31), (2, "bob", 25), @@ -171,6 +183,34 @@ async fn register_table(ctx: &SessionContext, pool: &Arc, + table: &str, +) { + let table_provider = CassandraTableFactory::new(Arc::clone(pool)) + .read_write_table_provider(TableReference::partial(KEYSPACE, table)) + .await + .expect("table provider should be created"); + ctx.register_table(table, table_provider) + .expect("table should be registered"); +} + +/// Runs an `INSERT` and returns the number of inserted rows that it reports. +async fn insert(ctx: &SessionContext, sql: &str) -> u64 { + let batches = query(ctx, sql).await; + array_value_to_string(batches[0].column(0), 0) + .unwrap() + .parse() + .unwrap() +} + +async fn query_to_string(ctx: &SessionContext, sql: &str) -> String { + arrow::util::pretty::pretty_format_batches(&query(ctx, sql).await) + .unwrap() + .to_string() +} + async fn query(ctx: &SessionContext, sql: &str) -> Vec { ctx.sql(sql) .await @@ -377,7 +417,15 @@ async fn test_cassandra_tables_and_schemas(pool: &Arc) .await .expect("tables should be listed"); tables.sort(); - assert_eq!(tables, vec!["all_types".to_string(), "users".to_string()]); + assert_eq!( + tables, + vec![ + "all_types".to_string(), + "all_types_written".to_string(), + "users".to_string(), + "users_written".to_string(), + ] + ); } async fn test_cassandra_missing_table(pool: &Arc) { @@ -452,3 +500,98 @@ async fn test_cassandra_connection_pool_with_schema(port: u16) { "carol" ); } + +/// Copying a table with every supported CQL type through `INSERT INTO ... SELECT` must give +/// back the same data. +async fn test_cassandra_insert_all_types(pool: &Arc) { + let ctx = SessionContext::new(); + register_table(&ctx, pool, "all_types").await; + register_writable_table(&ctx, pool, "all_types_written").await; + + let inserted = insert( + &ctx, + "INSERT INTO all_types_written SELECT * FROM all_types", + ) + .await; + assert_eq!(inserted, 2); + + let expected = query_to_string(&ctx, "SELECT * FROM all_types ORDER BY id").await; + let actual = query_to_string(&ctx, "SELECT * FROM all_types_written ORDER BY id").await; + assert_eq!(actual, expected); +} + +async fn test_cassandra_insert_values(pool: &Arc) { + let ctx = SessionContext::new(); + register_writable_table(&ctx, pool, "users_written").await; + + // Cassandra orders the columns as the primary key, then the other columns by name, so + // the columns of users_written are (id, age, name) + let inserted = insert( + &ctx, + "INSERT INTO users_written VALUES (1, 31, 'alice'), (2, NULL, 'bob')", + ) + .await; + assert_eq!(inserted, 2); + // the columns that are not given are null + let inserted = insert( + &ctx, + "INSERT INTO users_written (id, name) VALUES (3, 'carol')", + ) + .await; + assert_eq!(inserted, 1); + assert_eq!( + query_to_string(&ctx, "SELECT id, name, age FROM users_written ORDER BY id").await, + "+----+-------+-----+ +| id | name | age | ++----+-------+-----+ +| 1 | alice | 31 | +| 2 | bob | | +| 3 | carol | | ++----+-------+-----+" + ); + + // as in CQL, inserting an existing primary key replaces the row + let inserted = insert( + &ctx, + "INSERT INTO users_written (id, name, age) VALUES (1, 'alicia', 32)", + ) + .await; + assert_eq!(inserted, 1); + assert_eq!( + query_to_string(&ctx, "SELECT id, name, age FROM users_written WHERE id = 1").await, + "+----+--------+-----+ +| id | name | age | ++----+--------+-----+ +| 1 | alicia | 32 | ++----+--------+-----+" + ); + + // INSERT OVERWRITE replaces all rows + let inserted = insert( + &ctx, + "INSERT OVERWRITE users_written (id, name, age) VALUES (4, 'dave', 35)", + ) + .await; + assert_eq!(inserted, 1); + assert_eq!( + query_to_string(&ctx, "SELECT id, name, age FROM users_written").await, + "+----+------+-----+ +| id | name | age | ++----+------+-----+ +| 4 | dave | 35 | ++----+------+-----+" + ); + + // Cassandra rejects a null primary key + let error = ctx + .sql("INSERT INTO users_written (id, name) VALUES (NULL, 'nobody')") + .await + .expect("query should be planned") + .collect() + .await + .expect_err("a null primary key should be rejected"); + assert!( + error.to_string().contains("users_written"), + "the error should name the table: {error}" + ); +} diff --git a/crates/cassandra/Cargo.toml b/crates/cassandra/Cargo.toml index e4053a983..6f4e891b8 100644 --- a/crates/cassandra/Cargo.toml +++ b/crates/cassandra/Cargo.toml @@ -18,6 +18,7 @@ num-bigint = "0.4" scylla = { version = "1.9", features = ["bigdecimal-04", "num-bigint-04"] } secrecy = "0.10" snafu = "0.9" +uuid = "1" [dev-dependencies] tokio = { version = "1.52", features = ["macros", "rt-multi-thread"] } diff --git a/crates/cassandra/src/lib.rs b/crates/cassandra/src/lib.rs index b515830ce..0bf8c0cca 100644 --- a/crates/cassandra/src/lib.rs +++ b/crates/cassandra/src/lib.rs @@ -4,11 +4,15 @@ //! CQL is far more restrictive than SQL: it has no joins or aggregations, and `WHERE` clauses //! are limited to key columns. Scans therefore only push the projection and `LIMIT` down to //! Cassandra. Filters, aggregations and joins are evaluated by DataFusion. +//! +//! Tables created with [`CassandraTableFactory::read_write_table_provider`] can also be written +//! to with `INSERT INTO`; see [`CassandraTableWriter`]. pub mod conn; pub mod pool; mod table; pub mod types; +pub mod write; use std::sync::Arc; @@ -19,6 +23,7 @@ use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::As pub use crate::pool::CassandraConnectionPool; pub use crate::table::{CassandraExec, CassandraTable}; +pub use crate::write::CassandraTableWriter; /// Creates [`TableProvider`]s for Cassandra tables. #[derive(Debug)] @@ -54,4 +59,24 @@ impl CassandraTableFactory { schema, ))) } + + /// Creates a table provider for `table_reference` that also supports `INSERT INTO`. + /// + /// See [`CassandraTableWriter`] for how the rows are written. + /// + /// # Errors + /// + /// Returns an error if the schema of the table cannot be read. + pub async fn read_write_table_provider( + &self, + table_reference: TableReference, + ) -> Result, Box> + { + let read_provider = self.table_provider(table_reference.clone()).await?; + Ok(CassandraTableWriter::create( + read_provider, + table_reference, + Arc::clone(&self.pool), + )) + } } diff --git a/crates/cassandra/src/types.rs b/crates/cassandra/src/types.rs index ecf3f1420..aae36b9d2 100644 --- a/crates/cassandra/src/types.rs +++ b/crates/cassandra/src/types.rs @@ -1,5 +1,7 @@ //! Conversion between CQL and Arrow types and values. +use std::net::IpAddr; +use std::str::FromStr; use std::sync::Arc; use arrow::array::{ @@ -10,15 +12,22 @@ use arrow::array::{ Time64NanosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, }; +use arrow::array::{Array, AsArray}; use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::compute::cast; use arrow::datatypes::{ - DataType, Field, Fields, IntervalMonthDayNano, Schema, SchemaRef, TimeUnit, + DataType, Date32Type, Field, Fields, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, + Int8Type, IntervalMonthDayNano, IntervalMonthDayNanoType, Schema, SchemaRef, + Time64NanosecondType, TimeUnit, TimestampMillisecondType, }; use bigdecimal::BigDecimal; use num_bigint::BigInt; use scylla::frame::response::result::{CollectionType, ColumnSpec, ColumnType, NativeType}; -use scylla::value::{CqlValue, Row}; +use scylla::value::{ + CqlDate, CqlDecimal, CqlDuration, CqlTime, CqlTimestamp, CqlTimeuuid, CqlValue, CqlVarint, Row, +}; use snafu::prelude::*; +use uuid::Uuid; #[derive(Debug, Snafu)] pub enum Error { @@ -39,6 +48,17 @@ pub enum Error { #[snafu(display("Failed to build Arrow data: {source}"))] Arrow { source: arrow::error::ArrowError }, + + #[snafu(display("Cannot write column '{column_name}' to Cassandra: {message}"))] + InvalidValue { + column_name: String, + message: String, + }, + + #[snafu(display( + "Cannot write column '{column_name}' to Cassandra: CQL collections can't contain nulls" + ))] + NullInCollection { column_name: String }, } pub type Result = std::result::Result; @@ -592,6 +612,415 @@ fn nulls(values: &[Option<&CqlValue>]) -> Option { } } +/// Converts the values of `array` to CQL values of type `typ`, for binding them to a statement. +/// +/// Returns one value per row, `None` for nulls. The array is first cast to the Arrow type that +/// [`map_column_type`] maps `typ` to, so any Arrow type that can be cast to it is accepted, +/// e.g. `Int32` for a `bigint` column or `Utf8` for a `timestamp` column. `decimal`, `varint`, +/// `uuid`, `timeuuid` and `inet` values are parsed from their string representation; `uuid` and +/// `timeuuid` values are also accepted as 16 bytes of `FixedSizeBinary(16)`. +/// +/// # Errors +/// +/// Returns an error if the array cannot be converted to `typ`, or if a collection contains a +/// null: CQL collections can't contain nulls. +pub fn array_to_cql_values( + column_name: &str, + array: &ArrayRef, + typ: &ColumnType<'_>, +) -> Result>> { + match typ { + ColumnType::Native(native) => native_values(column_name, array, native, typ), + ColumnType::Collection { + typ: collection, .. + } => match collection { + CollectionType::List(element) => { + list_values(column_name, array, element, typ, CqlValue::List) + } + CollectionType::Set(element) => { + list_values(column_name, array, element, typ, CqlValue::Set) + } + CollectionType::Map(key, value) => map_values(column_name, array, key, value, typ), + _ => unsupported(column_name, typ), + }, + ColumnType::Vector { typ: element, .. } => { + list_values(column_name, array, element, typ, CqlValue::Vector) + } + ColumnType::UserDefinedType { definition, .. } => { + let struct_array = as_struct(column_name, array, typ)?; + if let Some(field) = struct_array.fields().iter().find(|field| { + !definition + .field_types + .iter() + .any(|(name, _)| name == field.name()) + }) { + return InvalidValueSnafu { + column_name, + message: format!( + "field '{}' is not a field of the user-defined type {}.{}", + field.name(), + definition.keyspace, + definition.name + ), + } + .fail(); + } + let fields = definition + .field_types + .iter() + .map(|(name, typ)| match struct_array.column_by_name(name) { + Some(column) => array_to_cql_values(column_name, column, typ), + None => Ok(vec![None; struct_array.len()]), + }) + .collect::>>()?; + Ok((0..struct_array.len()) + .map(|row| { + struct_array + .is_valid(row) + .then(|| CqlValue::UserDefinedType { + keyspace: definition.keyspace.to_string(), + name: definition.name.to_string(), + fields: definition + .field_types + .iter() + .zip(&fields) + .map(|((name, _), values)| (name.to_string(), values[row].clone())) + .collect(), + }) + }) + .collect()) + } + ColumnType::Tuple(types) => { + let struct_array = as_struct(column_name, array, typ)?; + ensure!( + struct_array.num_columns() == types.len(), + InvalidValueSnafu { + column_name, + message: format!( + "a tuple of {} elements can't be written from a struct with {} fields", + types.len(), + struct_array.num_columns() + ), + } + ); + let elements = struct_array + .columns() + .iter() + .zip(types) + .map(|(column, typ)| array_to_cql_values(column_name, column, typ)) + .collect::>>()?; + Ok((0..struct_array.len()) + .map(|row| { + struct_array.is_valid(row).then(|| { + CqlValue::Tuple(elements.iter().map(|values| values[row].clone()).collect()) + }) + }) + .collect()) + } + _ => unsupported(column_name, typ), + } +} + +fn native_values( + column_name: &str, + array: &ArrayRef, + native: &NativeType, + typ: &ColumnType<'_>, +) -> Result>> { + if matches!(native, NativeType::Uuid | NativeType::Timeuuid) { + if let DataType::FixedSizeBinary(16) = array.data_type() { + let bytes = array.as_fixed_size_binary(); + return each_value(array, |row| { + let uuid = Uuid::from_slice(bytes.value(row)).map_err(|e| e.to_string())?; + Ok(uuid_value(native, uuid)) + }) + .map_err(|message| invalid_value(column_name, message)); + } + } + + let target_type = match (native, array.data_type()) { + // Keep the time zone, so that the cast only changes the unit and not the instant + (NativeType::Timestamp, DataType::Timestamp(_, tz)) => { + DataType::Timestamp(TimeUnit::Millisecond, tz.clone()) + } + (NativeType::Timestamp, _) => DataType::Timestamp(TimeUnit::Millisecond, None), + (NativeType::Counter, _) => { + return InvalidValueSnafu { + column_name, + message: "counter columns can't be written with INSERT, only with UPDATE", + } + .fail() + } + _ => map_column_type(column_name, typ)?, + }; + let array = cast(array, &target_type).map_err(|e| { + invalid_value( + column_name, + format!("can't convert {} to CQL {typ:?}: {e}", array.data_type()), + ) + })?; + + let values = match native { + NativeType::Ascii => { + let strings = array.as_string::(); + each_value(&array, |row| { + Ok(CqlValue::Ascii(strings.value(row).to_string())) + }) + } + NativeType::Text => { + let strings = array.as_string::(); + each_value(&array, |row| { + Ok(CqlValue::Text(strings.value(row).to_string())) + }) + } + NativeType::Boolean => { + let booleans = array.as_boolean(); + each_value(&array, |row| Ok(CqlValue::Boolean(booleans.value(row)))) + } + NativeType::Blob => { + let blobs = array.as_binary::(); + each_value(&array, |row| Ok(CqlValue::Blob(blobs.value(row).to_vec()))) + } + NativeType::BigInt => { + let ints = array.as_primitive::(); + each_value(&array, |row| Ok(CqlValue::BigInt(ints.value(row)))) + } + NativeType::Int => { + let ints = array.as_primitive::(); + each_value(&array, |row| Ok(CqlValue::Int(ints.value(row)))) + } + NativeType::SmallInt => { + let ints = array.as_primitive::(); + each_value(&array, |row| Ok(CqlValue::SmallInt(ints.value(row)))) + } + NativeType::TinyInt => { + let ints = array.as_primitive::(); + each_value(&array, |row| Ok(CqlValue::TinyInt(ints.value(row)))) + } + NativeType::Double => { + let floats = array.as_primitive::(); + each_value(&array, |row| Ok(CqlValue::Double(floats.value(row)))) + } + NativeType::Float => { + let floats = array.as_primitive::(); + each_value(&array, |row| Ok(CqlValue::Float(floats.value(row)))) + } + NativeType::Date => { + let dates = array.as_primitive::(); + each_value(&array, |row| { + u32::try_from(i64::from(dates.value(row)) + CQL_DATE_UNIX_EPOCH_OFFSET) + .map(|days| CqlValue::Date(CqlDate(days))) + .map_err(|_| format!("date {} is out of range", dates.value(row))) + }) + } + NativeType::Timestamp => { + let timestamps = array.as_primitive::(); + each_value(&array, |row| { + Ok(CqlValue::Timestamp(CqlTimestamp(timestamps.value(row)))) + }) + } + NativeType::Time => { + let times = array.as_primitive::(); + each_value(&array, |row| Ok(CqlValue::Time(CqlTime(times.value(row))))) + } + NativeType::Duration => { + let intervals = array.as_primitive::(); + each_value(&array, |row| { + let interval = intervals.value(row); + Ok(CqlValue::Duration(CqlDuration { + months: interval.months, + days: interval.days, + nanoseconds: interval.nanoseconds, + })) + }) + } + NativeType::Decimal => { + let strings = array.as_string::(); + each_value(&array, |row| { + let decimal = BigDecimal::from_str(strings.value(row).trim()) + .map_err(|e| format!("'{}' is not a decimal: {e}", strings.value(row)))?; + CqlDecimal::try_from(decimal) + .map(CqlValue::Decimal) + .map_err(|_| format!("the scale of '{}' is too large", strings.value(row))) + }) + } + NativeType::Varint => { + let strings = array.as_string::(); + each_value(&array, |row| { + BigInt::from_str(strings.value(row).trim()) + .map(|v| CqlValue::Varint(CqlVarint::from(v))) + .map_err(|e| format!("'{}' is not an integer: {e}", strings.value(row))) + }) + } + NativeType::Uuid | NativeType::Timeuuid => { + let strings = array.as_string::(); + each_value(&array, |row| { + Uuid::parse_str(strings.value(row).trim()) + .map(|uuid| uuid_value(native, uuid)) + .map_err(|e| format!("'{}' is not a UUID: {e}", strings.value(row))) + }) + } + NativeType::Inet => { + let strings = array.as_string::(); + each_value(&array, |row| { + IpAddr::from_str(strings.value(row).trim()) + .map(CqlValue::Inet) + .map_err(|e| format!("'{}' is not an IP address: {e}", strings.value(row))) + }) + } + _ => return unsupported(column_name, typ), + }; + values.map_err(|message| invalid_value(column_name, message)) +} + +fn invalid_value(column_name: &str, message: String) -> Error { + Error::InvalidValue { + column_name: column_name.to_string(), + message, + } +} + +fn uuid_value(native: &NativeType, uuid: Uuid) -> CqlValue { + match native { + NativeType::Timeuuid => CqlValue::Timeuuid(CqlTimeuuid::from(uuid)), + _ => CqlValue::Uuid(uuid), + } +} + +/// Converts the non-null values of `array` with `f`, keeping nulls as `None`. +fn each_value( + array: &dyn Array, + mut f: impl FnMut(usize) -> Result, +) -> Result>, String> { + (0..array.len()) + .map(|row| { + if array.is_null(row) { + Ok(None) + } else { + f(row).map(Some) + } + }) + .collect() +} + +/// Converts a list-like array (`List`, `LargeList` or `FixedSizeList`) to CQL collections built +/// by `collection`. +fn list_values( + column_name: &str, + array: &ArrayRef, + element: &ColumnType<'_>, + typ: &ColumnType<'_>, + collection: fn(Vec) -> CqlValue, +) -> Result>> { + let (values, ranges): (&ArrayRef, Vec<(usize, usize)>) = match array.data_type() { + DataType::List(_) => { + let list = array.as_list::(); + let ranges = list + .value_offsets() + .windows(2) + .map(|w| (w[0] as usize, w[1] as usize)) + .collect(); + (list.values(), ranges) + } + DataType::LargeList(_) => { + let list = array.as_list::(); + let ranges = list + .value_offsets() + .windows(2) + .map(|w| (w[0] as usize, w[1] as usize)) + .collect(); + (list.values(), ranges) + } + DataType::FixedSizeList(_, _) => { + let list = array.as_fixed_size_list(); + let length = list.value_length() as usize; + let ranges = (0..list.len()) + .map(|row| { + let start = list.value_offset(row) as usize; + (start, start + length) + }) + .collect(); + (list.values(), ranges) + } + _ => return incompatible_array(column_name, array, typ), + }; + + let elements = array_to_cql_values(column_name, values, element)?; + ranges + .into_iter() + .enumerate() + .map(|(row, (start, end))| { + if array.is_null(row) { + return Ok(None); + } + elements[start..end] + .iter() + .map(|element| { + element + .clone() + .context(NullInCollectionSnafu { column_name }) + }) + .collect::>>() + .map(|elements| Some(collection(elements))) + }) + .collect() +} + +fn map_values( + column_name: &str, + array: &ArrayRef, + key: &ColumnType<'_>, + value: &ColumnType<'_>, + typ: &ColumnType<'_>, +) -> Result>> { + let DataType::Map(_, _) = array.data_type() else { + return incompatible_array(column_name, array, typ); + }; + let map = array.as_map(); + let keys = array_to_cql_values(column_name, map.keys(), key)?; + let values = array_to_cql_values(column_name, map.values(), value)?; + map.value_offsets() + .windows(2) + .enumerate() + .map(|(row, w)| { + if map.is_null(row) { + return Ok(None); + } + (w[0] as usize..w[1] as usize) + .map(|i| { + let key = keys[i] + .clone() + .context(NullInCollectionSnafu { column_name })?; + let value = values[i] + .clone() + .context(NullInCollectionSnafu { column_name })?; + Ok((key, value)) + }) + .collect::>>() + .map(|entries| Some(CqlValue::Map(entries))) + }) + .collect() +} + +fn as_struct<'a>( + column_name: &str, + array: &'a ArrayRef, + typ: &ColumnType<'_>, +) -> Result<&'a StructArray> { + match array.data_type() { + DataType::Struct(_) => Ok(array.as_struct()), + _ => incompatible_array(column_name, array, typ), + } +} + +fn incompatible_array(column_name: &str, array: &ArrayRef, typ: &ColumnType<'_>) -> Result { + InvalidValueSnafu { + column_name, + message: format!("can't convert {} to CQL {typ:?}", array.data_type()), + } + .fail() +} + #[cfg(test)] mod tests { use super::*; @@ -978,4 +1407,234 @@ mod tests { )])); assert!(rows_to_record_batch(&schema, &column_names, &rows).is_err()); } + + fn write(array: ArrayRef, typ: &ColumnType<'_>) -> Result>> { + array_to_cql_values("c", &array, typ) + } + + fn map_of(key: ColumnType<'static>, value: ColumnType<'static>) -> ColumnType<'static> { + ColumnType::Collection { + frozen: false, + typ: CollectionType::Map(Box::new(key), Box::new(value)), + } + } + + fn address_udt() -> ColumnType<'static> { + ColumnType::UserDefinedType { + frozen: true, + definition: Arc::new(UserDefinedType { + name: Cow::Borrowed("address"), + keyspace: Cow::Borrowed("ks"), + field_types: vec![ + (Cow::Borrowed("street"), native(NativeType::Text)), + (Cow::Borrowed("zip"), native(NativeType::Int)), + ], + }), + } + } + + /// Converting CQL values to Arrow and back must give the original values. + #[test] + fn writes_what_it_reads() { + let uuid = Uuid::parse_str("8f2bfa59-4e30-4b0a-8b9a-3c1e26f9a0b1").unwrap(); + let timeuuid = Uuid::parse_str("e9c0b2a0-6c1e-11ee-8c99-0242ac120002").unwrap(); + let cases = vec![ + (native(NativeType::Ascii), CqlValue::Ascii("abc".into())), + (native(NativeType::Text), CqlValue::Text("héllo".into())), + (native(NativeType::Boolean), CqlValue::Boolean(true)), + (native(NativeType::Blob), CqlValue::Blob(vec![1, 2, 3])), + (native(NativeType::BigInt), CqlValue::BigInt(i64::MIN)), + (native(NativeType::Int), CqlValue::Int(-7)), + (native(NativeType::SmallInt), CqlValue::SmallInt(300)), + (native(NativeType::TinyInt), CqlValue::TinyInt(-3)), + (native(NativeType::Double), CqlValue::Double(1.5)), + (native(NativeType::Float), CqlValue::Float(2.5)), + ( + native(NativeType::Date), + CqlValue::Date(CqlDate((1 << 31) + 19_000)), + ), + ( + native(NativeType::Timestamp), + CqlValue::Timestamp(CqlTimestamp(1_700_000_000_123)), + ), + (native(NativeType::Time), CqlValue::Time(CqlTime(42_000))), + ( + native(NativeType::Duration), + CqlValue::Duration(CqlDuration { + months: 1, + days: 2, + nanoseconds: 3, + }), + ), + ( + native(NativeType::Decimal), + CqlValue::Decimal(CqlDecimal::from_signed_be_bytes_slice_and_exponent( + &[0x30, 0x39], + 2, + )), + ), + ( + native(NativeType::Varint), + CqlValue::Varint(CqlVarint::from_signed_bytes_be_slice(&[0x01, 0x00, 0x00])), + ), + (native(NativeType::Uuid), CqlValue::Uuid(uuid)), + ( + native(NativeType::Timeuuid), + CqlValue::Timeuuid(CqlTimeuuid::from(timeuuid)), + ), + ( + native(NativeType::Inet), + CqlValue::Inet("2001:db8::1".parse().unwrap()), + ), + ( + list_of(native(NativeType::Int)), + CqlValue::List(vec![CqlValue::Int(1), CqlValue::Int(2)]), + ), + ( + ColumnType::Collection { + frozen: false, + typ: CollectionType::Set(Box::new(native(NativeType::Text))), + }, + CqlValue::Set(vec![CqlValue::Text("a".into())]), + ), + ( + map_of(native(NativeType::Text), native(NativeType::BigInt)), + CqlValue::Map(vec![(CqlValue::Text("a".into()), CqlValue::BigInt(1))]), + ), + ( + ColumnType::Vector { + typ: Box::new(native(NativeType::Float)), + dimensions: 2, + }, + CqlValue::Vector(vec![CqlValue::Float(1.0), CqlValue::Float(2.0)]), + ), + ( + address_udt(), + CqlValue::UserDefinedType { + keyspace: "ks".into(), + name: "address".into(), + fields: vec![ + ("street".into(), Some(CqlValue::Text("Main".into()))), + ("zip".into(), None), + ], + }, + ), + ( + ColumnType::Tuple(vec![native(NativeType::Int), native(NativeType::Text)]), + CqlValue::Tuple(vec![Some(CqlValue::Int(7)), None]), + ), + ]; + + for (typ, value) in cases { + let data_type = map_column_type("c", &typ).unwrap(); + let array = values_to_array(&data_type, &[Some(&value), None]).unwrap(); + let written = write(array, &typ).unwrap(); + assert_eq!(written, vec![Some(value), None], "CQL type {typ:?}"); + } + } + + #[test] + fn casts_to_the_column_type() { + let ints: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None])); + assert_eq!( + write(Arc::clone(&ints), &native(NativeType::BigInt)).unwrap(), + vec![Some(CqlValue::BigInt(1)), None] + ); + assert_eq!( + write(ints, &native(NativeType::Varint)).unwrap(), + vec![ + Some(CqlValue::Varint(CqlVarint::from(BigInt::from(1)))), + None + ] + ); + + let large_strings: ArrayRef = Arc::new(LargeStringArray::from(vec!["x"])); + assert_eq!( + write(large_strings, &native(NativeType::Text)).unwrap(), + vec![Some(CqlValue::Text("x".into()))] + ); + + let decimals: ArrayRef = Arc::new( + Decimal128Array::from(vec![12345]) + .with_precision_and_scale(10, 2) + .unwrap(), + ); + assert_eq!( + write(decimals, &native(NativeType::Decimal)).unwrap(), + vec![Some(CqlValue::Decimal( + CqlDecimal::from_signed_be_bytes_slice_and_exponent(&[0x30, 0x39], 2) + ))] + ); + } + + #[test] + fn keeps_the_instant_of_timestamps() { + let timestamps: ArrayRef = + Arc::new(TimestampMicrosecondArray::from(vec![1_500_000]).with_timezone("+02:00")); + assert_eq!( + write(timestamps, &native(NativeType::Timestamp)).unwrap(), + vec![Some(CqlValue::Timestamp(CqlTimestamp(1_500)))] + ); + } + + #[test] + fn writes_uuids_from_bytes() { + let uuid = Uuid::parse_str("8f2bfa59-4e30-4b0a-8b9a-3c1e26f9a0b1").unwrap(); + let bytes: ArrayRef = Arc::new( + arrow::array::FixedSizeBinaryArray::try_from_iter(std::iter::once(uuid.as_bytes())) + .unwrap(), + ); + assert_eq!( + write(bytes, &native(NativeType::Uuid)).unwrap(), + vec![Some(CqlValue::Uuid(uuid))] + ); + } + + #[test] + fn fails_for_invalid_values() { + let strings: ArrayRef = Arc::new(StringArray::from(vec!["not a uuid"])); + let err = write(strings, &native(NativeType::Uuid)).unwrap_err(); + assert!(err.to_string().contains("is not a UUID"), "{err}"); + + let ints: ArrayRef = Arc::new(Int64Array::from(vec![1])); + let err = write(ints, &native(NativeType::Counter)).unwrap_err(); + assert!(err.to_string().contains("counter"), "{err}"); + + let strings: ArrayRef = Arc::new(StringArray::from(vec!["a"])); + let err = write(strings, &list_of(native(NativeType::Text))).unwrap_err(); + assert!(err.to_string().contains("can't convert Utf8"), "{err}"); + } + + #[test] + fn fails_for_nulls_in_collections() { + let list = values_to_array( + &DataType::List(list_item_field(DataType::Int32)), + &[Some(&CqlValue::List(vec![CqlValue::Int(1)]))], + ) + .unwrap(); + // Replace the element with a null + let list = list.as_list::(); + let list: ArrayRef = Arc::new(ListArray::new( + list_item_field(DataType::Int32), + list.offsets().clone(), + Arc::new(Int32Array::from(vec![None::])), + None, + )); + let err = write(list, &list_of(native(NativeType::Int))).unwrap_err(); + assert!(matches!(err, Error::NullInCollection { .. }), "{err}"); + } + + #[test] + fn fails_for_struct_fields_missing_from_the_udt() { + let array: ArrayRef = Arc::new(StructArray::from(vec![( + Arc::new(Field::new("city", DataType::Utf8, true)), + Arc::new(StringArray::from(vec!["Sofia"])) as ArrayRef, + )])); + let err = write(array, &address_udt()).unwrap_err(); + assert!( + err.to_string() + .contains("'city' is not a field of the user-defined type ks.address"), + "{err}" + ); + } } diff --git a/crates/cassandra/src/write.rs b/crates/cassandra/src/write.rs new file mode 100644 index 000000000..0113df2b0 --- /dev/null +++ b/crates/cassandra/src/write.rs @@ -0,0 +1,236 @@ +use std::{fmt, sync::Arc}; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::catalog::Session; +use datafusion::common::TableReference; +use datafusion::datasource::sink::{DataSink, DataSinkExec}; +use datafusion::datasource::{TableProvider, TableType}; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::logical_expr::dml::InsertOp; +use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; +use futures::{StreamExt, TryStreamExt}; +use scylla::value::CqlValue; +use snafu::prelude::*; + +use crate::conn::{quote_identifier, quote_table_reference}; +use crate::pool::CassandraConnectionPool; +use crate::types::array_to_cql_values; + +/// The maximum number of `INSERT` statements in flight at a time. +const MAX_CONCURRENT_INSERTS: usize = 64; + +#[derive(Debug, Snafu)] +pub enum Error { + #[snafu(display("Failed to prepare the Cassandra statement {cql}: {source}"))] + Prepare { + cql: String, + source: scylla::errors::PrepareError, + }, + + #[snafu(display("Failed to write to the Cassandra table {table_reference}: {source}"))] + Write { + table_reference: TableReference, + source: scylla::errors::ExecutionError, + }, + + #[snafu(display("Failed to convert the data for Cassandra: {source}"))] + Conversion { source: crate::types::Error }, +} + +fn to_datafusion_error(e: Error) -> DataFusionError { + DataFusionError::External(Box::new(e)) +} + +/// A Cassandra table that can be written to with `INSERT INTO`. +/// +/// Reads are delegated to `read_provider`. Every row is written with its own CQL `INSERT`, so, +/// as in Cassandra, inserting a row whose primary key exists overwrites that row (an upsert). +/// `INSERT OVERWRITE` truncates the table first. +/// +/// Cassandra has no transactions: when a write fails, the rows written before the failure +/// are kept. +#[derive(Debug)] +pub struct CassandraTableWriter { + pub read_provider: Arc, + table_reference: TableReference, + pool: Arc, +} + +impl CassandraTableWriter { + pub fn create( + read_provider: Arc, + table_reference: TableReference, + pool: Arc, + ) -> Arc { + Arc::new(Self { + read_provider, + table_reference, + pool, + }) + } +} + +#[async_trait] +impl TableProvider for CassandraTableWriter { + fn schema(&self) -> SchemaRef { + self.read_provider.schema() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> DataFusionResult> { + self.read_provider.supports_filters_pushdown(filters) + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> DataFusionResult> { + self.read_provider + .scan(state, projection, filters, limit) + .await + } + + async fn insert_into( + &self, + _state: &dyn Session, + input: Arc, + op: InsertOp, + ) -> DataFusionResult> { + Ok(Arc::new(DataSinkExec::new( + input, + Arc::new(CassandraDataSink { + table_reference: self.table_reference.clone(), + pool: Arc::clone(&self.pool), + op, + schema: self.schema(), + }), + None, + )) as _) + } +} + +struct CassandraDataSink { + table_reference: TableReference, + pool: Arc, + op: InsertOp, + schema: SchemaRef, +} + +#[async_trait] +impl DataSink for CassandraDataSink { + fn metrics(&self) -> Option { + None + } + + fn schema(&self) -> &SchemaRef { + &self.schema + } + + async fn write_all( + &self, + mut data: SendableRecordBatchStream, + _context: &Arc, + ) -> DataFusionResult { + let session = self.pool.session(); + let table = quote_table_reference(&self.table_reference); + + let columns = self + .schema + .fields() + .iter() + .map(|f| quote_identifier(f.name())) + .collect::>(); + let cql = format!( + "INSERT INTO {table} ({}) VALUES ({})", + columns.join(", "), + vec!["?"; columns.len()].join(", ") + ); + let prepared = session + .prepare(cql.as_str()) + .await + .context(PrepareSnafu { cql }) + .map_err(to_datafusion_error)?; + let column_types = prepared + .get_variable_col_specs() + .iter() + .map(|spec| spec.typ().clone()) + .collect::>(); + + // Truncate only after the INSERT is prepared, so that an invalid INSERT (e.g. a + // column that doesn't exist) leaves the table untouched + if matches!(self.op, InsertOp::Overwrite) { + session + .query_unpaged(format!("TRUNCATE {table}"), ()) + .await + .context(WriteSnafu { + table_reference: self.table_reference.clone(), + }) + .map_err(to_datafusion_error)?; + } + + let mut num_rows = 0; + while let Some(batch) = data.next().await { + let batch = batch?; + if batch.num_rows() == 0 { + continue; + } + + let mut columns = batch + .columns() + .iter() + .zip(batch.schema().fields()) + .zip(&column_types) + .map(|((array, field), typ)| { + array_to_cql_values(field.name(), array, typ).map(Vec::into_iter) + }) + .collect::, _>>() + .context(ConversionSnafu) + .map_err(to_datafusion_error)?; + let rows = (0..batch.num_rows()).map(|_| { + columns + .iter_mut() + .map(|values| values.next().flatten()) + .collect::>>() + }); + + futures::stream::iter(rows) + .map(|row| session.execute_unpaged(&prepared, row)) + .buffer_unordered(MAX_CONCURRENT_INSERTS) + .try_for_each(|_| async { Ok(()) }) + .await + .context(WriteSnafu { + table_reference: self.table_reference.clone(), + }) + .map_err(to_datafusion_error)?; + + num_rows += batch.num_rows() as u64; + } + + Ok(num_rows) + } +} + +impl fmt::Debug for CassandraDataSink { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "CassandraDataSink table={}", self.table_reference) + } +} + +impl DisplayAs for CassandraDataSink { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "CassandraDataSink table={}", self.table_reference) + } +} From 828907b4d502403af674a69b9fd83fe58afb6912 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Sun, 18 Oct 2026 23:42:42 +0300 Subject: [PATCH 18/27] Add support for federation Signed-off-by: Martin Tzvetanov Grigorov --- Cargo.lock | 1 + README.md | 2 +- core/tests/cassandra/mod.rs | 65 +++++++ crates/cassandra/Cargo.toml | 5 + crates/cassandra/src/federation.rs | 288 +++++++++++++++++++++++++++++ crates/cassandra/src/lib.rs | 11 +- crates/cassandra/src/table.rs | 12 +- 7 files changed, 377 insertions(+), 7 deletions(-) create mode 100644 crates/cassandra/src/federation.rs diff --git a/Cargo.lock b/Cargo.lock index 13328f713..2194047ec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2505,6 +2505,7 @@ dependencies = [ "async-trait", "bigdecimal", "datafusion", + "datafusion-federation", "datafusion-table-providers-common", "futures", "num-bigint 0.4.8", diff --git a/README.md b/README.md index 6d725bd20..545047aff 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ Prefer a single-provider leaf crate when you only need one backend — for examp - Leaf crates enable `federation` by default (opt out with `default-features = false`). SQLite also defaults to the `bundled` rusqlite feature. - The existing `datafusion-table-providers` facade keeps the same module paths and provider feature names (`duckdb`, `postgres`, `sqlite`, etc.). Enabling a provider feature pulls in that leaf crate with its defaults (including federation). There is no separate `federation` / `*-federation` feature on the facade. - ClickHouse pool connections are exposed as `ClickHouseConnection` (wrapping the ClickHouse `Client`) under `sql::db_connection_pool::dbconnection::clickhouseconn`. -- Cassandra pool connections are exposed as `CassandraConnection` (wrapping a shared `scylla` driver `Session`) under `sql::db_connection_pool::dbconnection::cassandraconn`. The Cassandra crate has no `federation` feature: CQL cannot run the queries federation pushes down, so only projections and `LIMIT` are pushed down to Cassandra. +- Cassandra pool connections are exposed as `CassandraConnection` (wrapping a shared `scylla` driver `Session`) under `sql::db_connection_pool::dbconnection::cassandraconn`. Its `federation` feature works like the other crates', except that CQL cannot run the joins, aggregations and filters that federation pushes down: a federated Cassandra sub-plan is narrowed to its table scans, so only projections and `LIMIT` are pushed down to Cassandra. Existing examples continue to use the facade crate and its feature flags. diff --git a/core/tests/cassandra/mod.rs b/core/tests/cassandra/mod.rs index 435e6b1c0..2f498b0e0 100644 --- a/core/tests/cassandra/mod.rs +++ b/core/tests/cassandra/mod.rs @@ -86,6 +86,7 @@ async fn test_cassandra() { test_cassandra_projection_filter_and_limit(&pool).await; test_cassandra_count(&pool).await; test_cassandra_explain_shows_pushed_down_cql(&pool).await; + test_cassandra_federation(&pool).await; test_cassandra_tables_and_schemas(&pool).await; test_cassandra_missing_table(&pool).await; test_cassandra_connection_pool_with_schema(port).await; @@ -404,6 +405,70 @@ async fn test_cassandra_explain_shows_pushed_down_cql(pool: &Arc) { + let ctx = SessionContext::new_with_state(datafusion_federation::default_session_state()); + register_table(&ctx, pool, "users").await; + + assert_eq!( + query_to_string( + &ctx, + "SELECT name, age FROM users WHERE age > 30 ORDER BY age DESC" + ) + .await, + "+-------+-----+ +| name | age | ++-------+-----+ +| carol | 42 | +| dave | 35 | +| alice | 31 | ++-------+-----+" + ); + assert_eq!( + query_to_string( + &ctx, + "SELECT u.name, count(*) AS n FROM users u JOIN users v ON u.id = v.id \ + WHERE v.age < 30 GROUP BY u.name ORDER BY u.name" + ) + .await, + "+------+---+ +| name | n | ++------+---+ +| bob | 1 | +| eve | 1 | ++------+---+" + ); + + let batches = query(&ctx, "SELECT count(*) FROM users").await; + assert_eq!(array_value_to_string(batches[0].column(0), 0).unwrap(), "5"); + + // the filter must be applied after the limit + let batches = query( + &ctx, + "SELECT count(*) FROM (SELECT age FROM users LIMIT 3) WHERE age > 0", + ) + .await; + assert_eq!(array_value_to_string(batches[0].column(0), 0).unwrap(), "3"); + + let batches = query(&ctx, "SELECT id FROM users LIMIT 2").await; + let num_rows: usize = batches.iter().map(RecordBatch::num_rows).sum(); + assert_eq!(num_rows, 2); + + let batches = query( + &ctx, + "EXPLAIN SELECT name FROM users WHERE age > 30 LIMIT 3", + ) + .await; + let plan = arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string(); + assert!( + plan.contains(r#"CassandraExec cql=SELECT "age", "name" FROM "it_ks"."users""#), + "only the projection should be pushed down:\n{plan}" + ); +} + async fn test_cassandra_tables_and_schemas(pool: &Arc) { let conn = pool.connect().await.expect("connection should be created"); let conn = conn.as_async().expect("connection should be async"); diff --git a/crates/cassandra/Cargo.toml b/crates/cassandra/Cargo.toml index 6f4e891b8..20af7a622 100644 --- a/crates/cassandra/Cargo.toml +++ b/crates/cassandra/Cargo.toml @@ -12,6 +12,7 @@ arrow = { workspace = true } async-trait = "0.1" bigdecimal = "0.4" datafusion = { workspace = true, default-features = false, features = ["sql"] } +datafusion-federation = { workspace = true, features = ["sql"], optional = true } datafusion-table-providers-common = { workspace = true } futures = "0.3" num-bigint = "0.4" @@ -23,6 +24,10 @@ uuid = "1" [dev-dependencies] tokio = { version = "1.52", features = ["macros", "rt-multi-thread"] } +[features] +default = ["federation"] +federation = ["dep:datafusion-federation", "datafusion-table-providers-common/federation"] + [package.metadata.docs.rs] all-features = true rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/cassandra/src/federation.rs b/crates/cassandra/src/federation.rs new file mode 100644 index 000000000..5a0cb29d3 --- /dev/null +++ b/crates/cassandra/src/federation.rs @@ -0,0 +1,288 @@ +//! Integration with [`datafusion_federation`]. +//! +//! The federation optimizer hands a provider the largest sub-plan that only reads from it, to be +//! run remotely as one SQL query. CQL can't run most such queries: it has no joins or +//! aggregations, and `WHERE` clauses are limited to key columns. So the sub-plan is narrowed down +//! to its table scans, which keep their projection and `LIMIT` and run as [`CassandraExec`]s, +//! and the rest of it is evaluated by DataFusion. + +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::catalog::Session; +use datafusion::common::tree_node::Transformed; +use datafusion::common::TableReference; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::SendableRecordBatchStream; +use datafusion::logical_expr::{Extension, LogicalPlan, TableScan}; +use datafusion::optimizer::optimize_projections::OptimizeProjections; +use datafusion::optimizer::optimizer::Optimizer; +use datafusion::optimizer::push_down_limit::PushDownLimit; +use datafusion::optimizer::OptimizerContext; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::sql::unparser::dialect::{DefaultDialect, Dialect}; +use datafusion_federation::sql::{ + LogicalOptimizer, RemoteTableRef, SQLExecutor, SQLFederationProvider, SQLTableSource, +}; +use datafusion_federation::{ + FederatedPlanNode, FederatedTableProviderAdaptor, FederatedTableSource, FederationPlanner, +}; +use futures::TryStreamExt; + +use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::AsyncDbConnection; +use datafusion_table_providers_common::sql::db_connection_pool::{DbConnectionPool, JoinPushDown}; + +use crate::pool::CassandraConnectionPool; +use crate::{CassandraExec, CassandraTable}; + +impl CassandraTable { + fn create_federated_table_source( + self: Arc, + ) -> DataFusionResult> { + let table_reference = self.table_reference.clone(); + let schema = Arc::clone(&self.schema); + let fed_provider = Arc::new(SQLFederationProvider::new(self)); + Ok(Arc::new(SQLTableSource::new_with_schema( + fed_provider, + RemoteTableRef::from(table_reference), + schema, + ))) + } + + /// Wraps the table in a provider that takes part in federation. Outside of a federated + /// session, it scans the table like the `CassandraTable` itself. + /// + /// # Errors + /// + /// Returns an error if the federated table source cannot be created. + pub fn create_federated_table_provider( + self: Arc, + ) -> DataFusionResult { + let table_source = Self::create_federated_table_source(Arc::clone(&self))?; + Ok(FederatedTableProviderAdaptor::new_with_provider( + table_source, + self, + )) + } +} + +#[async_trait] +impl SQLExecutor for CassandraTable { + fn name(&self) -> &str { + "cassandra" + } + + fn compute_context(&self) -> Option { + match self.pool.join_push_down() { + JoinPushDown::Disallow => Some(format!("{}", std::ptr::from_ref(self) as usize)), + JoinPushDown::AllowedFor(s) => Some(s), + } + } + + fn dialect(&self) -> Arc { + Arc::new(DefaultDialect {}) + } + + fn logical_optimizer(&self) -> Option { + let planner: Arc = Arc::new(CassandraScanPlanner { + table_reference: self.table_reference.clone(), + pool: Arc::clone(&self.pool), + schema: Arc::clone(&self.schema), + }); + Some(Box::new(move |plan| { + federate_table_scans_only(plan, Arc::clone(&planner)) + })) + } + + /// Runs `query` as CQL. + /// + /// The federated scans are planned by [`CassandraScanPlanner`] instead, so no SQL + /// generated by the federation reaches Cassandra. + fn execute( + &self, + query: &str, + schema: SchemaRef, + _filters: &[Arc], + ) -> DataFusionResult { + let cql = query.to_string(); + let conn = self.pool.connect_direct(); + let stream_schema = Arc::clone(&schema); + let stream = futures::stream::once(async move { + conn.query_arrow(&cql, &[], Some(stream_schema)) + .await + .map_err(DataFusionError::External) + }) + .try_flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + + async fn table_names(&self) -> DataFusionResult> { + Err(DataFusionError::NotImplemented( + "table inference not implemented".to_string(), + )) + } + + async fn get_table_schema(&self, table_name: &str) -> DataFusionResult { + self.pool + .connect_direct() + .get_schema(&TableReference::from(table_name)) + .await + .map_err(|e| DataFusionError::External(Box::new(e))) + } +} + +/// Replaces a federated sub-plan with the same plan in which only the table scans are +/// federated, to be planned by `planner`, as CQL can run nothing more (see the module +/// documentation). +/// +/// The projections and limits are pushed down to the scans first: once federated, the scans +/// are opaque to the DataFusion optimizer. +fn federate_table_scans_only( + plan: LogicalPlan, + planner: Arc, +) -> DataFusionResult { + let LogicalPlan::Extension(Extension { node }) = &plan else { + return Ok(plan); + }; + let Some(federated) = node.as_any().downcast_ref::() else { + return Ok(plan); + }; + + let optimizer = Optimizer::with_rules(vec![ + Arc::new(OptimizeProjections::new()), + Arc::new(PushDownLimit::new()), + ]); + let inner = optimizer.optimize( + federated.plan().clone(), + &OptimizerContext::new(), + |_, _| {}, + )?; + + inner + .transform_up_with_subqueries(|plan| match plan { + LogicalPlan::TableScan(_) => Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(FederatedPlanNode::new(plan, Arc::clone(&planner))), + }))), + _ => Ok(Transformed::no(plan)), + }) + .map(|transformed| transformed.data) +} + +/// Plans a federated scan of a Cassandra table as a [`CassandraExec`]. +/// +/// The SQL federation would plan it as a scan that claims all the filters DataFusion pushes +/// down to it at runtime, but CQL can't run them. +struct CassandraScanPlanner { + table_reference: TableReference, + pool: Arc, + schema: SchemaRef, +} + +#[async_trait] +impl FederationPlanner for CassandraScanPlanner { + async fn plan_federation( + &self, + node: &FederatedPlanNode, + _session_state: &dyn Session, + ) -> DataFusionResult> { + let scan = match node.plan() { + LogicalPlan::TableScan(scan @ TableScan { filters, .. }) if filters.is_empty() => scan, + plan => { + return Err(DataFusionError::Plan(format!( + "Cassandra can only run plain table scans, not:\n{plan}" + ))) + } + }; + Ok(Arc::new(CassandraExec::try_new( + self.table_reference.clone(), + Arc::clone(&self.pool), + &self.schema, + scan.projection.as_ref(), + scan.fetch, + )?)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::logical_expr::{col, lit, LogicalPlanBuilder, LogicalTableSource}; + + struct NoPlanner; + + #[async_trait] + impl FederationPlanner for NoPlanner { + async fn plan_federation( + &self, + _node: &FederatedPlanNode, + _session_state: &dyn Session, + ) -> DataFusionResult> { + unimplemented!() + } + } + + fn federated(plan: LogicalPlan) -> LogicalPlan { + LogicalPlan::Extension(Extension { + node: Arc::new(FederatedPlanNode::new(plan, Arc::new(NoPlanner))), + }) + } + + #[test] + fn federates_only_the_scans_with_projection_and_limit() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("name", DataType::Utf8, true), + Field::new("age", DataType::Int32, true), + ])); + let scan = |name| { + LogicalPlanBuilder::scan( + TableReference::partial("ks", name), + Arc::new(LogicalTableSource::new(Arc::clone(&schema))), + None, + ) + .unwrap() + }; + let plan = scan("users") + .filter(col("age").gt(lit(30))) + .unwrap() + .join_on( + scan("pets").limit(0, Some(5)).unwrap().build().unwrap(), + datafusion::logical_expr::JoinType::Inner, + vec![col("ks.users.id").eq(col("ks.pets.id"))], + ) + .unwrap() + .project(vec![col("ks.users.name"), col("ks.pets.name")]) + .unwrap() + .build() + .unwrap(); + + let rewritten = + federate_table_scans_only(federated(plan.clone()), Arc::new(NoPlanner)).unwrap(); + assert_eq!(rewritten.schema(), plan.schema()); + assert_eq!( + rewritten.display_indent().to_string(), + "Projection: ks.users.name, ks.pets.name + Inner Join: Filter: ks.users.id = ks.pets.id + Projection: ks.users.id, ks.users.name + Filter: ks.users.age > Int32(30) + Federated + TableScan: ks.users projection=[id, name, age] + Limit: skip=0, fetch=5 + Federated + TableScan: ks.pets projection=[id, name], fetch=5" + ); + } + + #[test] + fn keeps_other_plans() { + let plan = LogicalPlanBuilder::empty(false).build().unwrap(); + assert_eq!( + federate_table_scans_only(plan.clone(), Arc::new(NoPlanner)).unwrap(), + plan + ); + } +} diff --git a/crates/cassandra/src/lib.rs b/crates/cassandra/src/lib.rs index 0bf8c0cca..fe83f7479 100644 --- a/crates/cassandra/src/lib.rs +++ b/crates/cassandra/src/lib.rs @@ -9,6 +9,8 @@ //! to with `INSERT INTO`; see [`CassandraTableWriter`]. pub mod conn; +#[cfg(feature = "federation")] +mod federation; pub mod pool; mod table; pub mod types; @@ -53,11 +55,16 @@ impl CassandraTableFactory { .get_schema(&table_reference) .await?; - Ok(Arc::new(CassandraTable::new( + let table_provider = Arc::new(CassandraTable::new( table_reference, Arc::clone(&self.pool), schema, - ))) + )); + + #[cfg(feature = "federation")] + let table_provider = Arc::new(table_provider.create_federated_table_provider()?); + + Ok(table_provider) } /// Creates a table provider for `table_reference` that also supports `INSERT INTO`. diff --git a/crates/cassandra/src/table.rs b/crates/cassandra/src/table.rs index eca07abd0..7312b6d14 100644 --- a/crates/cassandra/src/table.rs +++ b/crates/cassandra/src/table.rs @@ -27,9 +27,9 @@ use crate::pool::CassandraConnectionPool; /// /// Filters are not pushed down (see the crate documentation), so DataFusion applies them. pub struct CassandraTable { - table_reference: TableReference, - pool: Arc, - schema: SchemaRef, + pub(crate) table_reference: TableReference, + pub(crate) pool: Arc, + pub(crate) schema: SchemaRef, } impl fmt::Debug for CassandraTable { @@ -157,7 +157,11 @@ fn scan_columns( } } -fn build_cql(table_reference: &TableReference, columns: &[String], limit: Option) -> String { +pub(crate) fn build_cql( + table_reference: &TableReference, + columns: &[String], + limit: Option, +) -> String { let columns = columns .iter() .map(|c| quote_identifier(c)) From cc2c32b34744027e7a8ee0e2242ef05c83757a67 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Mon, 19 Oct 2026 09:04:03 +0300 Subject: [PATCH 19/27] Start the Cassandra Docker container in single-node setup, i.e. do not wait for cluster --- core/tests/cassandra/common.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/core/tests/cassandra/common.rs b/core/tests/cassandra/common.rs index eb2c52920..e3402abe0 100644 --- a/core/tests/cassandra/common.rs +++ b/core/tests/cassandra/common.rs @@ -30,6 +30,17 @@ pub(super) async fn start_cassandra_docker_container( // the contact point. By default that is the container's internal IP, which, combined // with the published host port, isn't reachable from the host. .add_env_var("CASSANDRA_BROADCAST_RPC_ADDRESS", "127.0.0.1") + // Single-node cluster: the node is its own seed (the image's default), so there is + // no one to gossip with. Skip the gossip-settle polling and the 30s ring delay. + .add_env_var( + "JVM_EXTRA_OPTS", + "-Dcassandra.skip_wait_for_gossip_to_settle=0 -Dcassandra.ring_delay_ms=0", + ) + .add_env_var( + "CASSANDRA_NUM_TOKENS", + &std::env::var("CASSANDRA_NUM_TOKENS").unwrap_or("1".to_string()), + ) + .add_env_var("CASSANDRA_CLUSTER_NAME", "df-t-p-cassandra") // Keep the JVM small so that the node starts quickly in CI .add_env_var("MAX_HEAP_SIZE", "512M") .add_env_var("HEAP_NEWSIZE", "128M") From 70be91c5f14791db821045ef953604f8937d1265 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Mon, 19 Oct 2026 10:16:38 +0300 Subject: [PATCH 20/27] Implement filter pushdown for key columns --- README.md | 4 +- core/examples/cassandra.rs | 15 +- core/tests/cassandra/mod.rs | 257 ++++++++++++ crates/cassandra/src/conn.rs | 47 ++- crates/cassandra/src/federation.rs | 63 ++- crates/cassandra/src/filter.rs | 605 +++++++++++++++++++++++++++++ crates/cassandra/src/lib.rs | 28 +- crates/cassandra/src/table.rs | 157 +++++++- crates/cassandra/src/types.rs | 2 +- 9 files changed, 1111 insertions(+), 67 deletions(-) create mode 100644 crates/cassandra/src/filter.rs diff --git a/README.md b/README.md index 545047aff..96b71f9c8 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ Prefer a single-provider leaf crate when you only need one backend — for examp - Leaf crates enable `federation` by default (opt out with `default-features = false`). SQLite also defaults to the `bundled` rusqlite feature. - The existing `datafusion-table-providers` facade keeps the same module paths and provider feature names (`duckdb`, `postgres`, `sqlite`, etc.). Enabling a provider feature pulls in that leaf crate with its defaults (including federation). There is no separate `federation` / `*-federation` feature on the facade. - ClickHouse pool connections are exposed as `ClickHouseConnection` (wrapping the ClickHouse `Client`) under `sql::db_connection_pool::dbconnection::clickhouseconn`. -- Cassandra pool connections are exposed as `CassandraConnection` (wrapping a shared `scylla` driver `Session`) under `sql::db_connection_pool::dbconnection::cassandraconn`. Its `federation` feature works like the other crates', except that CQL cannot run the joins, aggregations and filters that federation pushes down: a federated Cassandra sub-plan is narrowed to its table scans, so only projections and `LIMIT` are pushed down to Cassandra. +- Cassandra pool connections are exposed as `CassandraConnection` (wrapping a shared `scylla` driver `Session`) under `sql::db_connection_pool::dbconnection::cassandraconn`. Its `federation` feature works like the other crates', except that CQL cannot run the joins, aggregations and filters that federation pushes down: a federated Cassandra sub-plan is narrowed to its table scans, so only projections, filters on key columns and `LIMIT` are pushed down to Cassandra. Existing examples continue to use the facade crate and its feature flags. @@ -194,7 +194,7 @@ cargo run -p datafusion-table-providers --example mysql --features mysql ### Cassandra -The Cassandra provider is built on the [ScyllaDB Rust driver](https://crates.io/crates/scylla), so it works with ScyllaDB too. CQL has no joins or aggregations and can filter only on key columns, so only projections and `LIMIT` are pushed down to Cassandra; filters, aggregations and joins are evaluated by DataFusion. Tables created with `CassandraTableFactory::read_write_table_provider` also support `INSERT INTO`, which writes each row with a CQL `INSERT` (an upsert), and `INSERT OVERWRITE`, which truncates the table first. +The Cassandra provider is built on the [ScyllaDB Rust driver](https://crates.io/crates/scylla), so it works with ScyllaDB too. CQL has no joins or aggregations and can filter only on key columns, so only projections, `LIMIT` and the filters on key columns that make up a valid CQL `WHERE` clause (`=` or `IN` on every partition key column, then `=` on leading clustering columns and optionally a range on the next one) are pushed down to Cassandra; the other filters, aggregations and joins are evaluated by DataFusion. Tables created with `CassandraTableFactory::read_write_table_provider` also support `INSERT INTO`, which writes each row with a CQL `INSERT` (an upsert), and `INSERT OVERWRITE`, which truncates the table first. In order to run the Cassandra example, you need to have a Cassandra server running. You can use the following command to start a Cassandra server in a Docker container the example can use: diff --git a/core/examples/cassandra.rs b/core/examples/cassandra.rs index 23f85d5c9..0ea25aafe 100644 --- a/core/examples/cassandra.rs +++ b/core/examples/cassandra.rs @@ -12,7 +12,8 @@ use datafusion_table_providers::{ /// 2. Create and use CassandraTableFactory to generate TableProvider /// 3. Register TableProvider with DataFusion /// 4. Use SQL queries to access Cassandra table data -/// 5. Write data to a Cassandra table with `INSERT INTO` +/// 5. Push filters on key columns down to Cassandra +/// 6. Write data to a Cassandra table with `INSERT INTO` /// /// Prerequisites: /// Start a Cassandra server using Docker: @@ -93,4 +94,16 @@ async fn main() { .await .expect("select failed"); df.show().await.expect("show failed"); + + // Query Example 3: Filter on the partition key. + // CQL can filter on it, so the filter is pushed down to Cassandra and only the matching + // rows are read. The plan shows it in the CQL: `... WHERE "id" IN (1, 3)`. + let sql = "SELECT name FROM companies WHERE id IN (1, 3)"; + let df = ctx + .sql(&format!("EXPLAIN {sql}")) + .await + .expect("explain failed"); + df.show().await.expect("show failed"); + let df = ctx.sql(sql).await.expect("select failed"); + df.show().await.expect("show failed"); } diff --git a/core/tests/cassandra/mod.rs b/core/tests/cassandra/mod.rs index 2f498b0e0..b167c22ad 100644 --- a/core/tests/cassandra/mod.rs +++ b/core/tests/cassandra/mod.rs @@ -87,6 +87,9 @@ async fn test_cassandra() { test_cassandra_count(&pool).await; test_cassandra_explain_shows_pushed_down_cql(&pool).await; test_cassandra_federation(&pool).await; + test_cassandra_filter_pushdown(&pool).await; + test_cassandra_filter_pushdown_federated(&pool).await; + test_cassandra_scan_with_filters_not_pushed_down(&pool).await; test_cassandra_tables_and_schemas(&pool).await; test_cassandra_missing_table(&pool).await; test_cassandra_connection_pool_with_schema(port).await; @@ -160,6 +163,32 @@ async fn create_test_data(pool: &CassandraConnectionPool) { ), ) .await; + execute( + pool, + &format!( + "CREATE TABLE IF NOT EXISTS {KEYSPACE}.readings (sensor text, bucket int, ts bigint, \ + seq int, value int, PRIMARY KEY ((sensor, bucket), ts, seq))" + ), + ) + .await; + for (sensor, bucket, ts, seq, value) in [ + ("a", 1, 10, 1, 100), + ("a", 1, 10, 2, 101), + ("a", 1, 20, 1, 102), + ("a", 1, 30, 1, -103), + ("a", 2, 10, 1, 104), + ("b", 1, 10, 1, 105), + ("b", 1, 40, 1, 106), + ] { + execute( + pool, + &format!( + "INSERT INTO {KEYSPACE}.readings (sensor, bucket, ts, seq, value) \ + VALUES ('{sensor}', {bucket}, {ts}, {seq}, {value})" + ), + ) + .await; + } for (id, name, age) in [ (1, "alice", 31), (2, "bob", 25), @@ -469,6 +498,233 @@ async fn test_cassandra_federation(pool: &Arc) { ); } +/// Returns the `EXPLAIN` output of `sql`. +async fn explain(ctx: &SessionContext, sql: &str) -> String { + query_to_string(ctx, &format!("EXPLAIN {sql}")).await +} + +/// The filters on key columns that make up a valid CQL `WHERE` clause must be pushed down to +/// Cassandra, the others must be applied by DataFusion. +async fn test_cassandra_filter_pushdown(pool: &Arc) { + let ctx = SessionContext::new(); + register_table(&ctx, pool, "users").await; + register_table(&ctx, pool, "readings").await; + + let sql = "SELECT name FROM users WHERE id = 3"; + let plan = explain(&ctx, sql).await; + assert!( + plan.contains(r#"CassandraExec cql=SELECT "name" FROM "it_ks"."users" WHERE "id" = 3"#), + "the filter on the partition key should be pushed down:\n{plan}" + ); + assert!(!plan.contains("FilterExec"), "{plan}"); + assert_eq!( + query_to_string(&ctx, sql).await, + "+-------+ +| name | ++-------+ +| carol | ++-------+" + ); + + let sql = "SELECT name FROM users WHERE id IN (2, 4) AND age > 30"; + let plan = explain(&ctx, sql).await; + assert!( + plan.contains(r#"FROM "it_ks"."users" WHERE "id" IN (2, 4)"#), + "the IN list on the partition key should be pushed down:\n{plan}" + ); + assert!(plan.contains("FilterExec: age@"), "{plan}"); + assert_eq!( + query_to_string(&ctx, sql).await, + "+------+ +| name | ++------+ +| dave | ++------+" + ); + + // the whole partition key, the 1st clustering column by `=` and the 2nd by a range + let sql = "SELECT ts, seq, value FROM readings \ + WHERE sensor = 'a' AND bucket = 1 AND ts = 10 AND seq > 1 AND seq <= 5 \ + ORDER BY seq"; + let plan = explain(&ctx, sql).await; + assert!( + plan.contains( + r#"WHERE "sensor" = 'a' AND "bucket" = 1 AND "ts" = 10 AND "seq" > 1 AND "seq" <= 5"# + ), + "the filters on the key columns should be pushed down:\n{plan}" + ); + assert!(!plan.contains("FilterExec"), "{plan}"); + assert_eq!( + query_to_string(&ctx, sql).await, + "+----+-----+-------+ +| ts | seq | value | ++----+-----+-------+ +| 10 | 2 | 101 | ++----+-----+-------+" + ); + + // a range on the 1st clustering column, with a filter on a regular column + let sql = "SELECT ts, value FROM readings \ + WHERE sensor = 'a' AND bucket = 1 AND ts >= 20 AND value > 0"; + let plan = explain(&ctx, sql).await; + assert!( + plan.contains(r#"WHERE "sensor" = 'a' AND "bucket" = 1 AND "ts" >= 20"#), + "the filters on the key columns should be pushed down:\n{plan}" + ); + assert!(plan.contains("FilterExec: value@"), "{plan}"); + assert_eq!( + query_to_string(&ctx, sql).await, + "+----+-------+ +| ts | value | ++----+-------+ +| 20 | 102 | ++----+-------+" + ); + + // CQL can't restrict a part of the partition key, nor clustering columns without it + let sql = "SELECT count(*) FROM readings WHERE sensor = 'a' AND ts = 10"; + let plan = explain(&ctx, sql).await; + assert!( + !plan.contains("WHERE"), + "nothing should be pushed down:\n{plan}" + ); + let batches = query(&ctx, sql).await; + assert_eq!(array_value_to_string(batches[0].column(0), 0).unwrap(), "3"); + + // only one restriction of a column can be pushed down + let sql = "SELECT count(*) FROM readings WHERE sensor = 'a' AND sensor = 'b' AND bucket = 1"; + let batches = query(&ctx, sql).await; + assert_eq!(array_value_to_string(batches[0].column(0), 0).unwrap(), "0"); + + // the clustering column after a range can't be restricted + let sql = "SELECT count(*) FROM readings \ + WHERE sensor = 'a' AND bucket = 1 AND ts > 5 AND seq = 1"; + let plan = explain(&ctx, sql).await; + assert!( + plan.contains(r#"WHERE "sensor" = 'a' AND "bucket" = 1 AND "ts" > 5"#), + "the partition key and the range should be pushed down:\n{plan}" + ); + let batches = query(&ctx, sql).await; + assert_eq!(array_value_to_string(batches[0].column(0), 0).unwrap(), "3"); + + // the limit is pushed down along with the filters + let sql = "SELECT seq FROM readings WHERE sensor = 'a' AND bucket = 1 LIMIT 2"; + let plan = explain(&ctx, sql).await; + assert!( + plan.contains(r#"WHERE "sensor" = 'a' AND "bucket" = 1 LIMIT 2"#), + "the filters and the limit should be pushed down:\n{plan}" + ); + let batches = query(&ctx, sql).await; + let num_rows: usize = batches.iter().map(RecordBatch::num_rows).sum(); + assert_eq!(num_rows, 2); +} + +/// In a federated session, the filters on key columns must be pushed down too. +async fn test_cassandra_filter_pushdown_federated(pool: &Arc) { + let ctx = SessionContext::new_with_state(datafusion_federation::default_session_state()); + register_table(&ctx, pool, "users").await; + register_table(&ctx, pool, "readings").await; + + let sql = "SELECT name FROM users WHERE id = 3 AND age > 30"; + let plan = explain(&ctx, sql).await; + assert!( + plan.contains(r#"FROM "it_ks"."users" WHERE "id" = 3"#), + "the filter on the partition key should be pushed down:\n{plan}" + ); + assert_eq!( + query_to_string(&ctx, sql).await, + "+-------+ +| name | ++-------+ +| carol | ++-------+" + ); + + // The tables of a join are federated separately before DataFusion pushes the filters of + // the join down, so only the filters on the table itself reach the federated scan + let sql = "SELECT r.value FROM \ + (SELECT * FROM readings WHERE sensor = 'b' AND bucket = 1 AND ts < 40) r \ + JOIN users u ON r.bucket = u.id"; + let plan = explain(&ctx, sql).await; + assert!( + plan.contains(r#"WHERE "sensor" = 'b' AND "bucket" = 1 AND "ts" < 40"#), + "the filters on the key columns should be pushed down:\n{plan}" + ); + assert_eq!( + query_to_string(&ctx, sql).await, + "+-------+ +| value | ++-------+ +| 105 | ++-------+" + ); +} + +/// DataFusion may pass `scan` filters that don't make up a valid CQL `WHERE` clause together, +/// as it decides on them in separate optimizer passes: those not pushed down must be applied. +async fn test_cassandra_scan_with_filters_not_pushed_down(pool: &Arc) { + use datafusion::logical_expr::{col, lit}; + use datafusion::physical_plan::{collect, displayable}; + + let ctx = SessionContext::new(); + let provider = CassandraTableFactory::new(Arc::clone(pool)) + .table_provider(TableReference::partial(KEYSPACE, "users")) + .await + .expect("table provider should be created"); + + let num_rows = |filters: Vec, limit| { + let ctx = &ctx; + let provider = Arc::clone(&provider); + async move { + let plan = provider + .scan(&ctx.state(), Some(&vec![1]), &filters, limit) + .await + .expect("scan should be planned"); + let plan_string = displayable(plan.as_ref()).indent(false).to_string(); + let batches = collect(plan, ctx.task_ctx()) + .await + .expect("scan should be executed"); + assert!( + batches.iter().all(|b| b.num_columns() == 1), + "{plan_string}" + ); + ( + batches.iter().map(RecordBatch::num_rows).sum::(), + plan_string, + ) + } + }; + + let (rows, plan) = num_rows(vec![col("id").eq(lit(1)), col("id").eq(lit(2))], None).await; + assert_eq!(rows, 0, "{plan}"); + assert!(plan.contains(r#"WHERE "id" = 1"#), "{plan}"); + assert!(plan.contains("FilterExec: id@1 = 2"), "{plan}"); + + let (rows, plan) = num_rows( + vec![ + col("id").eq(lit(1)), + col("id").in_list(vec![lit(1), lit(2)], false), + ], + Some(5), + ) + .await; + assert_eq!(rows, 1, "{plan}"); + + let (rows, plan) = num_rows( + vec![ + col("id").in_list(vec![lit(1), lit(2), lit(3)], false), + col("age").gt(lit(26)), + ], + Some(1), + ) + .await; + assert_eq!(rows, 1, "{plan}"); + assert!( + !plan.contains("LIMIT"), + "the limit must be applied after the filter: {plan}" + ); +} + async fn test_cassandra_tables_and_schemas(pool: &Arc) { let conn = pool.connect().await.expect("connection should be created"); let conn = conn.as_async().expect("connection should be async"); @@ -487,6 +743,7 @@ async fn test_cassandra_tables_and_schemas(pool: &Arc) vec![ "all_types".to_string(), "all_types_written".to_string(), + "readings".to_string(), "users".to_string(), "users_written".to_string(), ] diff --git a/crates/cassandra/src/conn.rs b/crates/cassandra/src/conn.rs index 3cad432bf..c998f5cd6 100644 --- a/crates/cassandra/src/conn.rs +++ b/crates/cassandra/src/conn.rs @@ -13,6 +13,7 @@ use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::{ self, AsyncDbConnection, DbConnection, GenericError, Result, TableSchemas, }; +use crate::filter::{KeyColumn, PrimaryKey}; use crate::types::{column_specs_to_schema, rows_to_record_batch}; /// Number of rows converted to Arrow at a time; Cassandra pages are 5000 rows by default. @@ -49,7 +50,7 @@ pub enum Error { /// The driver's [`Session`] manages its own connection pool, so a connection is a cheap handle /// to a shared session. pub struct CassandraConnection { - pub session: Arc, + session: Arc, schemas: Arc, } @@ -75,6 +76,50 @@ impl CassandraConnection { .await .context(ReadRowsSnafu) } + + /// Returns the primary key of a table, given as `table` (in the session's keyspace) or + /// `keyspace.table`. It is empty if the table has no keyspace. + /// + /// # Errors + /// + /// Returns an error if the columns of the table cannot be read. + pub async fn primary_key(&self, table_reference: &TableReference) -> Result { + let keyspace = match table_reference.schema() { + Some(keyspace) => keyspace.to_string(), + None => match self.session.get_keyspace() { + Some(keyspace) => keyspace.to_string(), + None => return Ok(PrimaryKey::default()), + }, + }; + + let mut columns: Vec<(String, String, i32, String)> = self + .session + .query_iter( + "SELECT column_name, kind, position, type FROM system_schema.columns \ + WHERE keyspace_name = ? AND table_name = ?", + (keyspace, table_reference.table()), + ) + .await + .context(QuerySnafu)? + .rows_stream::<(String, String, i32, String)>() + .context(TypeCheckSnafu)? + .try_collect() + .await + .context(ReadRowsSnafu)?; + columns.sort_by_key(|(_, _, position, _)| *position); + + let key_columns = |kind: &str| { + columns + .iter() + .filter(|(_, column_kind, _, _)| column_kind == kind) + .map(|(name, _, _, cql_type)| KeyColumn::new(name, cql_type)) + .collect() + }; + Ok(PrimaryKey { + partition_key: key_columns("partition_key"), + clustering_key: key_columns("clustering"), + }) + } } impl DbConnection, CqlValue> for CassandraConnection { diff --git a/crates/cassandra/src/federation.rs b/crates/cassandra/src/federation.rs index 5a0cb29d3..3d728d76d 100644 --- a/crates/cassandra/src/federation.rs +++ b/crates/cassandra/src/federation.rs @@ -3,8 +3,8 @@ //! The federation optimizer hands a provider the largest sub-plan that only reads from it, to be //! run remotely as one SQL query. CQL can't run most such queries: it has no joins or //! aggregations, and `WHERE` clauses are limited to key columns. So the sub-plan is narrowed down -//! to its table scans, which keep their projection and `LIMIT` and run as [`CassandraExec`]s, -//! and the rest of it is evaluated by DataFusion. +//! to its table scans, which keep their projection, the filters they can push down and `LIMIT`, +//! and run as [`CassandraExec`]s. The rest of it is evaluated by DataFusion. use std::sync::Arc; @@ -15,10 +15,12 @@ use datafusion::common::tree_node::Transformed; use datafusion::common::TableReference; use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::SendableRecordBatchStream; -use datafusion::logical_expr::{Extension, LogicalPlan, TableScan}; +use datafusion::logical_expr::{Extension, LogicalPlan}; use datafusion::optimizer::optimize_projections::OptimizeProjections; use datafusion::optimizer::optimizer::Optimizer; +use datafusion::optimizer::push_down_filter::PushDownFilter; use datafusion::optimizer::push_down_limit::PushDownLimit; +use datafusion::optimizer::simplify_expressions::SimplifyExpressions; use datafusion::optimizer::OptimizerContext; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; @@ -35,8 +37,7 @@ use futures::TryStreamExt; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::AsyncDbConnection; use datafusion_table_providers_common::sql::db_connection_pool::{DbConnectionPool, JoinPushDown}; -use crate::pool::CassandraConnectionPool; -use crate::{CassandraExec, CassandraTable}; +use crate::CassandraTable; impl CassandraTable { fn create_federated_table_source( @@ -88,9 +89,7 @@ impl SQLExecutor for CassandraTable { fn logical_optimizer(&self) -> Option { let planner: Arc = Arc::new(CassandraScanPlanner { - table_reference: self.table_reference.clone(), - pool: Arc::clone(&self.pool), - schema: Arc::clone(&self.schema), + table: self.clone(), }); Some(Box::new(move |plan| { federate_table_scans_only(plan, Arc::clone(&planner)) @@ -138,8 +137,10 @@ impl SQLExecutor for CassandraTable { /// federated, to be planned by `planner`, as CQL can run nothing more (see the module /// documentation). /// -/// The projections and limits are pushed down to the scans first: once federated, the scans -/// are opaque to the DataFusion optimizer. +/// The projections, filters and limits are pushed down to the scans first: once federated, the +/// scans are opaque to the DataFusion optimizer. The federation runs before the DataFusion +/// optimizer simplifies the expressions, which the filter push down relies on (e.g. to turn +/// `CAST(id AS BIGINT) = 1` into `id = 1`), so they are simplified first. fn federate_table_scans_only( plan: LogicalPlan, planner: Arc, @@ -152,7 +153,9 @@ fn federate_table_scans_only( }; let optimizer = Optimizer::with_rules(vec![ + Arc::new(SimplifyExpressions::new()), Arc::new(OptimizeProjections::new()), + Arc::new(PushDownFilter::new()), Arc::new(PushDownLimit::new()), ]); let inner = optimizer.optimize( @@ -171,14 +174,14 @@ fn federate_table_scans_only( .map(|transformed| transformed.data) } -/// Plans a federated scan of a Cassandra table as a [`CassandraExec`]. +/// Plans a federated scan of a Cassandra table like [`CassandraTable::scan`]. /// /// The SQL federation would plan it as a scan that claims all the filters DataFusion pushes /// down to it at runtime, but CQL can't run them. +/// +/// [`CassandraTable::scan`]: datafusion::catalog::TableProvider::scan struct CassandraScanPlanner { - table_reference: TableReference, - pool: Arc, - schema: SchemaRef, + table: CassandraTable, } #[async_trait] @@ -186,23 +189,20 @@ impl FederationPlanner for CassandraScanPlanner { async fn plan_federation( &self, node: &FederatedPlanNode, - _session_state: &dyn Session, + session_state: &dyn Session, ) -> DataFusionResult> { - let scan = match node.plan() { - LogicalPlan::TableScan(scan @ TableScan { filters, .. }) if filters.is_empty() => scan, - plan => { - return Err(DataFusionError::Plan(format!( - "Cassandra can only run plain table scans, not:\n{plan}" - ))) - } + let LogicalPlan::TableScan(scan) = node.plan() else { + return Err(DataFusionError::Plan(format!( + "Cassandra can only run table scans, not:\n{}", + node.plan() + ))); }; - Ok(Arc::new(CassandraExec::try_new( - self.table_reference.clone(), - Arc::clone(&self.pool), - &self.schema, + self.table.plan_scan( + session_state, scan.projection.as_ref(), + &scan.filters, scan.fetch, - )?)) + ) } } @@ -231,8 +231,9 @@ mod tests { }) } + /// `LogicalTableSource` supports pushing down all filters exactly. #[test] - fn federates_only_the_scans_with_projection_and_limit() { + fn federates_only_the_scans_with_projection_filters_and_limit() { let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Int32, true), Field::new("name", DataType::Utf8, true), @@ -267,10 +268,8 @@ mod tests { rewritten.display_indent().to_string(), "Projection: ks.users.name, ks.pets.name Inner Join: Filter: ks.users.id = ks.pets.id - Projection: ks.users.id, ks.users.name - Filter: ks.users.age > Int32(30) - Federated - TableScan: ks.users projection=[id, name, age] + Federated + TableScan: ks.users projection=[id, name], full_filters=[ks.users.age > Int32(30)] Limit: skip=0, fetch=5 Federated TableScan: ks.pets projection=[id, name], fetch=5" diff --git a/crates/cassandra/src/filter.rs b/crates/cassandra/src/filter.rs new file mode 100644 index 000000000..c95999c40 --- /dev/null +++ b/crates/cassandra/src/filter.rs @@ -0,0 +1,605 @@ +//! Pushing filters on key columns down to Cassandra. +//! +//! Without `ALLOW FILTERING`, which makes Cassandra scan the whole table, a CQL `WHERE` clause +//! may only restrict key columns, and only in these ways: +//! +//! * either every partition key column is restricted by `=` or `IN`, or none is; +//! * the clustering columns may only be restricted if the partition key is, and only in order: +//! any number of them by `=`, then optionally the next one by a range (`<`, `<=`, `>`, `>=`), +//! with at most one lower and one upper bound. +//! +//! [`select_pushed_down_filters`] picks the filters that make up such a `WHERE` clause. Only +//! comparisons of a key column with a literal are considered, and only for the CQL types whose +//! comparison in Cassandra gives the same result as in DataFusion, so that the pushed down +//! filters are exact. All the other filters are evaluated by DataFusion. + +use datafusion::common::{Column, ScalarValue}; +use datafusion::logical_expr::utils::split_binary; +use datafusion::logical_expr::{BinaryExpr, Expr, Operator}; + +use crate::conn::quote_identifier; + +/// The primary key of a Cassandra table. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct PrimaryKey { + /// The partition key columns, in order. + pub partition_key: Vec, + /// The clustering columns, in order. + pub clustering_key: Vec, +} + +/// A primary key column. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct KeyColumn { + pub name: String, + /// The CQL type, as in `system_schema.columns`, e.g. `int` or `frozen>`. + pub cql_type: String, +} + +impl KeyColumn { + pub fn new(name: impl Into, cql_type: impl Into) -> Self { + Self { + name: name.into(), + cql_type: cql_type.into(), + } + } +} + +/// A restriction of a key column, with its values as CQL literals. +#[derive(Debug, Clone, PartialEq)] +enum Restriction { + Eq(String), + In(Vec), + /// `>` (exclusive) or `>=`. + Lower { + value: String, + inclusive: bool, + }, + /// `<` (exclusive) or `<=`. + Upper { + value: String, + inclusive: bool, + }, +} + +impl Restriction { + fn to_cql(&self, column: &str) -> String { + let column = quote_identifier(column); + match self { + Self::Eq(value) => format!("{column} = {value}"), + Self::In(values) => format!("{column} IN ({})", values.join(", ")), + Self::Lower { value, inclusive } => { + format!("{column} {} {value}", if *inclusive { ">=" } else { ">" }) + } + Self::Upper { value, inclusive } => { + format!("{column} {} {value}", if *inclusive { "<=" } else { "<" }) + } + } + } +} + +/// The filters to push down to Cassandra. +#[derive(Debug, Default, PartialEq)] +pub(crate) struct PushedDownFilters { + /// For each filter, whether it is pushed down. + pub pushed_down: Vec, + /// The relations of the CQL `WHERE` clause, to be joined with `AND`. + pub relations: Vec, +} + +/// Selects the filters that can be pushed down to Cassandra as a CQL `WHERE` clause, see the +/// module documentation. +pub(crate) fn select_pushed_down_filters( + primary_key: &PrimaryKey, + filters: &[&Expr], +) -> PushedDownFilters { + let mut pushed_down = vec![false; filters.len()]; + let mut relations = vec![]; + if primary_key.partition_key.is_empty() { + return PushedDownFilters { + pushed_down, + relations, + }; + } + + let restrictions: Vec> = filters + .iter() + .map(|filter| key_column_restriction(primary_key, filter)) + .collect(); + // Finds the first filter restricting `column` with a restriction accepted by `accept`. + // Further restrictions of the same kind are left to DataFusion: CQL rejects them. + let find = |column: &str, accept: fn(&Restriction) -> bool| { + restrictions.iter().position(|r| { + r.as_ref() + .is_some_and(|(name, restriction)| *name == column && accept(restriction)) + }) + }; + let mut push_down = |index: usize, column: &str| { + pushed_down[index] = true; + let (_, restriction) = restrictions[index].as_ref().expect("restriction exists"); + relations.push(restriction.to_cql(column)); + }; + + let partition_key_filters: Option> = primary_key + .partition_key + .iter() + .map(|column| { + find(&column.name, |r| { + matches!(r, Restriction::Eq(_) | Restriction::In(_)) + }) + }) + .collect(); + let Some(partition_key_filters) = partition_key_filters else { + // CQL can't restrict only some of the partition key columns, nor the clustering + // columns without the partition key + return PushedDownFilters { + pushed_down, + relations, + }; + }; + for (index, column) in partition_key_filters + .into_iter() + .zip(&primary_key.partition_key) + { + push_down(index, &column.name); + } + + for column in &primary_key.clustering_key { + if let Some(index) = find(&column.name, |r| matches!(r, Restriction::Eq(_))) { + push_down(index, &column.name); + continue; + } + // A range restriction ends the restricted clustering columns + let lower = find(&column.name, |r| matches!(r, Restriction::Lower { .. })); + let upper = find(&column.name, |r| matches!(r, Restriction::Upper { .. })); + for index in [lower, upper].into_iter().flatten() { + push_down(index, &column.name); + } + break; + } + + PushedDownFilters { + pushed_down, + relations, + } +} + +/// Returns the key column `filter` restricts and the restriction, if CQL can evaluate it +/// exactly like DataFusion does. +fn key_column_restriction<'a>( + primary_key: &'a PrimaryKey, + filter: &Expr, +) -> Option<(&'a str, Restriction)> { + let key_column = |column: &Column| { + let is_partition_key = primary_key + .partition_key + .iter() + .any(|c| c.name == column.name); + primary_key + .partition_key + .iter() + .chain(&primary_key.clustering_key) + .find(|c| c.name == column.name) + .map(|c| (c, is_partition_key)) + }; + + match filter { + // DataFusion rewrites short `IN` lists to `OR`s of `=` + Expr::BinaryExpr(BinaryExpr { + op: Operator::Or, .. + }) => { + let mut column = None; + let mut literals = vec![]; + for disjunct in split_binary(filter, Operator::Or) { + let (name, restriction) = key_column_restriction(primary_key, disjunct)?; + if *column.get_or_insert(name) != name { + return None; + } + match restriction { + Restriction::Eq(literal) => literals.push(literal), + Restriction::In(values) => literals.extend(values), + _ => return None, + } + } + let column = column?; + let is_partition_key = primary_key.partition_key.iter().any(|c| c.name == column); + is_partition_key.then_some((column, Restriction::In(literals))) + } + Expr::BinaryExpr(BinaryExpr { left, op, right }) => { + let (column, op, value) = match (left.as_ref(), right.as_ref()) { + (Expr::Column(column), Expr::Literal(value, _)) => (column, *op, value), + (Expr::Literal(value, _), Expr::Column(column)) => (column, op.swap()?, value), + _ => return None, + }; + let (key_column, is_partition_key) = key_column(column)?; + let literal = cql_literal(&key_column.cql_type, value, is_partition_key)?; + if op != Operator::Eq && !supports_range(&key_column.cql_type) { + return None; + } + let restriction = match op { + Operator::Eq => Restriction::Eq(literal), + Operator::Gt | Operator::GtEq => Restriction::Lower { + value: literal, + inclusive: op == Operator::GtEq, + }, + Operator::Lt | Operator::LtEq => Restriction::Upper { + value: literal, + inclusive: op == Operator::LtEq, + }, + _ => return None, + }; + Some((&key_column.name, restriction)) + } + Expr::InList(in_list) if !in_list.negated && !in_list.list.is_empty() => { + let Expr::Column(column) = in_list.expr.as_ref() else { + return None; + }; + let (key_column, is_partition_key) = key_column(column)?; + let literals = in_list + .list + .iter() + .map(|value| match value { + Expr::Literal(value, _) => { + cql_literal(&key_column.cql_type, value, is_partition_key) + } + _ => None, + }) + .collect::>>()?; + // CQL supports `IN` on clustering columns too, but with more caveats, so only + // the partition key is restricted with it + is_partition_key.then_some((key_column.name.as_str(), Restriction::In(literals))) + } + _ => None, + } +} + +/// Whether ordering values of `cql_type` in Cassandra gives the same order as in DataFusion. +/// +/// This isn't the case for UUIDs (`timeuuid`s are ordered by time, `uuid`s by version first), +/// whose values are strings in DataFusion. +fn supports_range(cql_type: &str) -> bool { + !matches!(cql_type, "uuid" | "timeuuid") +} + +/// Formats `value` as a CQL literal of `cql_type`. +/// +/// Returns `None` if the value isn't one of the column, as DataFusion produces it, or if +/// Cassandra would reject the query or compare the values differently than DataFusion: +/// in these cases DataFusion must evaluate the filter. +fn cql_literal(cql_type: &str, value: &ScalarValue, is_partition_key: bool) -> Option { + let literal = match (cql_type, value) { + ("tinyint", ScalarValue::Int8(Some(v))) => v.to_string(), + ("smallint", ScalarValue::Int16(Some(v))) => v.to_string(), + ("int", ScalarValue::Int32(Some(v))) => v.to_string(), + ("bigint", ScalarValue::Int64(Some(v))) => v.to_string(), + ("boolean", ScalarValue::Boolean(Some(v))) => v.to_string(), + // Cassandra rejects empty partition keys, while DataFusion just finds no rows + ("text" | "varchar", ScalarValue::Utf8(Some(s))) if !(is_partition_key && s.is_empty()) => { + quote_string(s) + } + // Cassandra rejects non-ASCII values for `ascii` columns + ("ascii", ScalarValue::Utf8(Some(s))) + if s.is_ascii() && !(is_partition_key && s.is_empty()) => + { + quote_string(s) + } + ("blob", ScalarValue::Binary(Some(bytes))) if !(is_partition_key && bytes.is_empty()) => { + let hex: String = bytes.iter().map(|b| format!("{b:02x}")).collect(); + format!("0x{hex}") + } + // An integer `date` literal is the number of days, with the epoch at 2^31 + ("date", ScalarValue::Date32(Some(days))) => (i64::from(*days) + (1 << 31)).to_string(), + // An integer `timestamp` literal is the number of milliseconds since the epoch + ("timestamp", ScalarValue::TimestampMillisecond(Some(millis), _)) => millis.to_string(), + // An integer `time` literal is the number of nanoseconds since midnight + ("time", ScalarValue::Time64Nanosecond(Some(nanos))) => nanos.to_string(), + // DataFusion compares UUIDs as strings, so only the canonical form, in which Cassandra + // returns them, matches the same rows. Cassandra rejects `timeuuid`s of other versions + ("uuid" | "timeuuid", ScalarValue::Utf8(Some(s))) => { + let uuid = uuid::Uuid::try_parse(s).ok()?; + if uuid.to_string() != *s || (cql_type == "timeuuid" && uuid.get_version_num() != 1) { + return None; + } + uuid.to_string() + } + _ => return None, + }; + Some(literal) +} + +fn quote_string(s: &str) -> String { + format!("'{}'", s.replace('\'', "''")) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::logical_expr::{col, lit}; + + /// `PRIMARY KEY ((p1, p2), c1, c2)` + fn primary_key() -> PrimaryKey { + PrimaryKey { + partition_key: vec![KeyColumn::new("p1", "int"), KeyColumn::new("p2", "text")], + clustering_key: vec![ + KeyColumn::new("c1", "bigint"), + KeyColumn::new("c2", "timeuuid"), + ], + } + } + + fn select(filters: &[Expr]) -> (Vec, String) { + let filters: Vec<&Expr> = filters.iter().collect(); + let selected = select_pushed_down_filters(&primary_key(), &filters); + (selected.pushed_down, selected.relations.join(" AND ")) + } + + #[test] + fn pushes_down_the_partition_key() { + assert_eq!( + select(&[ + col("p2").eq(lit("a'b")), + col("v").eq(lit(1)), + lit(1).eq(col("p1")), + ]), + ( + vec![true, false, true], + r#""p1" = 1 AND "p2" = 'a''b'"#.to_string() + ) + ); + } + + #[test] + fn pushes_down_in_lists_on_the_partition_key() { + assert_eq!( + select(&[ + col("p1").in_list(vec![lit(1), lit(2)], false), + col("p2").eq(lit("a")) + ]), + ( + vec![true, true], + r#""p1" IN (1, 2) AND "p2" = 'a'"#.to_string() + ) + ); + assert_eq!( + select(&[ + col("p1").in_list(vec![lit(1)], true), + col("p2").eq(lit("a")) + ]), + (vec![false, false], String::new()) + ); + } + + #[test] + fn pushes_down_disjunctions_on_a_partition_key_column_as_in_lists() { + assert_eq!( + select(&[ + col("p1") + .eq(lit(1)) + .or(lit(2).eq(col("p1"))) + .or(col("p1").in_list(vec![lit(3)], false)), + col("p2").eq(lit("a")), + ]), + ( + vec![true, true], + r#""p1" IN (1, 2, 3) AND "p2" = 'a'"#.to_string() + ) + ); + for filter in [ + col("p1").eq(lit(1)).or(col("p2").eq(lit("b"))), + col("p1").eq(lit(1)).or(col("p1").gt(lit(2))), + col("p1").eq(lit(1)).or(col("v").eq(lit(2))), + ] { + assert_eq!( + select(&[filter.clone(), col("p2").eq(lit("a"))]), + (vec![false, false], String::new()), + "{filter}" + ); + } + // only the partition key is restricted with `IN` + assert_eq!( + select(&[ + col("p1").eq(lit(1)), + col("p2").eq(lit("a")), + col("c1").eq(lit(1i64)).or(col("c1").eq(lit(2i64))), + ]), + ( + vec![true, true, false], + r#""p1" = 1 AND "p2" = 'a'"#.to_string() + ) + ); + } + + #[test] + fn keeps_partial_partition_keys() { + assert_eq!( + select(&[col("p1").eq(lit(1)), col("c1").eq(lit(1i64))]), + (vec![false, false], String::new()) + ); + } + + #[test] + fn keeps_unsupported_partition_key_filters() { + for filter in [ + col("p1").gt(lit(1)), + col("p1").not_eq(lit(1)), + col("p1").eq(lit(1i64)), // not the type of the column + col("p1").eq(lit(ScalarValue::Int32(None))), + col("p1").in_list(vec![lit(1), col("p1")], false), + ] { + assert_eq!( + select(&[filter.clone(), col("p2").eq(lit("a"))]), + (vec![false, false], String::new()), + "{filter}" + ); + } + assert_eq!( + select(&[col("p1").eq(lit(1)), col("p2").eq(lit(""))]), + (vec![false, false], String::new()) + ); + } + + #[test] + fn pushes_down_clustering_columns_in_order() { + assert_eq!( + select(&[ + col("p1").eq(lit(1)), + col("p2").eq(lit("a")), + col("c1").eq(lit(2i64)), + col("c2").eq(lit("50554d6e-29bb-11e5-b345-feff819cdc9f")), + ]), + ( + vec![true, true, true, true], + r#""p1" = 1 AND "p2" = 'a' AND "c1" = 2 AND "c2" = 50554d6e-29bb-11e5-b345-feff819cdc9f"# + .to_string() + ) + ); + // c2 can't be restricted without c1 + assert_eq!( + select(&[ + col("p1").eq(lit(1)), + col("p2").eq(lit("a")), + col("c2").eq(lit("50554d6e-29bb-11e5-b345-feff819cdc9f")), + ]), + ( + vec![true, true, false], + r#""p1" = 1 AND "p2" = 'a'"#.to_string() + ) + ); + } + + #[test] + fn pushes_down_one_range_on_a_clustering_column() { + assert_eq!( + select(&[ + col("p1").eq(lit(1)), + col("p2").eq(lit("a")), + col("c1").lt_eq(lit(9i64)), + col("c1").gt(lit(2i64)), + col("c1").gt_eq(lit(3i64)), + col("c2").eq(lit("50554d6e-29bb-11e5-b345-feff819cdc9f")), + ]), + ( + vec![true, true, true, true, false, false], + r#""p1" = 1 AND "p2" = 'a' AND "c1" > 2 AND "c1" <= 9"#.to_string() + ) + ); + assert_eq!( + select(&[ + col("p1").eq(lit(1)), + col("p2").eq(lit("a")), + lit(2i64).lt(col("c1")), + ]), + ( + vec![true, true, true], + r#""p1" = 1 AND "p2" = 'a' AND "c1" > 2"#.to_string() + ) + ); + } + + #[test] + fn keeps_duplicate_restrictions() { + assert_eq!( + select(&[ + col("p1").eq(lit(1)), + col("p1").eq(lit(2)), + col("p2").eq(lit("a")), + col("p2").in_list(vec![lit("a")], false), + ]), + ( + vec![true, false, true, false], + r#""p1" = 1 AND "p2" = 'a'"#.to_string() + ) + ); + } + + #[test] + fn keeps_ranges_and_non_canonical_values_of_uuids() { + let key = |c2| { + vec![ + col("p1").eq(lit(1)), + col("p2").eq(lit("a")), + col("c1").eq(lit(1i64)), + c2, + ] + }; + for c2 in [ + col("c2").gt(lit("50554d6e-29bb-11e5-b345-feff819cdc9f")), + col("c2").eq(lit("50554D6E-29BB-11E5-B345-FEFF819CDC9F")), + // a version 4 UUID + col("c2").eq(lit("123e4567-e89b-42d3-a456-426614174000")), + ] { + assert_eq!( + select(&key(c2.clone())).0, + vec![true, true, true, false], + "{c2}" + ); + } + } + + #[test] + fn pushes_down_nothing_without_a_primary_key() { + let filter = col("p1").eq(lit(1)); + assert_eq!( + select_pushed_down_filters(&PrimaryKey::default(), &[&filter]), + PushedDownFilters { + pushed_down: vec![false], + relations: vec![], + } + ); + } + + #[test] + fn formats_literals() { + let cases = [ + ("tinyint", ScalarValue::Int8(Some(-1)), "-1"), + ("smallint", ScalarValue::Int16(Some(2)), "2"), + ("boolean", ScalarValue::Boolean(Some(true)), "true"), + ("ascii", ScalarValue::Utf8(Some("x".into())), "'x'"), + ("varchar", ScalarValue::Utf8(Some("".into())), "''"), + ( + "blob", + ScalarValue::Binary(Some(vec![0xca, 0xfe])), + "0xcafe", + ), + ("date", ScalarValue::Date32(Some(0)), "2147483648"), + ("date", ScalarValue::Date32(Some(-1)), "2147483647"), + ( + "timestamp", + ScalarValue::TimestampMillisecond(Some(1_705_314_030_123), Some("UTC".into())), + "1705314030123", + ), + ("time", ScalarValue::Time64Nanosecond(Some(5)), "5"), + ( + "uuid", + ScalarValue::Utf8(Some("123e4567-e89b-12d3-a456-426614174000".into())), + "123e4567-e89b-12d3-a456-426614174000", + ), + ]; + for (cql_type, value, expected) in cases { + assert_eq!( + cql_literal(cql_type, &value, false).as_deref(), + Some(expected), + "{cql_type} {value}" + ); + } + + for (cql_type, value) in [ + ("ascii", ScalarValue::Utf8(Some("é".into()))), + ("double", ScalarValue::Float64(Some(1.0))), + ("decimal", ScalarValue::Utf8(Some("1.0".into()))), + ("inet", ScalarValue::Utf8(Some("127.0.0.1".into()))), + ("frozen>", ScalarValue::Int32(Some(1))), + ("int", ScalarValue::Utf8(Some("1".into()))), + ] { + assert_eq!( + cql_literal(cql_type, &value, false), + None, + "{cql_type} {value}" + ); + } + assert_eq!( + cql_literal("blob", &ScalarValue::Binary(Some(vec![])), true), + None + ); + } +} diff --git a/crates/cassandra/src/lib.rs b/crates/cassandra/src/lib.rs index fe83f7479..339191c47 100644 --- a/crates/cassandra/src/lib.rs +++ b/crates/cassandra/src/lib.rs @@ -2,8 +2,9 @@ //! [ScyllaDB Rust driver](https://docs.rs/scylla), so it works with ScyllaDB too. //! //! CQL is far more restrictive than SQL: it has no joins or aggregations, and `WHERE` clauses -//! are limited to key columns. Scans therefore only push the projection and `LIMIT` down to -//! Cassandra. Filters, aggregations and joins are evaluated by DataFusion. +//! are limited to key columns. Scans therefore push the projection, the filters on key columns +//! that CQL can evaluate (see [`filter`]) and `LIMIT` down to Cassandra. The other filters, +//! aggregations and joins are evaluated by DataFusion. //! //! Tables created with [`CassandraTableFactory::read_write_table_provider`] can also be written //! to with `INSERT INTO`; see [`CassandraTableWriter`]. @@ -11,6 +12,7 @@ pub mod conn; #[cfg(feature = "federation")] mod federation; +pub mod filter; pub mod pool; mod table; pub mod types; @@ -23,6 +25,7 @@ use datafusion::common::TableReference; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::AsyncDbConnection; +pub use crate::filter::{KeyColumn, PrimaryKey}; pub use crate::pool::CassandraConnectionPool; pub use crate::table::{CassandraExec, CassandraTable}; pub use crate::write::CassandraTableWriter; @@ -43,23 +46,20 @@ impl CassandraTableFactory { /// /// # Errors /// - /// Returns an error if the schema of the table cannot be read. + /// Returns an error if the schema or the primary key of the table cannot be read. pub async fn table_provider( &self, table_reference: TableReference, ) -> Result, Box> { - let schema = self - .pool - .connect_direct() - .get_schema(&table_reference) - .await?; + let conn = self.pool.connect_direct(); + let schema = conn.get_schema(&table_reference).await?; + let primary_key = conn.primary_key(&table_reference).await?; - let table_provider = Arc::new(CassandraTable::new( - table_reference, - Arc::clone(&self.pool), - schema, - )); + let table_provider = Arc::new( + CassandraTable::new(table_reference, Arc::clone(&self.pool), schema) + .with_primary_key(primary_key), + ); #[cfg(feature = "federation")] let table_provider = Arc::new(table_provider.create_federated_table_provider()?); @@ -73,7 +73,7 @@ impl CassandraTableFactory { /// /// # Errors /// - /// Returns an error if the schema of the table cannot be read. + /// Returns an error if the schema or the primary key of the table cannot be read. pub async fn read_write_table_provider( &self, table_reference: TableReference, diff --git a/crates/cassandra/src/table.rs b/crates/cassandra/src/table.rs index 7312b6d14..aeca92ef3 100644 --- a/crates/cassandra/src/table.rs +++ b/crates/cassandra/src/table.rs @@ -5,13 +5,16 @@ use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef; use datafusion::catalog::Session; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{project_schema, TableReference}; +use datafusion::common::{project_schema, DFSchema, TableReference}; use datafusion::datasource::TableType; use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; -use datafusion::logical_expr::Expr; +use datafusion::logical_expr::expr_rewriter::unnormalize_col; +use datafusion::logical_expr::utils::conjunction; +use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::filter::FilterExecBuilder; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, @@ -21,15 +24,19 @@ use futures::TryStreamExt; use datafusion_table_providers_common::sql::db_connection_pool::dbconnection::AsyncDbConnection; use crate::conn::{quote_identifier, quote_table_reference}; +use crate::filter::{select_pushed_down_filters, PrimaryKey}; use crate::pool::CassandraConnectionPool; /// A Cassandra table. /// -/// Filters are not pushed down (see the crate documentation), so DataFusion applies them. +/// The filters on key columns that CQL can evaluate are pushed down to Cassandra, see +/// [`crate::filter`]. DataFusion applies the others. +#[derive(Clone)] pub struct CassandraTable { pub(crate) table_reference: TableReference, pub(crate) pool: Arc, pub(crate) schema: SchemaRef, + pub(crate) primary_key: PrimaryKey, } impl fmt::Debug for CassandraTable { @@ -37,11 +44,14 @@ impl fmt::Debug for CassandraTable { f.debug_struct("CassandraTable") .field("table_reference", &self.table_reference) .field("schema", &self.schema) + .field("primary_key", &self.primary_key) .finish_non_exhaustive() } } impl CassandraTable { + /// Creates a table without a primary key, so no filters are pushed down to Cassandra. Use + /// [`Self::with_primary_key`] to push down the filters on key columns. #[must_use] pub fn new( table_reference: TableReference, @@ -52,8 +62,82 @@ impl CassandraTable { table_reference, pool, schema, + primary_key: PrimaryKey::default(), } } + + /// Sets the primary key of the table, which determines the filters pushed down to + /// Cassandra. + #[must_use] + pub fn with_primary_key(mut self, primary_key: PrimaryKey) -> Self { + self.primary_key = primary_key; + self + } + + /// Plans a scan of the table that returns the `projection` columns of the rows matching + /// all `filters`, at most `limit` of them. + /// + /// DataFusion only passes the filters for which [`supports_filters_pushdown`](datafusion::catalog::TableProvider::supports_filters_pushdown) + /// doesn't return `Unsupported` for, but it decides on them separately in each optimizer pass, + /// so together they may not make a valid CQL `WHERE` clause. The filters that aren't + /// pushed down to Cassandra are therefore applied by a `FilterExec` on top of the scan. + pub(crate) fn plan_scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> DataFusionResult> { + let filter_refs: Vec<&Expr> = filters.iter().collect(); + let pushed_down = select_pushed_down_filters(&self.primary_key, &filter_refs); + let remaining: Vec = filters + .iter() + .zip(&pushed_down.pushed_down) + .filter(|(_, pushed_down)| !**pushed_down) + .map(|(filter, _)| unnormalize_col(filter.clone())) + .collect(); + + let Some(predicate) = conjunction(remaining) else { + return Ok(Arc::new(CassandraExec::try_new( + self.table_reference.clone(), + Arc::clone(&self.pool), + &self.schema, + projection, + pushed_down.relations, + limit, + )?)); + }; + + // Scan the projected columns, then the other columns the predicate needs + let mut scan_projection: Vec = match projection { + Some(projection) => projection.clone(), + None => (0..self.schema.fields().len()).collect(), + }; + let output_projection: Vec = (0..scan_projection.len()).collect(); + for column in predicate.column_refs() { + let index = self.schema.index_of(&column.name)?; + if !scan_projection.contains(&index) { + scan_projection.push(index); + } + } + + let scan = Arc::new(CassandraExec::try_new( + self.table_reference.clone(), + Arc::clone(&self.pool), + &self.schema, + Some(&scan_projection), + pushed_down.relations, + None, + )?); + let scan_schema = DFSchema::try_from(scan.schema())?; + let predicate = state.create_physical_expr(predicate, &scan_schema)?; + Ok(Arc::new( + FilterExecBuilder::new(predicate, scan) + .apply_projection(Some(output_projection))? + .with_fetch(limit) + .build()?, + )) + } } #[async_trait] @@ -66,36 +150,52 @@ impl datafusion::catalog::TableProvider for CassandraTable { TableType::Base } + /// The filters on key columns that make up a valid CQL `WHERE` clause are `Exact`, the + /// others `Unsupported`, see [`crate::filter`]. + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> DataFusionResult> { + Ok(select_pushed_down_filters(&self.primary_key, filters) + .pushed_down + .into_iter() + .map(|pushed_down| { + if pushed_down { + TableProviderFilterPushDown::Exact + } else { + TableProviderFilterPushDown::Unsupported + } + }) + .collect()) + } + async fn scan( &self, - _state: &dyn Session, + state: &dyn Session, projection: Option<&Vec>, - _filters: &[Expr], + filters: &[Expr], limit: Option, ) -> DataFusionResult> { - Ok(Arc::new(CassandraExec::try_new( - self.table_reference.clone(), - Arc::clone(&self.pool), - &self.schema, - projection, - limit, - )?)) + self.plan_scan(state, projection, filters, limit) } } -/// Scans a Cassandra table with `SELECT FROM
[LIMIT n]`. +/// Scans a Cassandra table with `SELECT FROM
[WHERE ] [LIMIT n]`. pub struct CassandraExec { table_reference: TableReference, pool: Arc, /// The columns to select from Cassandra. columns: Vec, + /// The relations of the `WHERE` clause, joined with `AND`. + relations: Vec, projected_schema: SchemaRef, limit: Option, properties: Arc, } impl CassandraExec { - /// Creates a scan of the `projection` columns of `schema`. + /// Creates a scan of the `projection` columns of `schema`, of the rows matching the CQL + /// `relations`. /// /// # Errors /// @@ -105,6 +205,7 @@ impl CassandraExec { pool: Arc, schema: &SchemaRef, projection: Option<&Vec>, + relations: Vec, limit: Option, ) -> DataFusionResult { let projected_schema = project_schema(schema, projection)?; @@ -121,6 +222,7 @@ impl CassandraExec { table_reference, pool, columns, + relations, projected_schema, limit, properties, @@ -130,7 +232,12 @@ impl CassandraExec { /// The CQL statement this scan runs. #[must_use] pub fn cql(&self) -> String { - build_cql(&self.table_reference, &self.columns, self.limit) + build_cql( + &self.table_reference, + &self.columns, + &self.relations, + self.limit, + ) } } @@ -160,6 +267,7 @@ fn scan_columns( pub(crate) fn build_cql( table_reference: &TableReference, columns: &[String], + relations: &[String], limit: Option, ) -> String { let columns = columns @@ -171,6 +279,9 @@ pub(crate) fn build_cql( "SELECT {columns} FROM {}", quote_table_reference(table_reference) ); + if !relations.is_empty() { + cql.push_str(&format!(" WHERE {}", relations.join(" AND "))); + } // CQL requires a positive limit; `LIMIT 0` doesn't query Cassandra at all if let Some(limit) = limit.filter(|l| *l > 0) { cql.push_str(&format!(" LIMIT {limit}")); @@ -234,6 +345,7 @@ impl ExecutionPlan for CassandraExec { table_reference: self.table_reference.clone(), pool: Arc::clone(&self.pool), columns: self.columns.clone(), + relations: self.relations.clone(), projected_schema: Arc::clone(&self.projected_schema), limit, properties: Arc::clone(&self.properties), @@ -289,7 +401,7 @@ mod tests { let schema = schema(); let projected_schema = project_schema(&schema, projection).unwrap(); let columns = scan_columns(&table_reference, &schema, &projected_schema).unwrap(); - build_cql(&table_reference, &columns, limit) + build_cql(&table_reference, &columns, &[], limit) } #[test] @@ -324,6 +436,19 @@ mod tests { ); } + #[test] + fn adds_where_clause() { + assert_eq!( + build_cql( + &TableReference::bare("users"), + &["Name".to_string()], + &[r#""id" = 1"#.to_string(), r#""ts" > 2"#.to_string()], + Some(3), + ), + r#"SELECT "Name" FROM "users" WHERE "id" = 1 AND "ts" > 2 LIMIT 3"# + ); + } + #[test] fn fails_for_table_without_columns() { let schema = Arc::new(Schema::empty()); diff --git a/crates/cassandra/src/types.rs b/crates/cassandra/src/types.rs index aae36b9d2..0e63bf83a 100644 --- a/crates/cassandra/src/types.rs +++ b/crates/cassandra/src/types.rs @@ -61,7 +61,7 @@ pub enum Error { NullInCollection { column_name: String }, } -pub type Result = std::result::Result; +type Result = std::result::Result; /// Number of days between -5877641-06-23 (the CQL `date` epoch) and the Unix epoch. const CQL_DATE_UNIX_EPOCH_OFFSET: i64 = 1 << 31; From e7d06551f4366a88a07829e299fd72a13d3ced21 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Mon, 19 Oct 2026 11:51:44 +0300 Subject: [PATCH 21/27] Use constructors for PrimaryKey and KeyColumn --- crates/cassandra/src/conn.rs | 8 ++++---- crates/cassandra/src/filter.rs | 17 +++++++++++++---- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/crates/cassandra/src/conn.rs b/crates/cassandra/src/conn.rs index c998f5cd6..e3e77335d 100644 --- a/crates/cassandra/src/conn.rs +++ b/crates/cassandra/src/conn.rs @@ -115,10 +115,10 @@ impl CassandraConnection { .map(|(name, _, _, cql_type)| KeyColumn::new(name, cql_type)) .collect() }; - Ok(PrimaryKey { - partition_key: key_columns("partition_key"), - clustering_key: key_columns("clustering"), - }) + Ok(PrimaryKey::new( + key_columns("partition_key"), + key_columns("clustering"), + )) } } diff --git a/crates/cassandra/src/filter.rs b/crates/cassandra/src/filter.rs index c95999c40..36c202b59 100644 --- a/crates/cassandra/src/filter.rs +++ b/crates/cassandra/src/filter.rs @@ -23,17 +23,26 @@ use crate::conn::quote_identifier; #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct PrimaryKey { /// The partition key columns, in order. - pub partition_key: Vec, + partition_key: Vec, /// The clustering columns, in order. - pub clustering_key: Vec, + clustering_key: Vec, +} + +impl PrimaryKey { + pub fn new(partition_key: Vec, clustering_key: Vec) -> Self { + PrimaryKey { + partition_key, + clustering_key, + } + } } /// A primary key column. #[derive(Debug, Clone, PartialEq, Eq)] pub struct KeyColumn { - pub name: String, + name: String, /// The CQL type, as in `system_schema.columns`, e.g. `int` or `frozen>`. - pub cql_type: String, + cql_type: String, } impl KeyColumn { From 1385b2ce2e6e378b4ebe66ab2014c51ec9cb70c0 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Mon, 19 Oct 2026 12:03:51 +0300 Subject: [PATCH 22/27] Add support for ScyllaDB - IT tests and CI workflow --- .github/workflows/pr.yaml | 38 +++++++++++ core/tests/cassandra/common.rs | 115 ++++++++++++++++++++++----------- core/tests/cassandra/mod.rs | 3 +- core/tests/docker/mod.rs | 11 ++++ 4 files changed, 128 insertions(+), 39 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index cf3aa8227..55d6441f1 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -250,6 +250,44 @@ jobs: - name: Run Oracle integration tests run: cargo test -p datafusion-table-providers --test integration --no-default-features --features oracle -- --test-threads=1 --nocapture + scylla-integration-test: + name: ScyllaDB Integration Tests + runs-on: ubuntu-latest + timeout-minutes: 25 + + env: + # The Cassandra provider is tested against ScyllaDB with the Cassandra integration tests + CASSANDRA_DOCKER_IMAGE: scylladb/scylla:2026.3.3 + RUST_LOG: info + + steps: + - uses: actions/checkout@v7 + + - uses: dtolnay/rust-toolchain@stable + + - name: Cache Rust dependencies + uses: Swatinem/rust-cache@v2 + + - name: Cache ScyllaDB Docker image + id: scylla-docker-cache + uses: actions/cache@v6 + with: + path: /tmp/scylla-image + key: scylla-image-${{ env.CASSANDRA_DOCKER_IMAGE }} + + - name: Load or pull ScyllaDB image + run: | + mkdir -p /tmp/scylla-image + if [ -f /tmp/scylla-image/scylla.tar ]; then + docker load -i /tmp/scylla-image/scylla.tar + else + docker pull "${CASSANDRA_DOCKER_IMAGE}" + docker save -o /tmp/scylla-image/scylla.tar "${CASSANDRA_DOCKER_IMAGE}" + fi + + - name: Run ScyllaDB integration tests + run: cargo test -p datafusion-table-providers --test integration --no-default-features --features cassandra -- --nocapture cassandra + python-integration-test: name: Python Tests runs-on: ubuntu-latest diff --git a/core/tests/cassandra/common.rs b/core/tests/cassandra/common.rs index e3402abe0..6d70c3a55 100644 --- a/core/tests/cassandra/common.rs +++ b/core/tests/cassandra/common.rs @@ -22,48 +22,87 @@ pub(super) async fn start_cassandra_docker_container( let cassandra_docker_image = std::env::var("CASSANDRA_DOCKER_IMAGE") .unwrap_or_else(|_| format!("{}cassandra:5.0", container_registry())); - let running_container = - ContainerRunnerBuilder::new(format!("{CASSANDRA_DOCKER_CONTAINER}-{port}")) - .image(cassandra_docker_image) - .add_port_binding(9042, port) - // The driver connects to the address each node advertises in `system.local`, not to - // the contact point. By default that is the container's internal IP, which, combined - // with the published host port, isn't reachable from the host. - .add_env_var("CASSANDRA_BROADCAST_RPC_ADDRESS", "127.0.0.1") - // Single-node cluster: the node is its own seed (the image's default), so there is - // no one to gossip with. Skip the gossip-settle polling and the 30s ring delay. - .add_env_var( - "JVM_EXTRA_OPTS", - "-Dcassandra.skip_wait_for_gossip_to_settle=0 -Dcassandra.ring_delay_ms=0", - ) - .add_env_var( - "CASSANDRA_NUM_TOKENS", - &std::env::var("CASSANDRA_NUM_TOKENS").unwrap_or("1".to_string()), - ) - .add_env_var("CASSANDRA_CLUSTER_NAME", "df-t-p-cassandra") - // Keep the JVM small so that the node starts quickly in CI - .add_env_var("MAX_HEAP_SIZE", "512M") - .add_env_var("HEAP_NEWSIZE", "128M") - .healthcheck(HealthConfig { - test: Some(vec![ - "CMD-SHELL".to_string(), - "cqlsh -e 'DESCRIBE KEYSPACES' || exit 1".to_string(), - ]), - interval: Some(2_000_000_000), // 2s - timeout: Some(10_000_000_000), // 10s - retries: Some(90), - start_period: Some(20_000_000_000), // 20s - start_interval: None, - }) - // Cassandra takes much longer to boot than the other databases - .health_timeout(std::time::Duration::from_secs(300)) - .build()? - .run() - .await?; + // ScyllaDB is compatible with Cassandra, but its image is configured differently + let is_scylla = cassandra_docker_image.contains("scylla"); + + let builder = ContainerRunnerBuilder::new(format!("{CASSANDRA_DOCKER_CONTAINER}-{port}")) + .image(cassandra_docker_image) + .add_port_binding(9042, port); + let builder = if is_scylla { + configure_scylla(builder) + } else { + configure_cassandra(builder) + }; + + let running_container = builder + .healthcheck(HealthConfig { + test: Some(vec![ + "CMD-SHELL".to_string(), + "cqlsh -e 'DESCRIBE KEYSPACES' || exit 1".to_string(), + ]), + interval: Some(2_000_000_000), // 2s + timeout: Some(10_000_000_000), // 10s + retries: Some(90), + start_period: Some(20_000_000_000), // 20s + start_interval: None, + }) + // Cassandra takes much longer to boot than the other databases + .health_timeout(std::time::Duration::from_secs(300)) + .build()? + .run() + .await?; Ok(running_container) } +/// The `cassandra` image is configured with environment variables. +fn configure_cassandra(builder: ContainerRunnerBuilder<'_>) -> ContainerRunnerBuilder<'_> { + builder + // The driver connects to the address each node advertises in `system.local`, not to + // the contact point. By default that is the container's internal IP, which, combined + // with the published host port, isn't reachable from the host. + .add_env_var("CASSANDRA_BROADCAST_RPC_ADDRESS", "127.0.0.1") + // Single-node cluster: the node is its own seed (the image's default), so there is + // no one to gossip with. Skip the gossip-settle polling and the 30s ring delay. + .add_env_var( + "JVM_EXTRA_OPTS", + "-Dcassandra.skip_wait_for_gossip_to_settle=0 -Dcassandra.ring_delay_ms=0", + ) + .add_env_var( + "CASSANDRA_NUM_TOKENS", + &std::env::var("CASSANDRA_NUM_TOKENS").unwrap_or("1".to_string()), + ) + .add_env_var("CASSANDRA_CLUSTER_NAME", "df-t-p-cassandra") + // Keep the JVM small so that the node starts quickly in CI + .add_env_var("MAX_HEAP_SIZE", "512M") + .add_env_var("HEAP_NEWSIZE", "128M") +} + +/// The `scylladb/scylla` image ignores environment variables: it is configured with the +/// arguments of its entrypoint. +fn configure_scylla(builder: ContainerRunnerBuilder<'_>) -> ContainerRunnerBuilder<'_> { + builder.cmd( + [ + // See `configure_cassandra` + "--broadcast-rpc-address", + "127.0.0.1", + // By default Scylla uses all CPUs and most of the memory. With a single shard, the + // driver doesn't need the (unpublished) shard-aware port either. + "--smp", + "1", + "--memory", + "1G", + "--overprovisioned", + "1", + // Skips the checks of the host's I/O and kernel settings, which fail in Docker + "--developer-mode", + "1", + ] + .map(String::from) + .to_vec(), + ) +} + #[instrument] pub(super) async fn get_cassandra_connection_pool( port: u16, diff --git a/core/tests/cassandra/mod.rs b/core/tests/cassandra/mod.rs index b167c22ad..ca587dfbd 100644 --- a/core/tests/cassandra/mod.rs +++ b/core/tests/cassandra/mod.rs @@ -105,10 +105,11 @@ async fn execute(pool: &CassandraConnectionPool, cql: &str) { } async fn create_test_data(pool: &CassandraConnectionPool) { + // ScyllaDB creates keyspaces with tablets, which don't support `SimpleStrategy` execute( pool, &format!( - "CREATE KEYSPACE IF NOT EXISTS {KEYSPACE} WITH replication = {{'class': 'SimpleStrategy', 'replication_factor': 1}}" + "CREATE KEYSPACE IF NOT EXISTS {KEYSPACE} WITH replication = {{'class': 'NetworkTopologyStrategy', 'replication_factor': 1}}" ), ) .await; diff --git a/core/tests/docker/mod.rs b/core/tests/docker/mod.rs index 301ad04ce..0e831b185 100644 --- a/core/tests/docker/mod.rs +++ b/core/tests/docker/mod.rs @@ -58,6 +58,7 @@ pub struct ContainerRunnerBuilder<'a> { image: Option, port_bindings: Vec<(u16, u16)>, env_vars: Vec<(String, String)>, + cmd: Option>, healthcheck: Option, health_timeout: Option, } @@ -69,6 +70,7 @@ impl<'a> ContainerRunnerBuilder<'a> { image: None, port_bindings: Vec::new(), env_vars: Vec::new(), + cmd: None, healthcheck: None, health_timeout: None, } @@ -90,6 +92,12 @@ impl<'a> ContainerRunnerBuilder<'a> { self } + /// Overrides the image's default command, i.e. the arguments passed to its entrypoint. + pub fn cmd(mut self, cmd: Vec) -> Self { + self.cmd = Some(cmd); + self + } + pub fn healthcheck(mut self, healthcheck: HealthConfig) -> Self { self.healthcheck = Some(healthcheck); self @@ -112,6 +120,7 @@ impl<'a> ContainerRunnerBuilder<'a> { image, port_bindings: self.port_bindings, env_vars: self.env_vars, + cmd: self.cmd, healthcheck: self.healthcheck, health_timeout: self.health_timeout, }) @@ -124,6 +133,7 @@ pub struct ContainerRunner<'a> { image: String, port_bindings: Vec<(u16, u16)>, env_vars: Vec<(String, String)>, + cmd: Option>, healthcheck: Option, health_timeout: Option, } @@ -179,6 +189,7 @@ impl ContainerRunner<'_> { let config = ContainerCreateBody { image: Some(self.image.clone()), env: Some(env_vars), + cmd: self.cmd, exposed_ports, host_config, healthcheck: self.healthcheck, From 9db593441054c20a65cfe7fc2628aa1fddad1f86 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Mon, 19 Oct 2026 12:22:30 +0300 Subject: [PATCH 23/27] Update the README and the Cassandra example to use NetworkTopologyStrategy This way they are applicable for Scylla too --- README.md | 2 +- core/examples/cassandra.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 96b71f9c8..34305fd00 100644 --- a/README.md +++ b/README.md @@ -211,7 +211,7 @@ done # Create a table in the Cassandra server and insert some data docker exec -i cassandra cqlsh < Date: Mon, 19 Oct 2026 15:55:44 +0300 Subject: [PATCH 24/27] Run unit & IT tests for all features --- .github/workflows/pr.yaml | 4 ++-- Makefile | 6 ++---- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 55d6441f1..3e14e0744 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -115,7 +115,7 @@ jobs: - name: Install ODBC & Sqlite run: | - sudo apt-get install -y unixodbc-dev libsqlite3-dev + sudo apt-get install -y unixodbc-dev libsqliteodbc libsqlite3-dev - name: Run unit tests run: make test @@ -184,7 +184,7 @@ jobs: - name: Install ODBC & Sqlite run: | sudo apt-get update - sudo apt-get install -y unixodbc-dev libsqlite3-dev + sudo apt-get install -y unixodbc-dev libsqliteodbc libsqlite3-dev - name: Install ADBC drivers run: | diff --git a/Makefile b/Makefile index 75d4d58ca..9f2a210e2 100644 --- a/Makefile +++ b/Makefile @@ -3,9 +3,7 @@ all: .PHONY: test test: - cargo test --features clickhouse,duckdb,flight,mysql,postgres,sqlite,adbc -p datafusion-table-providers --lib - cargo test -p datafusion-table-providers-oracle - cargo test -p datafusion-table-providers-cassandra + RUST_LOG=$${RUST_LOG:-info} cargo test --all-features --lib -- --nocapture .PHONY: lint lint: @@ -13,4 +11,4 @@ lint: .PHONY: test-integration test-integration: - RUST_LOG=$${RUST_LOG:-info} cargo test -p datafusion-table-providers --test integration --no-default-features --features postgres,sqlite,mysql,flight,clickhouse,duckdb,adbc,cassandra -- --nocapture + RUST_LOG=$${RUST_LOG:-info} cargo test -p datafusion-table-providers --test integration --all-features -- --nocapture From 4801d2003adcdbcd6e60b493d2f4f6189ec35f33 Mon Sep 17 00:00:00 2001 From: Martin Grigorov Date: Mon, 19 Oct 2026 17:57:15 +0300 Subject: [PATCH 25/27] Update ODBC & Sqlite installation in workflow Removed 'libsqliteodbc' from ODBC & Sqlite installation commands. Signed-off-by: Martin Grigorov --- .github/workflows/pr.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 3e14e0744..55d6441f1 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -115,7 +115,7 @@ jobs: - name: Install ODBC & Sqlite run: | - sudo apt-get install -y unixodbc-dev libsqliteodbc libsqlite3-dev + sudo apt-get install -y unixodbc-dev libsqlite3-dev - name: Run unit tests run: make test @@ -184,7 +184,7 @@ jobs: - name: Install ODBC & Sqlite run: | sudo apt-get update - sudo apt-get install -y unixodbc-dev libsqliteodbc libsqlite3-dev + sudo apt-get install -y unixodbc-dev libsqlite3-dev - name: Install ADBC drivers run: | From 1d05667e1e717040a730285420738e950f9c78a1 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Mon, 19 Oct 2026 18:54:39 +0300 Subject: [PATCH 26/27] Install ADBC drivers for the unittests run --- .github/workflows/pr.yaml | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 55d6441f1..d24292f96 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -115,7 +115,12 @@ jobs: - name: Install ODBC & Sqlite run: | - sudo apt-get install -y unixodbc-dev libsqlite3-dev + sudo apt-get install -y unixodbc-dev libsqliteodbc libsqlite3-dev + curl -LsSf https://dbc.columnar.tech/install.sh | sh + echo "$HOME/.local/bin" >> "$GITHUB_PATH" + export PATH="$HOME/.local/bin:$PATH" + dbc install sqlite + dbc install duckdb - name: Run unit tests run: make test @@ -184,7 +189,7 @@ jobs: - name: Install ODBC & Sqlite run: | sudo apt-get update - sudo apt-get install -y unixodbc-dev libsqlite3-dev + sudo apt-get install -y unixodbc-dev libsqliteodbc libsqlite3-dev - name: Install ADBC drivers run: | From 05cfb0ea430211abdcd2259d845e36f0d7988241 Mon Sep 17 00:00:00 2001 From: Martin Tzvetanov Grigorov Date: Tue, 20 Oct 2026 08:31:44 +0300 Subject: [PATCH 27/27] Run Oracle iT tests with the rest IT test --- .github/workflows/pr.yaml | 45 +++++---------------------------------- 1 file changed, 5 insertions(+), 40 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index d24292f96..d8e30dda9 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -136,6 +136,7 @@ jobs: MYSQL_DOCKER_IMAGE: mysql:8.0 CLICKHOUSE_DOCKER_IMAGE: clickhouse/clickhouse-server:24.8 CASSANDRA_DOCKER_IMAGE: cassandra:5.0 + ORACLE_DOCKER_IMAGE: docker.io/gvenzl/oracle-free:latest # Keep integration test logs concise for faster CI runs. RUST_LOG: info # Serialize Docker-backed tests a bit to avoid stampeding many Postgres containers. @@ -163,7 +164,7 @@ jobs: uses: actions/cache@v6 with: path: /tmp/docker-images - key: docker-images-${{ env.PG_DOCKER_IMAGE }}-${{ env.MYSQL_DOCKER_IMAGE }}-${{ env.CLICKHOUSE_DOCKER_IMAGE }}-${{ env.CASSANDRA_DOCKER_IMAGE }} + key: docker-images-${{ env.PG_DOCKER_IMAGE }}-${{ env.MYSQL_DOCKER_IMAGE }}-${{ env.CLICKHOUSE_DOCKER_IMAGE }}-${{ env.CASSANDRA_DOCKER_IMAGE }}-${{ env.ORACLE_DOCKER_IMAGE }} - name: Load or pull Docker images run: | @@ -185,6 +186,7 @@ jobs: load_or_pull mysql "${{ env.MYSQL_DOCKER_IMAGE }}" load_or_pull clickhouse "${{ env.CLICKHOUSE_DOCKER_IMAGE }}" load_or_pull cassandra "${{ env.CASSANDRA_DOCKER_IMAGE }}" + load_or_pull oracle "${{ env.ORACLE_DOCKER_IMAGE }}" - name: Install ODBC & Sqlite run: | @@ -200,26 +202,6 @@ jobs: dbc install duckdb ls -la "${XDG_CONFIG_HOME:-$HOME/.config}/adbc/drivers" || true - - name: Run integration tests - run: make test-integration - - oracle-integration-test: - name: Oracle Integration Tests - runs-on: ubuntu-latest - timeout-minutes: 45 - - env: - ORACLE_DOCKER_IMAGE: docker.io/gvenzl/oracle-free:latest - RUST_LOG: info - - steps: - - uses: actions/checkout@v7 - - - uses: dtolnay/rust-toolchain@stable - - - name: Cache Rust dependencies - uses: Swatinem/rust-cache@v2 - - name: Install Oracle Instant Client run: | # `rust-oracle` uses ODPI-C and dlopens the Instant Client at runtime. @@ -235,25 +217,8 @@ jobs: fi echo "LD_LIBRARY_PATH=$IC_DIR" >> "$GITHUB_ENV" - - name: Cache Oracle Docker image - id: oracle-docker-cache - uses: actions/cache@v6 - with: - path: /tmp/oracle-image - key: oracle-image-${{ env.ORACLE_DOCKER_IMAGE }} - - - name: Load or pull Oracle image - run: | - mkdir -p /tmp/oracle-image - if [ -f /tmp/oracle-image/oracle.tar ]; then - docker load -i /tmp/oracle-image/oracle.tar - else - docker pull "${ORACLE_DOCKER_IMAGE}" - docker save -o /tmp/oracle-image/oracle.tar "${ORACLE_DOCKER_IMAGE}" - fi - - - name: Run Oracle integration tests - run: cargo test -p datafusion-table-providers --test integration --no-default-features --features oracle -- --test-threads=1 --nocapture + - name: Run integration tests + run: make test-integration scylla-integration-test: name: ScyllaDB Integration Tests