Add setup command & schema admin to users config (#387)

Author: levkk

integration/users.toml

@@ -8,6 +8,13 @@ name = "pgdog"
 database = "pgdog_sharded"
 password = "pgdog"
 
+[[users]]
+name = "pgdog_migrator"
+database = "pgdog_sharded"
+password = "pgdog"
+server_user = "pgdog"
+schema_admin = true
+
 [[users]]
 name = "pgdog"
 database = "failover"

pgdog/src/admin/show_pools.rs

@@ -41,6 +41,7 @@ impl Command for ShowPools {
             Field::numeric("out_of_sync"),
             Field::bool("online"),
             Field::text("replica_lag"),
+            Field::bool("schema_admin"),
         ]);
         let mut messages = vec![rd.message()?];
         for (user, cluster) in databases().all() {
@@ -71,7 +72,8 @@ impl Command for ShowPools {
                         .add(state.re_synced)
                         .add(state.out_of_sync)
                         .add(state.online)
-                        .add(state.replica_lag.simple_display());
+                        .add(state.replica_lag.simple_display())
+                        .add(cluster.schema_admin());
 
                     messages.push(row.message()?);
                 }

pgdog/src/backend/databases.rs

@@ -215,6 +215,17 @@ impl Databases {
         }
     }
 
+    /// Get the schema owner for this database.
+    pub fn schema_owner(&self, database: &str) -> Result<Cluster, Error> {
+        for (user, cluster) in &self.databases {
+            if cluster.schema_admin() && user.database == database {
+                return Ok(cluster.clone());
+            }
+        }
+
+        Err(Error::NoSchemaOwner(database.to_owned()))
+    }
+
     pub fn mirrors(&self, user: impl ToUser) -> Result<Option<&[Cluster]>, Error> {
         let user = user.to_user();
         if let Some(cluster) = self.databases.get(&user) {
@@ -443,6 +454,23 @@ pub fn from_config(config: &ConfigAndUsers) -> Databases {
         }
     }
 
+    // Duplicate schema owner check.
+    let mut dupl_schema_owners = HashMap::<String, usize>::new();
+    for (user, cluster) in &mut databases {
+        if cluster.schema_admin() {
+            let entry = dupl_schema_owners.entry(user.database.clone()).or_insert(0);
+            *entry += 1;
+
+            if *entry > 1 {
+                warn!(
+                    r#"database "{}" has duplicate schema owner "{}", ignoring setting"#,
+                    user.database, user.user
+                );
+                cluster.toggle_schema_admin(false);
+            }
+        }
+    }
+
     let mut mirrors = HashMap::new();
 
     for cluster in databases.values() {

pgdog/src/backend/error.rs

@@ -48,6 +48,9 @@ pub enum Error {
     #[error("no cluster connected")]
     NoCluster,
 
+    #[error("database \"{0}\" has no schema owner")]
+    NoSchemaOwner(String),
+
     #[error("{0}")]
     ScramAuth(#[from] crate::auth::scram::Error),
 

pgdog/src/backend/pool/cluster.rs

@@ -45,6 +45,7 @@ pub struct Cluster {
     multi_tenant: Option<MultiTenant>,
     rw_strategy: ReadWriteStrategy,
     rw_split: ReadWriteSplit,
+    schema_admin: bool,
 }
 
 /// Sharding configuration from the cluster.
@@ -81,6 +82,7 @@ pub struct ClusterConfig<'a> {
     pub multi_tenant: &'a Option<MultiTenant>,
     pub rw_strategy: ReadWriteStrategy,
     pub rw_split: ReadWriteSplit,
+    pub schema_admin: bool,
 }
 
 impl<'a> ClusterConfig<'a> {
@@ -105,6 +107,7 @@ impl<'a> ClusterConfig<'a> {
             multi_tenant,
             rw_strategy: general.read_write_strategy,
             rw_split: general.read_write_split,
+            schema_admin: user.schema_admin,
         }
     }
 }
@@ -125,6 +128,7 @@ impl Cluster {
             multi_tenant,
             rw_strategy,
             rw_split,
+            schema_admin,
         } = config;
 
         Self {
@@ -143,6 +147,7 @@ impl Cluster {
             multi_tenant: multi_tenant.clone(),
             rw_strategy,
             rw_split,
+            schema_admin,
         }
     }
 
@@ -193,6 +198,7 @@ impl Cluster {
             multi_tenant: self.multi_tenant.clone(),
             rw_strategy: self.rw_strategy,
             rw_split: self.rw_split,
+            schema_admin: self.schema_admin,
         }
     }
 
@@ -267,6 +273,16 @@ impl Cluster {
         true
     }
 
+    /// This database/user pair is responsible for schema management.
+    pub fn schema_admin(&self) -> bool {
+        self.schema_admin
+    }
+
+    /// Change schema owner attribute.
+    pub fn toggle_schema_admin(&mut self, owner: bool) {
+        self.schema_admin = owner;
+    }
+
     /// We'll need the query router to figure out
     /// where a query should go.
     pub fn router_needed(&self) -> bool {

pgdog/src/backend/schema/setup.sql

@@ -49,21 +49,25 @@ GRANT USAGE ON SEQUENCE pgdog.validator_bigint_id_seq TO PUBLIC;
 
 -- Generate a primary key from a sequence that will
 -- match the shard number this is ran on.
-CREATE OR REPLACE FUNCTION pgdog.next_id_auto() RETURNS BIGINT AS $body$
+CREATE OR REPLACE FUNCTION pgdog.next_id_seq(sequence_name regclass) RETURNS BIGINT AS $body$
 DECLARE next_value BIGINT;
 DECLARE seq_oid oid;
 DECLARE table_oid oid;
 DECLARE shards INTEGER;
 DECLARE shard INTEGER;
 BEGIN
-    SELECT 'pgdog.validator_bigint_id_seq'::regclass INTO seq_oid;
+    SELECT sequence_name INTO seq_oid;
     SELECT 'pgdog.validator_bigint'::regclass INTO table_oid;
     SELECT
         pgdog.config.shard,
         pgdog.config.shards
     INTO shard, shards
     FROM pgdog.config;
 
+    IF shards IS NULL OR shard IS NULL THEN
+        RAISE EXCEPTION 'pgdog.config not set';
+    END IF;
+
     LOOP
         -- This is atomic.
         SELECT nextval(seq_oid) INTO next_value;
@@ -75,6 +79,12 @@ BEGIN
 END;
 $body$ LANGUAGE plpgsql;
 
+CREATE OR REPLACE FUNCTION pgdog.next_id_auto() RETURNS BIGINT AS $body$
+BEGIN
+    RETURN pgdog.next_id_seq('pgdog.validator_bigint_id_seq'::regclass);
+END;
+$body$ LANGUAGE plpgsql;
+
 -- Generate a primary key from a sequence that will
 -- match the shard number this is ran on.
 CREATE OR REPLACE FUNCTION pgdog.next_uuid_auto() RETURNS UUID AS $body$

pgdog/src/backend/schema/sync/config.rs

@@ -0,0 +1,63 @@
+//! Shard configuration sync.
+
+use tracing::info;
+
+use crate::backend::pool::Request;
+use crate::backend::Cluster;
+use crate::backend::Error;
+use crate::backend::Pool;
+use crate::backend::Schema;
+
+pub struct ShardConfig {
+    shard: usize,
+    shards: usize,
+    pool: Pool,
+}
+
+impl ShardConfig {
+    /// Sync schema and set shard config.
+    pub async fn sync(&self) -> Result<(), Error> {
+        let mut conn = self.pool.get(&Request::default()).await?;
+
+        Schema::setup(&mut conn).await?;
+
+        conn.execute("BEGIN").await?;
+        conn.execute("DELETE FROM pgdog.config").await?;
+        conn.execute(format!(
+            "INSERT INTO pgdog.config (shard, shards) VALUES ({}, {})",
+            self.shard, self.shards
+        ))
+        .await?;
+        conn.execute("COMMIT").await?;
+
+        Ok(())
+    }
+
+    /// Sync all shards.
+    pub async fn sync_all(cluster: &Cluster) -> Result<(), Error> {
+        let shards = cluster.shards().len();
+
+        info!("setting up schema on {} shards", shards);
+
+        let shards: Vec<_> = cluster
+            .shards()
+            .iter()
+            .filter(|s| s.has_primary())
+            .map(|s| s.pools().first().unwrap().clone())
+            .enumerate()
+            .map(|(shard, pool)| ShardConfig {
+                pool,
+                shard,
+                shards,
+            })
+            .collect();
+
+        for shard in &shards {
+            shard.sync().await?;
+        }
+
+        info!("schema setup complete for {} shards", shards.len());
+
+        Ok(())
+    }
+}

pgdog/src/backend/schema/sync/mod.rs

@@ -1,3 +1,4 @@
+pub mod config;
 pub mod error;
 pub mod pg_dump;
 pub mod progress;

pgdog/src/cli.rs

@@ -7,6 +7,7 @@ use thiserror::Error;
 use tokio::{select, signal::ctrl_c};
 use tracing::error;
 
+use crate::backend::schema::sync::config::ShardConfig;
 use crate::backend::schema::sync::pg_dump::{PgDump, SyncState};
 use crate::backend::{databases::databases, replication::logical::Publisher};
 use crate::config::{Config, Users};
@@ -94,19 +95,13 @@ pub enum Commands {
         /// Source database name.
         #[arg(long)]
         from_database: String,
-        /// Source user name.
-        #[arg(long)]
-        from_user: String,
         /// Publication name.
         #[arg(long)]
         publication: String,
 
         /// Destination database.
         #[arg(long)]
         to_database: String,
-        /// Destination user name.
-        #[arg(long)]
-        to_user: String,
 
         /// Dry run. Print schema commands, don't actually execute them.
         #[arg(long)]
@@ -120,6 +115,14 @@ pub enum Commands {
         #[arg(long)]
         data_sync_complete: bool,
     },
+
+    /// Perform cluster configuration steps
+    /// required for sharded operations.
+    Setup {
+        /// Database name.
+        #[arg(long)]
+        database: String,
+    },
 }
 
 /// Fingerprint some queries.
@@ -251,17 +254,15 @@ pub async fn schema_sync(commands: Commands) -> Result<(), Box<dyn std::error::E
     let (source, destination, publication, dry_run, ignore_errors, data_sync_complete) =
         if let Commands::SchemaSync {
             from_database,
-            from_user,
             to_database,
-            to_user,
             publication,
             dry_run,
             ignore_errors,
             data_sync_complete,
         } = commands
         {
-            let source = databases().cluster((from_user.as_str(), from_database.as_str()))?;
-            let dest = databases().cluster((to_user.as_str(), to_database.as_str()))?;
+            let source = databases().schema_owner(&from_database)?;
+            let dest = databases().schema_owner(&to_database)?;
 
             (
                 source,
@@ -283,6 +284,10 @@ pub async fn schema_sync(commands: Commands) -> Result<(), Box<dyn std::error::E
         SyncState::PreData
     };
 
+    if state == SyncState::PreData {
+        ShardConfig::sync_all(&destination).await?;
+    }
+
     for output in output {
         if dry_run {
             let queries = output.statements(state)?;
@@ -296,3 +301,12 @@ pub async fn schema_sync(commands: Commands) -> Result<(), Box<dyn std::error::E
 
     Ok(())
 }
+
+pub async fn setup(database: &str) -> Result<(), Box<dyn std::error::Error>> {
+    let databases = databases();
+    let schema_owner = databases.schema_owner(database)?;
+
+    ShardConfig::sync_all(&schema_owner).await?;
+
+    Ok(())
+}

pgdog/src/config/mod.rs

@@ -954,6 +954,9 @@ pub struct User {
     pub idle_timeout: Option<u64>,
     /// Read-only mode.
     pub read_only: Option<bool>,
+    /// Schema owner.
+    #[serde(default)]
+    pub schema_admin: bool,
 }
 
 impl User {