change big brain to use postgres schemas (#37710)

GitOrigin-RevId: bfb8d89ed7a43d110433b5da975229e4465a4567

Author: gautamg795

crates/clusters/src/db_driver_tag.rs

@@ -7,6 +7,7 @@ use common::types::PersistenceVersion;
 pub enum DbDriverTag {
     Sqlite,
     Postgres(PersistenceVersion),
+    PostgresMultiSchema(PersistenceVersion),
     PostgresAwsIam(PersistenceVersion),
     MySql(PersistenceVersion),
     MySqlAwsIam(PersistenceVersion),
@@ -21,6 +22,7 @@ impl clap::ValueEnum for DbDriverTag {
             DbDriverTag::MySql(PersistenceVersion::V5),
             DbDriverTag::MySqlAwsIam(PersistenceVersion::V5),
             DbDriverTag::Postgres(PersistenceVersion::V5),
+            DbDriverTag::PostgresMultiSchema(PersistenceVersion::V5),
             DbDriverTag::PostgresAwsIam(PersistenceVersion::V5),
             #[cfg(any(test, feature = "testing"))]
             DbDriverTag::TestPersistence,
@@ -36,6 +38,7 @@ impl DbDriverTag {
     pub fn persistence_version(&self) -> anyhow::Result<PersistenceVersion> {
         match self {
             Self::Postgres(version)
+            | Self::PostgresMultiSchema(version)
             | Self::PostgresAwsIam(version)
             | Self::MySql(version)
             | Self::MySqlAwsIam(version) => Ok(*version),
@@ -53,6 +56,7 @@ impl DbDriverTag {
         match self {
             DbDriverTag::Sqlite => "sqlite",
             DbDriverTag::Postgres(PersistenceVersion::V5) => "postgres-v5",
+            DbDriverTag::PostgresMultiSchema(PersistenceVersion::V5) => "postgres-v5-multi-schema",
             DbDriverTag::PostgresAwsIam(PersistenceVersion::V5) => "postgres-v5-aws-iam",
             DbDriverTag::MySql(PersistenceVersion::V5) => "mysql-v5",
             DbDriverTag::MySqlAwsIam(PersistenceVersion::V5) => "mysql-v5-aws-iam",
@@ -69,6 +73,9 @@ impl FromStr for DbDriverTag {
         match s {
             "sqlite" => Ok(Self::Sqlite),
             "postgres-v5" => Ok(DbDriverTag::Postgres(PersistenceVersion::V5)),
+            "postgres-v5-multi-schema" => {
+                Ok(DbDriverTag::PostgresMultiSchema(PersistenceVersion::V5))
+            },
             "postgres-v5-aws-iam" => Ok(DbDriverTag::PostgresAwsIam(PersistenceVersion::V5)),
             "mysql-v5" => Ok(DbDriverTag::MySql(PersistenceVersion::V5)),
             "mysql-v5-aws-iam" => Ok(DbDriverTag::MySqlAwsIam(PersistenceVersion::V5)),

crates/clusters/src/lib.rs

@@ -33,14 +33,10 @@ pub fn persistence_args_from_cluster_url(
         "cluster url username must be set",
     );
     match driver {
-        DbDriverTag::Postgres(_) | DbDriverTag::PostgresAwsIam(_) => {
-            // NOTE: for PostgresAwsIam we do not set any database so we can
-            // reuse connections between databases
-            let schema = if matches!(driver, DbDriverTag::PostgresAwsIam(_)) {
-                // N.B.: unlike mysql we use the instance name as-is as a schema
-                // name (we don't change - to _)
-                Some(instance_name.to_string())
-            } else {
+        DbDriverTag::Postgres(_)
+        | DbDriverTag::PostgresMultiSchema(_)
+        | DbDriverTag::PostgresAwsIam(_) => {
+            let schema = if matches!(driver, DbDriverTag::Postgres(_)) {
                 // selfhosted case
                 let db_name = instance_name.replace('-', "_");
                 anyhow::ensure!(
@@ -50,6 +46,11 @@ pub fn persistence_args_from_cluster_url(
                 );
                 cluster_url.set_path(&db_name);
                 None
+            } else {
+                // NOTE: we do not set any database in this case
+                // N.B.: unlike mysql we use the instance name as-is as a schema
+                // name (we don't change - to _)
+                Some(instance_name.to_string())
             };
             if require_ssl {
                 cluster_url

crates/db_connection/src/lib.rs

@@ -43,6 +43,7 @@ pub async fn connect_persistence<RT: Runtime>(
             persistence
         },
         DbDriverTag::Postgres(version)
+        | DbDriverTag::PostgresMultiSchema(version)
         | DbDriverTag::PostgresAwsIam(version)
         | DbDriverTag::MySql(version)
         | DbDriverTag::MySqlAwsIam(version) => {
@@ -111,6 +112,7 @@ pub async fn connect_persistence_reader<RT: Runtime>(
     let persistence: Arc<dyn PersistenceReader> = match db {
         DbDriverTag::Sqlite => Arc::new(SqlitePersistence::new(db_spec, false)?),
         DbDriverTag::Postgres(version)
+        | DbDriverTag::PostgresMultiSchema(version)
         | DbDriverTag::PostgresAwsIam(version)
         | DbDriverTag::MySql(version)
         | DbDriverTag::MySqlAwsIam(version) => {