Spiceai mongo test

.github/workflows/pr.yaml

@@ -22,6 +22,10 @@ jobs:
         with:
           repo-token: ${{ secrets.GITHUB_TOKEN }}
 
+      - uses: Swatinem/rust-cache@v2
+        with:
+          key: clippy
+
       - run: cargo clippy --all-features -- -D warnings
 
   build:
@@ -33,6 +37,10 @@ jobs:
 
       - uses: dtolnay/rust-toolchain@stable
 
+      - uses: Swatinem/rust-cache@v2
+        with:
+          key: build
+
       # Putting this into a GitHub Actions matrix will run a separate job per matrix item, whereas in theory
       # this can re-use the existing build cache to go faster.
       - name: Build without default features
@@ -50,23 +58,32 @@ jobs:
       - name: Build with only mysql
         run: cargo check --no-default-features --features mysql
 
+      - name: Build with only mongodb
+        run: cargo check --no-default-features --features mongodb
+
   integration-test:
-    name: Tests
+    name: Tests postgres,postgres-federation,sqlite,sqlite-federation,mysql
     runs-on: ubuntu-latest
 
     env:
       PG_DOCKER_IMAGE: ghcr.io/cloudnative-pg/postgresql:16-bookworm
       MYSQL_DOCKER_IMAGE: public.ecr.aws/ubuntu/mysql:8.0-22.04_beta
+      MONGODB_DOCKER_IMAGE: public.ecr.aws/docker/library/mongo:7
 
     steps:
       - uses: actions/checkout@v4
 
       - uses: dtolnay/rust-toolchain@stable
 
+      - uses: Swatinem/rust-cache@v2
+        with:
+          key: test
+
       - name: Pull the Postgres/MySQL images
         run: |
           docker pull ${{ env.PG_DOCKER_IMAGE }}
           docker pull ${{ env.MYSQL_DOCKER_IMAGE }}
+          docker pull ${{ env.MONGODB_DOCKER_IMAGE }}
 
       - name: Free Disk Space
         run: |
@@ -103,4 +120,4 @@ jobs:
           sudo apt-get install -y libsqlite3-dev
 
       - name: Run tests
-        run: make test
+        run: cargo test --all-features

Cargo.toml

