feat: preliminary support odbc (datafusion-contrib#177)

* feat: preliminary support odbc

* feat: add sqlite odbc example

* chore: clippy lint

---------

Co-authored-by: Phillip LeBlanc <phillip@leblanc.tech>

Author: xhwhis

Cargo.toml

@@ -62,6 +62,7 @@ serde_json = "1.0.124"
 snafu = "0.8.3"
 time = "0.3.36"
 tokio = { version = "1.38.0", features = ["macros", "fs"] }
+tokio-util = "0.7.12"
 tokio-postgres = { version = "0.7.10", features = [
   "with-chrono-0_4",
   "with-uuid-1",
@@ -87,6 +88,9 @@ dyn-clone = { version = "1.0.17", optional = true }
 geo-types = "0.7.13"
 fundu = "2.0.1"
 dashmap = "6.1.0"
+odbc-api = { version = "9.0.0", optional = true }
+arrow-odbc = { version = "13.0.0", optional = true }
+sha2 = "0.10.8"
 
 [dev-dependencies]
 anyhow = "1.0.86"
@@ -135,8 +139,10 @@ flight = [
   "dep:serde",
   "dep:tonic",
 ]
+odbc = ["dep:odbc-api", "dep:arrow-odbc", "dep:async-stream", "dep:dyn-clone"]
 federation = ["dep:datafusion-federation"]
 duckdb-federation = ["duckdb", "federation"]
 sqlite-federation = ["sqlite", "federation"]
 postgres-federation = ["postgres", "federation"]
 mysql-federation = ["mysql", "federation"]
+odbc-federation = ["odbc", "federation"]

examples/odbc_sqlite.rs

@@ -0,0 +1,51 @@
+use std::{collections::HashMap, sync::Arc};
+
+use datafusion::prelude::SessionContext;
+use datafusion::sql::TableReference;
+use datafusion_table_providers::{
+    odbc::ODBCTableFactory, sql::db_connection_pool::odbcpool::ODBCPool,
+    util::secrets::to_secret_map,
+};
+
+/// This example demonstrates how to:
+/// 1. Create a SQLite ODBC connection pool
+/// 2. Create and use ODBCTableFactory to generate TableProvider
+/// 3. Register TableProvider with DataFusion
+/// 4. Use SQL queries to access SQLite ODBC table data
+#[tokio::main]
+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(),
+    )]));
+    let odbc_pool =
+        Arc::new(ODBCPool::new(params).expect("unable to create SQLite ODBC connection pool"));
+
+    // Create SQLite ODBC table provider factory
+    // Used to generate TableProvider instances that can read SQLite ODBC table data
+    let table_factory = ODBCTableFactory::new(odbc_pool.clone());
+
+    // Create DataFusion session context
+    let ctx = SessionContext::new();
+
+    // Demonstrate direct table provider registration
+    // This method registers the table in the default catalog
+    // Here we register the SQLite ODBC "companies" table as "companies_v2"
+    ctx.register_table(
+        "companies_v2",
+        table_factory
+            .table_provider(TableReference::bare("companies"), None)
+            .await
+            .expect("failed to register table provider"),
+    )
+    .expect("failed to register table");
+
+    // Query Example 1: Query the renamed table through default catalog
+    let df = ctx
+        .sql("SELECT * FROM datafusion.public.companies_v2")
+        .await
+        .expect("select failed");
+    df.show().await.expect("show failed");
+}

src/lib.rs

@@ -8,6 +8,8 @@ pub mod duckdb;
 pub mod flight;
 #[cfg(feature = "mysql")]
 pub mod mysql;
+#[cfg(feature = "odbc")]
+pub mod odbc;
 #[cfg(feature = "postgres")]
 pub mod postgres;
 #[cfg(feature = "sqlite")]

src/odbc.rs

