feat: add fk constraints to schema (#768)

- feat: load foreign key constraints in schema on startup
- feat: add `load_schema` general setting: on/off/auto

Author: levkk

integration/go/go_pgx/sharded_test.go

@@ -152,11 +152,11 @@ func TestShardedTwoPc(t *testing.T) {
 		assert.NoError(t, err)
 	}
 
-	// +3 is for schema sync
+	// +4 is for schema sync
 	assertShowField(t, "SHOW STATS", "total_xact_2pc_count", 200, "pgdog_2pc", "pgdog_sharded", 0, "primary")
 	assertShowField(t, "SHOW STATS", "total_xact_2pc_count", 200, "pgdog_2pc", "pgdog_sharded", 1, "primary")
-	assertShowField(t, "SHOW STATS", "total_xact_count", 401+3, "pgdog_2pc", "pgdog_sharded", 0, "primary") // PREPARE, COMMIT for each transaction + TRUNCATE
-	assertShowField(t, "SHOW STATS", "total_xact_count", 401+3, "pgdog_2pc", "pgdog_sharded", 1, "primary")
+	assertShowField(t, "SHOW STATS", "total_xact_count", 401+4, "pgdog_2pc", "pgdog_sharded", 0, "primary") // PREPARE, COMMIT for each transaction + TRUNCATE
+	assertShowField(t, "SHOW STATS", "total_xact_count", 401+4, "pgdog_2pc", "pgdog_sharded", 1, "primary")
 
 	for i := range 200 {
 		rows, err := conn.Query(

pgdog-config/src/general.rs

@@ -5,7 +5,7 @@ use std::path::PathBuf;
 use std::time::Duration;
 
 use crate::pooling::ConnectionRecovery;
-use crate::{CopyFormat, QueryParserEngine, QueryParserLevel, SystemCatalogsBehavior};
+use crate::{CopyFormat, LoadSchema, QueryParserEngine, QueryParserLevel, SystemCatalogsBehavior};
 
 use super::auth::{AuthType, PassthoughAuth};
 use super::database::{LoadBalancingStrategy, ReadWriteSplit, ReadWriteStrategy};
@@ -209,6 +209,9 @@ pub struct General {
     /// Trigger a schema reload on DDL like CREATE TABLE.
     #[serde(default = "General::reload_schema_on_ddl")]
     pub reload_schema_on_ddl: bool,
+    /// Load database schema.
+    #[serde(default = "General::load_schema")]
+    pub load_schema: LoadSchema,
 }
 
 impl Default for General {
@@ -282,6 +285,7 @@ impl Default for General {
             omnisharded_sticky: bool::default(),
             resharding_copy_format: CopyFormat::default(),
             reload_schema_on_ddl: Self::reload_schema_on_ddl(),
+            load_schema: Self::load_schema(),
         }
     }
 }
@@ -566,6 +570,10 @@ impl General {
         )
     }
 
+    fn load_schema() -> LoadSchema {
+        Self::env_enum_or_default("PGDOG_LOAD_SCHEMA")
+    }
+
     pub fn mirror_queue() -> usize {
         Self::env_or_default("PGDOG_MIRROR_QUEUE", 128)
     }

pgdog-config/src/sharding.rs

@@ -370,6 +370,27 @@ impl Display for CopyFormat {
     }
 }
 
