convert procedure result to JSON with consistent scalar representations

Cargo.toml

@@ -3,15 +3,16 @@ version = "0.0.3"
 
 [workspace]
 members = [
+  "crates/cli",
   "crates/configuration",
-  "crates/mongodb-connector",
-  "crates/mongodb-agent-common",
-  "crates/mongodb-support",
   "crates/dc-api",
-  "crates/dc-api-types",
   "crates/dc-api-test-helpers",
+  "crates/dc-api-types",
+  "crates/mongodb-agent-common",
+  "crates/mongodb-connector",
+  "crates/mongodb-support",
   "crates/ndc-test-helpers",
-  "crates/cli"
+  "crates/test-helpers",
 ]
 resolver = "2"
 

crates/cli/Cargo.toml

@@ -18,7 +18,8 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = { version = "1.0.113", features = ["raw_value"] }
 thiserror = "1.0.57"
 tokio = { version = "1.36.0", features = ["full"] }
-these = "2.0.0"
 
 [dev-dependencies]
+test-helpers = { path = "../test-helpers" }
+
 proptest = "1"

crates/cli/src/introspection/mod.rs

@@ -1,6 +1,7 @@
 pub mod sampling;
-pub mod validation_schema;
 pub mod type_unification;
+pub mod validation_schema;
 
+pub use sampling::{sample_schema_from_db, type_from_bson};
 pub use validation_schema::get_metadata_from_validation_schema;
-pub use sampling::sample_schema_from_db;
+

crates/cli/src/introspection/sampling.rs

@@ -114,6 +114,15 @@ fn make_object_field(
     (collected_otds, object_field)
 }
 
