diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 6aa81cd7..c6effbf1 100644 --- a/.github/workflows/pr.yaml +++ b/.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,6 +58,9 @@ 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 runs-on: ubuntu-latest @@ -57,16 +68,22 @@ jobs: 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: | diff --git a/Cargo.toml b/Cargo.toml index df6cd1c9..8e47214e 100644 --- a/Cargo.toml +++ b/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 \ No newline at end of file +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"] diff --git a/Makefile b/Makefile index 96b6f735..fc993126 100644 --- a/Makefile +++ b/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 diff --git a/README.md b/README.md index 5d514492..7da1fb03 100644 --- a/README.md +++ b/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 < }, + + #[snafu(display("Unable to get schema: {source}"))] + UnableToGetSchema { source: Box }, + + #[snafu(display("Unable to get schemas: {source}"))] + UnableToGetSchemas { source: Box }, + + #[snafu(display("Failed to execute MongoDB query: {source}"))] + QueryError { source: Box }, + + #[snafu(display("Failed to convert MongoDB documents to Arrow: {source}"))] + ConversionError { source: Box }, + + #[snafu(display("Invalid decimal parameters: {source}"))] + InvalidDecimalError { source: ArrowError }, + + #[snafu(display("Authentication failed. Verify username and password."))] + InvalidUsernameOrPassword, +} + +type Result = std::result::Result; + +pub struct MongoDBTableFactory { + pool: Arc, +} + +impl MongoDBTableFactory { + #[must_use] + pub fn new(pool: Arc) -> Self { + Self { pool } + } + + pub async fn table_provider( + &self, + table_reference: TableReference, + ) -> Result, Box> { + 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)?, + ); + + Ok(table_provider) + } +} diff --git a/src/mongodb/connection.rs b/src/mongodb/connection.rs new file mode 100644 index 00000000..9640786f --- /dev/null +++ b/src/mongodb/connection.rs @@ -0,0 +1,122 @@ +use std::sync::Arc; +use async_stream::stream; +use datafusion::arrow::datatypes::{Schema, SchemaRef}; +use datafusion::execution::SendableRecordBatchStream; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::sql::TableReference; +use futures::TryStreamExt; +use futures::StreamExt; +use mongodb::{bson::{Document, doc}, Client, Collection}; +use snafu::prelude::*; + +use crate::mongodb::utils::arrow::mongo_docs_to_arrow; +use crate::mongodb::utils::schema::infer_arrow_schema_from_documents; +use crate::mongodb::{Error, QuerySnafu, Result, UnableToGetSchemaSnafu}; + +const NUM_DOCUMENTS_TO_INFER_SCHEMA: i64 = 20; + + +pub struct MongoDBConnection { + pub client: Arc, + pub db_name: String, +} + +impl MongoDBConnection { + pub fn new(client: Arc, db_name: String) -> Self { + MongoDBConnection { client, db_name } + } + + fn get_collection(&self, collection: &str) -> Collection { + self.client.database(&self.db_name).collection(collection) + } + + pub async fn get_schema( + &self, + table_reference: &TableReference, + ) -> Result { + let collection_name = table_reference.table(); + let coll = self.get_collection(collection_name); + + let sample = coll + .find(doc! {}) + .limit(NUM_DOCUMENTS_TO_INFER_SCHEMA) + .await + .boxed() + .context(UnableToGetSchemaSnafu)?; + + let docs: Vec = sample.try_collect().await.boxed().context(UnableToGetSchemaSnafu)?; + + infer_arrow_schema_from_documents(&docs) + .boxed() + .context(UnableToGetSchemaSnafu) + } + + pub async fn query_arrow( + &self, + table_reference: &Arc, + projected_schema: &SchemaRef, + filters_doc: &Document, + limit: Option, + ) -> Result { + let collection_name = table_reference.table(); + let coll = self.get_collection(collection_name); + + let mut find = coll + .find(filters_doc.clone()) + .projection(schema_to_mongo_projection(projected_schema)); + + if let Some(l) = limit { + find = find.limit(l.into()); + } + + let cursor = find.await.boxed().context(QuerySnafu)?; + let chunked_stream = cursor.try_chunks(4_000); + let projected_schema_clone = Arc::clone(projected_schema); + + let mut batch_stream = Box::pin(stream! { + for await chunk in chunked_stream { + match chunk { + Ok(docs) => { + let batch = mongo_docs_to_arrow(&docs, Arc::clone(&projected_schema_clone))?; + yield Ok(batch); + } + Err(e) => yield Err(Error::QueryError { source: Box::new(e) }), + } + } + }); + + let Some(first_batch_result) = batch_stream.next().await else { + return Ok(Box::pin(RecordBatchStreamAdapter::new( + Arc::new(Schema::empty()), + futures::stream::empty().boxed(), + ))); + }; + + let first_batch = first_batch_result?; + let schema = first_batch.schema(); + + let full_stream = Box::pin(stream! { + yield Ok(first_batch); + while let Some(batch_result) = batch_stream.next().await { + yield batch_result.map_err(|e| datafusion::error::DataFusionError::Execution(e.to_string())); + } + }); + + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, full_stream))) + } + +} + +pub fn schema_to_mongo_projection(projected_schema: &SchemaRef) -> Document { + let mut projection = Document::new(); + + if projected_schema.fields().is_empty() { + return projection; + } + + for field in projected_schema.fields() { + projection.insert(field.name(), 1); + } + + projection +} diff --git a/src/mongodb/connection_pool.rs b/src/mongodb/connection_pool.rs new file mode 100644 index 00000000..65e302bf --- /dev/null +++ b/src/mongodb/connection_pool.rs @@ -0,0 +1,552 @@ +use std::{collections::HashMap, path::PathBuf, sync::Arc}; +use mongodb::{ + error::ErrorKind, + bson::doc, + options::{ClientOptions, Tls, TlsOptions}, + Client, +}; +use secrecy::{ExposeSecret, SecretString}; +use snafu::ResultExt; +use crate::mongodb::{connection::MongoDBConnection, ConnectionFailedSnafu, Error, InvalidUriSnafu, Result}; +#[derive(Clone, Debug)] +pub struct MongoDBConnectionPool { + client: Arc, + db_name: String, +} + +const DEFAULT_HOST: &str = "localhost"; +const DEFAULT_PORT: &str = "27017"; +const DEFAULT_DATABASE : &str = "default"; +const DEFAULT_MIN_POOL_SIZE: u32 = 10; +const DEFAULT_MAX_POOL_SIZE: u32 = 100; +const DEFAULT_SSL_MODE: &str = "required"; + +impl MongoDBConnectionPool { + pub async fn new(params: HashMap) -> Result { + let params = crate::util::remove_prefix_from_hashmap_keys(params, "mongodb_"); + + let (uri, explicit_db_name) = build_connection_uri(¶ms)?; + + let mut client_options = ClientOptions::parse(&uri) + .await + .context(InvalidUriSnafu)?; + + configure_pool_size(&mut client_options, ¶ms)?; + configure_tls(&mut client_options, ¶ms)?; + + let db_name = explicit_db_name + .or(client_options.default_database.clone()) + .unwrap_or(DEFAULT_DATABASE.to_string()); + + let client = Client::with_options(client_options).context(ConnectionFailedSnafu)?; + + test_connection(&client, &db_name).await?; + + Ok(Self { + client: Arc::new(client), + db_name, + }) + } + + pub async fn connect(&self) -> Result> { + Ok(Box::new(MongoDBConnection::new( + Arc::clone(&self.client), + self.db_name.clone(), + ))) + } +} + +fn build_connection_uri(params: &HashMap) -> Result<(String, Option)> { + if let Some(uri) = params.get("connection_string") { + return Ok((uri.expose_secret().to_string(), None)); + } + + let db_name = get_param_or_default(params, "db", DEFAULT_DATABASE); + let host = get_param_or_default(params, "host", DEFAULT_HOST); + let port = get_param_or_default(params, "port", DEFAULT_PORT); + + let auth = match (params.get("user"), params.get("pass")) { + (Some(user), Some(pass)) => { + format!("{}:{}@", user.expose_secret(), pass.expose_secret()) + } + (Some(_), None) => { + return Err(Error::InvalidParameter {parameter_name: "pass".to_string(),}); + } + (None, Some(_)) => { + return Err(Error::InvalidParameter {parameter_name: "user".to_string()}); + } + (None, None) => String::new(), + }; + + let uri = format!("mongodb://{}{}:{}/{}", auth, host, port, db_name); + Ok((uri, Some(db_name.to_string()))) +} + +fn configure_pool_size(client_options: &mut ClientOptions, params: &HashMap) -> Result<()> { + let pool_min = parse_u32_param(params, "pool_min", DEFAULT_MIN_POOL_SIZE)?; + let pool_max = parse_u32_param(params, "pool_max", DEFAULT_MAX_POOL_SIZE)?; + + if pool_min > pool_max { + return Err(Error::InvalidParameter { + parameter_name: "pool_min/pool_max".to_string(), + }); + } + + client_options.min_pool_size = Some(pool_min); + client_options.max_pool_size = Some(pool_max); + + Ok(()) +} + +fn configure_tls(client_options: &mut ClientOptions, params: &HashMap) -> Result<()> { + let has_explicit_tls_params = params.contains_key("sslmode") || params.contains_key("sslrootcert"); + + if client_options.tls.is_some() && !has_explicit_tls_params { + return Ok(()); + } + + let ssl_mode = get_param_or_default(params, "sslmode", DEFAULT_SSL_MODE); + + match ssl_mode.to_lowercase().as_str() { + "disabled" | "required" | "preferred" => {}, + _ => { + return Err(Error::InvalidParameter { + parameter_name: "sslmode".to_string(), + }); + } + } + + let ssl_rootcert_path = if let Some(cert_path) = params.get("sslrootcert") { + let path = PathBuf::from(cert_path.expose_secret()); + if !path.exists() { + return Err(Error::InvalidRootCertPath { + path: cert_path.expose_secret().to_string(), + }); + } + Some(path) + } else { + None + }; + + client_options.tls = build_tls_options(&ssl_mode, ssl_rootcert_path); + Ok(()) +} + +fn build_tls_options(ssl_mode: &str, rootcert_path: Option) -> Option { + if ssl_mode == "disabled" { + return Some(Tls::Disabled); + } + + let tls_options = match (rootcert_path, ssl_mode) { + // Root cert + preferred + (Some(path), "preferred") => TlsOptions::builder() + .ca_file_path(Some(path)) + .allow_invalid_certificates(Some(true)) + .allow_invalid_hostnames(Some(true)) + .build(), + + // Root cert + required + (Some(path), _) => TlsOptions::builder() + .ca_file_path(Some(path)) + .build(), + + // No root cert + preferred + (None, "preferred") => TlsOptions::builder() + .allow_invalid_certificates(Some(true)) + .allow_invalid_hostnames(Some(true)) + .build(), + + // No root cert + required + (None, _) => TlsOptions::builder().build(), + }; + + Some(Tls::Enabled(tls_options)) +} + +async fn test_connection(client: &Client, db_name: &str) -> Result<()> { + client + .database(db_name) + .run_command(doc! { "ping": 1 }) + .await + .map_err(|err| match *err.kind { + ErrorKind::Authentication { .. } => Error::InvalidUsernameOrPassword {}, + _ => Error::ConnectionFailed { source: err }, + })?; + Ok(()) +} + +fn get_param_or_default(params: &HashMap, key: &str, default: &str) -> String { + params + .get(key) + .map(|s| s.expose_secret().to_string()) + .unwrap_or_else(|| default.to_string()) +} + +fn parse_u32_param(params: &HashMap, key: &str, default: u32) -> Result { + params + .get(key) + .map(|s| s.expose_secret().parse::()) + .transpose() + .map_err(|_| Error::InvalidParameter { + parameter_name: key.to_string(), + }) + .map(|opt| opt.unwrap_or(default)) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + use secrecy::SecretString; + use mongodb::options::{Tls, TlsOptions}; + + fn create_secret_string(value: &str) -> SecretString { + SecretString::new(value.to_string().into_boxed_str()) + } + + fn create_params(pairs: Vec<(&str, &str)>) -> HashMap { + pairs + .into_iter() + .map(|(k, v)| (k.to_string(), create_secret_string(v))) + .collect() + } + + #[test] + fn test_build_connection_uri_with_connection_string() { + let params = create_params(vec![ + ("connection_string", "mongodb://user:pass@example.com:27017/testdb"), + ]); + + let result = build_connection_uri(¶ms).unwrap(); + assert_eq!(result.0, "mongodb://user:pass@example.com:27017/testdb"); + assert_eq!(result.1, None); + } + + #[test] + fn test_build_connection_uri_with_individual_params() { + let params = create_params(vec![ + ("db", "mydb"), + ("host", "example.com"), + ("port", "27018"), + ("user", "testuser"), + ("pass", "testpass"), + ]); + + let result = build_connection_uri(¶ms).unwrap(); + assert_eq!(result.0, "mongodb://testuser:testpass@example.com:27018/mydb"); + assert_eq!(result.1, Some("mydb".to_string())); + } + + #[test] + fn test_build_connection_uri_with_defaults() { + let params = HashMap::new(); + + let result = build_connection_uri(¶ms).unwrap(); + assert_eq!(result.0, "mongodb://localhost:27017/default"); + assert_eq!(result.1, Some("default".to_string())); + } + + #[test] + fn test_build_connection_uri_without_auth() { + let params = create_params(vec![ + ("db", "testdb"), + ("host", "localhost"), + ("port", "27017"), + ]); + + let result = build_connection_uri(¶ms).unwrap(); + assert_eq!(result.0, "mongodb://localhost:27017/testdb"); + assert_eq!(result.1, Some("testdb".to_string())); + } + + #[test] + fn test_build_connection_uri_user_without_password() { + let params = create_params(vec![ + ("user", "testuser"), + ]); + + let result = build_connection_uri(¶ms); + assert!(result.is_err()); + if let Err(Error::InvalidParameter { parameter_name }) = result { + assert_eq!(parameter_name, "pass"); + } else { + panic!("Expected InvalidParameter error for pass"); + } + } + + #[test] + fn test_build_connection_uri_password_without_user() { + let params = create_params(vec![ + ("pass", "testpass"), + ]); + + let result = build_connection_uri(¶ms); + assert!(result.is_err()); + if let Err(Error::InvalidParameter { parameter_name }) = result { + assert_eq!(parameter_name, "user"); + } else { + panic!("Expected InvalidParameter error for user"); + } + } + + #[test] + fn test_configure_pool_size_with_valid_params() { + let mut client_options = ClientOptions::default(); + let params = create_params(vec![ + ("pool_min", "5"), + ("pool_max", "50"), + ]); + + let result = configure_pool_size(&mut client_options, ¶ms); + assert!(result.is_ok()); + assert_eq!(client_options.min_pool_size, Some(5)); + assert_eq!(client_options.max_pool_size, Some(50)); + } + + #[test] + fn test_configure_pool_size_with_defaults() { + let mut client_options = ClientOptions::default(); + let params = HashMap::new(); + + let result = configure_pool_size(&mut client_options, ¶ms); + assert!(result.is_ok()); + assert_eq!(client_options.min_pool_size, Some(DEFAULT_MIN_POOL_SIZE)); + assert_eq!(client_options.max_pool_size, Some(DEFAULT_MAX_POOL_SIZE)); + } + + #[test] + fn test_configure_pool_size_min_greater_than_max() { + let mut client_options = ClientOptions::default(); + let params = create_params(vec![ + ("pool_min", "100"), + ("pool_max", "50"), + ]); + + let result = configure_pool_size(&mut client_options, ¶ms); + assert!(result.is_err()); + if let Err(Error::InvalidParameter { parameter_name }) = result { + assert_eq!(parameter_name, "pool_min/pool_max"); + } else { + panic!("Expected InvalidParameter error for pool_min/pool_max"); + } + } + + #[test] + fn test_configure_tls_skips_when_already_configured() { + let mut client_options = ClientOptions::default(); + client_options.tls = Some(Tls::Enabled(TlsOptions::builder().build())); + + let params = HashMap::new(); + + let result = configure_tls(&mut client_options, ¶ms); + assert!(result.is_ok()); + + assert!(client_options.tls.is_some()); + if let Some(Tls::Enabled(_)) = client_options.tls {} else { + panic!("Expected TLS to remain enabled"); + } + } + + #[test] + fn test_configure_tls_overrides_when_explicit_params() { + let mut client_options = ClientOptions::default(); + client_options.tls = Some(Tls::Enabled(TlsOptions::builder().build())); + + let params = create_params(vec![ + ("sslmode", "disabled"), + ]); + + let result = configure_tls(&mut client_options, ¶ms); + assert!(result.is_ok()); + + if let Some(Tls::Disabled) = client_options.tls {} else { + panic!("Expected TLS to be disabled"); + } + } + + #[test] + fn test_configure_tls_with_required_mode() { + let mut client_options = ClientOptions::default(); + let params = create_params(vec![ + ("sslmode", "required"), + ]); + + let result = configure_tls(&mut client_options, ¶ms); + assert!(result.is_ok()); + + if let Some(Tls::Enabled(_)) = client_options.tls {} else { + panic!("Expected TLS to be enabled"); + } + } + + #[test] + fn test_configure_tls_with_preferred_mode() { + let mut client_options = ClientOptions::default(); + let params = create_params(vec![ + ("sslmode", "preferred"), + ]); + + let result = configure_tls(&mut client_options, ¶ms); + assert!(result.is_ok()); + + if let Some(Tls::Enabled(_)) = client_options.tls {} else { + panic!("Expected TLS to be enabled"); + } + } + + #[test] + fn test_configure_tls_with_disabled_mode() { + let mut client_options = ClientOptions::default(); + let params = create_params(vec![ + ("sslmode", "disabled"), + ]); + + let result = configure_tls(&mut client_options, ¶ms); + assert!(result.is_ok()); + + if let Some(Tls::Disabled) = client_options.tls {} else { + panic!("Expected TLS to be disabled"); + } + } + + #[test] + fn test_configure_tls_with_invalid_mode() { + let mut client_options = ClientOptions::default(); + let params = create_params(vec![ + ("sslmode", "invalid_mode"), + ]); + + let result = configure_tls(&mut client_options, ¶ms); + assert!(result.is_err()); + if let Err(Error::InvalidParameter { parameter_name }) = result { + assert_eq!(parameter_name, "sslmode"); + } else { + panic!("Expected InvalidParameter error for sslmode"); + } + } + + #[test] + fn test_configure_tls_with_nonexistent_cert_path() { + let mut client_options = ClientOptions::default(); + let params = create_params(vec![ + ("sslmode", "required"), + ("sslrootcert", "/nonexistent/path/cert.pem"), + ]); + + let result = configure_tls(&mut client_options, ¶ms); + assert!(result.is_err()); + if let Err(Error::InvalidRootCertPath { path }) = result { + assert_eq!(path, "/nonexistent/path/cert.pem"); + } else { + panic!("Expected InvalidRootCertPath error"); + } + } + + #[test] + fn test_build_tls_options_disabled() { + let result = build_tls_options("disabled", None); + + if let Some(Tls::Disabled) = result {} else { + panic!("Expected TLS to be disabled"); + } + } + + #[test] + fn test_build_tls_options_required_without_cert() { + let result = build_tls_options("required", None); + + if let Some(Tls::Enabled(_)) = result {} else { + panic!("Expected TLS to be enabled"); + } + } + + #[test] + fn test_build_tls_options_preferred_without_cert() { + let result = build_tls_options("preferred", None); + + if let Some(Tls::Enabled(_)) = result {} else { + panic!("Expected TLS to be enabled"); + } + } + + #[test] + fn test_build_tls_options_with_cert_path() { + let temp_dir = std::env::temp_dir(); + let cert_path = temp_dir.join("test_cert.pem"); + std::fs::write(&cert_path, "dummy cert content").unwrap(); + + let result = build_tls_options("required", Some(cert_path.clone())); + std::fs::remove_file(&cert_path).ok(); + + if let Some(Tls::Enabled(_)) = result {} else { + panic!("Expected TLS to be enabled with certificate"); + } + } + + #[test] + fn test_build_tls_options_preferred_with_cert_path() { + let temp_dir = std::env::temp_dir(); + let cert_path = temp_dir.join("test_cert_preferred.pem"); + std::fs::write(&cert_path, "dummy cert content").unwrap(); + + let result = build_tls_options("preferred", Some(cert_path.clone())); + std::fs::remove_file(&cert_path).ok(); + + if let Some(Tls::Enabled(_)) = result {} else { + panic!("Expected TLS to be enabled with certificate in preferred mode"); + } + } + + #[test] + fn test_configure_tls_case_insensitive_ssl_mode() { + let mut client_options = ClientOptions::default(); + let params = create_params(vec![ + ("sslmode", "REQUIRED"), + ]); + + let result = configure_tls(&mut client_options, ¶ms); + assert!(result.is_ok()); + + if let Some(Tls::Enabled(_)) = client_options.tls {} else { + panic!("Expected TLS to be enabled with uppercase mode"); + } + } + + #[test] + fn test_configure_tls_mixed_case_ssl_mode() { + let mut client_options = ClientOptions::default(); + let params = create_params(vec![ + ("sslmode", "Preferred"), + ]); + + let result = configure_tls(&mut client_options, ¶ms); + assert!(result.is_ok()); + + if let Some(Tls::Enabled(_)) = client_options.tls {} else { + panic!("Expected TLS to be enabled with mixed case mode"); + } + } + + #[test] + fn test_configure_tls_with_rootcert_param_triggers_override() { + let mut client_options = ClientOptions::default(); + client_options.tls = Some(Tls::Enabled(TlsOptions::builder().build())); + + let temp_dir = std::env::temp_dir(); + let cert_path = temp_dir.join("test_override_cert.pem"); + std::fs::write(&cert_path, "dummy cert content").unwrap(); + + let params = create_params(vec![ + ("sslrootcert", cert_path.to_str().unwrap()), + ]); + + let result = configure_tls(&mut client_options, ¶ms); + std::fs::remove_file(&cert_path).ok(); + + assert!(result.is_ok()); + + if let Some(Tls::Enabled(_)) = client_options.tls {} else { + panic!("Expected TLS to be enabled after override"); + } + } +} \ No newline at end of file diff --git a/src/mongodb/table.rs b/src/mongodb/table.rs new file mode 100644 index 00000000..07ac17db --- /dev/null +++ b/src/mongodb/table.rs @@ -0,0 +1,231 @@ +use crate::mongodb::connection_pool::MongoDBConnectionPool; +use crate::mongodb::utils::expression::{combine_exprs_with_and, expr_to_mongo_filter}; +use crate::mongodb::Error; +use async_trait::async_trait; +use datafusion::common::project_schema; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::catalog::{Session, TableProvider}; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::{TaskContext}; +use datafusion::logical_expr::{Expr, TableType}; +use datafusion::physical_expr::EquivalenceProperties; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, SendableRecordBatchStream +}; +use datafusion::sql::TableReference; +use mongodb::bson::Document; +use std::{any::Any, fmt, sync::Arc}; +use serde_json; +use futures::TryStreamExt; + +#[derive(Debug)] +pub struct MongoDBTable { + pool: Arc, + schema: SchemaRef, + table_reference: Arc, +} + +impl MongoDBTable { + pub async fn new( + pool: &Arc, + table_reference: impl Into, + ) -> Result { + + let table_reference = table_reference.into(); + let schema= pool + .connect() + .await? + .get_schema(&table_reference) + .await?; + + Ok(Self { + pool: Arc::clone(pool), + schema, + table_reference: Arc::new(table_reference), + }) + } + + +} + +#[async_trait] +impl TableProvider for MongoDBTable { + fn as_any(&self) -> &dyn Any { + self + } + + 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(MongoDBExec::new( + Arc::clone(&self.table_reference), + Arc::clone(&self.pool), + Arc::clone(&self.schema), + projection, + filters, + limit, + )?)) + } +} + +#[derive(Debug)] +struct MongoDBExec { + table_reference: Arc, + pool: Arc, + projected_schema: SchemaRef, + filters_doc: Document, + limit: Option, + properties: PlanProperties, +} + + +impl MongoDBExec { + pub fn new( + table_reference: Arc, + pool: Arc, + schema: SchemaRef, + projections: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> DataFusionResult { + + let projected_schema = project_schema(&schema, projections)?; + let limit = limit + .map(|u| { + let Ok(u) = u32::try_from(u) else { + return Err(DataFusionError::Execution( + "Value is too large to fit in a u32".to_string(), + )); + }; + if let Ok(u) = i32::try_from(u) { + Ok(u) + } else { + Err(DataFusionError::Execution( + "Value is too large to fit in an i32".to_string(), + )) + } + }) + .transpose()?; + + let combined_exprs = combine_exprs_with_and(filters); + + let mongo_filters_doc = match combined_exprs { + Some(e) => expr_to_mongo_filter(&e) + .ok_or(DataFusionError::Execution("Failed to convert expressions".to_string()))?, + None => Document::new(), + }; + + Ok(Self { + table_reference: Arc::clone(&table_reference), + pool, + projected_schema: Arc::clone(&projected_schema), + filters_doc: mongo_filters_doc, + limit, + properties: PlanProperties::new( + EquivalenceProperties::new(projected_schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + ), + }) + } +} + +impl DisplayAs for MongoDBExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> std::fmt::Result { + let columns = self + .projected_schema + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect::>(); + + let filters = serde_json::to_string(&self.filters_doc) + .map_err(|_| fmt::Error)?; + + write!( + f, + "MongoDBExec projection=[{}] filters=[{}]", + columns.join(", "), + filters, + ) + } +} + +impl ExecutionPlan for MongoDBExec { + fn name(&self) -> &'static str { + "MongoDBExec" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.projected_schema) + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DataFusionResult> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> DataFusionResult { + let schema = self.schema(); + + let table_reference = Arc::clone(&self.table_reference); + let pool = Arc::clone(&self.pool); + let projected_schema = Arc::clone(&self.projected_schema); + let filters_doc = self.filters_doc.clone(); + let limit = self.limit; + + let stream = futures::stream::once(async move { + let conn = pool + .connect() + .await + .map_err(to_execution_error)?; + + conn.query_arrow(&table_reference, &projected_schema, &filters_doc, limit) + .await + .map_err(to_execution_error) + }) + .try_flatten(); + + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + + +} + +#[allow(clippy::needless_pass_by_value)] +pub fn to_execution_error(e: impl Into>) -> DataFusionError { + DataFusionError::Execution(format!("{}", e.into()).to_string()) +} diff --git a/src/mongodb/utils/arrow.rs b/src/mongodb/utils/arrow.rs new file mode 100644 index 00000000..e84d84cd --- /dev/null +++ b/src/mongodb/utils/arrow.rs @@ -0,0 +1,1355 @@ +use std::sync::Arc; +use std::collections::HashMap; +use std::str::FromStr; +use arrow::array::{ + ArrayRef, BooleanBuilder, Float64Builder, Int32Builder, Int64Builder, + StringBuilder, TimestampMillisecondBuilder, BinaryBuilder, ListBuilder, + NullBuilder, Decimal128Builder, RecordBatch +}; +use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit}; +use mongodb::bson::{Bson, Document}; +use rust_decimal::Decimal; +use snafu::prelude::*; +use num_traits::ToPrimitive; +use crate::mongodb::{Error, InvalidDecimalSnafu, Result}; + + +pub fn mongo_docs_to_arrow( + docs: &[Document], + projected_schema: SchemaRef, +) -> Result { + if docs.is_empty() { + // Return empty batch with correct schema + let empty_arrays: Vec = projected_schema + .fields() + .iter() + .map(|field| create_empty_array(field.data_type())) + .collect(); + + return RecordBatch::try_new(projected_schema, empty_arrays) + .map_err(|e| Error::ConversionError { + source: Box::new(e) + }); + } + + let mut builders = create_builders(&projected_schema, docs.len())?; + + for doc in docs { + append_document_to_builders(doc, &projected_schema, &mut builders)?; + } + + let arrays = finish_builders(builders, &projected_schema)?; + + RecordBatch::try_new(projected_schema, arrays) + .map_err(|e| Error::ConversionError { + source: Box::new(e) + }) +} + +fn create_empty_array(data_type: &DataType) -> ArrayRef { + match data_type { + DataType::Boolean => Arc::new(BooleanBuilder::new().finish()), + DataType::Int32 => Arc::new(Int32Builder::new().finish()), + DataType::Int64 => Arc::new(Int64Builder::new().finish()), + DataType::Float64 => Arc::new(Float64Builder::new().finish()), + DataType::Utf8 => Arc::new(StringBuilder::new().finish()), + DataType::Binary => Arc::new(BinaryBuilder::new().finish()), + DataType::Timestamp(TimeUnit::Millisecond, None) => { + Arc::new(TimestampMillisecondBuilder::new().finish()) + } + DataType::Decimal128(_, _) => { + Arc::new(Decimal128Builder::new().finish()) + } + DataType::List(_) => { + let values_builder = StringBuilder::new(); + Arc::new(ListBuilder::new(values_builder).finish()) + } + DataType::Null => Arc::new(NullBuilder::new().finish()), + _ => { + // Fallback to string for unsupported types + Arc::new(StringBuilder::new().finish()) + } + } +} + +type BuilderMap = HashMap>; + +trait ArrayBuilderTrait { + fn append_bson(&mut self, value: Option<&Bson>) -> Result<(), Error>; + fn finish_builder(self: Box) -> Result; +} + +fn create_builders(schema: &SchemaRef, capacity: usize) -> Result { + let mut builders: BuilderMap = HashMap::new(); + + for field in schema.fields() { + let builder: Box = match field.data_type() { + DataType::Boolean => Box::new(BooleanArrayBuilder::new(capacity)), + DataType::Int32 => Box::new(Int32ArrayBuilder::new(capacity)), + DataType::Int64 => Box::new(Int64ArrayBuilder::new(capacity)), + DataType::Float64 => Box::new(Float64ArrayBuilder::new(capacity)), + DataType::Utf8 => Box::new(StringArrayBuilder::new(capacity)), + DataType::Binary => Box::new(BinaryArrayBuilder::new(capacity)), + DataType::Timestamp(TimeUnit::Millisecond, None) => { + Box::new(TimestampArrayBuilder::new(capacity)) + } + DataType::Decimal128(precision, scale) => { + Box::new(Decimal128ArrayBuilder::new(capacity, *precision, *scale)?) + } + DataType::List(_) => Box::new(ListArrayBuilder::new(capacity)), + DataType::Null => Box::new(NullArrayBuilder::new()), + _ => { + // Fallback to string for unsupported types + Box::new(StringArrayBuilder::new(capacity)) + } + }; + + builders.insert(field.name().clone(), builder); + } + + Ok(builders) +} + +fn append_document_to_builders( + doc: &Document, + schema: &SchemaRef, + builders: &mut BuilderMap, +) -> Result<(), Error> { + for field in schema.fields() { + let field_name = field.name(); + let value = doc.get(field_name); + + if let Some(builder) = builders.get_mut(field_name) { + builder.append_bson(value)?; + } + } + Ok(()) +} + +fn finish_builders( + mut builders: BuilderMap, + schema: &SchemaRef, +) -> Result, Error> { + let mut arrays = Vec::new(); + + for field in schema.fields() { + let field_name = field.name(); + if let Some(builder) = builders.remove(field_name) { + arrays.push(builder.finish_builder()?); + } else { + return Err(Error::ConversionError { + source: Box::new(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Missing builder for field: {}", field_name) + )) + }); + } + } + + Ok(arrays) +} + +struct BooleanArrayBuilder(BooleanBuilder); +struct Int32ArrayBuilder(Int32Builder); +struct Int64ArrayBuilder(Int64Builder); +struct Float64ArrayBuilder(Float64Builder); +struct StringArrayBuilder(StringBuilder); +struct BinaryArrayBuilder(BinaryBuilder); +struct TimestampArrayBuilder(TimestampMillisecondBuilder); +pub struct Decimal128ArrayBuilder { + builder: Decimal128Builder, + scale: i8, +} +struct ListArrayBuilder(ListBuilder); +struct NullArrayBuilder(NullBuilder); + +impl BooleanArrayBuilder { + fn new(capacity: usize) -> Self { + Self(BooleanBuilder::with_capacity(capacity)) + } +} + +impl ArrayBuilderTrait for BooleanArrayBuilder { + fn append_bson(&mut self, value: Option<&Bson>) -> Result<(), Error> { + match value { + Some(Bson::Boolean(b)) => self.0.append_value(*b), + Some(_) => self.0.append_null(), + None => self.0.append_null(), + } + Ok(()) + } + + fn finish_builder(mut self: Box) -> Result { + Ok(Arc::new(self.0.finish())) + } +} + +impl Int32ArrayBuilder { + fn new(capacity: usize) -> Self { + Self(Int32Builder::with_capacity(capacity)) + } +} + +impl ArrayBuilderTrait for Int32ArrayBuilder { + fn append_bson(&mut self, value: Option<&Bson>) -> Result<(), Error> { + match value { + Some(Bson::Int32(i)) => self.0.append_value(*i), + Some(Bson::Int64(i)) if *i >= i32::MIN as i64 && *i <= i32::MAX as i64 => { + self.0.append_value(*i as i32) + } + Some(_) => self.0.append_null(), + None => self.0.append_null(), + } + Ok(()) + } + + fn finish_builder(mut self: Box) -> Result { + Ok(Arc::new(self.0.finish())) + } +} + +impl Int64ArrayBuilder { + fn new(capacity: usize) -> Self { + Self(Int64Builder::with_capacity(capacity)) + } +} + +impl ArrayBuilderTrait for Int64ArrayBuilder { + fn append_bson(&mut self, value: Option<&Bson>) -> Result<(), Error> { + match value { + Some(Bson::Int32(i)) => self.0.append_value(*i as i64), + Some(Bson::Int64(i)) => self.0.append_value(*i), + Some(_) => self.0.append_null(), + None => self.0.append_null(), + } + Ok(()) + } + + fn finish_builder(mut self: Box) -> Result { + Ok(Arc::new(self.0.finish())) + } +} + +impl Float64ArrayBuilder { + fn new(capacity: usize) -> Self { + Self(Float64Builder::with_capacity(capacity)) + } +} + +impl ArrayBuilderTrait for Float64ArrayBuilder { + fn append_bson(&mut self, value: Option<&Bson>) -> Result<(), Error> { + match value { + Some(Bson::Double(d)) => self.0.append_value(*d), + Some(Bson::Int32(i)) => self.0.append_value(*i as f64), + Some(Bson::Int64(i)) => self.0.append_value(*i as f64), + Some(_) => self.0.append_null(), + None => self.0.append_null(), + } + Ok(()) + } + + fn finish_builder(mut self: Box) -> Result { + Ok(Arc::new(self.0.finish())) + } +} + +impl StringArrayBuilder { + fn new(capacity: usize) -> Self { + Self(StringBuilder::with_capacity(capacity, 1024)) + } +} + +impl ArrayBuilderTrait for StringArrayBuilder { + fn append_bson(&mut self, value: Option<&Bson>) -> Result<(), Error> { + match value { + Some(Bson::String(s)) => self.0.append_value(s), + Some(Bson::ObjectId(oid)) => self.0.append_value(oid.to_hex()), + Some(Bson::Document(doc)) => { + // Convert document to JSON string. Maybe later add support for nested documents + let json_str = serde_json::to_string(doc) + .map_err(|e| Error::ConversionError { source: Box::new(e) })?; + self.0.append_value(&json_str); + } + Some(Bson::Null) => self.0.append_null(), + Some(other) => { + self.0.append_value(format!("{}", other)); + } + None => self.0.append_null(), + } + Ok(()) + } + + fn finish_builder(mut self: Box) -> Result { + Ok(Arc::new(self.0.finish())) + } +} + +impl BinaryArrayBuilder { + fn new(capacity: usize) -> Self { + Self(BinaryBuilder::with_capacity(capacity, 1024)) + } +} + +impl ArrayBuilderTrait for BinaryArrayBuilder { + fn append_bson(&mut self, value: Option<&Bson>) -> Result<(), Error> { + match value { + Some(Bson::Binary(binary)) => self.0.append_value(&binary.bytes), + Some(_) => self.0.append_null(), + None => self.0.append_null(), + } + Ok(()) + } + + fn finish_builder(mut self: Box) -> Result { + Ok(Arc::new(self.0.finish())) + } +} + +impl TimestampArrayBuilder { + fn new(capacity: usize) -> Self { + Self(TimestampMillisecondBuilder::with_capacity(capacity)) + } +} + +impl ArrayBuilderTrait for TimestampArrayBuilder { + fn append_bson(&mut self, value: Option<&Bson>) -> Result<(), Error> { + match value { + Some(Bson::DateTime(dt)) => { + self.0.append_value(dt.timestamp_millis()) + } + Some(Bson::Timestamp(ts)) => { + self.0.append_value((ts.time as i64) * 1000) + } + Some(_) => self.0.append_null(), + None => self.0.append_null(), + } + Ok(()) + } + + fn finish_builder(mut self: Box) -> Result { + Ok(Arc::new(self.0.finish())) + } +} + +impl Decimal128ArrayBuilder { + fn new(capacity: usize, precision: u8, scale: i8) -> Result { + let builder = Decimal128Builder::with_capacity(capacity) + .with_precision_and_scale(precision, scale) + .context(InvalidDecimalSnafu)?; + Ok(Self { builder, scale } ) + } +} + +impl ArrayBuilderTrait for Decimal128ArrayBuilder { + fn append_bson(&mut self, value: Option<&Bson>) -> Result<(), Error> { + match value { + Some(Bson::Decimal128(decimal)) => { + let parsed_decimal = rust_decimal::Decimal::from_str(&decimal.to_string()) + .map_err(|e| Error::ConversionError { source: Box::new(e) })?; + + let scaling_factor: Decimal = if self.scale >= 0 { + ten_pow_decimal(self.scale as u32) + .map_err(|_| Error::ConversionError { + source: Box::new(std::io::Error::new( + std::io::ErrorKind::InvalidData,"overflow in scaling factor")) + })? + } else { + let abs_scale = (-(self.scale as i32)) as u32; + if abs_scale > 28 { + return Err(Error::ConversionError { + source: Box::new(std::io::Error::new( + std::io::ErrorKind::InvalidData,"Negative scale too large for rust_decimal")) + }); + } + rust_decimal::Decimal::new(1, abs_scale) + }; + + let scaled_decimal = parsed_decimal + .checked_mul(scaling_factor) + .ok_or_else(|| Error::ConversionError { + source: Box::new(std::io::Error::new( + std::io::ErrorKind::InvalidData,"overflow during decimal conversion")) + })?; + + let rounded_decimal = scaled_decimal.round(); + + let value = rounded_decimal + .to_i128() + .ok_or_else(|| Error::ConversionError { + source: Box::new(std::io::Error::new( + std::io::ErrorKind::InvalidData,"overflow during decimal conversion")) + })?; + + self.builder.append_value(value); + } + Some(_) => self.builder.append_null(), + None => self.builder.append_null(), + } + Ok(()) + } + + fn finish_builder(mut self: Box) -> Result { + Ok(Arc::new(self.builder.finish())) + } +} + +fn ten_pow_decimal(exp: u32) -> Result { + let mut result = Decimal::ONE; + for _ in 0..exp { + result = result.checked_mul(Decimal::TEN) + .ok_or_else(|| Error::ConversionError { + source: Box::new(std::io::Error::new( + std::io::ErrorKind::InvalidData,"Multiplication overflow during decimal conversion")) + })?; + } + Ok(result) +} + +impl ListArrayBuilder { + fn new(capacity: usize) -> Self { + let values_builder = StringBuilder::with_capacity(capacity * 4, 256); + Self(ListBuilder::new(values_builder)) + } +} + +impl ArrayBuilderTrait for ListArrayBuilder { + fn append_bson(&mut self, value: Option<&Bson>) -> Result<(), Error> { + match value { + Some(Bson::Array(arr)) => { + for item in arr { + match item { + Bson::String(s) => self.0.values().append_value(s), + other => self.0.values().append_value(format!("{}", other)), + } + } + self.0.append(true); + } + Some(_) => self.0.append_null(), + None => self.0.append_null(), + } + Ok(()) + } + + fn finish_builder(mut self: Box) -> Result { + Ok(Arc::new(self.0.finish())) + } +} + +impl NullArrayBuilder { + fn new() -> Self { + Self(NullBuilder::new()) + } +} + +impl ArrayBuilderTrait for NullArrayBuilder { + fn append_bson(&mut self, _value: Option<&Bson>) -> Result<(), Error> { + self.0.append_null(); + Ok(()) + } + + fn finish_builder(mut self: Box) -> Result { + Ok(Arc::new(self.0.finish())) + } +} + + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::*; + use arrow::datatypes::{Schema, Field, DataType, TimeUnit}; + use mongodb::bson::{doc, Bson, Document, oid::ObjectId, DateTime, Timestamp, Binary, spec::BinarySubtype}; + + #[test] + fn test_empty_documents() { + let docs: Vec = vec![]; + let schema = Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, true), + Field::new("age", DataType::Int32, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema.clone()).unwrap(); + + assert_eq!(result.num_rows(), 0); + assert_eq!(result.num_columns(), 2); + assert_eq!(result.schema(), schema); + } + + #[test] + fn test_single_document_basic_types() { + let doc = doc! { + "name": "Alice", + "age": 30_i32, + "height": 5.6_f64, + "is_active": true + }; + let docs = vec![doc]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, true), + Field::new("age", DataType::Int32, true), + Field::new("height", DataType::Float64, true), + Field::new("is_active", DataType::Boolean, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + assert_eq!(result.num_rows(), 1); + assert_eq!(result.num_columns(), 4); + + // Check string value + let name_array = result.column_by_name("name").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(name_array.value(0), "Alice"); + + // Check int32 value + let age_array = result.column_by_name("age").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(age_array.value(0), 30); + + // Check float64 value + let height_array = result.column_by_name("height").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(height_array.value(0), 5.6); + + // Check boolean value + let active_array = result.column_by_name("is_active").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(active_array.value(0), true); + } + + #[test] + fn test_multiple_documents() { + let docs = vec![ + doc! { "name": "Alice", "age": 30_i32 }, + doc! { "name": "Bob", "age": 25_i32 }, + doc! { "name": "Charlie", "age": 35_i32 }, + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, true), + Field::new("age", DataType::Int32, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + assert_eq!(result.num_rows(), 3); + + let name_array = result.column_by_name("name").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(name_array.value(0), "Alice"); + assert_eq!(name_array.value(1), "Bob"); + assert_eq!(name_array.value(2), "Charlie"); + + let age_array = result.column_by_name("age").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(age_array.value(0), 30); + assert_eq!(age_array.value(1), 25); + assert_eq!(age_array.value(2), 35); + } + + #[test] + fn test_missing_fields() { + let docs = vec![ + doc! { "name": "Alice", "age": 30_i32 }, + doc! { "name": "Bob" }, // Missing age + doc! { "age": 25_i32 }, // Missing name + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, true), + Field::new("age", DataType::Int32, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + assert_eq!(result.num_rows(), 3); + + let name_array = result.column_by_name("name").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(name_array.value(0), "Alice"); + assert_eq!(name_array.value(1), "Bob"); + assert!(name_array.is_null(2)); // Missing name + + let age_array = result.column_by_name("age").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(age_array.value(0), 30); + assert!(age_array.is_null(1)); // Missing age + assert_eq!(age_array.value(2), 25); + } + + #[test] + fn test_mongodb_specific_types() { + let test_oid = ObjectId::new(); + let test_datetime = DateTime::now(); + let test_timestamp = Timestamp { time: 1234567890, increment: 1 }; + let test_binary_data = vec![1, 2, 3]; + + let doc = doc! { + "id": test_oid, + "created_at": test_datetime, + "timestamp": test_timestamp, + "binary_data": Binary { + subtype: BinarySubtype::Generic, + bytes: test_binary_data.clone() + }, + }; + let docs = vec![doc]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, true), + Field::new("created_at", DataType::Timestamp(TimeUnit::Millisecond, None), true), + Field::new("timestamp", DataType::Timestamp(TimeUnit::Millisecond, None), true), + Field::new("binary_data", DataType::Binary, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + // Check ObjectId conversion + let id_array = result.column_by_name("id").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(id_array.value(0), test_oid.to_hex()); + + // Check DateTime conversion + let datetime_array = result.column_by_name("created_at").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(datetime_array.value(0), test_datetime.timestamp_millis()); + + // Check Timestamp conversion + let timestamp_array = result.column_by_name("timestamp").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(timestamp_array.value(0), (test_timestamp.time as i64) * 1000); + + // Check Binary conversion + let binary_array = result.column_by_name("binary_data").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(binary_array.value(0), test_binary_data); + } + + #[test] + fn test_numeric_type_coercion() { + let docs = vec![ + doc! { + "int32_to_int64": 100_i32, + "int32_to_float": 50_i32, + "int64_to_float": 75_i64 + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("int32_to_int64", DataType::Int64, true), + Field::new("int32_to_float", DataType::Float64, true), + Field::new("int64_to_float", DataType::Float64, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + // Int32 -> Int64 + let int64_array = result.column_by_name("int32_to_int64").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(int64_array.value(0), 100_i64); + + // Int32 -> Float64 + let float_array1 = result.column_by_name("int32_to_float").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(float_array1.value(0), 50.0); + + // Int64 -> Float64 + let float_array2 = result.column_by_name("int64_to_float").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(float_array2.value(0), 75.0); + } + + #[test] + fn test_array_conversion() { + let docs = vec![ + doc! { + "string_array": ["a", "b", "c"], + "mixed_array": ["text", 42_i32, true], + "empty_array": [] + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("string_array", DataType::List( + Arc::new(Field::new("item", DataType::Utf8, true)) + ), true), + Field::new("mixed_array", DataType::List( + Arc::new(Field::new("item", DataType::Utf8, true)) + ), true), + Field::new("empty_array", DataType::List( + Arc::new(Field::new("item", DataType::Utf8, true)) + ), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + // Check string array + let string_list = result.column_by_name("string_array").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(string_list.len(), 1); + + let string_array_ref = string_list.value(0); + let string_values = string_array_ref + .as_any().downcast_ref::().unwrap(); + assert_eq!(string_values.len(), 3); + assert_eq!(string_values.value(0), "a"); + assert_eq!(string_values.value(1), "b"); + assert_eq!(string_values.value(2), "c"); + + // Check mixed array (all converted to strings) + let mixed_list = result.column_by_name("mixed_array").unwrap() + .as_any().downcast_ref::().unwrap(); + let mixed_array_ref = mixed_list.value(0); + let mixed_values = mixed_array_ref + .as_any().downcast_ref::().unwrap(); + assert_eq!(mixed_values.len(), 3); + assert_eq!(mixed_values.value(0), "text"); + assert_eq!(mixed_values.value(1), "42"); + assert_eq!(mixed_values.value(2), "true"); + + // Check empty array + let empty_list = result.column_by_name("empty_array").unwrap() + .as_any().downcast_ref::().unwrap(); + let empty_array_ref = empty_list.value(0); + let empty_values = empty_array_ref + .as_any().downcast_ref::().unwrap(); + assert_eq!(empty_values.len(), 0); + } + + #[test] + fn test_nested_document_conversion() { + let docs = vec![ + doc! { + "user": { + "name": "Alice", + "age": 30_i32 + }, + "metadata": {} + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("user", DataType::Utf8, true), + Field::new("metadata", DataType::Utf8, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let user_array = result.column_by_name("user").unwrap() + .as_any().downcast_ref::().unwrap(); + let user_json = user_array.value(0); + + // Should be valid JSON + let parsed: serde_json::Value = serde_json::from_str(user_json).unwrap(); + assert_eq!(parsed["name"], "Alice"); + assert_eq!(parsed["age"], 30); + + let metadata_array = result.column_by_name("metadata").unwrap() + .as_any().downcast_ref::().unwrap(); + let metadata_json = metadata_array.value(0); + assert_eq!(metadata_json, "{}"); + } + + #[test] + fn test_null_values() { + let docs = vec![ + doc! { + "nullable_string": Bson::Null, + "nullable_int": Bson::Null, + "nullable_bool": Bson::Null, + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("nullable_string", DataType::Utf8, true), + Field::new("nullable_int", DataType::Int32, true), + Field::new("nullable_bool", DataType::Boolean, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let string_array = result.column_by_name("nullable_string").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(string_array.len(), 1); + + let int_array = result.column_by_name("nullable_int").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(int_array.len(), 1); + + let bool_array = result.column_by_name("nullable_bool").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(bool_array.len(), 1); + } + + #[test] + fn test_null_array_type() { + let docs = vec![ + doc! { "null_field": Bson::Null } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("null_field", DataType::Null, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let null_array = result.column_by_name("null_field").unwrap() + .as_any().downcast_ref::().unwrap(); + + assert_eq!(null_array.len(), 1); + } + + #[test] + fn test_type_mismatch_fallback() { + let docs = vec![ + doc! { + "wrong_type_string": 42_i32, // Int32 in string field + "wrong_type_int": "not_a_number", // String in int field + "wrong_type_bool": 3.14_f64, // Float in bool field + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("wrong_type_string", DataType::Utf8, true), + Field::new("wrong_type_int", DataType::Int32, true), + Field::new("wrong_type_bool", DataType::Boolean, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + // String builder should convert int to string + let string_array = result.column_by_name("wrong_type_string").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(string_array.value(0), "42"); + + // Int builder should null out non-int values + let int_array = result.column_by_name("wrong_type_int").unwrap() + .as_any().downcast_ref::().unwrap(); + assert!(int_array.is_null(0)); + + // Bool builder should null out non-bool values + let bool_array = result.column_by_name("wrong_type_bool").unwrap() + .as_any().downcast_ref::().unwrap(); + assert!(bool_array.is_null(0)); + } + + #[test] + fn test_extreme_values() { + let docs = vec![ + doc! { + "max_int32": i32::MAX, + "min_int32": i32::MIN, + "max_int64": i64::MAX, + "min_int64": i64::MIN, + "infinity": f64::INFINITY, + "neg_infinity": f64::NEG_INFINITY, + "nan": f64::NAN, + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("max_int32", DataType::Int32, true), + Field::new("min_int32", DataType::Int32, true), + Field::new("max_int64", DataType::Int64, true), + Field::new("min_int64", DataType::Int64, true), + Field::new("infinity", DataType::Float64, true), + Field::new("neg_infinity", DataType::Float64, true), + Field::new("nan", DataType::Float64, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + // Check extreme integers + let max_int32_array = result.column_by_name("max_int32").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(max_int32_array.value(0), i32::MAX); + + let min_int64_array = result.column_by_name("min_int64").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(min_int64_array.value(0), i64::MIN); + + // Check special float values + let inf_array = result.column_by_name("infinity").unwrap() + .as_any().downcast_ref::().unwrap(); + assert!(inf_array.value(0).is_infinite()); + assert!(inf_array.value(0).is_sign_positive()); + + let nan_array = result.column_by_name("nan").unwrap() + .as_any().downcast_ref::().unwrap(); + assert!(nan_array.value(0).is_nan()); + } + + #[test] + fn test_large_binary_data() { + let large_data = vec![0u8; 10000]; // 10KB of zeros + let docs = vec![ + doc! { + "large_binary": Binary { + subtype: BinarySubtype::Generic, + bytes: large_data.clone() + } + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("large_binary", DataType::Binary, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let binary_array = result.column_by_name("large_binary").unwrap() + .as_any().downcast_ref::().unwrap(); + let retrieved_data = binary_array.value(0); + + assert_eq!(retrieved_data.len(), 10000); + assert_eq!(retrieved_data, large_data); + } + + #[test] + fn test_unicode_strings() { + let docs = vec![ + doc! { + "unicode": "Hello 世界 🌍 مرحبا Здравствуй", + "emoji": "🚀🎉💯", + "complex": "𝕳𝖊𝖑𝖑𝖔", + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("unicode", DataType::Utf8, true), + Field::new("emoji", DataType::Utf8, true), + Field::new("complex", DataType::Utf8, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let unicode_array = result.column_by_name("unicode").unwrap() + .as_any().downcast_ref::().unwrap(); + let unicode_value = unicode_array.value(0); + assert!(unicode_value.contains("世界")); + assert!(unicode_value.contains("🌍")); + assert!(unicode_value.contains("مرحبا")); + + let emoji_array = result.column_by_name("emoji").unwrap() + .as_any().downcast_ref::().unwrap(); + assert_eq!(emoji_array.value(0), "🚀🎉💯"); + } + + #[test] + fn test_bson_null_string_is_real_null() { + use mongodb::bson::{doc, Bson}; + use arrow::datatypes::{Field, Schema, DataType}; + use arrow::array::StringArray; + + let docs = vec![ + doc! { + "name": Bson::Null, + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let name_array = result + .column_by_name("name") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(name_array.len(), 1); + assert!(name_array.is_null(0), "Expected Arrow null, got {:?}", name_array.value(0)); + } + + #[test] + fn test_schema_field_order_preservation() { + let docs = vec![ + doc! { + "z_field": "last", + "a_field": "first", + "m_field": "middle", + } + ]; + + // Schema with specific field order + let schema = Arc::new(Schema::new(vec![ + Field::new("a_field", DataType::Utf8, true), + Field::new("m_field", DataType::Utf8, true), + Field::new("z_field", DataType::Utf8, true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema.clone()).unwrap(); + + // Verify field order matches schema order + assert_eq!(result.schema(), schema); + + // Verify data is in correct positions + let a_array = result.column(0).as_any().downcast_ref::().unwrap(); + assert_eq!(a_array.value(0), "first"); + + let m_array = result.column(1).as_any().downcast_ref::().unwrap(); + assert_eq!(m_array.value(0), "middle"); + + let z_array = result.column(2).as_any().downcast_ref::().unwrap(); + assert_eq!(z_array.value(0), "last"); + } +} + +#[cfg(test)] +mod decimal_tests { + use super::*; + use arrow::array::*; + use arrow::datatypes::{Schema, Field, DataType}; + use mongodb::bson::{doc, Decimal128 as BsonDecimal128}; + use std::str::FromStr; + + #[test] + fn test_decimal_basic_conversion() { + let docs = vec![ + doc! { + "price": BsonDecimal128::from_str("123.45").unwrap(), + "tax": BsonDecimal128::from_str("9.99").unwrap(), + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("price", DataType::Decimal128(10, 2), true), + Field::new("tax", DataType::Decimal128(10, 2), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let price_array = result.column_by_name("price").unwrap() + .as_any().downcast_ref::().unwrap(); + let tax_array = result.column_by_name("tax").unwrap() + .as_any().downcast_ref::().unwrap(); + + assert_eq!(price_array.value(0), 12345); + assert_eq!(tax_array.value(0), 999); + } + + #[test] + fn test_decimal_zero_scale() { + let docs = vec![ + doc! { + "whole_number": BsonDecimal128::from_str("123").unwrap(), + "decimal_truncated": BsonDecimal128::from_str("123.99").unwrap(), + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("whole_number", DataType::Decimal128(10, 0), true), + Field::new("decimal_truncated", DataType::Decimal128(10, 0), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let whole_array = result.column_by_name("whole_number").unwrap() + .as_any().downcast_ref::().unwrap(); + let truncated_array = result.column_by_name("decimal_truncated").unwrap() + .as_any().downcast_ref::().unwrap(); + + assert_eq!(whole_array.value(0), 123); + assert_eq!(truncated_array.value(0), 124); + } + + #[test] + fn test_decimal_negative_scale() { + let docs = vec![ + doc! { + "large_number": BsonDecimal128::from_str("123456").unwrap(), + "scientific": BsonDecimal128::from_str("1.23456").unwrap(), + } + ]; + + // Negative scale means division by power of 10 + let schema = Arc::new(Schema::new(vec![ + Field::new("large_number", DataType::Decimal128(10, -2), true), + Field::new("scientific", DataType::Decimal128(10, -4), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let large_array = result.column_by_name("large_number").unwrap() + .as_any().downcast_ref::().unwrap(); + let scientific_array = result.column_by_name("scientific").unwrap() + .as_any().downcast_ref::().unwrap(); + + // 123456 with scale -2 = 123456 / 100 = 1234.56 rounded = 1235 + assert_eq!(large_array.value(0), 1235); + // 1.23456 with scale -4 = 1.23456 / 10000 = 0.000123456 rounded = 0 + assert_eq!(scientific_array.value(0), 0); + } + + #[test] + fn test_decimal_rounding() { + let docs = vec![ + doc! { + "round_up": BsonDecimal128::from_str("123.456").unwrap(), + "round_down": BsonDecimal128::from_str("123.454").unwrap(), + "round_half": BsonDecimal128::from_str("123.455").unwrap(), + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("round_up", DataType::Decimal128(10, 2), true), + Field::new("round_down", DataType::Decimal128(10, 2), true), + Field::new("round_half", DataType::Decimal128(10, 2), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let up_array = result.column_by_name("round_up").unwrap() + .as_any().downcast_ref::().unwrap(); + let down_array = result.column_by_name("round_down").unwrap() + .as_any().downcast_ref::().unwrap(); + let half_array = result.column_by_name("round_half").unwrap() + .as_any().downcast_ref::().unwrap(); + + // 123.456 rounded to 2 decimals = 123.46 = 12346 + assert_eq!(up_array.value(0), 12346); + // 123.454 rounded to 2 decimals = 123.45 = 12345 + assert_eq!(down_array.value(0), 12345); + // 123.455 rounded to 2 decimals = 123.46 = 12346 (banker's rounding may vary) + assert!(half_array.value(0) == 12345 || half_array.value(0) == 12346); + } + + #[test] + fn test_decimal_negative_numbers() { + let docs = vec![ + doc! { + "negative": BsonDecimal128::from_str("-123.45").unwrap(), + "negative_zero": BsonDecimal128::from_str("-0.00").unwrap(), + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("negative", DataType::Decimal128(10, 2), true), + Field::new("negative_zero", DataType::Decimal128(10, 2), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let neg_array = result.column_by_name("negative").unwrap() + .as_any().downcast_ref::().unwrap(); + let neg_zero_array = result.column_by_name("negative_zero").unwrap() + .as_any().downcast_ref::().unwrap(); + + // -123.45 with scale 2 = -12345 + assert_eq!(neg_array.value(0), -12345); + // -0.00 with scale 2 = 0 + assert_eq!(neg_zero_array.value(0), 0); + } + + #[test] + fn test_decimal_large_numbers() { + let docs = vec![ + doc! { + "billion": BsonDecimal128::from_str("1000000000.00").unwrap(), + "trillion": BsonDecimal128::from_str("1000000000000.00").unwrap(), + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("billion", DataType::Decimal128(15, 2), true), + Field::new("trillion", DataType::Decimal128(20, 2), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let billion_array = result.column_by_name("billion").unwrap() + .as_any().downcast_ref::().unwrap(); + let trillion_array = result.column_by_name("trillion").unwrap() + .as_any().downcast_ref::().unwrap(); + + // 1000000000.00 with scale 2 = 100000000000 + assert_eq!(billion_array.value(0), 100000000000); + // 1000000000000.00 with scale 2 = 100000000000000 + assert_eq!(trillion_array.value(0), 100000000000000); + } + + #[test] + fn test_decimal_null_values() { + let docs = vec![ + doc! { + "decimal_null": Bson::Null, + "decimal_value": BsonDecimal128::from_str("123.45").unwrap(), + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("decimal_null", DataType::Decimal128(10, 2), true), + Field::new("decimal_value", DataType::Decimal128(10, 2), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let null_array = result.column_by_name("decimal_null").unwrap() + .as_any().downcast_ref::().unwrap(); + let value_array = result.column_by_name("decimal_value").unwrap() + .as_any().downcast_ref::().unwrap(); + + assert!(null_array.is_null(0)); + assert!(!value_array.is_null(0)); + assert_eq!(value_array.value(0), 12345); + } + + #[test] + fn test_decimal_wrong_type_fallback() { + let docs = vec![ + doc! { + "not_decimal": "not a decimal", + "int_as_decimal": 42_i32, + "float_as_decimal": 3.14_f64, + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("not_decimal", DataType::Decimal128(10, 2), true), + Field::new("int_as_decimal", DataType::Decimal128(10, 2), true), + Field::new("float_as_decimal", DataType::Decimal128(10, 2), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let not_decimal_array = result.column_by_name("not_decimal").unwrap() + .as_any().downcast_ref::().unwrap(); + let int_array = result.column_by_name("int_as_decimal").unwrap() + .as_any().downcast_ref::().unwrap(); + let float_array = result.column_by_name("float_as_decimal").unwrap() + .as_any().downcast_ref::().unwrap(); + + // Non-decimal types should result in null values + assert!(not_decimal_array.is_null(0)); + assert!(int_array.is_null(0)); + assert!(float_array.is_null(0)); + } + + #[test] + fn test_decimal_scale_edge_cases() { + // Test maximum and minimum practical scales + let docs = vec![ + doc! { + "max_scale": BsonDecimal128::from_str("1.23456789012345678901234567").unwrap(), + "min_scale": BsonDecimal128::from_str("12345678901234567890").unwrap(), + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("max_scale", DataType::Decimal128(38, 28), true), + Field::new("min_scale", DataType::Decimal128(38, -10), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema); + + // This should succeed for max scale within rust_decimal limits + assert!(result.is_ok()); + } + + #[test] + fn test_decimal_precision_loss() { + let docs = vec![ + doc! { + "high_precision": BsonDecimal128::from_str("123.123456789").unwrap(), + } + ]; + + // Schema with lower precision than the input + let schema = Arc::new(Schema::new(vec![ + Field::new("high_precision", DataType::Decimal128(10, 4), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let array = result.column_by_name("high_precision").unwrap() + .as_any().downcast_ref::().unwrap(); + + // 123.123456789 rounded to 4 decimal places = 123.1235 = 1231235 + assert_eq!(array.value(0), 1231235); + } + + #[test] + fn test_decimal_scientific_notation() { + let docs = vec![ + doc! { + "scientific": BsonDecimal128::from_str("1.23E+2").unwrap(), // 123 + "small_scientific": BsonDecimal128::from_str("1.23E-2").unwrap(), // 0.0123 + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("scientific", DataType::Decimal128(10, 2), true), + Field::new("small_scientific", DataType::Decimal128(10, 4), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let sci_array = result.column_by_name("scientific").unwrap() + .as_any().downcast_ref::().unwrap(); + let small_array = result.column_by_name("small_scientific").unwrap() + .as_any().downcast_ref::().unwrap(); + + // 123.00 with scale 2 = 12300 + assert_eq!(sci_array.value(0), 12300); + // 0.0123 with scale 4 = 123 + assert_eq!(small_array.value(0), 123); + } + + #[test] + fn test_decimal_multiple_documents() { + let docs = vec![ + doc! { "amount": BsonDecimal128::from_str("100.50").unwrap() }, + doc! { "amount": BsonDecimal128::from_str("200.75").unwrap() }, + doc! { "amount": BsonDecimal128::from_str("-50.25").unwrap() }, + doc! { "amount": Bson::Null }, + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("amount", DataType::Decimal128(10, 2), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema).unwrap(); + + let array = result.column_by_name("amount").unwrap() + .as_any().downcast_ref::().unwrap(); + + assert_eq!(array.len(), 4); + assert_eq!(array.value(0), 10050); // 100.50 + assert_eq!(array.value(1), 20075); // 200.75 + assert_eq!(array.value(2), -5025); // -50.25 + assert!(array.is_null(3)); // null + } + + #[test] + fn test_decimal_invalid_scale_too_negative() { + let docs = vec![ + doc! { + "invalid": BsonDecimal128::from_str("123.45").unwrap(), + } + ]; + + // Scale of -29 should be too large for rust_decimal (max is 28) + let schema = Arc::new(Schema::new(vec![ + Field::new("invalid", DataType::Decimal128(38, -29), true), + ])); + + let result = mongo_docs_to_arrow(&docs, schema); + + // Should return an error due to invalid scale + assert!(result.is_err()); + } + + #[test] + fn test_decimal_builder_creation_invalid_precision_scale() { + // Test the builder creation directly with invalid parameters + let result = Decimal128ArrayBuilder::new(10, 39, 0); // precision > 38 + assert!(result.is_err()); + + let result = Decimal128ArrayBuilder::new(10, 10, 11); // scale > precision + assert!(result.is_err()); + } +} \ No newline at end of file diff --git a/src/mongodb/utils/expression.rs b/src/mongodb/utils/expression.rs new file mode 100644 index 00000000..8daab615 --- /dev/null +++ b/src/mongodb/utils/expression.rs @@ -0,0 +1,374 @@ +use datafusion::{logical_expr::{Expr, Operator}, scalar::ScalarValue}; +use mongodb::bson::{doc, Bson, Document}; + +pub fn combine_exprs_with_and(exprs: &[Expr]) -> Option { + let mut iter = exprs.iter(); + let first = iter.next()?.clone(); + Some(iter.fold(first, |acc, e| acc.and(e.clone()))) +} + +pub fn expr_to_mongo_filter(expr: &Expr) -> Option { + match expr { + Expr::BinaryExpr(binary) => { + match binary.op { + Operator::And => { + let l = expr_to_mongo_filter(&binary.left)?; + let r = expr_to_mongo_filter(&binary.right)?; + Some(doc! { "$and": [l, r] }) + } + Operator::Or => { + let l = expr_to_mongo_filter(&binary.left)?; + let r = expr_to_mongo_filter(&binary.right)?; + Some(doc! { "$or": [l, r] }) + } + Operator::Eq => { + let field = extract_column_name(&binary.left); + let value = extract_literal_value(&binary.right); + let field = field?; + let value = value?; + Some(doc! { field: value }) + } + Operator::Gt => { + let field = extract_column_name(&binary.left); + let value = extract_literal_value(&binary.right); + let field = field?; + let value = value?; + Some(doc! { field: { "$gt": value } }) + } + Operator::Lt => { + let field = extract_column_name(&binary.left)?; + let value = extract_literal_value(&binary.right)?; + Some(doc! { field: { "$lt": value } }) + } + Operator::GtEq => { + let field = extract_column_name(&binary.left)?; + let value = extract_literal_value(&binary.right)?; + Some(doc! { field: { "$gte": value } }) + } + Operator::LtEq => { + let field = extract_column_name(&binary.left)?; + let value = extract_literal_value(&binary.right)?; + Some(doc! { field: { "$lte": value } }) + } + Operator::NotEq => { + let field = extract_column_name(&binary.left)?; + let value = extract_literal_value(&binary.right)?; + Some(doc! { field: { "$ne": value } }) + } + _ => { + println!("Unsupported operator: {:?}", binary.op); + None + } + } + } + _ => { + println!("Non-binary expr: {:?}", expr); + None + } + } +} + +fn extract_column_name(expr: &Expr) -> Option { + match expr { + Expr::Column(col) => Some(col.name.clone()), + _ => None, + } +} + +fn extract_literal_value(expr: &Expr) -> Option { + match expr { + Expr::Literal(scalar) => match scalar { + ScalarValue::Utf8(Some(s)) => Some(Bson::String(s.clone())), + ScalarValue::Utf8(None) => Some(Bson::Null), + ScalarValue::Int32(Some(i)) => Some(Bson::Int32(*i)), + ScalarValue::Int32(None) => Some(Bson::Null), + ScalarValue::Int64(Some(i)) => Some(Bson::Int64(*i)), + ScalarValue::Int64(None) => Some(Bson::Null), + ScalarValue::Float32(Some(f)) => Some(Bson::Double(*f as f64)), + ScalarValue::Float32(None) => Some(Bson::Null), + ScalarValue::Float64(Some(f)) => Some(Bson::Double(*f)), + ScalarValue::Float64(None) => Some(Bson::Null), + ScalarValue::Boolean(Some(b)) => Some(Bson::Boolean(*b)), + ScalarValue::Boolean(None) => Some(Bson::Null), + + ScalarValue::UInt8(Some(i)) => Some(Bson::Int32(*i as i32)), + ScalarValue::UInt16(Some(i)) => Some(Bson::Int32(*i as i32)), + ScalarValue::UInt32(Some(i)) => Some(Bson::Int64(*i as i64)), + ScalarValue::UInt64(Some(i)) => Some(Bson::Int64(*i as i64)), + ScalarValue::Int8(Some(i)) => Some(Bson::Int32(*i as i32)), + ScalarValue::Int16(Some(i)) => Some(Bson::Int32(*i as i32)), + + ScalarValue::UInt8(None) | ScalarValue::UInt16(None) | ScalarValue::UInt32(None) | + ScalarValue::UInt64(None) | ScalarValue::Int8(None) | ScalarValue::Int16(None) => Some(Bson::Null), + + _ => None, + }, + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::prelude::{col, lit}; + use datafusion::logical_expr::BinaryExpr; + + #[test] + fn test_combine_exprs_with_and() { + let exprs = vec![ + col("age").gt(lit(30)), + col("active").eq(lit(true)), + ]; + + let combined = combine_exprs_with_and(&exprs).unwrap(); + + // Should produce (age > 30) AND (active = true) + if let Expr::BinaryExpr(bin) = &combined { + assert_eq!(bin.op, Operator::And); + } else { + panic!("Expected BinaryExpr with AND operator"); + } + } + + #[test] + fn test_simple_eq_filter() { + let expr = col("name").eq(lit("Alice")); + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { "name": "Alice" }; + assert_eq!(filter, expected); + } + + #[test] + fn test_gt_and_eq_filter() { + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(col("age").gt(lit(21))), + op: Operator::And, + right: Box::new(col("status").eq(lit("active"))), + }); + + + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { + "$and": [ + { "age": { "$gt": 21 } }, + { "status": "active" } + ] + }; + + assert_eq!(filter, expected); + } + + #[test] + fn test_not_eq_filter() { + let expr = col("role").not_eq(lit("guest")); + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { "role": { "$ne": "guest" } }; + assert_eq!(filter, expected); + } + + #[test] + fn test_unsupported_expr_returns_none() { + let expr = col("salary").in_list(vec![lit(100), lit(200)], false); // unsupported for now + assert!(expr_to_mongo_filter(&expr).is_none()); + } + + #[test] + fn test_lt_filter() { + let expr = col("age").lt(lit(65)); + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { "age": { "$lt": 65 } }; + assert_eq!(filter, expected); + } + + #[test] + fn test_gte_filter() { + let expr = col("score").gt_eq(lit(85)); + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { "score": { "$gte": 85 } }; + assert_eq!(filter, expected); + } + + #[test] + fn test_lte_filter() { + let expr = col("temperature").lt_eq(lit(100)); + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { "temperature": { "$lte": 100 } }; + assert_eq!(filter, expected); + } + + #[test] + fn test_or_filter() { + let expr = col("department").eq(lit("sales")).or(col("department").eq(lit("marketing"))); + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { + "$or": [ + { "department": "sales" }, + { "department": "marketing" } + ] + }; + assert_eq!(filter, expected); + } + + #[test] + fn test_complex_and_or_filter() { + let age_and_status = col("age").gt(lit(25)).and(col("status").eq(lit("active"))); + let expr = age_and_status.or(col("priority").eq(lit("high"))); + + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { + "$or": [ + { + "$and": [ + { "age": { "$gt": 25 } }, + { "status": "active" } + ] + }, + { "priority": "high" } + ] + }; + assert_eq!(filter, expected); + } + + #[test] + fn test_nested_and_filters() { + let age_range = col("age").gt(lit(18)).and(col("age").lt(lit(65))); + let expr = age_range.and(col("country").eq(lit("US"))); + + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { + "$and": [ + { + "$and": [ + { "age": { "$gt": 18 } }, + { "age": { "$lt": 65 } } + ] + }, + { "country": "US" } + ] + }; + assert_eq!(filter, expected); + } + + #[test] + fn test_boolean_filter() { + let expr = col("is_verified").eq(lit(true)); + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { "is_verified": true }; + assert_eq!(filter, expected); + } + + #[test] + fn test_float_filter() { + let expr = col("price").gt(lit(99.99)); + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { "price": { "$gt": 99.99 } }; + assert_eq!(filter, expected); + } + + #[test] + fn test_string_comparison_filters() { + let expr = col("name").gt(lit("M")); + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { "name": { "$gt": "M" } }; + assert_eq!(filter, expected); + } + + #[test] + fn test_mixed_type_and_filter() { + let expr = col("age").gt(lit(21)).and(col("name").eq(lit("John"))); + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { + "$and": [ + { "age": { "$gt": 21 } }, + { "name": "John" } + ] + }; + assert_eq!(filter, expected); + } + + #[test] + fn test_single_expr_combine() { + let exprs = vec![col("status").eq(lit("active"))]; + let combined = combine_exprs_with_and(&exprs).unwrap(); + + if let Expr::BinaryExpr(bin) = &combined { + assert_eq!(bin.op, Operator::Eq); + } else { + panic!("Expected BinaryExpr with EQ operator"); + } + } + + #[test] + fn test_empty_exprs_combine() { + let exprs: Vec = vec![]; + let result = combine_exprs_with_and(&exprs); + assert!(result.is_none()); + } + + #[test] + fn test_three_exprs_combine() { + let exprs = vec![ + col("age").gt(lit(25)), + col("status").eq(lit("active")), + col("department").eq(lit("engineering")), + ]; + let combined = combine_exprs_with_and(&exprs).unwrap(); + + if let Expr::BinaryExpr(bin) = &combined { + assert_eq!(bin.op, Operator::And); + if let Expr::BinaryExpr(left_bin) = &*bin.left { + assert_eq!(left_bin.op, Operator::And); + } else { + panic!("Expected nested AND on left side"); + } + } else { + panic!("Expected BinaryExpr with AND operator"); + } + } + + #[test] + fn test_null_literal_filter() { + let expr = col("optional_field").eq(lit(ScalarValue::Utf8(None))); + let filter = expr_to_mongo_filter(&expr); + + if let Some(doc) = filter { + let expected = doc! { "optional_field": mongodb::bson::Bson::Null }; + assert_eq!(doc, expected); + } + } + + #[test] + fn test_wrong_operand_order_returns_none() { + use datafusion::logical_expr::Expr; + + let expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(lit("Alice")), + op: Operator::Eq, + right: Box::new(col("name")), + }); + + let filter = expr_to_mongo_filter(&expr); + assert!(filter.is_none(), "Should return None for unsupported operand order"); + } + + #[test] + fn test_multiple_or_conditions() { + let expr = col("status").eq(lit("active")) + .or(col("status").eq(lit("pending"))) + .or(col("status").eq(lit("review"))); + + let filter = expr_to_mongo_filter(&expr).unwrap(); + let expected = doc! { + "$or": [ + { + "$or": [ + { "status": "active" }, + { "status": "pending" } + ] + }, + { "status": "review" } + ] + }; + assert_eq!(filter, expected); + } +} diff --git a/src/mongodb/utils/mod.rs b/src/mongodb/utils/mod.rs new file mode 100644 index 00000000..ae1aa170 --- /dev/null +++ b/src/mongodb/utils/mod.rs @@ -0,0 +1,3 @@ +pub mod arrow; +pub mod expression; +pub mod schema; diff --git a/src/mongodb/utils/schema.rs b/src/mongodb/utils/schema.rs new file mode 100644 index 00000000..8b34d385 --- /dev/null +++ b/src/mongodb/utils/schema.rs @@ -0,0 +1,361 @@ +use std::sync::Arc; +use std::collections::HashMap; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use mongodb::bson::{Bson, Document}; + +use crate::mongodb::{Result, Error}; + +pub fn infer_arrow_schema_from_documents(docs: &[Document]) -> Result { + if docs.is_empty() { + return Ok(Arc::new(Schema::empty())); + } + + let mut field_types: HashMap = HashMap::new(); + + for doc in docs { + analyze_document(doc, &mut field_types); + } + + let fields: Vec = field_types + .into_iter() + .map(|(name, data_type)| Field::new(name, data_type, true)) + .collect(); + + Ok(Arc::new(Schema::new(fields))) +} + +fn analyze_document(doc: &Document, field_types: &mut HashMap) { + for (key, value) in doc { + let inferred_type = infer_bson_type(value); + + match field_types.get(key) { + Some(existing_type) => { + // Ue the most general type + let unified_type = unify_types(existing_type, &inferred_type); + field_types.insert(key.clone(), unified_type); + } + None => { + field_types.insert(key.clone(), inferred_type); + } + } + } +} + +fn infer_bson_type(value: &Bson) -> DataType { + match value { + Bson::Double(_) => DataType::Float64, + Bson::String(_) => DataType::Utf8, + Bson::Array(_) => { + // MongoDB arrays can be heterogeneous [1, "foo", true] + // Arrow arrays must be homogeneous - use strings to preserve all data + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))) + } + Bson::Document(_) => { + // Represent nested documents as JSON strings + // Maybe consider to recursively infer nested schemas in the future + DataType::Utf8 + } + Bson::Boolean(_) => DataType::Boolean, + Bson::Null => DataType::Null, + Bson::RegularExpression(_) => DataType::Utf8, + Bson::JavaScriptCode(_) => DataType::Utf8, + Bson::JavaScriptCodeWithScope(_) => DataType::Utf8, + Bson::Int32(_) => DataType::Int32, + Bson::Int64(_) => DataType::Int64, + Bson::Timestamp(_) => DataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + Bson::Binary(_) => DataType::Binary, + Bson::ObjectId(_) => DataType::Utf8, + Bson::DateTime(_) => DataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + Bson::Symbol(_) => DataType::Utf8, + Bson::Decimal128(_) => DataType::Decimal128(38, 10), + Bson::Undefined => DataType::Null, + Bson::MaxKey => DataType::Utf8, + Bson::MinKey => DataType::Utf8, + Bson::DbPointer(_) => DataType::Utf8, + } +} + +fn unify_types(type1: &DataType, type2: &DataType) -> DataType { + match (type1, type2) { + (a, b) if a == b => a.clone(), + (DataType::Null, other) | (other, DataType::Null) => other.clone(), + + // Numeric type promotion + (DataType::Int32, DataType::Int64) | (DataType::Int64, DataType::Int32) => DataType::Int64, + (DataType::Int32, DataType::Float64) | (DataType::Float64, DataType::Int32) => DataType::Float64, + (DataType::Int64, DataType::Float64) | (DataType::Float64, DataType::Int64) => DataType::Float64, + + // Otherwise use string + _ => DataType::Utf8, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use mongodb::bson::{doc, Bson, Document}; + use datafusion::arrow::datatypes::{DataType, TimeUnit}; + use std::str::FromStr; + + #[test] + fn test_empty_documents() { + let docs: Vec = vec![]; + let schema = infer_arrow_schema_from_documents(&docs).unwrap(); + assert_eq!(schema.fields().len(), 0); + } + + #[test] + fn test_single_document_simple_types() { + let doc = doc! { + "name": "Alice", + "age": 30_i32, + "height": 5.6_f64, + "is_active": true + }; + let docs = vec![doc]; + + let schema = infer_arrow_schema_from_documents(&docs).unwrap(); + + // Check field count + assert_eq!(schema.fields().len(), 4); + + // Check each field type (order may vary due to HashMap) + let field_map: HashMap = schema.fields() + .iter() + .map(|f| (f.name().clone(), f.data_type())) + .collect(); + + assert_eq!(field_map.get("name"), Some(&&DataType::Utf8)); + assert_eq!(field_map.get("age"), Some(&&DataType::Int32)); + assert_eq!(field_map.get("height"), Some(&&DataType::Float64)); + assert_eq!(field_map.get("is_active"), Some(&&DataType::Boolean)); + + // Check all fields are nullable + for field in schema.fields() { + assert!(field.is_nullable()); + } + } + + #[test] + fn test_mongodb_specific_types() { + let doc = doc! { + "id": mongodb::bson::oid::ObjectId::new(), + "created_at": mongodb::bson::DateTime::now(), + "timestamp": mongodb::bson::Timestamp { time: 1234567890, increment: 1 }, + "binary_data": mongodb::bson::Binary { subtype: mongodb::bson::spec::BinarySubtype::Generic, bytes: vec![1, 2, 3] }, + "decimal": mongodb::bson::Decimal128::from_str("123.456").unwrap(), + }; + let docs = vec![doc]; + + let schema = infer_arrow_schema_from_documents(&docs).unwrap(); + let field_map: HashMap = schema.fields() + .iter() + .map(|f| (f.name().clone(), f.data_type())) + .collect(); + + assert_eq!(field_map.get("id"), Some(&&DataType::Utf8)); // ObjectId as string + assert_eq!(field_map.get("created_at"), Some(&&DataType::Timestamp(TimeUnit::Millisecond, None))); + assert_eq!(field_map.get("timestamp"), Some(&&DataType::Timestamp(TimeUnit::Millisecond, None))); + assert_eq!(field_map.get("binary_data"), Some(&&DataType::Binary)); + assert_eq!(field_map.get("decimal"), Some(&&DataType::Decimal128(38, 10))); + } + + #[test] + fn test_array_types() { + let doc = doc! { + "empty_array": [], + "string_array": ["a", "b", "c"], + "number_array": [1_i32, 2_i32, 3_i32], + "mixed_array": ["text", 42_i32, true], // Should infer from first non-null + "null_array": [Bson::Null, Bson::Null, "finally_text"] + }; + let docs = vec![doc]; + + let schema = infer_arrow_schema_from_documents(&docs).unwrap(); + let field_map: HashMap = schema.fields() + .iter() + .map(|f| (f.name().clone(), f.data_type())) + .collect(); + + assert!(matches!(field_map.get("empty_array"), Some(DataType::List(_)))); + + if let Some(DataType::List(field)) = field_map.get("string_array") { + assert_eq!(field.data_type(), &DataType::Utf8); + } else { + panic!("Expected List type for string_array"); + } + + if let Some(DataType::List(field)) = field_map.get("number_array") { + assert_eq!(field.data_type(), &DataType::Utf8); + } else { + panic!("Expected List type for number_array"); + } + + if let Some(DataType::List(field)) = field_map.get("mixed_array") { + assert_eq!(field.data_type(), &DataType::Utf8); + } else { + panic!("Expected List type for mixed_array"); + } + + if let Some(DataType::List(field)) = field_map.get("null_array") { + assert_eq!(field.data_type(), &DataType::Utf8); + } else { + panic!("Expected List type for null_array"); + } + } + + #[test] + fn test_nested_document() { + let doc = doc! { + "user": { + "name": "Alice", + "age": 30_i32 + }, + "metadata": {} + }; + let docs = vec![doc]; + + let schema = infer_arrow_schema_from_documents(&docs).unwrap(); + let field_map: HashMap = schema.fields() + .iter() + .map(|f| (f.name().clone(), f.data_type())) + .collect(); + + // Nested documents should be treated as strings (JSON) + assert_eq!(field_map.get("user"), Some(&&DataType::Utf8)); + assert_eq!(field_map.get("metadata"), Some(&&DataType::Utf8)); + } + + #[test] + fn test_type_unification_numeric_promotion() { + let docs = vec![ + doc! { "value": 10_i32 }, // Int32 + doc! { "value": 20_i64 }, // Int64 -> should promote to Int64 + ]; + + let schema = infer_arrow_schema_from_documents(&docs).unwrap(); + let field = schema.field_with_name("value").unwrap(); + assert_eq!(field.data_type(), &DataType::Int64); + } + + #[test] + fn test_type_unification_to_float() { + let docs = vec![ + doc! { "value": 10_i32 }, // Int32 + doc! { "value": 3.14_f64 }, // Float64 -> should promote to Float64 + ]; + + let schema = infer_arrow_schema_from_documents(&docs).unwrap(); + let field = schema.field_with_name("value").unwrap(); + assert_eq!(field.data_type(), &DataType::Float64); + } + + #[test] + fn test_type_unification_to_string_fallback() { + let docs = vec![ + doc! { "value": 10_i32 }, // Int32 + doc! { "value": "text" }, // String -> should fallback to String + ]; + + let schema = infer_arrow_schema_from_documents(&docs).unwrap(); + let field = schema.field_with_name("value").unwrap(); + assert_eq!(field.data_type(), &DataType::Utf8); + } + + #[test] + fn test_null_unification() { + let docs = vec![ + doc! { "value": Bson::Null }, // Null + doc! { "value": "text" }, // String -> should be String + ]; + + let schema = infer_arrow_schema_from_documents(&docs).unwrap(); + let field = schema.field_with_name("value").unwrap(); + assert_eq!(field.data_type(), &DataType::Utf8); + } + + #[test] + fn test_only_null_values() { + let docs = vec![ + doc! { "value": Bson::Null }, + doc! { "value": Bson::Null }, + ]; + + let schema = infer_arrow_schema_from_documents(&docs).unwrap(); + let field = schema.field_with_name("value").unwrap(); + assert_eq!(field.data_type(), &DataType::Null); + } + + #[test] + fn test_missing_fields_across_documents() { + let docs = vec![ + doc! { "name": "Alice", "age": 30_i32 }, + doc! { "name": "Bob", "city": "NYC" }, + doc! { "age": 25_i32, "country": "US" }, + ]; + + let schema = infer_arrow_schema_from_documents(&docs).unwrap(); + + // Should have all unique fields + assert_eq!(schema.fields().len(), 4); + + let field_names: std::collections::HashSet<&str> = schema.fields() + .iter() + .map(|f| f.name().as_str()) + .collect(); + + assert!(field_names.contains("name")); + assert!(field_names.contains("age")); + assert!(field_names.contains("city")); + assert!(field_names.contains("country")); + + // All fields should be nullable since they're missing in some docs + for field in schema.fields() { + assert!(field.is_nullable()); + } + } + + #[test] + fn test_large_document_set() { + let mut docs = Vec::new(); + + // Generate 100 documents with varying schemas + for i in 0..100 { + let mut doc = Document::new(); + doc.insert("id", i as i32); + doc.insert("name", format!("user_{}", i)); + + // Add optional fields for some documents + if i % 2 == 0 { + doc.insert("age", (20 + i % 50) as i32); + } + if i % 3 == 0 { + doc.insert("city", "NYC"); + } + if i % 5 == 0 { + doc.insert("score", (i as f64) / 10.0); + } + + docs.push(doc); + } + + let schema = infer_arrow_schema_from_documents(&docs).unwrap(); + + // Should have all the fields + let field_names: std::collections::HashSet<&str> = schema.fields() + .iter() + .map(|f| f.name().as_str()) + .collect(); + + assert!(field_names.contains("id")); + assert!(field_names.contains("name")); + assert!(field_names.contains("age")); + assert!(field_names.contains("city")); + assert!(field_names.contains("score")); + + // All fields should be nullable + for field in schema.fields() { + assert!(field.is_nullable()); + } + } +} diff --git a/tests/integration.rs b/tests/integration.rs index a9fa0e11..e1bb8903 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -6,6 +6,8 @@ mod docker; mod duckdb; #[cfg(feature = "flight")] mod flight; +#[cfg(feature = "mongodb")] +mod mongodb; #[cfg(feature = "mysql")] mod mysql; #[cfg(feature = "postgres")] diff --git a/tests/mongodb/common.rs b/tests/mongodb/common.rs new file mode 100644 index 00000000..ed7cf152 --- /dev/null +++ b/tests/mongodb/common.rs @@ -0,0 +1,139 @@ +use bollard::secret::HealthConfig; +use datafusion_table_providers::mongodb::connection_pool::MongoDBConnectionPool; +use mongodb::{Client, options::ClientOptions}; +use secrecy::SecretString; +use std::collections::HashMap; +use tracing::instrument; + +use crate::{ + container_registry, + docker::{ContainerRunnerBuilder, RunningContainer}, +}; + +const MONGODB_DOCKER_CONTAINER: &str = "runtime-integration-test-mongodb"; + +pub(super) fn get_mongodb_params(port: usize) -> HashMap { + let mut params = HashMap::new(); + params.insert( + "mongodb_host".to_string(), + SecretString::from("localhost".to_string()), + ); + params.insert( + "mongodb_port".to_string(), + SecretString::from(port.to_string()), + ); + params.insert( + "mongodb_database".to_string(), + SecretString::from("testdb".to_string()), + ); + params.insert( + "mongodb_username".to_string(), + SecretString::from("root".to_string()), + ); + params.insert( + "mongodb_password".to_string(), + SecretString::from("integration-test-pw".to_string()), + ); + params.insert( + "mongodb_auth_source".to_string(), + SecretString::from("admin".to_string()), + ); + params.insert( + "mongodb_connection_string".to_string(), + SecretString::from(format!("mongodb://root:integration-test-pw@localhost:{port}/testdb?authSource=admin")), + ); + params.insert( + "mongodb_pool_min".to_string(), + SecretString::from("1".to_string()), + ); + params.insert( + "mongodb_pool_max".to_string(), + SecretString::from("10".to_string()), + ); + params.insert( + "mongodb_sslmode".to_string(), + SecretString::from("disabled".to_string()), + ); + params +} + +#[instrument] +pub async fn start_mongodb_docker_container(port: usize) -> Result { + let container_name = format!("{MONGODB_DOCKER_CONTAINER}-{port}"); + + let port = port.try_into().unwrap_or(27017); + + let mongodb_docker_image = std::env::var("MONGODB_DOCKER_IMAGE") + .unwrap_or_else(|_| format!("{}mongo:7", container_registry())); + + let running_container = ContainerRunnerBuilder::new(container_name) + .image(mongodb_docker_image) + .add_port_binding(27017, port) + .add_env_var("MONGO_INITDB_ROOT_USERNAME", "root") + .add_env_var("MONGO_INITDB_ROOT_PASSWORD", "integration-test-pw") + .add_env_var("MONGO_INITDB_DATABASE", "testdb") + .healthcheck(HealthConfig { + test: Some(vec![ + "CMD".to_string(), + "mongosh".to_string(), + "mongodb://root:integration-test-pw@localhost:27017/testdb?authSource=admin".to_string(), + "--quiet".to_string(), + "--eval".to_string(), + "db.runCommand({ ping: 1 }).ok".to_string(), + ]), + interval: Some(500_000_000), + timeout: Some(500_000_000), + retries: Some(10), + start_period: Some(3_000_000_000), + start_interval: None, + }) + .build()? + .run() + .await?; + + tokio::time::sleep(std::time::Duration::from_secs(10)).await; + Ok(running_container) +} + +#[instrument] +pub(super) async fn get_mongodb_connection_pool( + port: usize, +) -> Result { + let mongodb_pool = MongoDBConnectionPool::new(get_mongodb_params(port)) + .await + .expect("Failed to create MongoDB Connection Pool"); + + Ok(mongodb_pool) +} + +#[instrument] +pub(super) async fn get_mongodb_client(port: usize) -> Result { + let connection_string = format!("mongodb://root:integration-test-pw@localhost:{port}/testdb?authSource=admin"); + + let client_options = ClientOptions::parse(connection_string) + .await + .expect("Failed to parse MongoDB connection string"); + + let client = Client::with_options(client_options) + .expect("Failed to create MongoDB client"); + + let mut retries = 10; + let mut last_err = None; + while retries > 0 { + match client + .database("testdb") + .run_command(mongodb::bson::doc! { "ping": 1 }) + .await + { + Ok(_) => {println!("Client created"); return Ok(client)}, + Err(e) => { + last_err = Some(e); + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + retries -= 1; + println!("Ping failed"); + } + } + } + + Ok(client) +} diff --git a/tests/mongodb/mod.rs b/tests/mongodb/mod.rs new file mode 100644 index 00000000..03e470ff --- /dev/null +++ b/tests/mongodb/mod.rs @@ -0,0 +1,707 @@ +use std::time::SystemTime; +use std::sync::Arc; +use chrono::{DateTime, Utc}; +use datafusion::{error::DataFusionError, execution::context::SessionContext}; +use datafusion_table_providers::mongodb::table::MongoDBTable; +use mongodb::bson::{doc, Document, Bson, DateTime as BsonDateTime, Decimal128}; +use std::str::FromStr; +use rstest::rstest; + +use arrow::{ + array::*, + datatypes::{DataType, Field, Schema, TimeUnit}, +}; + +use crate::docker::RunningContainer; + +mod common; + +async fn test_mongodb_datetime_types(port: usize) { + let ts0 = DateTime::parse_from_rfc3339("2024-09-12T10:00:00Z").unwrap().with_timezone(&Utc); + let ts1 = DateTime::parse_from_rfc3339("2024-09-12T10:00:00.1Z").unwrap().with_timezone(&Utc); + let ts2 = DateTime::parse_from_rfc3339("2024-09-12T10:00:00.12Z").unwrap().with_timezone(&Utc); + let ts3 = DateTime::parse_from_rfc3339("2024-09-12T10:00:00.123Z").unwrap().with_timezone(&Utc); + + let test_docs = vec![ + doc! { + "timestamp_field": Bson::DateTime(BsonDateTime::from(SystemTime::from(ts0))), + "timestamp_one_fraction": Bson::DateTime(BsonDateTime::from(SystemTime::from(ts1))), + "timestamp_two_fraction": Bson::DateTime(BsonDateTime::from(SystemTime::from(ts2))), + "timestamp_three_fraction": Bson::DateTime(BsonDateTime::from(SystemTime::from(ts3))), + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("timestamp_field", DataType::Timestamp(TimeUnit::Millisecond, None), true), + Field::new("timestamp_one_fraction", DataType::Timestamp(TimeUnit::Millisecond, None), true), + Field::new("timestamp_two_fraction", DataType::Timestamp(TimeUnit::Millisecond, None), true), + Field::new("timestamp_three_fraction", DataType::Timestamp(TimeUnit::Millisecond, None), true), + ])); + + let expected_record = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![1_726_135_200_000])), + Arc::new(TimestampMillisecondArray::from(vec![1_726_135_200_100])), + Arc::new(TimestampMillisecondArray::from(vec![1_726_135_200_120])), + Arc::new(TimestampMillisecondArray::from(vec![1_726_135_200_123])), + ], + ) + .expect("Failed to create arrow record batch"); + + arrow_mongodb_one_way( + port, + "timestamp_collection", + test_docs, + expected_record, + ) + .await; +} + +async fn test_mongodb_numeric_types(port: usize) { + + let decimal = Decimal128::from_str("123.456").unwrap(); + + let test_docs = vec![ + doc! { + "int32_field": 2147483647i32, + "int64_field": 9223372036854775807i64, + "double_field": 3.14159265359, + "decimal_field": Bson::Decimal128(decimal), + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("int32_field", DataType::Int32, true), + Field::new("int64_field", DataType::Int64, true), + Field::new("double_field", DataType::Float64, true), + Field::new("decimal_field", DataType::Decimal128(38, 10), true), + ])); + + let expected_record = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2147483647i32])), + Arc::new(Int64Array::from(vec![9223372036854775807i64])), + Arc::new(Float64Array::from(vec![3.14159265359])), + Arc::new( + Decimal128Array::from(vec![Some(1234560000000i128)]) + .with_precision_and_scale(38, 10) + .unwrap(), + ), + ], + ) + .expect("Failed to create arrow record batch"); + + let array = expected_record + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + + println!("Decimal as i128: {:?}", array.value(0)); + + arrow_mongodb_one_way( + port, + "numeric_collection", + test_docs, + expected_record, + ) + .await; +} + +async fn test_mongodb_string_types(port: usize) { + let test_docs = vec![ + doc! { + "name": "Alice", + "description": "Software Engineer", + "notes": Bson::Null, + }, + doc! { + "name": "Bob", + "description": "Data Scientist", + "notes": "Likes MongoDB", + }, + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, true), + Field::new("description", DataType::Utf8, true), + Field::new("notes", DataType::Utf8, true), + ])); + + let expected_record = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec!["Alice", "Bob"])), + Arc::new(StringArray::from(vec!["Software Engineer", "Data Scientist"])), + Arc::new(StringArray::from(vec![None, Some("Likes MongoDB")])), + ], + ) + .expect("Failed to create arrow record batch"); + + arrow_mongodb_one_way( + port, + "string_collection", + test_docs, + expected_record, + ) + .await; +} + +async fn test_mongodb_boolean_types(port: usize) { + let test_docs = vec![ + doc! { + "is_active": true, + "is_verified": false, + "is_premium": Bson::Null, + }, + doc! { + "is_active": false, + "is_verified": true, + "is_premium": true, + }, + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("is_active", DataType::Boolean, true), + Field::new("is_verified", DataType::Boolean, true), + Field::new("is_premium", DataType::Boolean, true), + ])); + + let expected_record = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(BooleanArray::from(vec![true, false])), + Arc::new(BooleanArray::from(vec![false, true])), + Arc::new(BooleanArray::from(vec![None, Some(true)])), + ], + ) + .expect("Failed to create arrow record batch"); + + arrow_mongodb_one_way( + port, + "boolean_collection", + test_docs, + expected_record, + ) + .await; +} + +async fn test_mongodb_binary_types(port: usize) { + let test_docs = vec![ + doc! { + "binary_data": Bson::Binary(mongodb::bson::Binary { + subtype: mongodb::bson::spec::BinarySubtype::Generic, + bytes: b"hello world".to_vec(), + }), + "file_content": Bson::Binary(mongodb::bson::Binary { + subtype: mongodb::bson::spec::BinarySubtype::Generic, + bytes: b"binary file content".to_vec(), + }), + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("binary_data", DataType::Binary, true), + Field::new("file_content", DataType::Binary, true), + ])); + + let expected_record = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(BinaryArray::from_vec(vec![b"hello world"])), + Arc::new(BinaryArray::from_vec(vec![b"binary file content"])), + ], + ) + .expect("Failed to create arrow record batch"); + + arrow_mongodb_one_way( + port, + "binary_collection", + test_docs, + expected_record, + ) + .await; +} + +async fn test_mongodb_object_id_types(port: usize) { + let oid1 = mongodb::bson::oid::ObjectId::new(); + let oid2 = mongodb::bson::oid::ObjectId::new(); + + let test_docs = vec![ + doc! { + "_id": oid1, + "ref_id": oid2, + } + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("_id", DataType::Utf8, true), // ObjectId typically converted to string + Field::new("ref_id", DataType::Utf8, true), + ])); + + let expected_record = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec![oid1.to_hex()])), + Arc::new(StringArray::from(vec![oid2.to_hex()])), + ], + ) + .expect("Failed to create arrow record batch"); + + arrow_mongodb_one_way( + port, + "objectid_collection", + test_docs, + expected_record, + ) + .await; +} + +async fn test_mongodb_array_types(port: usize) { + let test_docs = vec![ + doc! { + "string_tags": ["rust", "mongodb", "arrow"], + "mixed_array": ["text", 42, true, 3.14], + "empty_array": [], + "numbers_as_strings": [1, 2, 3], + }, + doc! { + "string_tags": ["python", "sql"], + "mixed_array": ["another", false, 99], + "empty_array": [], + "numbers_as_strings": [4, 5], + }, + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new( + "string_tags", + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + true, + ), + Field::new( + "mixed_array", + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + true, + ), + Field::new( + "empty_array", + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + true, + ), + Field::new( + "numbers_as_strings", + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + true, + ), + ])); + + // Create the expected ListArrays manually + let string_tags_builder = ListBuilder::new(StringBuilder::new()); + let mut string_tags_list = string_tags_builder; + + // First document string_tags: ["rust", "mongodb", "arrow"] + string_tags_list.values().append_value("rust"); + string_tags_list.values().append_value("mongodb"); + string_tags_list.values().append_value("arrow"); + string_tags_list.append(true); + + // Second document string_tags: ["python", "sql"] + string_tags_list.values().append_value("python"); + string_tags_list.values().append_value("sql"); + string_tags_list.append(true); + + let string_tags_array = Arc::new(string_tags_list.finish()); + + // Mixed array (all converted to strings) + let mixed_array_builder = ListBuilder::new(StringBuilder::new()); + let mut mixed_array_list = mixed_array_builder; + + // First document mixed_array: ["text", "42", "true", "3.14"] + mixed_array_list.values().append_value("text"); + mixed_array_list.values().append_value("42"); + mixed_array_list.values().append_value("true"); + mixed_array_list.values().append_value("3.14"); + mixed_array_list.append(true); + + // Second document mixed_array: ["another", "false", "99"] + mixed_array_list.values().append_value("another"); + mixed_array_list.values().append_value("false"); + mixed_array_list.values().append_value("99"); + mixed_array_list.append(true); + + let mixed_array_array = Arc::new(mixed_array_list.finish()); + + // Empty arrays + let empty_array_builder = ListBuilder::new(StringBuilder::new()); + let mut empty_array_list = empty_array_builder; + + // First document: empty array + empty_array_list.append(true); + // Second document: empty array + empty_array_list.append(true); + + let empty_array_array = Arc::new(empty_array_list.finish()); + + // Numbers as strings array + let numbers_builder = ListBuilder::new(StringBuilder::new()); + let mut numbers_list = numbers_builder; + + // First document: [1, 2, 3] -> ["1", "2", "3"] + numbers_list.values().append_value("1"); + numbers_list.values().append_value("2"); + numbers_list.values().append_value("3"); + numbers_list.append(true); + + // Second document: [4, 5] -> ["4", "5"] + numbers_list.values().append_value("4"); + numbers_list.values().append_value("5"); + numbers_list.append(true); + + let numbers_array = Arc::new(numbers_list.finish()); + + let expected_record = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + string_tags_array, + mixed_array_array, + empty_array_array, + numbers_array, + ], + ) + .expect("Failed to create arrow record batch"); + + arrow_mongodb_one_way( + port, + "array_collection", + test_docs, + expected_record, + ) + .await; +} + +async fn test_mongodb_nested_object_types(port: usize) { + let test_docs = vec![ + doc! { + "user": { + "name": "Alice", + "age": 30, + "contact": { + "email": "alice@example.com", + "phone": "555-1234" + } + }, + "metadata": { + "created_at": "2024-01-01", + "tags": ["important", "user"], + "settings": { + "theme": "dark", + "notifications": true + } + }, + "empty_object": {}, + "simple_string": "not an object" + }, + doc! { + "user": { + "name": "Bob", + "age": 25, + "contact": { + "email": "bob@example.com" + } + }, + "metadata": { + "created_at": "2024-01-02", + "tags": ["user"], + "settings": { + "theme": "light", + "notifications": false + } + }, + "empty_object": {}, + "simple_string": "also not an object" + }, + ]; + + // We'll test the content, not the exact JSON string format + let ctx = SessionContext::new(); + let client = common::get_mongodb_client(port) + .await + .expect("MongoDB client should be created"); + + // Insert test data into MongoDB collection + let db = client.database("testdb"); + let collection = db.collection::("nested_object_collection"); + + // Drop collection if it exists + let _ = collection.drop().await; + + // Insert test documents + collection + .insert_many(test_docs) + .await + .expect("MongoDB documents should be inserted"); + + let expected_user1 = serde_json::json!({ + "name": "Alice", + "age": 30, + "contact": { + "email": "alice@example.com", + "phone": "555-1234" + } + }); + + let expected_user2 = serde_json::json!({ + "name": "Bob", + "age": 25, + "contact": { + "email": "bob@example.com" + } + }); + + let expected_metadata1 = serde_json::json!({ + "created_at": "2024-01-01", + "tags": ["important", "user"], + "settings": { + "theme": "dark", + "notifications": true + } + }); + + let expected_metadata2 = serde_json::json!({ + "created_at": "2024-01-02", + "tags": ["user"], + "settings": { + "theme": "light", + "notifications": false + } + }); + + let expected_empty = serde_json::json!({}); + + // Register DataFusion table + let mongo_conn_pool = common::get_mongodb_connection_pool(port) + .await + .expect("MongoDB connection pool should be created"); + + let table = MongoDBTable::new(&Arc::new(mongo_conn_pool), "nested_object_collection") + .await + .expect("Table should be created"); + + ctx.register_table("nested_object_collection", Arc::new(table)) + .expect("Table should be registered"); + + // Query the data + let sql = r#"SELECT "user", "metadata", "empty_object", "simple_string" FROM nested_object_collection"#; + let df = ctx + .sql(&sql) + .await + .expect("DataFrame should be created from query"); + + let record_batches = df.collect().await.expect("RecordBatch should be collected"); + assert_eq!(record_batches.len(), 1); + + let batch = &record_batches[0]; + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch.num_columns(), 4); + + // Verify the JSON content by parsing and comparing structure + let user_array = batch.column_by_name("user").unwrap() + .as_any().downcast_ref::().unwrap(); + let metadata_array = batch.column_by_name("metadata").unwrap() + .as_any().downcast_ref::().unwrap(); + let empty_array = batch.column_by_name("empty_object").unwrap() + .as_any().downcast_ref::().unwrap(); + let string_array = batch.column_by_name("simple_string").unwrap() + .as_any().downcast_ref::().unwrap(); + + // Parse actual JSON strings and compare with expected JSON objects + let actual_user1: serde_json::Value = serde_json::from_str(user_array.value(0)).unwrap(); + let actual_user2: serde_json::Value = serde_json::from_str(user_array.value(1)).unwrap(); + let actual_metadata1: serde_json::Value = serde_json::from_str(metadata_array.value(0)).unwrap(); + let actual_metadata2: serde_json::Value = serde_json::from_str(metadata_array.value(1)).unwrap(); + let actual_empty1: serde_json::Value = serde_json::from_str(empty_array.value(0)).unwrap(); + let actual_empty2: serde_json::Value = serde_json::from_str(empty_array.value(1)).unwrap(); + + // Direct JSON comparison - order doesn't matter! + assert_eq!(actual_user1, expected_user1); + assert_eq!(actual_user2, expected_user2); + assert_eq!(actual_metadata1, expected_metadata1); + assert_eq!(actual_metadata2, expected_metadata2); + assert_eq!(actual_empty1, expected_empty); + assert_eq!(actual_empty2, expected_empty); + + // String values remain simple + assert_eq!(string_array.value(0), "not an object"); + assert_eq!(string_array.value(1), "also not an object"); +} + +async fn test_mongodb_null_and_missing_fields(port: usize) { + let test_docs = vec![ + doc! { + "name": "Alice", + "age": 30, + "email": "alice@example.com", + }, + doc! { + "name": "Bob", + "age": Bson::Null, + "phone": "555-1234", + }, + doc! { + "name": "Charlie", + "age": 25, + // email and phone missing + }, + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, true), + Field::new("age", DataType::Int32, true), + Field::new("email", DataType::Utf8, true), + Field::new("phone", DataType::Utf8, true), + ])); + + let expected_record = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec!["Alice", "Bob", "Charlie"])), + Arc::new(Int32Array::from(vec![Some(30), None, Some(25)])), + Arc::new(StringArray::from(vec![Some("alice@example.com"), None, None])), + Arc::new(StringArray::from(vec![None, Some("555-1234"), None])), + ], + ) + .expect("Failed to create arrow record batch"); + + arrow_mongodb_one_way( + port, + "null_fields_collection", + test_docs, + expected_record, + ) + .await; +} + +async fn arrow_mongodb_one_way( + port: usize, + collection_name: &str, + test_docs: Vec, + expected_record: RecordBatch, +) -> Vec { + tracing::debug!("Running MongoDB tests on {collection_name}"); + + let ctx = SessionContext::new(); + let client = common::get_mongodb_client(port) + .await + .expect("MongoDB client should be created"); + + // Insert test data into MongoDB collection + let db = client.database("testdb"); + let collection = db.collection::(collection_name); + + // Drop collection if it exists + let _ = collection.drop().await; + + // Insert test documents + if !test_docs.is_empty() { + collection + .insert_many(test_docs) + .await + .expect("MongoDB documents should be inserted"); + } + + // Register DataFusion table + let mongo_conn_pool = common::get_mongodb_connection_pool(port) + .await + .expect("MongoDB connection pool should be created"); + + let table = MongoDBTable::new(&Arc::new(mongo_conn_pool), collection_name) + .await + .expect("Table should be created"); + + ctx.register_table(collection_name, Arc::new(table)) + .expect("Table should be registered"); + + // Extract expected columns (excluding _id) + let schema_ref = expected_record.schema(); + let expected_fields: Vec<&str> = schema_ref + .fields() + .iter() + .map(|f| f.name().as_str()) + .filter(|name| *name != "_id") + .collect(); + + // Build SELECT query with correct projection + let projection = expected_fields + .iter() + .map(|c| format!("\"{c}\"")) + .collect::>() + .join(", "); + let sql = format!("SELECT {projection} FROM {collection_name}"); + + let df = ctx + .sql(&sql) + .await + .expect("DataFrame should be created from query"); + + let record_batches = df.collect().await.expect("RecordBatch should be collected"); + assert_eq!(record_batches.len(), 1); + + // Normalize actual and expected + let actual_projected = + project_record_batch(&record_batches[0], &expected_fields).expect("Project actual"); + let expected_projected = + project_record_batch(&expected_record, &expected_fields).expect("Project expected"); + + assert_eq!(actual_projected, expected_projected); + + record_batches +} + +use datafusion::common::Result as DFResult; +fn project_record_batch(batch: &RecordBatch, columns: &[&str]) -> DFResult { + let schema = batch.schema(); + let indices: Vec = columns + .iter() + .map(|col| schema.index_of(col).expect("Column not found")) + .collect(); + let arrays = indices.iter().map(|&i| batch.column(i).clone()).collect(); + let fields = indices + .iter() + .map(|&i| schema.field(i).clone()) + .collect::>(); + let projected_schema = Arc::new(arrow::datatypes::Schema::new(fields)); + RecordBatch::try_new(projected_schema, arrays) + .map_err(|e| DataFusionError::ArrowError(e, None)) +} + +async fn start_mongodb_container(port: usize) -> RunningContainer { + let running_container = common::start_mongodb_docker_container(port) + .await + .expect("MongoDB container to start"); + + tracing::debug!("MongoDB Container started"); + + running_container +} + +#[rstest] +#[test_log::test(tokio::test)] +async fn test_mongodb_arrow_oneway() { + let port = crate::get_random_port(); + let mongodb_container = start_mongodb_container(port).await; + + test_mongodb_datetime_types(port).await; + test_mongodb_numeric_types(port).await; + test_mongodb_string_types(port).await; + test_mongodb_boolean_types(port).await; + test_mongodb_binary_types(port).await; + test_mongodb_object_id_types(port).await; + test_mongodb_array_types(port).await; + test_mongodb_nested_object_types(port).await; + test_mongodb_null_and_missing_fields(port).await; + + mongodb_container.remove().await.expect("container to stop"); +} diff --git a/tests/postgres/mod.rs b/tests/postgres/mod.rs index 5083625e..22dfae20 100644 --- a/tests/postgres/mod.rs +++ b/tests/postgres/mod.rs @@ -1,6 +1,6 @@ use crate::{arrow_record_batch_gen::*, docker::RunningContainer}; use arrow::{ - array::{Decimal128Array, RecordBatch}, + array::{Decimal128Array, RecordBatch, Array, StringArray}, datatypes::{DataType, Field, Schema, SchemaRef}, }; use datafusion::logical_expr::CreateExternalTable; @@ -18,7 +18,7 @@ use datafusion_table_providers::{ UnsupportedTypeAction, }; use rstest::{fixture, rstest}; -use serde_json::Value; +use serde_json::{Value, from_str}; use std::collections::HashMap; use std::sync::Arc; use tokio::sync::{Mutex, MutexGuard}; @@ -207,6 +207,7 @@ async fn test_postgres_enum_type(port: usize) { extra_stmt, expected_record, UnsupportedTypeAction::default(), + true, ) .await; } @@ -264,6 +265,7 @@ async fn test_postgres_numeric_type(port: usize) { extra_stmt, expected_record, UnsupportedTypeAction::default(), + true, ) .await; } @@ -285,39 +287,49 @@ async fn test_postgres_jsonb_type(port: usize) { let schema = Arc::new(Schema::new(vec![Field::new("data", DataType::Utf8, true)])); - // Parse and re-serialize the JSON to ensure consistent ordering let expected_values = vec![ - serde_json::from_str::(r#"{"name":"John","age":30}"#) - .unwrap() - .to_string(), - serde_json::from_str::(r#"{"name":"Jane","age":25}"#) - .unwrap() - .to_string(), - serde_json::from_str::("[1,2,3]") - .unwrap() - .to_string(), - serde_json::from_str::("null").unwrap().to_string(), - serde_json::from_str::(r#"{"nested":{"key":"value"}}"#) - .unwrap() - .to_string(), + r#"{"name": "John", "age": 30}"#, + r#"{"name": "Jane", "age": 25}"#, + "[1, 2, 3]", + "null", + r#"{"nested": {"key": "value"}}"#, ]; + let expected_json: Vec = expected_values + .iter() + .map(|s| from_str(s).unwrap()) + .collect(); + let expected_record = RecordBatch::try_new( Arc::clone(&schema), vec![Arc::new(arrow::array::StringArray::from(expected_values))], ) .expect("Failed to create arrow record batch"); - arrow_postgres_one_way( + let actual_record_batch = arrow_postgres_one_way( port, "jsonb_values", create_table_stmt, insert_table_stmt, None, - expected_record, + expected_record.clone(), UnsupportedTypeAction::String, + false, ) .await; + + let actual_data_column = actual_record_batch[0] + .column_by_name("data") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + + let actual_json: Vec = (0..actual_data_column.len()) + .map(|i| from_str(actual_data_column.value(i)).unwrap()) + .collect(); + + assert_eq!(actual_json, expected_json); } async fn arrow_postgres_one_way( @@ -328,7 +340,8 @@ async fn arrow_postgres_one_way( extra_stmt: Option<&str>, expected_record: RecordBatch, unsupported_type_action: UnsupportedTypeAction, -) { + perform_check: bool, +) -> Vec { tracing::debug!("Running tests on {table_name}"); let ctx = SessionContext::new(); @@ -377,5 +390,9 @@ async fn arrow_postgres_one_way( let record_batch = df.collect().await.expect("RecordBatch should be collected"); - assert_eq!(record_batch[0], expected_record); + if perform_check { + assert_eq!(record_batch[0], expected_record); + } + + record_batch }