feat: Implement Python ODBC, MySQL, Postgres, Flight table providers (#281)

* Implement ODBC python table provider

* Update ODBC documentation

* Update API

* Remove unused imports

* Update comments

* Update comments

* Fix a typo

* Implement mysql and postgresql, and refactor tokio runtime

* Enable flight

* Fix styling and documentation

* Update core/src/sql/db_connection_pool/runtime.rs

Co-authored-by: Phillip LeBlanc <[email protected]>

* Update python/python/datafusion_table_providers/odbc.py

Co-authored-by: Phillip LeBlanc <[email protected]>

* Fix formatting and address comments

* Update core/src/sql/db_connection_pool/runtime.rs

---------

Co-authored-by: Phillip LeBlanc <[email protected]>

Author: crystalxyz

Cargo.lock

@@ -27,7 +27,7 @@ let ctx = SessionContext::with_state(state);
 - Flight SQL
 - ODBC
 
-## Examples
+## Examples (in Rust)
 
 Run the included examples to see how to use the table providers:
 
@@ -45,6 +45,7 @@ cargo run --example duckdb_function --features duckdb
 ### SQLite
 
 ```bash
+# Run from repo folder
 cargo run --example sqlite --features sqlite
 ```
 
@@ -80,7 +81,9 @@ EOF
 ```
 
 ```bash
-cargo run --example postgres --features postgres
+# Run from repo folder
+cargo run -p datafusion-table-providers --example postgres --features postgres
+
 ```
 
 ### MySQL
@@ -104,7 +107,8 @@ EOF
 ```
 
 ```bash
-cargo run --example mysql --features mysql
+# Run from repo folder
+cargo run -p datafusion-table-providers --example mysql --features mysql
 ```
 
 ### Flight SQL
@@ -115,16 +119,37 @@ brew install roapi
 # cargo install --locked --git https://github.com/roapi/roapi --branch main --bins roapi
 roapi -t taxi=https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2024-01.parquet &
 
-cargo run --example flight-sql --features flight
+# Run from repo folder
+cargo run -p datafusion-table-providers --example flight-sql --features flight
 ```
 
 ### ODBC
-
 ```bash
 apt-get install unixodbc-dev libsqliteodbc
 # or
 # brew install unixodbc & brew install sqliteodbc
-# If you use ARM Mac, please see https://github.com/pacman82/odbc-api#os-x-arm--mac-m1
 
 cargo run --example odbc_sqlite --features odbc
 ```
+
+#### ARM Mac
+
+Please see https://github.com/pacman82/odbc-api#os-x-arm--mac-m1 for reference.
+
+Steps:
+1. Install unixodbc and sqliteodbc by `brew install unixodbc sqliteodbc`.
+2. Find local sqliteodbc driver path by running `brew info sqliteodbc`. The path might look like `/opt/homebrew/Cellar/sqliteodbc/0.99991`.
+3. Set up odbc config file at `~/.odbcinst.ini` with your local sqliteodbc path.
+Example config file:
+```
+[SQLite3]
+Description = SQLite3 ODBC Driver
+Driver      = /opt/homebrew/Cellar/sqliteodbc/0.99991/lib/libsqlite3odbc.dylib
+```
+4. Test configuration by running `odbcinst -q -d -n SQLite3`. If the path is printed out correctly, then you are all set.
+
+## Examples (in Python)
+1. Start a Python venv
+2. Enter into venv
+3. Inside python/ folder, run `maturin develop`.
+4. Inside python/examples/ folder, run the corresponding test using `python3 [file_name]`.

README.md

@@ -19,7 +19,7 @@ async fn main() {
     // Opening in ReadOnly mode allows multiple reader processes to access
     // the database at the same time
     let duckdb_pool = Arc::new(
-        DuckDbConnectionPool::new_file("examples/duckdb_example.db", &AccessMode::ReadOnly)
+        DuckDbConnectionPool::new_file("core/examples/duckdb_example.db", &AccessMode::ReadOnly)
             .expect("unable to create DuckDB connection pool"),
     );
 

core/examples/duckdb.rs

@@ -17,7 +17,7 @@ async fn main() {
     // Create SQLite ODBC connection pool
     let params = to_secret_map(HashMap::from([(
         "connection_string".to_owned(),
-        "driver=SQLite3;database=examples/sqlite_example.db;".to_owned(),
+        "driver=SQLite3;database=core/examples/sqlite_example.db;".to_owned(),
     )]));
     let odbc_pool =
         Arc::new(ODBCPool::new(params).expect("unable to create SQLite ODBC connection pool"));

core/examples/odbc_sqlite.rs

@@ -20,7 +20,7 @@ async fn main() {
     // - arg3: Connection timeout duration
     let sqlite_pool = Arc::new(
         SqliteConnectionPoolFactory::new(
-            "examples/sqlite_example.db",
+            "core/examples/sqlite_example.db",
             Mode::File,
             Duration::from_millis(5000),
         )

core/examples/sqlite.rs

@@ -24,6 +24,7 @@ use std::str::FromStr;
 use std::sync::Arc;
 
 use crate::flight::{flight_channel, to_df_err, FlightMetadata, FlightProperties, SizeLimits};
+use crate::sql::db_connection_pool::runtime::run_async_with_tokio;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::{FlightClient, FlightEndpoint, Ticket};
@@ -190,7 +191,8 @@ async fn flight_stream(
 ) -> Result<SendableRecordBatchStream> {
     let mut errors: Vec<Box<dyn Error + Send + Sync>> = vec![];
     for loc in partition.locations.iter() {
-        let client = flight_client(loc, grpc_headers.as_ref(), &size_limits).await?;
+        let get_client = || async { flight_client(loc, grpc_headers.as_ref(), &size_limits).await };
+        let client = run_async_with_tokio(get_client).await?;
         match try_fetch_stream(client, &partition.ticket, schema.clone()).await {
             Ok(stream) => return Ok(stream),
             Err(e) => errors.push(Box::new(e)),

core/src/flight/exec.rs

@@ -1,7 +1,6 @@
 use std::any::Any;
-use std::collections::HashSet;
 use std::sync::Arc;
-use std::sync::OnceLock;
+use std::collections::HashSet;
 
 use arrow::array::RecordBatch;
 use arrow_schema::{DataType, Field};
@@ -20,9 +19,9 @@ use duckdb::{Connection, DuckdbConnectionManager};
 use dyn_clone::DynClone;
 use rand::distr::{Alphanumeric, SampleString};
 use snafu::{prelude::*, ResultExt};
-use tokio::runtime::{Handle, Runtime};
 use tokio::sync::mpsc::Sender;
 
+use crate::sql::db_connection_pool::runtime::run_sync_with_tokio;
 use crate::util::schema::SchemaValidator;
 use crate::UnsupportedTypeAction;
 
@@ -283,13 +282,6 @@ impl DbConnection<r2d2::PooledConnection<DuckdbConnectionManager>, DuckDBParamet
     }
 }
 
-fn get_tokio_runtime() -> &'static Runtime {
-    // TODO: this function is a repetition of python/src/utils.rs::get_tokio_runtime.
-    // Think about how to refactor it
-    static RUNTIME: OnceLock<Runtime> = OnceLock::new();
-    RUNTIME.get_or_init(|| Runtime::new().expect("Failed to create Tokio runtime"))
-}
-
 impl SyncDbConnection<r2d2::PooledConnection<DuckdbConnectionManager>, DuckDBParameter>
     for DuckDbConnection
 {
@@ -449,13 +441,7 @@ impl SyncDbConnection<r2d2::PooledConnection<DuckdbConnectionManager>, DuckDBPar
             )))
         };
 
-        // If calling directly from Rust, there is already tokio runtime so no
-        // additional work is needed. If calling from Python FFI, there's no existing
-        // tokio runtime, so we need to start a new one.
-        match Handle::try_current() {
-            Ok(_) => create_stream(),
-            Err(_) => get_tokio_runtime().block_on(async { create_stream() }),
-        }
+        run_sync_with_tokio(create_stream)
     }
 
     fn execute(&self, sql: &str, params: &[DuckDBParameter]) -> Result<u64> {

core/src/sql/db_connection_pool/dbconnection/duckdbconn.rs

@@ -20,6 +20,7 @@ use std::sync::Arc;
 
 use crate::sql::db_connection_pool::{
     dbconnection::{self, AsyncDbConnection, DbConnection, GenericError},
+    runtime::run_async_with_tokio,
     DbConnectionPool,
 };
 use arrow_odbc::arrow_schema_from;
@@ -42,7 +43,7 @@ use odbc_api::handles::StatementImpl;
 use odbc_api::parameter::InputParameter;
 use odbc_api::Cursor;
 use odbc_api::CursorImpl;
-use secrecy::{SecretBox, ExposeSecret, SecretString};
+use secrecy::{ExposeSecret, SecretBox, SecretString};
 use snafu::prelude::*;
 use snafu::Snafu;
 use tokio::runtime::Handle;
@@ -184,69 +185,71 @@ where
         let params = params.iter().map(dyn_clone::clone).collect::<Vec<_>>();
         let secrets = Arc::clone(&self.params);
 
-        let join_handle = tokio::task::spawn_blocking(move || {
-            let handle = Handle::current();
-            let cxn = handle.block_on(async { conn.lock().await });
+        let create_stream = async || -> Result<SendableRecordBatchStream> {
+            let join_handle = tokio::task::spawn_blocking(move || {
+                let handle = Handle::current();
+                let cxn = handle.block_on(async { conn.lock().await });
 
-            let mut prepared = cxn.prepare(&sql)?;
-            let schema = Arc::new(arrow_schema_from(&mut prepared, false)?);
-            blocking_channel_send(&schema_tx, Arc::clone(&schema))?;
+                let mut prepared = cxn.prepare(&sql)?;
+                let schema = Arc::new(arrow_schema_from(&mut prepared, false)?);
+                blocking_channel_send(&schema_tx, Arc::clone(&schema))?;
 
-            let mut statement = prepared.into_statement();
+                let mut statement = prepared.into_statement();
 
-            bind_parameters(&mut statement, &params)?;
+                bind_parameters(&mut statement, &params)?;
 
-            // StatementImpl<'_>::execute is unsafe, CursorImpl<_>::new is unsafe
-            let cursor = unsafe {
-                if let SqlResult::Error { function } = statement.execute() {
-                    return Err(Error::ODBCAPIErrorNoSource {
-                        message: function.to_string(),
+                // StatementImpl<'_>::execute is unsafe, CursorImpl<_>::new is unsafe
+                let cursor = unsafe {
+                    if let SqlResult::Error { function } = statement.execute() {
+                        return Err(Error::ODBCAPIErrorNoSource {
+                            message: function.to_string(),
+                        }
+                        .into());
                     }
-                    .into());
-                }
 
-                Ok::<_, GenericError>(CursorImpl::new(statement.as_stmt_ref()))
-            }?;
+                    Ok::<_, GenericError>(CursorImpl::new(statement.as_stmt_ref()))
+                }?;
 
-            let reader = build_odbc_reader(cursor, &schema, &secrets)?;
-            for batch in reader {
-                blocking_channel_send(&batch_tx, batch.context(ArrowSnafu)?)?;
-            }
+                let reader = build_odbc_reader(cursor, &schema, &secrets)?;
+                for batch in reader {
+                    blocking_channel_send(&batch_tx, batch.context(ArrowSnafu)?)?;
+                }
 
-            Ok::<_, GenericError>(())
-        });
+                Ok::<_, GenericError>(())
+            });
 
-        // we need to wait for the schema first before we can build our RecordBatchStreamAdapter
-        let Some(schema) = schema_rx.recv().await else {
-            // if the channel drops, the task errored
-            if !join_handle.is_finished() {
-                unreachable!("Schema channel should not have dropped before the task finished");
-            }
+            // we need to wait for the schema first before we can build our RecordBatchStreamAdapter
+            let Some(schema) = schema_rx.recv().await else {
+                // if the channel drops, the task errored
+                if !join_handle.is_finished() {
+                    unreachable!("Schema channel should not have dropped before the task finished");
+                }
 
-            let result = join_handle.await?;
-            let Err(err) = result else {
-                unreachable!("Task should have errored");
+                let result = join_handle.await?;
+                let Err(err) = result else {
+                    unreachable!("Task should have errored");
+                };
+
+                return Err(err);
             };
 
-            return Err(err);
-        };
+            let output_stream = stream! {
+                while let Some(batch) = batch_rx.recv().await {
+                    yield Ok(batch);
+                }
 
-        let output_stream = stream! {
-            while let Some(batch) = batch_rx.recv().await {
-                yield Ok(batch);
-            }
+                if let Err(e) = join_handle.await {
+                    yield Err(DataFusionError::Execution(format!(
+                        "Failed to execute ODBC query: {e}"
+                    )))
+                }
+            };
 
-            if let Err(e) = join_handle.await {
-                yield Err(DataFusionError::Execution(format!(
-                    "Failed to execute ODBC query: {e}"
-                )))
-            }
+            let result: SendableRecordBatchStream =
+                Box::pin(RecordBatchStreamAdapter::new(schema, output_stream));
+            Ok(result)
         };
-
-        Ok(Box::pin(RecordBatchStreamAdapter::new(
-            schema,
-            output_stream,
-        )))
+        run_async_with_tokio(create_stream).await
     }
 
     async fn execute(&self, query: &str, params: &[ODBCParameter]) -> Result<u64> {

core/src/sql/db_connection_pool/dbconnection/odbcconn.rs

@@ -360,7 +360,6 @@ mod test {
 
     use super::*;
     use crate::sql::db_connection_pool::DbConnectionPool;
-    use std::sync::Arc;
 
     fn random_db_name() -> String {
         let mut rng = rand::rng();

core/src/sql/db_connection_pool/duckdbpool.rs

@@ -11,6 +11,7 @@ pub mod mysqlpool;
 pub mod odbcpool;
 #[cfg(feature = "postgres")]
 pub mod postgrespool;
+pub mod runtime;
 #[cfg(feature = "sqlite")]
 pub mod sqlitepool;