preparatory support for multi-file schema introspection (#4847)

psl/parser-database/src/files.rs

@@ -54,6 +54,12 @@ impl Files {
 
         String::from_utf8(out).unwrap()
     }
+
+    /// Returns the number of files.
+    #[allow(clippy::len_without_is_empty)]
+    pub fn len(&self) -> usize {
+        self.0.len()
+    }
 }
 
 impl Index<crate::FileId> for Files {

psl/parser-database/src/lib.rs

@@ -205,6 +205,11 @@ impl ParserDatabase {
         self.types.model_attributes.len()
     }
 
+    /// The total number of files for the schema. This is O(1).
+    pub fn files_count(&self) -> usize {
+        self.asts.len()
+    }
+
     /// The source file contents. This methods asserts that there is a single prisma schema file.
     /// As multi-file schemas are implemented, calls to this methods should be replaced with
     /// `ParserDatabase::source()` and `ParserDatabase::iter_sources()`.

psl/parser-database/src/walkers.rs

@@ -79,6 +79,15 @@ impl crate::ParserDatabase {
             .map(|model_id| self.walk(model_id))
     }
 
+    /// Find a composite type by name.
+    pub fn find_composite_type<'db>(&'db self, name: &str) -> Option<CompositeTypeWalker<'db>> {
+        self.interner
+            .lookup(name)
+            .and_then(|name_id| self.names.tops.get(&name_id))
+            .and_then(|(file_id, top_id)| top_id.as_composite_type_id().map(|id| (*file_id, id)))
+            .map(|ct_id| self.walk(ct_id))
+    }
+
     /// Traverse a schema element by id.
     pub fn walk<I>(&self, id: I) -> Walker<'_, I> {
         Walker { db: self, id }

psl/parser-database/src/walkers/composite_type.rs

@@ -28,6 +28,11 @@ impl<'db> CompositeTypeWalker<'db> {
         self.id
     }
 
+    /// The ID of the file containing the composite type.
+    pub fn file_id(self) -> FileId {
+        self.id.0
+    }
+
     /// The composite type node in the AST.
     pub fn ast_composite_type(self) -> &'db ast::CompositeType {
         &self.db.asts[self.id]

psl/parser-database/src/walkers/model.rs

@@ -25,6 +25,11 @@ impl<'db> ModelWalker<'db> {
         self.ast_model().name()
     }
 
+    /// The ID of the file containing the model.
+    pub fn file_id(self) -> FileId {
+        self.id.0
+    }
+
     /// Traverse the fields of the models in the order they were defined.
     pub fn fields(self) -> impl ExactSizeIterator<Item = FieldWalker<'db>> + Clone {
         self.ast_model()

psl/psl-core/src/configuration/configuration_struct.rs

