hasura · hallettj · Apr 19, 2024 · Apr 5, 2024 · Apr 6, 2024 · Apr 10, 2024
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -16,6 +16,12 @@ members = [
 ]
 resolver = "2"
 
+# The tag or rev of ndc-models must match the locked tag or rev of the
+# ndc-models dependency of ndc-sdk
+[workspace.dependencies]
+ndc-sdk = { git = "https://github.com/hasura/ndc-sdk-rs.git" }
+ndc-models = { git = "http://github.com/hasura/ndc-spec.git", tag = "v0.1.2" }
+
 # We have a fork of the mongodb driver with a fix for reading metadata from time
 # series collections.
 # See the upstream PR: https://github.com/mongodb/mongo-rust-driver/pull/1003

diff --git a/crates/cli/src/introspection/sampling.rs b/crates/cli/src/introspection/sampling.rs
@@ -7,7 +7,7 @@ use configuration::{
 };
 use futures_util::TryStreamExt;
 use mongodb::bson::{doc, Bson, Document};
-use mongodb_agent_common::interface_types::MongoConfig;
+use mongodb_agent_common::state::ConnectorState;
 use mongodb_support::BsonScalarType::{self, *};
 
 type ObjectField = WithName<schema::ObjectField>;
@@ -19,18 +19,18 @@ type ObjectType = WithName<schema::ObjectType>;
 /// are not unifiable.
 pub async fn sample_schema_from_db(
     sample_size: u32,
-    config: &MongoConfig,
+    state: &ConnectorState,
     existing_schemas: &HashSet<std::string::String>,
 ) -> anyhow::Result<BTreeMap<std::string::String, Schema>> {
     let mut schemas = BTreeMap::new();
-    let db = config.client.database(&config.database);
+    let db = state.database();
     let mut collections_cursor = db.list_collections(None, None).await?;
 
     while let Some(collection_spec) = collections_cursor.try_next().await? {
         let collection_name = collection_spec.name;
         if !existing_schemas.contains(&collection_name) {
             let collection_schema =
-                sample_schema_from_collection(&collection_name, sample_size, config).await?;
+                sample_schema_from_collection(&collection_name, sample_size, state).await?;
             schemas.insert(collection_name, collection_schema);
         }
     }
@@ -40,9 +40,9 @@ pub async fn sample_schema_from_db(
 async fn sample_schema_from_collection(
     collection_name: &str,
     sample_size: u32,
-    config: &MongoConfig,
+    state: &ConnectorState,
 ) -> anyhow::Result<Schema> {
-    let db = config.client.database(&config.database);
+    let db = state.database();
     let options = None;
     let mut cursor = db
         .collection::<Document>(collection_name)

diff --git a/crates/cli/src/introspection/validation_schema.rs b/crates/cli/src/introspection/validation_schema.rs
@@ -6,19 +6,22 @@ use configuration::{
 };
 use futures_util::TryStreamExt;
 use mongodb::bson::from_bson;
-use mongodb_agent_common::schema::{get_property_description, Property, ValidatorSchema};
+use mongodb_agent_common::{
+    schema::{get_property_description, Property, ValidatorSchema},
+    state::ConnectorState,
+};
 use mongodb_support::BsonScalarType;
 
-use mongodb_agent_common::interface_types::{MongoAgentError, MongoConfig};
+use mongodb_agent_common::interface_types::MongoAgentError;
 
 type Collection = WithName<schema::Collection>;
 type ObjectType = WithName<schema::ObjectType>;
 type ObjectField = WithName<schema::ObjectField>;
 
 pub async fn get_metadata_from_validation_schema(
-    config: &MongoConfig,
+    state: &ConnectorState,
 ) -> Result<BTreeMap<String, Schema>, MongoAgentError> {
-    let db = config.client.database(&config.database);
+    let db = state.database();
     let mut collections_cursor = db.list_collections(None, None).await?;
 
     let mut schemas: Vec<WithName<Schema>> = vec![];

diff --git a/crates/cli/src/lib.rs b/crates/cli/src/lib.rs
@@ -6,10 +6,9 @@ use std::path::PathBuf;
 
 use clap::{Parser, Subcommand};
 
-use mongodb_agent_common::interface_types::MongoConfig;
-
 // Exported for use in tests
 pub use introspection::type_from_bson;
+use mongodb_agent_common::state::ConnectorState;
 
 #[derive(Debug, Clone, Parser)]
 pub struct UpdateArgs {
@@ -29,7 +28,7 @@ pub enum Command {
 
 pub struct Context {
     pub path: PathBuf,
-    pub mongo_config: MongoConfig,
+    pub connector_state: ConnectorState,
 }
 
 /// Run a command in a given directory.
@@ -44,14 +43,14 @@ pub async fn run(command: Command, context: &Context) -> anyhow::Result<()> {
 async fn update(context: &Context, args: &UpdateArgs) -> anyhow::Result<()> {
     if !args.no_validator_schema {
         let schemas_from_json_validation =
-            introspection::get_metadata_from_validation_schema(&context.mongo_config).await?;
+            introspection::get_metadata_from_validation_schema(&context.connector_state).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,
+        &context.connector_state,
         &existing_schemas,
     )
     .await?;

diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs
@@ -45,10 +45,13 @@ pub async fn main() -> anyhow::Result<()> {
         Some(path) => path,
         None => env::current_dir()?,
     };
-    let mongo_config = try_init_state_from_uri(&args.connection_uri, &Default::default())
+    let connector_state = try_init_state_from_uri(&args.connection_uri)
         .await
         .map_err(|e| anyhow!("Error initializing MongoDB state {}", e))?;
-    let context = Context { path, mongo_config };
+    let context = Context {
+        path,
+        connector_state,
+    };
     run(args.subcommand, &context).await?;
     Ok(())
 }
diff --git a/crates/configuration/Cargo.toml b/crates/configuration/Cargo.toml
@@ -9,6 +9,7 @@ futures = "^0.3"
 itertools = "^0.12"
 mongodb = "2.8"
 mongodb-support = { path = "../mongodb-support" }
+ndc-models = { workspace = true }
 schemars = "^0.8.12"
 serde = { version = "1", features = ["derive"] }
 serde_json = { version = "1" }