SQL-2288: Create mongosqltranslate integration tests

constants/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "constants"
-version = "0.0.0"
+version = "2.0.0"
 authors = [
     "Ryan Chipman <[email protected]>",
     "Natacha Bagnard <[email protected]>",

constants/src/lib.rs

@@ -5,6 +5,7 @@ pub const DRIVER_NAME: &str = "MongoDB Atlas SQL ODBC Driver";
 pub const DBMS_NAME: &str = "MongoDB Atlas";
 pub const ODBC_VERSION: &str = "03.80";
 pub const DRIVER_SHORT_NAME: &str = "mongodb-odbc";
+pub const SQL_SCHEMAS_COLLECTION: &str = "__sql_schemas";
 
 lazy_static! {
     pub static ref DRIVER_METRICS_VERSION: String = format!(

core/src/col_metadata.rs

@@ -1,15 +1,15 @@
-use std::collections::HashMap;
-
 use crate::{
     json_schema::{
         simplified::{Atomic, ObjectSchema, Schema},
         BsonTypeName,
     },
     BsonTypeInfo, Error, Result, TypeMode,
 };
+use bson::{Bson, Document};
 use definitions::{Nullability, SqlCode, SqlDataType};
 use itertools::Itertools;
 use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
 
 // Metadata information for a column of the result set.
 // The information is to be used when reporting columns information from
@@ -173,10 +173,21 @@ pub struct VersionedJsonSchema {
 // Struct representing the ResultSetSchema.
 #[derive(Serialize, Deserialize, PartialEq, Debug, Clone, Default)]
 pub struct ResultSetSchema {
+    #[serde(rename = "result_set_schema")]
     pub schema: crate::json_schema::Schema,
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub select_order: Option<Vec<Vec<String>>>,
 }
 
+impl ResultSetSchema {
+    pub fn from_sql_schemas_document(doc: &Document) -> bson::de::Result<Self> {
+        let as_bson = Bson::Document(doc.clone());
+        let deserializer = bson::Deserializer::new(as_bson);
+        let deserializer = serde_stacker::Deserializer::new(deserializer);
+        Deserialize::deserialize(deserializer)
+    }
+}
+
 impl From<SqlGetSchemaResponse> for ResultSetSchema {
     fn from(sql_get_schema_response: SqlGetSchemaResponse) -> Self {
         Self {

core/src/collections.rs

@@ -123,7 +123,13 @@ impl MongoCollections {
                     .unwrap()
                     .iter()
                     // MHOUSE-7119 - admin database and empty strings are showing in list_database_names
-                    .filter(|&db_name| !db_name.is_empty() && !db_name.eq("admin"))
+                    .filter(|&db_name| {
+                        !db_name.is_empty()
+                            && !db_name.eq("admin")
+                            && !db_name.eq("config")
+                            && !db_name.eq("local")
+                            && !db_name.eq("system")
+                    })
                     .filter(|&db_name| is_match(db_name, db_name_filter, accept_search_patterns))
                     .map(|val| async move {
                         CollectionsForDb {

core/src/conn.rs

@@ -7,7 +7,7 @@ use crate::mongosqltranslate::{
 use crate::odbc_uri::UserOptions;
 use crate::{err::Result, Error};
 use crate::{MongoQuery, TypeMode};
-use constants::DRIVER_ODBC_VERSION;
+use constants::DRIVER_METRICS_VERSION;
 use lazy_static::lazy_static;
 use mongodb::{
     bson::{doc, Bson, UuidRepresentation},
@@ -154,7 +154,7 @@ impl MongoConnection {
     }
 
     fn is_libmongosqltranslate_compatible_with_driver_version() -> Result<bool> {
-        let command = CheckDriverVersion::new(DRIVER_ODBC_VERSION.clone());
+        let command = CheckDriverVersion::new(DRIVER_METRICS_VERSION.clone());
 
         let command_response = libmongosqltranslate_run_command(command)?;
 
@@ -196,7 +196,7 @@ impl MongoConnection {
 
         let uuid_repr = user_options.uuid_representation;
 
-        load_mongosqltranslate_library();
+        load_mongosqltranslate_library(false);
 
         let (is_libmongosqltranslate_compatible_with_driver_version, libmongosqltranslate_version) =
             if get_mongosqltranslate_library().is_some() {
@@ -243,7 +243,7 @@ impl MongoConnection {
                     .is_some_and(|is_compatible| is_compatible)
                 {
                     return Err(Error::LibmongosqltranslateLibraryIsIncompatible(
-                        &DRIVER_ODBC_VERSION,
+                        &DRIVER_METRICS_VERSION,
                         libmongosqltranslate_version
                             .ok_or(Error::EmptyLibmongosqltranslateVersion)?,
                     ));

core/src/databases.rs

@@ -221,7 +221,13 @@ impl MongoDatabases {
             })
             .unwrap()
             .iter()
-            .filter(|&db_name| !db_name.is_empty() && !db_name.eq("admin"))
+            .filter(|&db_name| {
+                !db_name.is_empty()
+                    && !db_name.eq("admin")
+                    && !db_name.eq("config")
+                    && !db_name.eq("local")
+                    && !db_name.eq("system")
+            })
             .map(|s| s.to_string())
             .collect();
 

core/src/err.rs

@@ -58,7 +58,7 @@ pub enum Error {
         "The ODBC driver version `{0}` is incompatible with libmongosqltranslate version `{1}`"
     )]
     LibmongosqltranslateLibraryIsIncompatible(&'static str, String),
-    #[error("The following collections were not found in the `__sql_schemas` collection: {0:?}")]
+    #[error("The following collection(s) were not found in the `__sql_schemas` collection: {0:?}")]
     SchemaDocumentNotFoundInSchemaCollection(Vec<String>),
     #[error(
         "The libmongosqltranslate command `{0}` failed. Error message: `{1}`. Error is internal: {2}"

core/src/fields.rs

@@ -1,4 +1,5 @@
 use crate::{
+    cluster_type::MongoClusterType,
     col_metadata::{MongoColMetadata, ResultSetSchema, SqlGetSchemaResponse},
     collections::MongoODBCCollectionSpecification,
     conn::MongoConnection,
@@ -7,9 +8,10 @@ use crate::{
     util::to_name_regex,
     BsonTypeInfo, TypeMode,
 };
+use constants::SQL_SCHEMAS_COLLECTION;
 use definitions::{Nullability, SqlDataType};
 use mongodb::{
-    bson::{doc, Bson},
+    bson::{doc, Bson, Document},
     results::CollectionType,
 };
 use once_cell::sync::OnceCell;
@@ -484,7 +486,13 @@ impl MongoFields {
                     .unwrap()
                     // MHOUSE-7119 - admin database and empty strings are showing in list_database_names
                     .iter()
-                    .filter(|&db_name| !db_name.is_empty() && !db_name.eq("admin"))
+                    .filter(|&db_name| {
+                        !db_name.is_empty()
+                            && !db_name.eq("admin")
+                            && !db_name.eq("config")
+                            && !db_name.eq("local")
+                            && !db_name.eq("system")
+                    })
                     .map(|s| s.to_string())
                     .collect()
             },
@@ -542,23 +550,66 @@ impl MongoFields {
                             // The collection does not match the filter, moving to the next one
                             continue;
                         }
-                        let get_schema_cmd = doc! {"sqlGetSchema": collection_name.clone()};
 
                         let db = mongo_connection.client.database(&self.current_db_name);
-                        let current_col_metadata_response: Result<SqlGetSchemaResponse> =
-                            mongodb::bson::from_document(
-                                db.run_command(get_schema_cmd).await.unwrap(),
-                            )
-                            .map_err(|e| {
-                                Error::CollectionDeserialization(collection_name.clone(), e)
-                            });
-                        if let Err(error) = current_col_metadata_response {
-                            // If there is an Error while deserializing the schema, we won't show any columns for it
-                            warnings.push(error);
-                            continue;
-                        }
-                        let current_col_metadata_response: ResultSetSchema =
-                            current_col_metadata_response.unwrap().into();
+
+                        let current_col_metadata_response: ResultSetSchema = if mongo_connection
+                            .cluster_type
+                            == MongoClusterType::AtlasDataFederation
+                        {
+                            let get_schema_cmd = doc! {"sqlGetSchema": &collection_name};
+
+                            let sql_get_schema_response: Result<SqlGetSchemaResponse> =
+                                mongodb::bson::from_document(
+                                    db.run_command(get_schema_cmd).await.unwrap(),
+                                )
+                                .map_err(|e| {
+                                    Error::CollectionDeserialization(collection_name.clone(), e)
+                                });
+
+                            if let Err(error) = sql_get_schema_response {
+                                // If there is an Error while deserializing the schema, we won't show any columns for it
+                                warnings.push(error);
+                                continue;
+                            }
+
+                            sql_get_schema_response.unwrap().into()
+                        } else if mongo_connection.cluster_type == MongoClusterType::Enterprise {
+                            let schema_collection =
+                                db.collection::<Document>(SQL_SCHEMAS_COLLECTION);
+
+                            let schema_doc: Document = mongo_connection
+                                .runtime
+                                .block_on(async {
+                                    schema_collection
+                                        .find_one(doc! {
+                                            "_id": &collection_name
+                                        })
+                                        .await
+                                        .map_err(Error::QueryExecutionFailed)
+                                })?
+                                .ok_or(Error::SchemaDocumentNotFoundInSchemaCollection(vec![
+                                    collection_name.clone(),
+                                ]))?;
+
+                            let result_set_schema: Result<ResultSetSchema> =
+                                ResultSetSchema::from_sql_schemas_document(&schema_doc).map_err(
+                                    |e| {
+                                        Error::CollectionDeserialization(collection_name.clone(), e)
+                                    },
+                                );
+
+                            if let Err(error) = result_set_schema {
+                                // If there is an Error while deserializing the schema, we won't show any columns for it
+                                warnings.push(error);
+                                continue;
+                            }
+
+                            result_set_schema.unwrap()
+                        } else {
+                            unreachable!()
+                        };
+
                         match current_col_metadata_response.process_collection_metadata(
                             &self.current_db_name,
                             collection_name.as_str(),

core/src/mongosqltranslate.rs

@@ -42,8 +42,8 @@ fn get_library_name(library_type: &str) -> String {
     }
 }
 
-fn get_library_path(library_type: &str) -> PathBuf {
-    let lib_name = get_library_name(library_type);
+fn get_library_path() -> PathBuf {
+    let lib_name = get_library_name(LIBRARY_NAME);
     let mut path = PathBuf::from(LIBRARY_INSTALL_PATH);
     path.push(lib_name);
     path
@@ -63,12 +63,12 @@ fn get_mock_library_path() -> PathBuf {
 // and is responsible for determining the library name and path.
 // The library name and path are determined based on the operating system and architecture.
 // It is stored in a static variable to ensure that it is only loaded once.
-pub fn load_mongosqltranslate_library() {
+pub fn load_mongosqltranslate_library(mock_library: bool) {
     INIT.call_once(|| {
-        let library_path = if cfg!(test) {
+        let library_path = if mock_library {
             get_mock_library_path()
         } else {
-            get_library_path(LIBRARY_NAME)
+            get_library_path()
         };
 
         match unsafe { Library::new(library_path.clone()) } {
@@ -221,6 +221,7 @@ impl CheckDriverVersion {
 }
 
 #[derive(Debug, Serialize, Deserialize)]
+#[serde(tag = "command_type")]
 pub enum CommandResponse {
     Translate(TranslateCommandResponse),
     GetNamespaces(GetNamespacesCommandResponse),
@@ -308,7 +309,7 @@ pub(crate) fn libmongosqltranslate_run_command<T: CommandName + Serialize>(
 
     let decomposed_returned_doc = unsafe { run_command_function(libmongosqltranslate_command) };
 
-    let command_response_doc: Document = unsafe {
+    let mut command_response_doc: Document = unsafe {
         bson::from_slice(
             Vec::from_raw_parts(
                 decomposed_returned_doc.data.cast_mut(),
@@ -320,6 +321,20 @@ pub(crate) fn libmongosqltranslate_run_command<T: CommandName + Serialize>(
         .map_err(Error::LibmongosqltranslateDeserialization)?
     };
 
+    let command_type = if command_response_doc.get_str("error").is_ok() {
+        "Error"
+    } else {
+        match T::command_name() {
+            "getNamespaces" => "GetNamespaces",
+            "translate" => "Translate",
+            "getMongosqlTranslateVersion" => "GetMongosqlTranslateVersion",
+            "checkDriverVersion" => "CheckDriverVersion",
+            _ => unreachable!(),
+        }
+    };
+
+    command_response_doc.insert("command_type", command_type);
+
     let command_response = CommandResponse::from_document(&command_response_doc)?;
 
     if let CommandResponse::Error(error_response) = command_response {
@@ -340,7 +355,7 @@ mod unit {
 
     #[test]
     fn library_load_and_run_command_test() {
-        load_mongosqltranslate_library();
+        load_mongosqltranslate_library(true);
         assert!(get_mongosqltranslate_library().is_some());
 
         let run_command = get_run_command_fn_ptr().expect("Failed to load runCommand symbol");