@@ -14,10 +14,22 @@ pub struct Configuration {
 }
 
 impl Configuration {
-    pub fn extend(&mut self, configuration: Configuration) {
-        self.generators.extend(configuration.generators);
-        self.datasources.extend(configuration.datasources);
-        self.warnings.extend(configuration.warnings);
+    pub fn new(
+        generators: Vec<Generator>,
+        datasources: Vec<Datasource>,
+        warnings: Vec<diagnostics::DatamodelWarning>,
+    ) -> Self {
+        Self {
+            generators,
+            datasources,
+            warnings,
+        }
+    }
+
+    pub fn extend(&mut self, other: Configuration) {
+        self.generators.extend(other.generators);
+        self.datasources.extend(other.datasources);
+        self.warnings.extend(other.warnings);
     }
 
     pub fn validate_that_one_datasource_is_provided(&self) -> Result<(), Diagnostics> {
@@ -167,4 +179,8 @@ impl Configuration {
 
         Ok(())
     }
+
+    pub fn first_datasource(&self) -> &Datasource {
+        self.datasources.first().expect("Expected a datasource to exist.")
+    }
 }

psl/psl-core/src/configuration/datasource.rs

@@ -1,3 +1,5 @@
+use schema_ast::ast::WithSpan;
+
 use crate::{
     configuration::StringFromEnvVar,
     datamodel_connector::{Connector, ConnectorCapabilities, RelationMode},
@@ -274,6 +276,12 @@ impl Datasource {
     }
 }
 
+impl WithSpan for Datasource {
+    fn span(&self) -> Span {
+        self.span
+    }
+}
+
 pub(crate) fn from_url<F>(url: &StringFromEnvVar, env: F) -> Result<String, UrlValidationError>
 where
     F: Fn(&str) -> Option<String>,

psl/psl-core/src/configuration/generator.rs

@@ -1,6 +1,8 @@
 use crate::{configuration::StringFromEnvVar, PreviewFeature};
+use diagnostics::Span;
 use enumflags2::BitFlags;
 use parser_database::ast::Expression;
+use schema_ast::ast::WithSpan;
 use serde::{ser::SerializeSeq, Serialize, Serializer};
 use std::collections::HashMap;
 
@@ -45,6 +47,15 @@ pub struct Generator {
 
     #[serde(skip_serializing_if = "Option::is_none")]
     pub documentation: Option<String>,
+
+    #[serde(skip)]
+    pub span: Span,
+}
+
+impl WithSpan for Generator {
+    fn span(&self) -> Span {
+        self.span
+    }
 }
 
 pub fn mcf_preview_features<S>(feats: &Option<BitFlags<PreviewFeature>>, s: S) -> Result<S::Ok, S::Error>

psl/psl-core/src/lib.rs

@@ -91,9 +91,7 @@ pub fn validate_multi_file(files: Vec<(String, SourceFile)>, connectors: Connect
     for ast in db.iter_asts() {
         let new_config = validate_configuration(ast, &mut diagnostics, connectors);
 
-        configuration.datasources.extend(new_config.datasources.into_iter());
-        configuration.generators.extend(new_config.generators.into_iter());
-        configuration.warnings.extend(new_config.warnings.into_iter());
+        configuration.extend(new_config);
     }
 
     let datasources = &configuration.datasources;
@@ -164,12 +162,7 @@ fn validate_configuration(
     connectors: ConnectorRegistry<'_>,
 ) -> Configuration {
     let generators = generator_loader::load_generators_from_ast(schema_ast, diagnostics);
-
     let datasources = datasource_loader::load_datasources_from_ast(schema_ast, diagnostics, connectors);
 
-    Configuration {
-        generators,
-        datasources,
-        warnings: diagnostics.warnings().to_owned(),
-    }
+    Configuration::new(generators, datasources, diagnostics.warnings().to_owned())
 }

psl/psl-core/src/validate/datasource_loader.rs

@@ -32,7 +32,7 @@ pub(crate) fn load_datasources_from_ast(
 
     for src in ast_schema.sources() {
         if let Some(source) = lift_datasource(src, diagnostics, connectors) {
-            sources.push(source)
+            sources.push(source);
         }
     }
 

psl/psl-core/src/validate/generator_loader.rs

@@ -26,7 +26,7 @@ pub(crate) fn load_generators_from_ast(ast_schema: &ast::SchemaAst, diagnostics:
 
     for gen in ast_schema.generators() {
         if let Some(generator) = lift_generator(gen, diagnostics) {
-            generators.push(generator)
+            generators.push(generator);
         }
     }
 
@@ -123,6 +123,7 @@ fn lift_generator(ast_generator: &ast::GeneratorConfig, diagnostics: &mut Diagno
         preview_features,
         config: properties,
         documentation: ast_generator.documentation().map(String::from),
+        span: ast_generator.span,
     })
 }
 

psl/schema-ast/src/ast/find_at_position.rs

@@ -352,7 +352,6 @@ impl<'ast> SourcePosition<'ast> {
 
 #[derive(Debug)]
 pub enum PropertyPosition<'ast> {
-    /// prop
     Property,
     Value(&'ast str),
     FunctionValue(&'ast str),

psl/schema-ast/src/parser.rs

@@ -18,4 +18,5 @@ pub use parse_schema::parse_schema;
 // It is more convenient if this enum is directly available here.
 #[derive(pest_derive::Parser)]
 #[grammar = "parser/datamodel.pest"]
+#[allow(clippy::empty_docs)]
 pub(crate) struct PrismaDatamodelParser;

schema-engine/connectors/mongodb-schema-connector/Cargo.toml

@@ -31,3 +31,5 @@ once_cell = "1.8.0"
 url.workspace = true
 expect-test = "1"
 names = { version = "0.12", default-features = false }
+itertools.workspace = true
+indoc.workspace = true

schema-engine/connectors/mongodb-schema-connector/src/sampler.rs

@@ -11,6 +11,7 @@ use mongodb::{
 use mongodb_schema_describer::MongoSchema;
 use schema_connector::{warnings::Model, IntrospectionContext, IntrospectionResult, Warnings};
 use statistics::*;
+use std::borrow::Cow;
 
 /// From the given database, lists all collections as models, and samples
 /// maximum of SAMPLE_SIZE documents for their fields with the following rules:
@@ -61,14 +62,25 @@ pub(super) async fn sample(
     }
 
     let mut data_model = render::Datamodel::default();
-    statistics.render(ctx.datasource(), &mut data_model, &mut warnings);
 
-    let psl_string = if ctx.render_config {
-        let config = render::Configuration::from_psl(ctx.configuration(), None);
-        format!("{config}\n{data_model}")
-    } else {
-        data_model.to_string()
-    };
+    // Ensures that all previous files are present in the new datamodel, even when empty after re-introspection.
+    for file_id in ctx.previous_schema().db.iter_file_ids() {
+        let file_name = ctx.previous_schema().db.file_name(file_id);
+
+        data_model.create_empty_file(Cow::Borrowed(file_name));
+    }
+
+    statistics.render(ctx, &mut data_model, &mut warnings);
+
+    let is_empty = data_model.is_empty();
+
+    if ctx.render_config {
+        let config = render::Configuration::from_psl(ctx.configuration(), ctx.previous_schema(), None);
+
+        data_model.set_configuration(config);
+    }
+
+    let sources = data_model.render();
 
     let warnings = if !warnings.is_empty() {
         Some(warnings.to_string())
@@ -77,8 +89,8 @@ pub(super) async fn sample(
     };
 
     Ok(IntrospectionResult {
-        data_model: psl::reformat(&psl_string, 2).unwrap(),
-        is_empty: data_model.is_empty(),
+        datamodels: psl::reformat_multiple(sources, 2),
+        is_empty,
         warnings,
         views: None,
     })

schema-engine/connectors/mongodb-schema-connector/src/sampler/statistics.rs

@@ -8,7 +8,7 @@ use renderer::{
 };
 use schema_connector::{
     warnings::{ModelAndField, ModelAndFieldAndType, TypeAndField, TypeAndFieldAndType},
-    CompositeTypeDepth, Warnings,
+    CompositeTypeDepth, IntrospectionContext, Warnings,
 };
 
 use super::field_type::FieldType;
@@ -20,6 +20,7 @@ use once_cell::sync::Lazy;
 use psl::datamodel_connector::constraint_names::ConstraintNames;
 use regex::Regex;
 use std::{
+    borrow::Cow,
     cmp::Ordering,
     collections::{BTreeMap, HashMap, HashSet},
     fmt,
@@ -121,7 +122,7 @@ impl<'a> Statistics<'a> {
 
     pub(super) fn render(
         &'a self,
-        datasource: &'a psl::Datasource,
+        ctx: &'a IntrospectionContext,
         rendered: &mut renderer::Datamodel<'a>,
         warnings: &mut Warnings,
     ) {
@@ -164,7 +165,7 @@ impl<'a> Statistics<'a> {
                     let mut field = renderer::datamodel::Field::new("id", "String");
 
                     field.map("_id");
-                    field.native_type(&datasource.name, "ObjectId", Vec::new());
+                    field.native_type(&ctx.datasource().name, "ObjectId", Vec::new());
                     field.default(renderer::datamodel::DefaultValue::function(Function::new("auto")));
                     field.id(IdFieldDefinition::new());
 
@@ -285,7 +286,7 @@ impl<'a> Statistics<'a> {
             }
 
             if let Some(native_type) = field_type.native_type() {
-                field.native_type(&datasource.name, native_type.to_string(), Vec::new());
+                field.native_type(&ctx.datasource().name, native_type.to_string(), Vec::new());
             }
 
             if field_type.is_array() {
@@ -410,12 +411,22 @@ impl<'a> Statistics<'a> {
             }
         }
 
-        for (_, r#type) in types {
-            rendered.push_composite_type(r#type);
+        for (ct_name, r#type) in types {
+            let file_name = match ctx.previous_schema().db.find_composite_type(ct_name) {
+                Some(walker) => ctx.previous_schema().db.file_name(walker.file_id()),
+                None => ctx.introspection_file_name(),
+            };
+
+            rendered.push_composite_type(Cow::Borrowed(file_name), r#type);
         }
 
-        for (_, model) in models.into_iter() {
-            rendered.push_model(model);
+        for (model_name, model) in models.into_iter() {
+            let file_name = match ctx.previous_schema().db.find_model(model_name) {
+                Some(walker) => ctx.previous_schema().db.file_name(walker.file_id()),
+                None => ctx.introspection_file_name(),
+            };
+
+            rendered.push_model(Cow::Borrowed(file_name), model);
         }
     }
 

schema-engine/connectors/mongodb-schema-connector/tests/introspection/mod.rs

@@ -4,6 +4,7 @@ mod basic;
 mod dirty_data;
 mod index;
 mod model_renames;
+mod multi_file;
 mod remapping_names;
 mod types;
 mod views;