prune superfluous inferred object types from native query configuration

crates/cli/Cargo.toml

@@ -21,6 +21,7 @@ ndc-models = { workspace = true }
 nom = "^7.1.3"
 nonempty = "^0.10.0"
 ref-cast = { workspace = true }
+regex = "^1.11.1"
 serde = { workspace = true }
 serde_json = { workspace = true }
 thiserror = "1.0.57"

crates/cli/src/native_query/helpers.rs

@@ -1,5 +1,8 @@
+use std::{borrow::Cow, collections::BTreeMap};
+
 use configuration::Configuration;
 use ndc_models::{CollectionInfo, CollectionName, ObjectTypeName};
+use regex::Regex;
 
 use super::error::{Error, Result};
 
@@ -24,3 +27,32 @@ pub fn find_collection_object_type(
     let collection = find_collection(configuration, collection_name)?;
     Ok(collection.collection_type.clone())
 }
+
+pub fn unique_type_name<A, B>(
+    object_types: &BTreeMap<ObjectTypeName, A>,
+    added_object_types: &BTreeMap<ObjectTypeName, B>,
+    desired_type_name: &str,
+) -> ObjectTypeName {
+    let (name, mut counter) = parse_counter_suffix(desired_type_name);
+    let mut type_name: ObjectTypeName = name.as_ref().into();
+    while object_types.contains_key(&type_name) || added_object_types.contains_key(&type_name) {
+        counter += 1;
+        type_name = format!("{desired_type_name}_{counter}").into();
+    }
+    type_name
+}
+
+/// [unique_type_name] adds a `_n` numeric suffix where necessary. There are cases where we go
+/// through multiple layers of unique names. Instead of accumulating multiple suffixes, we can
+/// increment the existing suffix. If there is no suffix then the count starts at zero.
+pub fn parse_counter_suffix(name: &str) -> (Cow<'_, str>, u32) {
+    let re = Regex::new(r"^(.*?)_(\d+)$").unwrap();
+    let Some(captures) = re.captures(name) else {
+        return (Cow::Borrowed(name), 0);
+    };
+    let prefix = captures.get(1).unwrap().as_str();
+    let Some(count) = captures.get(2).and_then(|s| s.as_str().parse().ok()) else {
+        return (Cow::Borrowed(name), 0);
+    };
+    (Cow::Owned(prefix.to_string()), count)
+}

crates/cli/src/native_query/mod.rs

@@ -3,6 +3,7 @@ pub mod error;
 mod helpers;
 mod pipeline;
 mod pipeline_type_context;
+mod prune_object_types;
 mod reference_shorthand;
 mod type_constraint;
 mod type_solver;
@@ -206,7 +207,7 @@ mod tests {
             pipeline.clone(),
         )?;
 
-        let expected_document_type_name: ObjectTypeName = "selected_title_documents_2".into();
+        let expected_document_type_name: ObjectTypeName = "selected_title_documents".into();
 
         let expected_object_types = [(
             expected_document_type_name.clone(),

crates/cli/src/native_query/pipeline/mod.rs

@@ -329,7 +329,7 @@ mod tests {
         let config = mflix_config();
         let pipeline_types = infer_pipeline_types(&config, "documents", None, &pipeline).unwrap();
         let expected = [(
-            "documents_documents_2".into(),
+            "documents_documents".into(),
             ObjectType {
                 fields: [
                     (

crates/cli/src/native_query/pipeline_type_context.rs

@@ -14,6 +14,8 @@ use ndc_models::{ArgumentName, ObjectTypeName};
 
 use super::{
     error::{Error, Result},
+    helpers::unique_type_name,
+    prune_object_types::prune_object_types,
     type_constraint::{ObjectTypeConstraint, TypeConstraint, TypeVariable, Variance},
     type_solver::unify,
 };
@@ -106,17 +108,12 @@ impl PipelineTypeContext<'_> {
             e => e,
         })?;
 
-        let result_document_type = variable_types
+        let mut result_document_type = variable_types
             .get(&result_document_type_variable)
-            .expect("missing result type variable is missing");
-        let result_document_type_name = match result_document_type {
-            Type::Object(type_name) => type_name.clone().into(),
-            t => Err(Error::ExpectedObject {
-                actual_type: t.clone(),
-            })?,
-        };
+            .expect("missing result type variable is missing")
+            .clone();
 
-        let parameter_types = self
+        let mut parameter_types: BTreeMap<ArgumentName, Type> = self
             .parameter_types
             .into_iter()
             .map(|(parameter_name, type_variable)| {
@@ -127,10 +124,31 @@ impl PipelineTypeContext<'_> {
             })
             .collect();
 
+        // Prune added object types to remove types that are not referenced by the return type or
+        // by parameter types, and therefore don't need to be included in the native query
+        // configuration.
+        let object_types = {
+            let mut reference_types = std::iter::once(&mut result_document_type)
+                .chain(parameter_types.values_mut())
+                .collect_vec();
+            prune_object_types(
+                &mut reference_types,
+                &self.configuration.object_types,
+                added_object_types,
+            )?
+        };
+
+        let result_document_type_name = match result_document_type {
+            Type::Object(type_name) => type_name.clone().into(),
+            t => Err(Error::ExpectedObject {
+                actual_type: t.clone(),
+            })?,
+        };
+
         Ok(PipelineTypes {
             result_document_type: result_document_type_name,
             parameter_types,
-            object_types: added_object_types,
+            object_types,
             warnings: self.warnings,
         })
     }
@@ -185,15 +203,11 @@ impl PipelineTypeContext<'_> {
     }
 
     pub fn unique_type_name(&self, desired_type_name: &str) -> ObjectTypeName {
-        let mut counter = 0;
-        let mut type_name: ObjectTypeName = desired_type_name.into();
-        while self.configuration.object_types.contains_key(&type_name)
-            || self.object_types.contains_key(&type_name)
-        {
-            counter += 1;
-            type_name = format!("{desired_type_name}_{counter}").into();
-        }
-        type_name
+        unique_type_name(
+            &self.configuration.object_types,
+            &self.object_types,
+            desired_type_name,
+        )
     }
 
     pub fn set_stage_doc_type(&mut self, doc_type: TypeConstraint) {