Add "sensible" default behaviour for the CLI update command

CHANGELOG.md

@@ -3,7 +3,13 @@ This changelog documents the changes between release versions.
 
 ## [Unreleased]
 - Use separate schema files for each collection
-- Don't sample from collections that already have a schema
+- Changes to `update` CLI command:
+  - new default behaviour:
+    - attempt to use validator schema if available
+    - if no validator schema then sample documents from the collection
+  - don't sample from collections that already have a schema
+  - if no --sample-size given on command line, default sample size is 10
+  - new option --no-validator-schema to disable attempting to use validator schema
 
 ## [0.0.2] - 2024-03-26
 - Rename CLI plugin to ndc-mongodb ([PR #13](https://github.com/hasura/ndc-mongodb/pull/13))

crates/cli/src/introspection/validation_schema.rs

@@ -4,11 +4,10 @@ use configuration::{
     schema::{self, Type},
     Schema, WithName,
 };
-use futures_util::{StreamExt, TryStreamExt};
-use indexmap::IndexMap;
+use futures_util::TryStreamExt;
 use mongodb::bson::from_bson;
 use mongodb_agent_common::schema::{get_property_description, Property, ValidatorSchema};
-use mongodb_support::{BsonScalarType, BsonType};
+use mongodb_support::BsonScalarType;
 
 use mongodb_agent_common::interface_types::{MongoAgentError, MongoConfig};
 
@@ -20,42 +19,34 @@ pub async fn get_metadata_from_validation_schema(
     config: &MongoConfig,
 ) -> Result<BTreeMap<String, Schema>, MongoAgentError> {
     let db = config.client.database(&config.database);
-    let collections_cursor = db.list_collections(None, None).await?;
-
-    let schemas: Vec<WithName<Schema>> = collections_cursor
-        .into_stream()
-        .map(
-            |collection_spec| -> Result<WithName<Schema>, MongoAgentError> {
-                let collection_spec_value = collection_spec?;
-                let name = &collection_spec_value.name;
-                let schema_bson_option = collection_spec_value
-                    .options
-                    .validator
-                    .as_ref()
-                    .and_then(|x| x.get("$jsonSchema"));
-
-                match schema_bson_option {
-                    Some(schema_bson) => {
-                        from_bson::<ValidatorSchema>(schema_bson.clone()).map_err(|err| {
-                            MongoAgentError::BadCollectionSchema(
-                                name.to_owned(),
-                                schema_bson.clone(),
-                                err,
-                            )
-                        })
-                    }
-                    None => Ok(ValidatorSchema {
-                        bson_type: BsonType::Object,
-                        description: None,
-                        required: Vec::new(),
-                        properties: IndexMap::new(),
-                    }),
-                }
-                .map(|validator_schema| make_collection_schema(name, &validator_schema))
-            },
-        )
-        .try_collect::<Vec<WithName<Schema>>>()
-        .await?;
+    let mut collections_cursor = db.list_collections(None, None).await?;
+
+    let mut schemas: Vec<WithName<Schema>> = vec![];
+
+    while let Some(collection_spec) = collections_cursor.try_next().await? {
+        let name = &collection_spec.name;
+        let schema_bson_option = collection_spec
+            .options
+            .validator
+            .as_ref()
+            .and_then(|x| x.get("$jsonSchema"));
+
+        match schema_bson_option {
+            Some(schema_bson) => {
+                let validator_schema =
+                    from_bson::<ValidatorSchema>(schema_bson.clone()).map_err(|err| {
+                        MongoAgentError::BadCollectionSchema(
+                            name.to_owned(),
+                            schema_bson.clone(),
+                            err,
+                        )
+                    })?;
+                let collection_schema = make_collection_schema(name, &validator_schema);
+                schemas.push(collection_schema);
+            }
+            None => {}
+        }
+    }
 
     Ok(WithName::into_map(schemas))
 }

crates/cli/src/lib.rs

@@ -10,8 +10,11 @@ use mongodb_agent_common::interface_types::MongoConfig;
 
 #[derive(Debug, Clone, Parser)]
 pub struct UpdateArgs {
-    #[arg(long = "sample-size", value_name = "N")]
-    sample_size: Option<u32>,
+    #[arg(long = "sample-size", value_name = "N", default_value_t = 10)]
+    sample_size: u32,
+
+    #[arg(long = "no-validator-schema", default_value_t = false)]
+    no_validator_schema: bool,
 }
 
 /// The command invoked by the user.
@@ -36,19 +39,18 @@ pub async fn run(command: Command, context: &Context) -> anyhow::Result<()> {
 
 /// Update the configuration in the current directory by introspecting the database.
 async fn update(context: &Context, args: &UpdateArgs) -> anyhow::Result<()> {
-    let schemas = match args.sample_size {
-        None => introspection::get_metadata_from_validation_schema(&context.mongo_config).await?,
-        Some(sample_size) => {
-            let existing_schemas = configuration::list_existing_schemas(&context.path).await?;
-            introspection::sample_schema_from_db(
-                sample_size,
-                &context.mongo_config,
-                &existing_schemas,
-            )
-            .await?
-        }
-    };
-    configuration::write_schema_directory(&context.path, schemas).await?;
-
-    Ok(())
+    if !args.no_validator_schema {
+        let schemas_from_json_validation =
+            introspection::get_metadata_from_validation_schema(&context.mongo_config).await?;
+        configuration::write_schema_directory(&context.path, schemas_from_json_validation).await?;
+    }
+
+    let existing_schemas = configuration::list_existing_schemas(&context.path).await?;
+    let schemas_from_sampling = introspection::sample_schema_from_db(
+        args.sample_size,
+        &context.mongo_config,
+        &existing_schemas,
+    )
+    .await?;
+    configuration::write_schema_directory(&context.path, schemas_from_sampling).await
 }