MongoDB: upgrade crate, add SRV support, re-enable Python bindings (#602)

* MongoDB: upgrade crate, add SRV support, re-enable Python bindings

- Upgrade mongodb crate from 3.2.2 to 3.5.2
- Add mongodb+srv:// SRV connection format support via 'srv' parameter
- Add tables() method to MongoDBConnection for listing collections
- Re-enable MongoDB Python bindings (feature flag, PyO3 module, wrapper class)
- Add MongoDB Python example and demo script
- Add unit tests for SRV connection URI building

* Update Rust toolchain to 1.92.0 to match DataFusion v52

* CI: authenticate to ghcr.io to avoid Docker pull rate limits

Author: lukekim

.github/workflows/pr.yaml

@@ -75,8 +75,13 @@ jobs:
       - name: Cache Rust dependencies
         uses: Swatinem/rust-cache@v2
 
+      - name: Log in to container registries
+        run: |
+          echo "${{ secrets.GITHUB_TOKEN }}" | docker login ghcr.io -u ${{ github.actor }} --password-stdin
+
       - name: Pull the Postgres/MySQL images
         run: |
+          : work around spurious network errors in curl 8.0
           docker pull ${{ env.PG_DOCKER_IMAGE }}
           docker pull ${{ env.MYSQL_DOCKER_IMAGE }}
 

Cargo.lock

@@ -50,7 +50,7 @@ fundu = { workspace = true }
 futures = { workspace = true }
 geo-types = { workspace = true }
 itertools = { workspace = true }
-mongodb = { version = "3.2.2", features = ["openssl-tls"], optional = true }
+mongodb = { version = "3.5.2", features = ["openssl-tls"], optional = true }
 mysql_async = { version = "0.36", features = [
   "native-tls-tls",
   "chrono",

core/Cargo.toml

@@ -15,7 +15,7 @@ use std::sync::Arc;
 use crate::mongodb::utils::arrow::mongo_docs_to_arrow;
 use crate::mongodb::utils::schema::infer_arrow_schema_from_documents;
 use crate::mongodb::utils::unnest::{unnest_bson_documents, UnnestBehavior, UnnestParameters};
-use crate::mongodb::{Error, QuerySnafu, Result, UnableToGetSchemaSnafu};
+use crate::mongodb::{Error, QuerySnafu, Result, UnableToGetSchemaSnafu, UnableToGetTablesSnafu};
 
 pub struct MongoDBConnection {
     pub client: Arc<Client>,
@@ -46,6 +46,16 @@ impl MongoDBConnection {
         self.client.database(&self.db_name).collection(collection)
     }
 
+    pub async fn tables(&self) -> Result<Vec<String>, Error> {
+        let db = self.client.database(&self.db_name);
+        let collections = db
+            .list_collection_names()
+            .await
+            .boxed()
+            .context(UnableToGetTablesSnafu)?;
+        Ok(collections)
+    }
+
     pub async fn get_schema(&self, table_reference: &TableReference) -> Result<SchemaRef, Error> {
         let collection_name = table_reference.table();
         let coll = self.get_collection(collection_name);

core/src/mongodb/connection.rs

@@ -133,7 +133,17 @@ fn build_connection_uri(
         format!("?{}", query_params.join("&"))
     };
 
-    let uri = format!("mongodb://{auth}{host}:{port}/{db_name}{query_string}");
+    let use_srv = params
+        .get("srv")
+        .map(|s| s.expose_secret().eq_ignore_ascii_case("true"))
+        .unwrap_or(false);
+
+    let uri = if use_srv {
+        // mongodb+srv:// uses DNS SRV records for discovery; port is not allowed
+        format!("mongodb+srv://{auth}{host}/{db_name}{query_string}")
+    } else {
+        format!("mongodb://{auth}{host}:{port}/{db_name}{query_string}")
+    };
     Ok((uri, Some(db_name.to_string())))
 }
 
@@ -327,6 +337,70 @@ mod tests {
         assert_eq!(result.1, Some("testdb".to_string()));
     }
 
+    #[test]
+    fn test_build_connection_uri_with_srv() {
+        let params = create_params(vec![
+            ("db", "mydb"),
+            ("host", "cluster0.example.mongodb.net"),
+            ("user", "testuser"),
+            ("pass", "testpass"),
+            ("srv", "true"),
+        ]);
+
+        let result = build_connection_uri(&params).unwrap();
+        assert_eq!(
+            result.0,
+            "mongodb+srv://testuser:testpass@cluster0.example.mongodb.net/mydb"
+        );
+        assert_eq!(result.1, Some("mydb".to_string()));
+    }
+
+    #[test]
+    fn test_build_connection_uri_with_srv_and_query_params() {
+        let params = create_params(vec![
+            ("db", "mydb"),
+            ("host", "cluster0.example.mongodb.net"),
+            ("srv", "true"),
+            ("auth_source", "admin"),
+        ]);
+
+        let result = build_connection_uri(&params).unwrap();
+        assert_eq!(
+            result.0,
+            "mongodb+srv://cluster0.example.mongodb.net/mydb?authSource=admin"
+        );
+        assert_eq!(result.1, Some("mydb".to_string()));
+    }
+
+    #[test]
+    fn test_build_connection_uri_with_srv_false() {
+        let params = create_params(vec![
+            ("db", "mydb"),
+            ("host", "localhost"),
+            ("port", "27017"),
+            ("srv", "false"),
+        ]);
+
+        let result = build_connection_uri(&params).unwrap();
+        assert_eq!(result.0, "mongodb://localhost:27017/mydb");
+        assert_eq!(result.1, Some("mydb".to_string()));
+    }
+
+    #[test]
+    fn test_build_connection_uri_with_connection_string_srv() {
+        let params = create_params(vec![(
+            "connection_string",
+            "mongodb+srv://user:pass@cluster0.example.mongodb.net/testdb",
+        )]);
+
+        let result = build_connection_uri(&params).unwrap();
+        assert_eq!(
+            result.0,
+            "mongodb+srv://user:pass@cluster0.example.mongodb.net/testdb"
+        );
+        assert_eq!(result.1, None);
+    }
+
     #[test]
     fn test_build_connection_uri_user_without_password() {
         let params = create_params(vec![("user", "testuser")]);

core/src/mongodb/connection_pool.rs

@@ -36,7 +36,8 @@ default = [
     "mysql",
     "postgres",
     "odbc",
-    "flight",    
+    "flight",
+    "mongodb",
 ]
 clickhouse = ["datafusion-table-providers/clickhouse-federation"]
 duckdb = ["dep:duckdb", "datafusion-table-providers/duckdb-federation"]
@@ -45,3 +46,4 @@ mysql = ["datafusion-table-providers/mysql-federation"]
 postgres = ["datafusion-table-providers/postgres-federation"]
 odbc = ["datafusion-table-providers/odbc-federation"]
 flight = ["dep:arrow-flight", "datafusion-table-providers/flight"]
+mongodb = ["datafusion-table-providers/mongodb"]

python/Cargo.toml

@@ -0,0 +1,75 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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
+#
+#   http://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.
+"""Example demonstrating MongoDB table provider usage.
+
+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
+```
+"""
+
+import datafusion
+from datafusion_table_providers.mongodb import MongoDBTableFactory
+
+
+def main():
+    # Create MongoDB connection parameters
+    mongodb_params = {
+        "connection_string": "mongodb://root:password@localhost:27017/mongo_db?authSource=admin&tls=false"
+    }
+
+    # Create MongoDB table factory
+    factory = MongoDBTableFactory(mongodb_params)
+
+    # List all tables
+    tables = factory.tables()
+    print(f"Tables: {tables}")
+
+    # Get table provider for 'companies' table
+    table = factory.get_table("companies")
+
+    # Create DataFusion context and register the table
+    ctx = datafusion.SessionContext()
+    ctx.register_table_provider("companies", table)
+
+    # Query the table
+    df = ctx.sql("SELECT * FROM companies")
+    df.show()
+
+
+if __name__ == "__main__":
+    main()

python/examples/mongodb_demo.py

@@ -24,7 +24,7 @@ classifier = [
     "Programming Language :: Python",
     "Programming Language :: Rust",
 ]
-dependencies = ["datafusion>=45.0.0"]
+dependencies = ["datafusion>=45.0.0,<52.0.0"]
 
 [project.urls]
 repository = "https://github.com/datafusion-contrib/datafusion-table-providers"

python/pyproject.toml

@@ -0,0 +1,39 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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
+#
+#   http://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.
+"""Python interface for MongoDB table provider."""
+
+from typing import Any, List
+from . import _internal
+
+class MongoDBTableFactory:
+    """MongoDB table factory."""
+
+    def __init__(self, params: dict) -> None:
+        """Create a MongoDB table factory."""
+        self._raw = _internal.mongodb.RawMongoDBTableFactory(params)
+
+    def tables(self) -> List[str]:
+        """Get all the table names."""
+        return self._raw.tables()
+
+    def get_table(self, table_reference: str) -> Any:
+        """Return the table provider for table named `table_reference`.
+
+        Args:
+            table_reference (str): table name
+        """
+        return self._raw.get_table(table_reference)

python/python/datafusion_table_providers/mongodb.py

@@ -49,6 +49,8 @@ pub mod clickhouse;
 pub mod duckdb;
 #[cfg(feature = "flight")]
 pub mod flight;
+#[cfg(feature = "mongodb")]
+pub mod mongodb;
 #[cfg(feature = "mysql")]
 pub mod mysql;
 #[cfg(feature = "odbc")]
@@ -113,5 +115,12 @@ fn _internal(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
         m.add_submodule(&clickhouse)?;
     }
 
+    #[cfg(feature = "mongodb")]
+    {
+        let mongodb = PyModule::new(py, "mongodb")?;
+        mongodb::init_module(&mongodb)?;
+        m.add_submodule(&mongodb)?;
+    }
+
     Ok(())
 }