me: implement multi-schema reset on postgres

libs/test-cli/src/main.rs

@@ -236,8 +236,7 @@ async fn main() -> anyhow::Result<()> {
             let schema = read_datamodel_from_file(&cmd.schema_path).context("Error reading the schema from file")?;
             let api = migration_core::migration_api(Some(schema), None)?;
 
-            // TODO(MultiSchema): Perhaps read namespaces from the schema somehow.
-            api.reset(None).await?;
+            api.reset().await?;
         }
         Command::CreateDatabase(cmd) => {
             let schema = read_datamodel_from_file(&cmd.schema_path).context("Error reading the schema from file")?;

migration-engine/connectors/migration-connector/src/namespaces.rs

@@ -28,3 +28,12 @@ impl Namespaces {
         }
     }
 }
+
+impl IntoIterator for Namespaces {
+    type Item = String;
+    type IntoIter = std::iter::Chain<std::iter::Once<String>, <Vec<String> as IntoIterator>::IntoIter>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        std::iter::once(self.0).chain(self.1.into_iter())
+    }
+}

migration-engine/connectors/sql-migration-connector/src/flavour.rs

@@ -241,7 +241,7 @@ pub(crate) trait SqlFlavour:
     fn raw_cmd<'a>(&'a mut self, sql: &'a str) -> BoxFuture<'a, ConnectorResult<()>>;
 
     /// Drop the database and recreate it empty.
-    fn reset(&mut self) -> BoxFuture<'_, ConnectorResult<()>>;
+    fn reset(&mut self, namespaces: Option<Namespaces>) -> BoxFuture<'_, ConnectorResult<()>>;
 
     /// Optionally scan a migration script that could have been altered by users and emit warnings.
     fn scan_migration_script(&self, _script: &str) {}

migration-engine/connectors/sql-migration-connector/src/flavour/mssql.rs

@@ -216,7 +216,7 @@ impl SqlFlavour for MssqlFlavour {
         })
     }
 