@@ -0,0 +1,89 @@
+/*
+Copyright 2024 The Spice.ai OSS Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+use crate::sql::db_connection_pool::dbconnection::odbcconn::ODBCDbConnectionPool;
+use crate::sql::{
+    db_connection_pool as db_connection_pool_datafusion,
+    sql_provider_datafusion::{Engine, SqlTable},
+};
+use arrow::datatypes::SchemaRef;
+use datafusion::error::DataFusionError;
+use datafusion::{datasource::TableProvider, sql::TableReference};
+use snafu::prelude::*;
+use std::sync::Arc;
+
+#[derive(Debug, Snafu)]
+pub enum Error {
+    #[snafu(display("DbConnectionError: {source}"))]
+    DbConnectionError {
+        source: db_connection_pool_datafusion::dbconnection::GenericError,
+    },
+    #[snafu(display("The table '{table_name}' doesn't exist in the Postgres server"))]
+    TableDoesntExist { table_name: String },
+
+    #[snafu(display("Unable to get a DB connection from the pool: {source}"))]
+    UnableToGetConnectionFromPool {
+        source: db_connection_pool_datafusion::Error,
+    },
+
+    #[snafu(display("Unable to get schema: {source}"))]
+    UnableToGetSchema {
+        source: db_connection_pool_datafusion::dbconnection::Error,
+    },
+
+    #[snafu(display("Unable to generate SQL: {source}"))]
+    UnableToGenerateSQL { source: DataFusionError },
+}
+
+type Result<T, E = Error> = std::result::Result<T, E>;
+
+pub struct ODBCTableFactory<'a> {
+    pool: Arc<ODBCDbConnectionPool<'a>>,
+}
+
+impl<'a> ODBCTableFactory<'a>
+where
+    'a: 'static,
+{
+    #[must_use]
+    pub fn new(pool: Arc<ODBCDbConnectionPool<'a>>) -> Self {
+        Self { pool }
+    }
+
+    pub async fn table_provider(
+        &self,
+        table_reference: TableReference,
+        _schema: Option<SchemaRef>,
+    ) -> Result<Arc<dyn TableProvider + 'static>, Box<dyn std::error::Error + Send + Sync>> {
+        let pool = Arc::clone(&self.pool);
+        let dyn_pool: Arc<ODBCDbConnectionPool<'a>> = pool;
+
+        let table = SqlTable::new("odbc", &dyn_pool, table_reference, Some(Engine::ODBC))
+            .await
+            .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
+
+        let table_provider = Arc::new(table);
+
+        #[cfg(feature = "odbc-federation")]
+        let table_provider = Arc::new(
+            table_provider
+                .create_federated_table_provider()
+                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?,
+        );
+
+        Ok(table_provider)
+    }
+}

src/sql/arrow_sql_gen/postgres/composite.rs

@@ -80,7 +80,7 @@ fn composite_type_fields(type_: &Type) -> &[Field] {
     }
 }
 
-impl<'a> CompositeType<'a> {
+impl CompositeType<'_> {
     /// Returns information about the fields of the composite type.
     #[must_use]
     pub fn fields(&self) -> &[Field] {
@@ -180,7 +180,7 @@ impl<'a> CompositeTypeRanges<'a> {
 }
 
 #[allow(clippy::cast_sign_loss)]
-impl<'a> FallibleIterator for CompositeTypeRanges<'a> {
+impl FallibleIterator for CompositeTypeRanges<'_> {
     type Item = Option<std::ops::Range<usize>>;
     type Error = std::io::Error;
 

src/sql/db_connection_pool/dbconnection.rs

@@ -9,6 +9,8 @@ use snafu::prelude::*;
 pub mod duckdbconn;
 #[cfg(feature = "mysql")]
 pub mod mysqlconn;
+#[cfg(feature = "odbc")]
+pub mod odbcconn;
 #[cfg(feature = "postgres")]
 pub mod postgresconn;
 #[cfg(feature = "sqlite")]