+// Exported for use in tests
+pub fn type_from_bson(
+    object_type_name: &str,
+    value: &Bson,
+) -> (BTreeMap<std::string::String, schema::ObjectType>, Type) {
+    let (object_types, t) = make_field_type(object_type_name, value);
+    (WithName::into_map(object_types), t)
+}
+
 fn make_field_type(
     object_type_name: &str,
     field_value: &Bson,

crates/cli/src/introspection/type_unification.rs

@@ -8,10 +8,7 @@ use configuration::{
 };
 use indexmap::IndexMap;
 use itertools::Itertools as _;
-use mongodb_support::{
-    align::align,
-    BsonScalarType::*,
-};
+use mongodb_support::{align::align, BsonScalarType::*};
 use std::string::String;
 
 type ObjectField = WithName<schema::ObjectField>;
@@ -131,10 +128,7 @@ fn unify_object_type(object_type_a: ObjectType, object_type_b: ObjectType) -> Ob
 
 /// Unify the types of two `ObjectField`s.
 /// If the types are not unifiable then return an error.
-fn unify_object_field(
-    object_field_a: ObjectField,
-    object_field_b: ObjectField,
-) -> ObjectField {
+fn unify_object_field(object_field_a: ObjectField, object_field_b: ObjectField) -> ObjectField {
     WithName::named(
         object_field_a.name,
         schema::ObjectField {
@@ -185,6 +179,7 @@ mod tests {
     };
     use mongodb_support::BsonScalarType;
     use proptest::{collection::hash_map, prelude::*};
+    use test_helpers::arb_type;
 
     #[test]
     fn test_unify_scalar() -> Result<(), anyhow::Error> {
@@ -208,45 +203,11 @@ mod tests {
         Ok(())
     }
 
-    fn arb_bson_scalar_type() -> impl Strategy<Value = BsonScalarType> {
-        prop_oneof![
-            Just(BsonScalarType::Double),
-            Just(BsonScalarType::Decimal),
-            Just(BsonScalarType::Int),
-            Just(BsonScalarType::Long),
-            Just(BsonScalarType::String),
-            Just(BsonScalarType::Date),
-            Just(BsonScalarType::Timestamp),
-            Just(BsonScalarType::BinData),
-            Just(BsonScalarType::ObjectId),
-            Just(BsonScalarType::Bool),
-            Just(BsonScalarType::Null),
-            Just(BsonScalarType::Regex),
-            Just(BsonScalarType::Javascript),
-            Just(BsonScalarType::JavascriptWithScope),
-            Just(BsonScalarType::MinKey),
-            Just(BsonScalarType::MaxKey),
-            Just(BsonScalarType::Undefined),
-            Just(BsonScalarType::DbPointer),
-            Just(BsonScalarType::Symbol),
-        ]
-    }
-
-    fn arb_type() -> impl Strategy<Value = Type> {
-        let leaf = prop_oneof![
-            arb_bson_scalar_type().prop_map(Type::Scalar),
-            any::<String>().prop_map(Type::Object)
-        ];
-        leaf.prop_recursive(3, 10, 10, |inner| {
-            prop_oneof![
-                inner.clone().prop_map(|t| Type::ArrayOf(Box::new(t))),
-                inner.prop_map(|t| Type::Nullable(Box::new(t)))
-            ]
-        })
-    }
-
     fn is_nullable(t: &Type) -> bool {
-        matches!(t, Type::Scalar(BsonScalarType::Null) | Type::Nullable(_) | Type::Any)
+        matches!(
+            t,
+            Type::Scalar(BsonScalarType::Null) | Type::Nullable(_) | Type::Any
+        )
     }
 
     proptest! {

crates/cli/src/lib.rs

@@ -8,6 +8,9 @@ use clap::{Parser, Subcommand};
 
 use mongodb_agent_common::interface_types::MongoConfig;
 
+// Exported for use in tests
+pub use introspection::type_from_bson;
+
 #[derive(Debug, Clone, Parser)]
 pub struct UpdateArgs {
     #[arg(long = "sample-size", value_name = "N", default_value_t = 10)]

crates/configuration/src/schema/database.rs

@@ -36,6 +36,10 @@ pub enum Type {
 }
 
 impl Type {
+    pub fn is_nullable(&self) -> bool {
+        matches!(self, Type::Any | Type::Nullable(_) | Type::Scalar(BsonScalarType::Null))
+    }
+
     pub fn normalize_type(self) -> Type {
         match self {
             Type::Any => Type::Any,

crates/configuration/src/with_name.rs

@@ -52,7 +52,7 @@ impl<T> From<(String, T)> for WithName<T> {
     }
 }
 
-#[derive(Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
+#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
 pub struct WithNameRef<'a, T> {
     pub name: &'a str,
     pub value: &'a T,

crates/mongodb-agent-common/Cargo.toml

@@ -5,22 +5,23 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+configuration = { path = "../configuration" }
+dc-api = { path = "../dc-api" }
+dc-api-types = { path = "../dc-api-types" }
+mongodb-support = { path = "../mongodb-support" }
+
 anyhow = "1.0.71"
 async-trait = "^0.1"
 axum = { version = "0.6", features = ["headers"] }
 bytes = "^1"
-configuration = { path = "../configuration" }
-dc-api = { path = "../dc-api" }
-dc-api-types = { path = "../dc-api-types" }
-enum-iterator = "1.4.1"
+enum-iterator = "^2.0.0"
 futures = "0.3.28"
 futures-util = "0.3.28"
 http = "^0.2"
 indexmap = { version = "1", features = ["serde"] } # must match the version that ndc-client uses
 indent = "^0.1"
 itertools = "^0.10"
 mongodb = "2.8"
-mongodb-support = { path = "../mongodb-support" }
 once_cell = "1"
 regex = "1"
 schemars = { version = "^0.8.12", features = ["smol_str"] }
@@ -32,6 +33,10 @@ time = { version = "0.3.29", features = ["formatting", "parsing", "serde"] }
 tracing = "0.1"
 
 [dev-dependencies]
+mongodb-cli-plugin = { path = "../cli" }
+test-helpers = { path = "../test-helpers" }
+
 mockall = "0.11.4"
 pretty_assertions = "1"
+proptest = "1"
 tokio = { version = "1", features = ["full"] }

crates/mongodb-agent-common/proptest-regressions/query/serialization/tests.txt

@@ -0,0 +1,10 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc 2efdea7f185f2f38ae643782b3523014ab7b8236e36a79cc6b7a7cac74b06f79 # shrinks to bytes = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 238, 161, 0]
+cc 26e2543468ab6d4ffa34f9f8a2c920801ef38a35337557a8f4e74c92cf57e344 # shrinks to bson = Document({" ": Document({"¡": DateTime(1970-01-01 0:00:00.001 +00:00:00)})})
+cc 7d760e540b56fedac7dd58e5bdb5bb9613b9b0bc6a88acfab3fc9c2de8bf026d # shrinks to bson = Document({"A": Array([Null, Undefined])})
+cc 21360610045c5a616b371fb8d5492eb0c22065d62e54d9c8a8761872e2e192f3 # shrinks to bson = Array([Document({}), Document({" ": Null})])

crates/mongodb-agent-common/src/procedure/mod.rs

@@ -5,7 +5,7 @@ use std::borrow::Cow;
 use std::collections::BTreeMap;
 
 use configuration::native_queries::NativeQuery;
-use configuration::schema::{ObjectField, ObjectType};
+use configuration::schema::{ObjectField, ObjectType, Type};
 use mongodb::options::SelectionCriteria;
 use mongodb::{bson, Database};
 
@@ -19,43 +19,47 @@ pub use self::interpolated_command::interpolated_command;
 pub struct Procedure<'a> {
     arguments: BTreeMap<String, serde_json::Value>,
     command: Cow<'a, bson::Document>,
-    object_types: Cow<'a, BTreeMap<String, ObjectType>>,
     parameters: Cow<'a, BTreeMap<String, ObjectField>>,
+    result_type: Type,
     selection_criteria: Option<Cow<'a, SelectionCriteria>>,
 }
 
 impl<'a> Procedure<'a> {
-    /// Note: the `object_types` argument here is not the object types from the native query - it
-    /// should be the set of *all* object types collected from schema and native query definitions.
     pub fn from_native_query(
         native_query: &'a NativeQuery,
-        object_types: &'a BTreeMap<String, ObjectType>,
         arguments: BTreeMap<String, serde_json::Value>,
     ) -> Self {
         Procedure {
             arguments,
             command: Cow::Borrowed(&native_query.command),
-            object_types: Cow::Borrowed(object_types),
             parameters: Cow::Borrowed(&native_query.arguments),
+            result_type: native_query.result_type.clone(),
             selection_criteria: native_query.selection_criteria.as_ref().map(Cow::Borrowed),
         }
     }
 
-    pub async fn execute(self, database: Database) -> Result<bson::Document, ProcedureError> {
+    pub async fn execute(
+        self,
+        object_types: &BTreeMap<String, ObjectType>,
+        database: Database,
+    ) -> Result<(bson::Document, Type), ProcedureError> {
         let selection_criteria = self.selection_criteria.map(Cow::into_owned);
         let command = interpolate(
-            &self.object_types,
+            object_types,
             &self.parameters,
             self.arguments,
             &self.command,
         )?;
         let result = database.run_command(command, selection_criteria).await?;
-        Ok(result)
+        Ok((result, self.result_type))
     }
 
-    pub fn interpolated_command(self) -> Result<bson::Document, ProcedureError> {
+    pub fn interpolated_command(
+        self,
+        object_types: &BTreeMap<String, ObjectType>,
+    ) -> Result<bson::Document, ProcedureError> {
         interpolate(
-            &self.object_types,
+            object_types,
             &self.parameters,
             self.arguments,
             &self.command,

crates/mongodb-agent-common/src/query/arguments/mod.rs → crates/mongodb-agent-common/src/query/arguments.rs

@@ -1,5 +1,3 @@
-mod json_to_bson;
-
 use std::collections::BTreeMap;
 
 use configuration::schema::{ObjectField, ObjectType, Type};
@@ -9,9 +7,7 @@ use mongodb::bson::Bson;
 use serde_json::Value;
 use thiserror::Error;
 
-use self::json_to_bson::json_to_bson;
-
-pub use self::json_to_bson::{json_to_bson_scalar, JsonToBsonError};
+use super::serialization::{json_to_bson, JsonToBsonError};
 
 #[derive(Debug, Error)]
 pub enum ArgumentError {

crates/mongodb-agent-common/src/query/make_selector.rs

@@ -10,7 +10,7 @@ use mongodb_support::BsonScalarType;
 
 use crate::{
     comparison_function::ComparisonFunction, interface_types::MongoAgentError,
-    query::arguments::json_to_bson_scalar, query::column_ref::column_ref,
+    query::serialization::json_to_bson_scalar, query::column_ref::column_ref,
 };
 
 use BinaryArrayComparisonOperator as ArrOp;

crates/mongodb-agent-common/src/query/mod.rs

@@ -8,6 +8,7 @@ mod make_selector;
 mod make_sort;
 mod pipeline;
 mod relations;
+pub mod serialization;
 
 use dc_api::JsonResponse;
 use dc_api_types::{QueryRequest, QueryResponse, Target};