-    fn reset(&mut self) -> BoxFuture<'_, ConnectorResult<()>> {
+    fn reset(&mut self, _namespaces: Option<Namespaces>) -> BoxFuture<'_, ConnectorResult<()>> {
         with_connection(&mut self.state, move |params, connection| async move {
             let schema_name = params.url.schema();
 
@@ -402,7 +402,7 @@ impl SqlFlavour for MssqlFlavour {
                 shadow_database.set_params(shadow_db_params)?;
                 shadow_database.ensure_connection_validity().await?;
 
-                if shadow_database.reset().await.is_err() {
+                if shadow_database.reset(None).await.is_err() {
                     crate::best_effort_reset(&mut shadow_database, namespaces).await?;
                 }
 

migration-engine/connectors/sql-migration-connector/src/flavour/mysql.rs

@@ -231,7 +231,7 @@ impl SqlFlavour for MysqlFlavour {
         })
     }
 
-    fn reset(&mut self) -> BoxFuture<'_, ConnectorResult<()>> {
+    fn reset(&mut self, _namespaces: Option<Namespaces>) -> BoxFuture<'_, ConnectorResult<()>> {
         with_connection(&mut self.state, move |params, circumstances, connection| async move {
             if circumstances.contains(Circumstances::IsVitess) {
                 return Err(ConnectorError::from_msg(
@@ -307,7 +307,7 @@ impl SqlFlavour for MysqlFlavour {
                 shadow_database.ensure_connection_validity().await?;
 
                 tracing::info!("Connecting to user-provided shadow database.");
-                if shadow_database.reset().await.is_err() {
+                if shadow_database.reset(None).await.is_err() {
                     crate::best_effort_reset(&mut shadow_database, namespaces).await?;
                 }
 

migration-engine/connectors/sql-migration-connector/src/flavour/postgres.rs

@@ -10,7 +10,7 @@ use migration_connector::{
 };
 use quaint::connector::PostgresUrl;
 use sql_schema_describer::SqlSchema;
-use std::{collections::HashMap, future, time};
+use std::{borrow::Cow, collections::HashMap, future, time};
 use url::Url;
 use user_facing_errors::{
     common::{DatabaseAccessDenied, DatabaseDoesNotExist},
@@ -264,14 +264,27 @@ impl SqlFlavour for PostgresFlavour {
         })
     }
 
-    fn reset(&mut self) -> BoxFuture<'_, ConnectorResult<()>> {
+    fn reset(&mut self, namespaces: Option<Namespaces>) -> BoxFuture<'_, ConnectorResult<()>> {
         with_connection(self, move |params, _circumstances, conn| async move {
-            let schema_name = params.url.schema();
+            let schemas_to_reset = match namespaces {
+                Some(ns) => ns.into_iter().map(Cow::Owned).collect(),
+                None => vec![Cow::Borrowed(params.url.schema())],
+            };
 
-            conn.raw_cmd(&format!("DROP SCHEMA \"{}\" CASCADE", schema_name), &params.url)
-                .await?;
-            conn.raw_cmd(&format!("CREATE SCHEMA \"{}\"", schema_name), &params.url)
-                .await?;
+            tracing::info!(?schemas_to_reset, "Resetting schema(s)");
+
+            for schema_name in schemas_to_reset {
+                conn.raw_cmd(&format!("DROP SCHEMA \"{}\" CASCADE", schema_name), &params.url)
+                    .await?;
+                conn.raw_cmd(&format!("CREATE SCHEMA \"{}\"", schema_name), &params.url)
+                    .await?;
+            }
+
+            // Drop the migrations table in the main schema, otherwise migrate dev will not
+            // perceive that as a reset, since migrations are still marked as applied.
+            //
+            // We don't care if this fails.
+            conn.raw_cmd("DROP TABLE _prisma_migrations", &params.url).await.ok();
 
             Ok(())
         })
@@ -342,7 +355,7 @@ impl SqlFlavour for PostgresFlavour {
 
                 tracing::info!("Connecting to user-provided shadow database.");
 
-                if shadow_database.reset().await.is_err() {
+                if shadow_database.reset(namespaces.clone()).await.is_err() {
                     crate::best_effort_reset(&mut shadow_database, namespaces.clone()).await?;
                 }
 

migration-engine/connectors/sql-migration-connector/src/flavour/sqlite.rs

@@ -245,7 +245,7 @@ impl SqlFlavour for SqliteFlavour {
         ready(with_connection(&mut self.state, |_, conn| conn.raw_cmd(sql)))
     }
 
-    fn reset(&mut self) -> BoxFuture<'_, ConnectorResult<()>> {
+    fn reset(&mut self, _namespaces: Option<Namespaces>) -> BoxFuture<'_, ConnectorResult<()>> {
         ready(with_connection(&mut self.state, move |params, connection| {
             let file_path = &params.file_path;
 

migration-engine/connectors/sql-migration-connector/src/lib.rs

@@ -262,7 +262,7 @@ impl MigrationConnector for SqlMigrationConnector {
 
     fn reset(&mut self, soft: bool, namespaces: Option<Namespaces>) -> BoxFuture<'_, ConnectorResult<()>> {
         Box::pin(async move {
-            if soft || self.flavour.reset().await.is_err() {
+            if soft || self.flavour.reset(namespaces.clone()).await.is_err() {
                 best_effort_reset(self.flavour.as_mut(), namespaces).await?;
             }
 

migration-engine/core/src/api.rs

@@ -1,7 +1,5 @@
 //! The external facing programmatic API to the migration engine.
 
-use migration_connector::Namespaces;
-
 use crate::{commands, json_rpc::types::*, CoreResult};
 
 /// The programmatic, generic, fantastic migration engine API.
@@ -70,7 +68,7 @@ pub trait GenericApi: Send + Sync + 'static {
     ) -> CoreResult<MarkMigrationRolledBackOutput>;
 
     /// Reset a database to an empty state (no data, no schema).
-    async fn reset(&self, namespaces: Option<Namespaces>) -> CoreResult<()>;
+    async fn reset(&self) -> CoreResult<()>;
 
     /// The command behind `prisma db push`.
     async fn schema_push(&self, input: SchemaPushInput) -> CoreResult<SchemaPushOutput>;

migration-engine/core/src/rpc.rs

@@ -3,9 +3,9 @@ use jsonrpc_core::{types::error::Error as JsonRpcError, IoHandler, Params};
 use std::sync::Arc;
 
 /// Initialize a JSON-RPC ready migration engine API.
-pub fn rpc_api(datamodel: Option<String>, host: Arc<dyn migration_connector::ConnectorHost>) -> IoHandler {
+pub fn rpc_api(prisma_schema: Option<String>, host: Arc<dyn migration_connector::ConnectorHost>) -> IoHandler {
     let mut io_handler = IoHandler::default();
-    let api = Arc::new(crate::state::EngineState::new(datamodel, Some(host)));
+    let api = Arc::new(crate::state::EngineState::new(prisma_schema, Some(host)));
 
     for cmd in METHOD_NAMES {
         let api = api.clone();
@@ -41,7 +41,7 @@ async fn run_command(
         MARK_MIGRATION_APPLIED => render(executor.mark_migration_applied(params.parse()?).await),
         MARK_MIGRATION_ROLLED_BACK => render(executor.mark_migration_rolled_back(params.parse()?).await),
         // TODO(MultiSchema): we probably need to grab the namespaces from the params
-        RESET => render(executor.reset(None).await),
+        RESET => render(executor.reset().await),
         SCHEMA_PUSH => render(executor.schema_push(params.parse()?).await),
         other => unreachable!("Unknown command {}", other),
     }

migration-engine/core/src/state.rs

@@ -19,7 +19,7 @@ use tracing_futures::Instrument;
 /// channels. That ensures that each connector is handling requests one at a time to avoid
 /// synchronization issues. You can think of it in terms of the actor model.
 pub(crate) struct EngineState {
-    initial_datamodel: Option<String>,
+    initial_datamodel: Option<psl::ValidatedSchema>,
     host: Arc<dyn ConnectorHost>,
     // A map from either:
     //
@@ -46,7 +46,7 @@ type ErasedConnectorRequest = Box<
 impl EngineState {
     pub(crate) fn new(initial_datamodel: Option<String>, host: Option<Arc<dyn ConnectorHost>>) -> Self {
         EngineState {
-            initial_datamodel,
+            initial_datamodel: initial_datamodel.map(|s| psl::validate(s.into())),
             host: host.unwrap_or_else(|| Arc::new(migration_connector::EmptyHost)),
             connectors: Default::default(),
         }
@@ -170,7 +170,7 @@ impl EngineState {
             return Err(ConnectorError::from_msg("Missing --datamodel".to_owned()));
         };
 
-        self.with_connector_for_schema(schema, None, f).await
+        self.with_connector_for_schema(schema.db.source(), None, f).await
     }
 }
 
@@ -366,9 +366,18 @@ impl GenericApi for EngineState {
         .await
     }
 
-    async fn reset(&self, namespaces: Option<Namespaces>) -> CoreResult<()> {
+    async fn reset(&self) -> CoreResult<()> {
         tracing::debug!("Resetting the database.");
 
+        let namespaces: Option<Namespaces> = self
+            .initial_datamodel
+            .as_ref()
+            .and_then(|schema| schema.configuration.datasources.first())
+            .and_then(|ds| {
+                let mut names = ds.namespaces.iter().map(|(ns, _)| ns.to_owned()).collect();
+                Namespaces::from_vec(&mut names)
+            });
+
         self.with_default_connector(Box::new(move |connector| {
             Box::pin(MigrationConnector::reset(connector, false, namespaces).instrument(tracing::info_span!("Reset")))
         }))

migration-engine/migration-engine-tests/src/assertions.rs

@@ -71,6 +71,7 @@ impl SchemaAssertion {
         }
     }
 
+    #[track_caller]
     fn assert_error(&self, table_name: &str, positive: bool) -> ! {
         let method = if positive {
             "assert_table"

migration-engine/migration-engine-tests/src/test_api.rs

@@ -1,7 +1,10 @@
 pub use crate::assertions::{MigrationsAssertions, ResultSetExt, SchemaAssertion};
 pub use expect_test::expect;
-pub use migration_core::json_rpc::types::{
-    DbExecuteDatasourceType, DbExecuteParams, DiffParams, DiffResult, SchemaContainer, UrlContainer,
+pub use migration_core::{
+    json_rpc::types::{
+        DbExecuteDatasourceType, DbExecuteParams, DiffParams, DiffResult, SchemaContainer, UrlContainer,
+    },
+    migration_connector::Namespaces,
 };
 pub use test_macros::test_connector;
 pub use test_setup::{runtime::run_with_thread_local_runtime as tok, BitFlags, Capabilities, Tags};
@@ -10,7 +13,7 @@ use crate::{commands::*, multi_engine_test_api::TestApi as RootTestApi};
 use migration_core::{
     commands::diff,
     migration_connector::{
-        BoxFuture, ConnectorHost, ConnectorResult, DiffTarget, MigrationConnector, MigrationPersistence, Namespaces,
+        BoxFuture, ConnectorHost, ConnectorResult, DiffTarget, MigrationConnector, MigrationPersistence,
     },
 };
 use psl::parser_database::SourceFile;

migration-engine/migration-engine-tests/tests/migrations/reset_tests.rs

@@ -93,3 +93,63 @@ fn reset_then_diagnostics_with_migrations_directory_works(api: TestApi) {
         .assert_has_table("Cat")
         .assert_has_table("_prisma_migrations");
 }
+
+#[test_connector(
+    tags(Postgres),
+    exclude(CockroachDb),
+    namespaces("felines", "rodents"),
+    preview_features("multiSchema")
+)]
+fn multi_schema_reset(mut api: TestApi) {
+    let prisma_schema = format! {
+        r#"
+            {}
+
+            generator js {{
+                provider = "prisma-client-js"
+                previewFeatures = ["multiSchema"]
+            }}
+
+            model Manul {{
+                id Int @id
+                @@schema("felines")
+            }}
+
+            model Capybara {{
+                id Int @id
+                @@schema("rodents")
+            }}
+        "#, api.datasource_block_with(&[("schemas", r#"["felines", "rodents"]"#)])
+    };
+    api.schema_push(&prisma_schema)
+        .send()
+        .assert_green()
+        .assert_has_executed_steps();
+
+    let migrations_dir = api.create_migrations_directory();
+    api.apply_migrations(&migrations_dir).send_sync();
+    api.raw_cmd("CREATE TABLE randomTable (id INTEGER PRIMARY KEY);");
+
+    let all_namespaces = Namespaces::from_vec(&mut vec!["felines".into(), "rodents".into(), api.schema_name()]);
+    let namespaces_in_psl = Namespaces::from_vec(&mut vec!["felines".into(), "rodents".into()]);
+
+    api.assert_schema_with_namespaces(all_namespaces.clone())
+        .assert_has_table("randomtable")
+        .assert_has_table("_prisma_migrations")
+        .assert_has_table("Manul")
+        .assert_has_table("Capybara");
+
+    api.reset().send_sync(namespaces_in_psl);
+
+    api.assert_schema_with_namespaces(all_namespaces)
+        .assert_has_table("randomtable") // we do not want to wipe the schema from search_path
+        .assert_has_no_table("_prisma_migrations")
+        .assert_has_no_table("Manul")
+        .assert_has_no_table("Capybara");
+
+    // Check that we can migrate from there.
+    api.schema_push(&prisma_schema)
+        .send()
+        .assert_green()
+        .assert_has_executed_steps();
+}

migration-engine/qe-setup/src/mssql.rs

@@ -22,7 +22,7 @@ pub(crate) async fn mssql_setup(url: String, prisma_schema: &str, db_schemas: &[
         conn.raw_cmd(&format!("USE [{db_name}];")).await.unwrap();
     } else {
         let api = migration_core::migration_api(Some(prisma_schema.to_owned()), None)?;
-        api.reset(None).await.ok();
+        api.reset().await.ok();
         // Without these, our poor connection gets deadlocks if other schemas
         // are modified while we introspect.
         let allow_snapshot_isolation = format!(