+#[derive(Serialize, Deserialize, Debug, Copy, Clone, PartialEq, Eq, Hash, Default)]
+#[serde(rename_all = "snake_case", deny_unknown_fields)]
+pub enum LoadSchema {
+    On,
+    Off,
+    #[default]
+    Auto,
+}
+
+impl FromStr for LoadSchema {
+    type Err = ();
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Ok(match s.to_lowercase().as_str() {
+            "on" => Self::On,
+            "auto" => Self::Auto,
+            "off" => Self::Off,
+            _ => return Err(()),
+        })
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

pgdog-stats/src/schema.rs

@@ -5,6 +5,40 @@ use std::{collections::HashMap, hash::Hash, ops::Deref, sync::Arc};
 /// Schema name -> Table name -> Relation
 pub type Relations = HashMap<String, HashMap<String, Relation>>;
 
+/// The action to take when a referenced row is deleted or updated.
+#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq, Hash)]
+pub enum ForeignKeyAction {
+    #[default]
+    NoAction,
+    Restrict,
+    Cascade,
+    SetNull,
+    SetDefault,
+}
+
+impl ForeignKeyAction {
+    /// Parse from PostgreSQL's information_schema representation.
+    pub fn from_pg_string(s: &str) -> Self {
+        match s {
+            "CASCADE" => Self::Cascade,
+            "SET NULL" => Self::SetNull,
+            "SET DEFAULT" => Self::SetDefault,
+            "RESTRICT" => Self::Restrict,
+            "NO ACTION" | _ => Self::NoAction,
+        }
+    }
+}
+
+/// A foreign key reference to another table's column.
+#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Hash)]
+pub struct ForeignKey {
+    pub schema: String,
+    pub table: String,
+    pub column: String,
+    pub on_delete: ForeignKeyAction,
+    pub on_update: ForeignKeyAction,
+}
+
 #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Hash)]
 pub struct Column {
     pub table_catalog: String,
@@ -16,6 +50,7 @@ pub struct Column {
     pub data_type: String,
     pub ordinal_position: i32,
     pub is_primary_key: bool,
+    pub foreign_keys: Vec<ForeignKey>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]

pgdog/src/backend/pool/cluster.rs

@@ -1,7 +1,9 @@
 //! A collection of replicas and a primary.
 
 use parking_lot::Mutex;
-use pgdog_config::{PreparedStatements, QueryParserEngine, QueryParserLevel, Rewrite, RewriteMode};
+use pgdog_config::{
+    LoadSchema, PreparedStatements, QueryParserEngine, QueryParserLevel, Rewrite, RewriteMode,
+};
 use std::{
     sync::{
         atomic::{AtomicBool, AtomicUsize, Ordering},
@@ -75,6 +77,7 @@ pub struct Cluster {
     client_connection_recovery: ConnectionRecovery,
     query_parser_engine: QueryParserEngine,
     reload_schema_on_ddl: bool,
+    load_schema: LoadSchema,
 }
 
 /// Sharding configuration from the cluster.
@@ -149,6 +152,7 @@ pub struct ClusterConfig<'a> {
     pub client_connection_recovery: ConnectionRecovery,
     pub lsn_check_interval: Duration,
     pub reload_schema_on_ddl: bool,
+    pub load_schema: LoadSchema,
 }
 
 impl<'a> ClusterConfig<'a> {
@@ -198,6 +202,7 @@ impl<'a> ClusterConfig<'a> {
             client_connection_recovery: general.client_connection_recovery,
             lsn_check_interval: Duration::from_millis(general.lsn_check_interval),
             reload_schema_on_ddl: general.reload_schema_on_ddl,
+            load_schema: general.load_schema,
         }
     }
 }
@@ -233,6 +238,7 @@ impl Cluster {
             lsn_check_interval,
             query_parser_engine,
             reload_schema_on_ddl,
+            load_schema,
         } = config;
 
         let identifier = Arc::new(DatabaseUser {
@@ -280,6 +286,7 @@ impl Cluster {
             client_connection_recovery,
             query_parser_engine,
             reload_schema_on_ddl,
+            load_schema,
         }
     }
 
@@ -486,10 +493,11 @@ impl Cluster {
     }
 
     fn load_schema(&self) -> bool {
-        self.shards.len() > 1
-            && self.sharded_schemas.is_empty()
-            && !self.sharded_tables.tables().is_empty()
-            || self.multi_tenant().is_some()
+        match self.load_schema {
+            LoadSchema::On => true,
+            LoadSchema::Off => false,
+            LoadSchema::Auto => self.shards.len() > 1 || self.multi_tenant().is_some(),
+        }
     }
 
     /// Get currently loaded schema from shard 0.
@@ -549,12 +557,11 @@ impl Cluster {
 
                     info!("loaded schema from {}/{} shards", done + 1, shards);
 
+                    schema_changed_hook(&shard.schema(), &identifier, &shard);
+
                     // Loaded schema on all shards.
                     if done >= shards - 1 {
                         readiness.schemas_ready.notify_waiters();
-                        // We assume the schema is the same on all shards.
-                        // TODO: check that this is the case and raise a stink if its not.
-                        schema_changed_hook(&shard.schema(), &identifier);
                     }
                 });
             }
@@ -745,7 +752,8 @@ mod test {
         let config = ConfigAndUsers::default();
         let cluster = Cluster::new_test(&config);
 
-        assert!(!cluster.load_schema());
+        // In Auto mode with multiple shards, load_schema returns true
+        assert!(cluster.load_schema());
     }
 
     #[test]
@@ -755,7 +763,9 @@ mod test {
         cluster.sharded_schemas = ShardedSchemas::default();
         cluster.sharded_tables = ShardedTables::default();
 
-        assert!(!cluster.load_schema());
+        // In Auto mode with multiple shards, load_schema returns true
+        // (sharded_schemas and sharded_tables no longer affect this decision)
+        assert!(cluster.load_schema());
     }
 
     #[test]
@@ -839,9 +849,9 @@ mod test {
     #[tokio::test]
     async fn test_wait_schema_loaded_returns_immediately_when_not_needed() {
         let config = ConfigAndUsers::default();
-        let cluster = Cluster::new_test(&config);
+        let cluster = Cluster::new_test_single_shard(&config);
 
-        // load_schema() returns false because sharded_schemas is not empty
+        // load_schema() returns false for single shard without multi_tenant
         assert!(!cluster.load_schema());
 
         // Should return immediately without waiting

pgdog/src/backend/pool/ee/mod.rs

@@ -1,3 +1,3 @@
-use crate::backend::{databases::User, Schema};
+use crate::backend::{databases::User, Schema, Shard};
 
-pub(crate) fn schema_changed_hook(_schema: &Schema, _user: &User) {}
+pub(crate) fn schema_changed_hook(_schema: &Schema, _user: &User, _shard: &Shard) {}