@@ -35,9 +35,12 @@ duckdb = { version = "1.1.3", features = [
 fallible-iterator = "0.3.0"
 futures = "0.3.30"
 mysql_async = { version = "0.35.1", features = ["native-tls-tls", "chrono", "hdrhistogram", "bigdecimal", "time"], optional = true }
+mongodb = { version = "3.2.2", features = ["openssl-tls"], optional = true }
+num-traits = { version = "0.2", optional = true }
 prost = { version = "0.13.2", optional = true }
 rand = "0.8.5"
 r2d2 = { version = "0.8.10", optional = true }
+rust_decimal = { version = "1.32", optional = true }
 rusqlite = { version = "0.31.0", optional = true }
 sea-query = { git = "https://github.com/spiceai/sea-query.git", rev = "213b6b876068f58159ebdd5852604a021afaebf9", features = ["backend-sqlite", "backend-postgres", "postgres-array", "with-rust_decimal", "with-bigdecimal", "with-time", "with-chrono"] }
 secrecy = "0.10.3"
@@ -100,6 +103,13 @@ flight = [
 duckdb-federation = ["duckdb"]
 sqlite-federation = ["sqlite"]
 postgres-federation = ["postgres"]
+mongodb = [
+  "dep:mongodb", 
+  "dep:async-stream",
+  "dep:arrow-schema",
+  "dep:rust_decimal",
+  "dep:num-traits",
+]
 
 [patch.crates-io]
 datafusion-federation = { git = "https://github.com/spiceai/datafusion-federation.git", rev = "9db74a4b360df6be1bb554c59a474a2fd4bfb7e9" } # spiceai-47
@@ -109,4 +119,54 @@ datafusion = { git = "https://github.com/spiceai/datafusion.git", rev = "b5c62f2
 datafusion-expr = { git = "https://github.com/spiceai/datafusion.git", rev = "b5c62f29d2c70c5331ff50015b67b5e1cafcd578" }          # spiceai-47
 datafusion-physical-expr = { git = "https://github.com/spiceai/datafusion.git", rev = "b5c62f29d2c70c5331ff50015b67b5e1cafcd578" } # spiceai-47
 datafusion-physical-plan = { git = "https://github.com/spiceai/datafusion.git", rev = "b5c62f29d2c70c5331ff50015b67b5e1cafcd578" } # spiceai-47
-datafusion-proto = { git = "https://github.com/spiceai/datafusion.git", rev = "b5c62f29d2c70c5331ff50015b67b5e1cafcd578" }         # spiceai-47
+datafusion-proto = { git = "https://github.com/spiceai/datafusion.git", rev = "b5c62f29d2c70c5331ff50015b67b5e1cafcd578" }         # spiceai-47
+
+[[example]]
+name = "odbc_sqlite"
+path = "examples/odbc_sqlite.rs"
+required-features = ["sqlite", "odbc"]
+
+[[example]]
+name = "duckdb"
+path = "examples/duckdb.rs"
+required-features = ["duckdb"]
+
+[[example]]
+name = "duckdb_external_table"
+path = "examples/duckdb_external_table.rs"
+required-features = ["duckdb"]
+
+[[example]]
+name = "duckdb_function"
+path = "examples/duckdb_function.rs"
+required-features = ["duckdb"]
+
+[[example]]
+name = "flight-sql"
+path = "examples/flight-sql.rs"
+required-features = ["flight"]
+
+[[example]]
+name = "sqlite"
+path = "examples/sqlite.rs"
+required-features = ["sqlite"]
+
+[[example]]
+name = "clickhouse"
+path = "examples/clickhouse.rs"
+required-features = ["clickhouse"]
+
+[[example]]
+name = "mysql"
+path = "examples/mysql.rs"
+required-features = ["mysql"]
+
+[[example]]
+name = "postgres"
+path = "examples/postgres.rs"
+required-features = ["postgres"]
+
+[[example]]
+name = "mongodb"
+path = "examples/mongodb.rs"
+required-features = ["mongodb"]

Makefile

@@ -11,4 +11,4 @@ lint:
 
 .PHONY: test-integration
 test-integration:
-	RUST_LOG=debug cargo test --test integration --no-default-features --features postgres,sqlite,mysql -- --nocapture
+	RUST_LOG=debug cargo test --test integration --no-default-features --features postgres,sqlite,mysql,mongodb -- --nocapture

README.md

@@ -13,6 +13,7 @@ Many of the table providers in this repo are for querying data from other databa
 - SQLite
 - DuckDB
 - Flight SQL
+- MongoDB
 
 ## Examples
 
@@ -104,3 +105,30 @@ roapi -t taxi=https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_20
 
 cargo run --example flight-sql --features flight
 ```
+
+### MongoDB
+
+In order to run the MongoDB example, you need to have a MongoDB server running. You can use the following command to start a MongoDB server in a Docker container the example can use:
+
+```bash
+docker run --name mongodb \
+  -e MONGO_INITDB_ROOT_USERNAME=root \
+  -e MONGO_INITDB_ROOT_PASSWORD=password \
+  -e MONGO_INITDB_DATABASE=mongo_db \
+  -p 27017:27017 \
+  -d mongo:7.0
+# Wait for the MongoDB server to start
+sleep 30
+
+# Create a table in the MongoDB server and insert some data
+docker exec -i mongodb mongosh -u root -p password --authenticationDatabase admin <<EOF
+use mongo_db;
+db.companies.insertOne({
+  id: 1,
+  name: "Acme Corporation"
+});
+EOF
+
+# Run from repo folder
+cargo run -p datafusion-table-providers --example mongodb --features mongodb
+```

examples/mongodb.rs

@@ -0,0 +1,80 @@
+use std::{collections::HashMap, sync::Arc};
+
+use datafusion::prelude::SessionContext;
+use datafusion::sql::TableReference;
+use datafusion_table_providers::{
+    mongodb::{connection_pool::MongoDBConnectionPool, MongoDBTableFactory}, util::secrets::to_secret_map,
+};
+
+/// This example demonstrates how to:
+/// 1. Create a MongoDB connection pool
+/// 2. Create and use MongoDBTableFactory to generate TableProvider
+/// 3. Use SQL queries to access MongoDB table data
+///
+/// Prerequisites:
+/// Start a MongoDB server using Docker:
+/// ```bash
+/// docker run --name mongodb \
+/// -e MONGO_INITDB_ROOT_USERNAME=root \
+/// -e MONGO_INITDB_ROOT_PASSWORD=password \
+/// -e MONGO_INITDB_DATABASE=mongo_db \
+/// -p 27017:27017 \
+/// -d mongo:7.0
+/// # Wait for the MongoDB server to start
+/// sleep 30
+///
+/// # Create a table in the MongoDB server and insert some data
+/// docker exec -i mongodb mongosh -u root -p password --authenticationDatabase admin <<EOF
+/// use mongo_db;
+///
+/// db.companies.insertOne({
+/// id: 1,
+/// name: "Acme Corporation"
+/// });
+/// EOF
+/// ```
+#[tokio::main]
+async fn main(){
+
+    // Create MongoDB connection parameters
+    // Including connection string and SSL mode settings
+    let mongodb_params = to_secret_map(HashMap::from([
+        (
+            "connection_string".to_string(),
+            "mongodb://root:password@localhost:27017/mongo_db?authSource=admin&tls=true".to_string(),
+        ),
+    ]));
+
+    // Create MongoDB connection pool
+    let mongodb_pool = Arc::new(
+        MongoDBConnectionPool::new(mongodb_params)
+            .await
+            .expect("unable to create MongoDB connection pool"),
+    );
+
+    // Create MongoDB table provider factory
+    // Used to generate TableProvider instances that can read MongoDB table data
+    let table_factory = MongoDBTableFactory::new(mongodb_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 MongoDB "companies" table as "companies_v2"
+    ctx.register_table(
+        "companies_v2",
+        table_factory
+            .table_provider(TableReference::bare("companies"))
+            .await
+            .expect("failed to register table provider"),
+    )
+    .expect("failed to register table");
+
+    // Query Example: 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/duckdb/creator.rs

@@ -297,6 +297,7 @@ impl TableManager {
 
     /// Inserts data from this table into the target table.
     #[tracing::instrument(level = "debug", skip_all)]
+    #[allow(dead_code)]
     pub(crate) fn insert_into(
         &self,
         table: &TableManager,

src/lib.rs

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

src/mongodb.rs

@@ -0,0 +1,75 @@
+pub mod connection;
+pub mod connection_pool;
+pub mod table;
+pub mod utils;
+
+use crate::mongodb::table::MongoDBTable;
+use crate::mongodb::connection_pool::MongoDBConnectionPool;
+use arrow_schema::ArrowError;
+use datafusion::datasource::TableProvider;
+use datafusion::sql::TableReference;
+use snafu::prelude::*;
+use std::sync::Arc;
+
+#[derive(Debug, Snafu)]
+pub enum Error {
+    #[snafu(display("Invalid MongoDB URI: {source}"))]
+    InvalidUri { source: mongodb::error::Error },
+
+    #[snafu(display("Invalid value for parameter {parameter_name}\nEnsure the value is valid for parameter {parameter_name}"))]
+    InvalidParameter { parameter_name: String },
+
+    #[snafu(display("TLS root certificate path is invalid: {path}"))]
+    InvalidRootCertPath { path: String },
+
+    #[snafu(display("Failed to connect to MongoDB: {source}"))]
+    ConnectionFailed { source: mongodb::error::Error },
+
+    #[snafu(display("Unable to get tables: {source}"))]
+    UnableToGetTables { source: Box<dyn std::error::Error + std::marker::Send + Sync> },
+
+    #[snafu(display("Unable to get schema: {source}"))]
+    UnableToGetSchema { source: Box<dyn std::error::Error + std::marker::Send + Sync> },
+
+    #[snafu(display("Unable to get schemas: {source}"))]
+    UnableToGetSchemas { source: Box<dyn std::error::Error + std::marker::Send + Sync> },
+
+    #[snafu(display("Failed to execute MongoDB query: {source}"))]
+    QueryError { source: Box<dyn std::error::Error + std::marker::Send + Sync> },
+
+    #[snafu(display("Failed to convert MongoDB documents to Arrow: {source}"))]
+    ConversionError { source: Box<dyn std::error::Error + std::marker::Send + Sync> },
+
+    #[snafu(display("Invalid decimal parameters: {source}"))]
+    InvalidDecimalError { source: ArrowError },
+
+    #[snafu(display("Authentication failed. Verify username and password."))]
+    InvalidUsernameOrPassword,
+}
+
+type Result<T, E = Error> = std::result::Result<T, E>;
+
+pub struct MongoDBTableFactory {
+    pool: Arc<MongoDBConnectionPool>,
+}
+
+impl MongoDBTableFactory {
+    #[must_use]
+    pub fn new(pool: Arc<MongoDBConnectionPool>) -> Self {
+        Self { pool }
+    }
+
+    pub async fn table_provider(
+        &self,
+        table_reference: TableReference,
+    ) -> Result<Arc<dyn TableProvider + 'static>, Box<dyn std::error::Error + Send + Sync>> {
+        let pool = Arc::clone(&self.pool);
+        let table_provider = Arc::new(
+            MongoDBTable::new(&pool, table_reference)
+                .await
+                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?,
+        );
+
+        Ok(table_provider)
+    }
+}