feat: add configurable write_functions

.schema/pgdog.schema.json

@@ -266,6 +266,14 @@
           "type": "null"
         }
       ]
+    },
+    "write_functions": {
+      "description": "Per-database functions treated as writes by the query parser.",
+      "type": "array",
+      "default": [],
+      "items": {
+        "$ref": "#/$defs/WriteFunctions"
+      }
     }
   },
   "additionalProperties": false,
@@ -2343,6 +2351,30 @@
       "required": [
         "values"
       ]
+    },
+    "WriteFunctions": {
+      "description": "Per-database write function classification.",
+      "type": "object",
+      "properties": {
+        "database": {
+          "description": "Database name.",
+          "type": "string"
+        },
+        "name": {
+          "description": "Function that should be treated as write-only.",
+          "type": "string"
+        },
+        "schema": {
+          "description": "Schema containing the function.",
+          "type": "string"
+        }
+      },
+      "additionalProperties": false,
+      "required": [
+        "database",
+        "schema",
+        "name"
+      ]
     }
   }
 }

example.pgdog.toml

@@ -282,6 +282,19 @@ cross_shard_disabled = false
 #
 dns_ttl = 5_000
 
+#
+# Per-database function that should always route to the primary.
+#
+# Matching follows PostgreSQL identifier semantics:
+# - unquoted identifiers are folded to lowercase
+# - quoted identifiers preserve case
+#
+# [[write_functions]]
+# database = "pgdog"
+# schema = "partman"
+# name = "create_partition"
+#
+
 #
 # Admin database used for stats and system admin.
 #

integration/load_balancer/pgdog.toml

@@ -72,6 +72,11 @@ database_name = "postgres"
 role = "auto"
 port = 45002
 
+[[write_functions]]
+database = "postgres"
+schema = "public"
+name = "Lb_Write_Fn"
+
 
 [[plugins]]
 name = "pgdog_primary_only_tables"

integration/load_balancer/pgx/read_write_split_test.go

@@ -5,6 +5,8 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"os"
+	"strings"
 	"testing"
 	"time"
 
@@ -151,6 +153,39 @@ func TestWriteFunctions(t *testing.T) {
 	assert.Equal(t, int64(25), calls.Calls)
 }
 
+func TestConfiguredWriteFunctionsRouteToPrimary(t *testing.T) {
+	if !strings.Contains(os.Getenv("PGDOG_PLUGIN_FEATURES"), "new_parser") {
+		t.Skip("write_functions routing is implemented in the new parser path only")
+	}
+
+	pool := GetPool()
+	defer pool.Close()
+
+	_, err := pool.Exec(context.Background(), `
+		CREATE OR REPLACE FUNCTION lb_write_fn(val bigint)
+		RETURNS bigint
+		LANGUAGE sql
+		AS $$ SELECT $1; $$`)
+	assert.NoError(t, err)
+
+	// DDL replication can lag briefly.
+	time.Sleep(2 * time.Second)
+	ResetStats()
+
+	for i := range 20 {
+		_, err = pool.Exec(context.Background(), "SELECT public.lb_write_fn($1)", int64(i))
+		assert.NoError(t, err)
+	}
+
+	primaryCalls := LoadStatsForPrimary("SELECT public.lb_write_fn")
+	assert.Equal(t, int64(20), primaryCalls.Calls)
+
+	replicaCalls := LoadStatsForReplicas("SELECT public.lb_write_fn")
+	for _, call := range replicaCalls {
+		assert.Equal(t, int64(0), call.Calls)
+	}
+}
+
 func withTransaction(t *testing.T, pool *pgxpool.Pool, f func(t pgx.Tx) error) error {
 	tx, err := pool.Begin(context.Background())
 	assert.NoError(t, err)

pgdog-config/src/core.rs

@@ -11,7 +11,7 @@ use crate::util::random_string;
 use crate::{
     EnumeratedDatabase, Memory, OmnishardedTable, PassthroughAuth, PreparedStatements, QueryParser,
     QueryParserEngine, QueryParserLevel, ReadWriteSplit, RewriteMode, Role, ShardedMappingKey,
-    ShardedTableConfig, SystemCatalogsBehavior, system_catalogs,
+    ShardedTableConfig, SystemCatalogsBehavior, WriteFunctions, system_catalogs,
 };
 
 use super::database::Database;
@@ -292,6 +292,10 @@ pub struct Config {
     /// Query parser levels per-database.
     #[serde(default)]
     pub query_parsers: Vec<QueryParser>,
+
+    /// Per-database functions treated as writes by the query parser.
+    #[serde(default)]
+    pub write_functions: Vec<WriteFunctions>,
 }
 
 impl Config {

pgdog-config/src/sharding.rs

@@ -597,6 +597,18 @@ pub struct QueryParser {
     pub engine: QueryParserEngine,
 }
 
+/// Per-database write function classification.
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash, Default, JsonSchema)]
+#[serde(rename_all = "snake_case", deny_unknown_fields)]
+pub struct WriteFunctions {
+    /// Database name.
+    pub database: String,
+    /// Schema containing the function.
+    pub schema: String,
+    /// Function that should be treated as write-only.
+    pub name: String,
+}
+
 #[cfg(test)]
 mod tests {
     use crate::Config;
@@ -620,4 +632,38 @@ database = "production"
             QueryParserEngine::PgQueryProtobuf
         );
     }
+
+    #[test]
+    fn write_functions_reads_from_config() {
+        let source = r#"
+[[write_functions]]
+database = "production"
+schema = "partman"
+name = "create_partition"
+"#;
+
+        let config: Config = toml::from_str(source).unwrap();
+
+        assert_eq!(config.write_functions.len(), 1);
+        assert_eq!(config.write_functions[0].database, "production");
+        assert_eq!(config.write_functions[0].schema, "partman");
+        assert_eq!(config.write_functions[0].name, "create_partition");
+    }
+
+    #[test]
+    fn write_functions_requires_schema_and_name() {
+        let missing_schema = r#"
+[[write_functions]]
+database = "production"
+name = "create_partition"
+"#;
+        assert!(toml::from_str::<Config>(missing_schema).is_err());
+
+        let missing_name = r#"
+[[write_functions]]
+database = "production"
+schema = "partman"
+"#;
+        assert!(toml::from_str::<Config>(missing_name).is_err());
+    }
 }

pgdog/src/backend/pool/cluster.rs

@@ -6,7 +6,7 @@ use pgdog_config::{
     LoadSchema, PreparedStatements, QueryParser, QueryParserEngine, QueryParserLevel, Rewrite,
     RewriteMode, users::PasswordKind,
 };
-use std::{sync::Arc, time::Duration};
+use std::{collections::HashSet, sync::Arc, time::Duration};
 
 use crate::backend::schema::SchemaCache;
 use crate::frontend::router::sharding::ShardedTable;
@@ -52,6 +52,7 @@ pub struct Cluster {
     multi_tenant: Option<MultiTenant>,
     rw_strategy: ReadWriteStrategy,
     rw_split: ReadWriteSplit,
+    write_functions: HashSet<WriteFunction>,
     schema_admin: bool,
     stats: Arc<Mutex<MirrorStats>>,
     cross_shard_disabled: bool,
@@ -91,6 +92,7 @@ pub struct ShardingSchema {
     pub schemas: ShardedSchemas,
     /// Rewrite config.
     pub rewrite: Rewrite,
+    pub write_functions: HashSet<WriteFunction>,
     /// Query parser engine.
     pub query_parser_engine: QueryParserEngine,
     pub log_min_duration_parse: Option<Duration>,
@@ -103,6 +105,32 @@ impl ShardingSchema {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct WriteFunction {
+    pub schema: String,
+    pub name: String,
+}
+
+impl WriteFunction {
+    /// Normalize one SQL identifier using PostgreSQL rules:
+    /// - unquoted: folded to lowercase
+    /// - quoted: preserve case and unescape doubled quotes
+    fn normalize_identifier(identifier: &str) -> String {
+        if identifier.len() >= 2 && identifier.starts_with('"') && identifier.ends_with('"') {
+            identifier[1..identifier.len() - 1].replace("\"\"", "\"")
+        } else {
+            identifier.to_ascii_lowercase()
+        }
+    }
+
+    fn from_config(schema: &str, name: &str) -> Self {
+        Self {
+            schema: Self::normalize_identifier(schema),
+            name: Self::normalize_identifier(name),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct ClusterShardConfig {
     pub primary: Option<PoolConfig>,
@@ -144,6 +172,7 @@ pub struct ClusterConfig<'a> {
     two_pc_auto: bool,
     sharded_schemas: ShardedSchemas,
     rewrite: &'a Rewrite,
+    write_functions: HashSet<WriteFunction>,
     prepared_statements: &'a PreparedStatements,
     dry_run: bool,
     expanded_explain: bool,
@@ -203,6 +232,12 @@ impl<'a> ClusterConfig<'a> {
             multi_tenant,
             rw_strategy: general.read_write_strategy,
             rw_split: general.read_write_split,
+            write_functions: config
+                .write_functions
+                .iter()
+                .filter(|entry| entry.database == user.database)
+                .map(|entry| WriteFunction::from_config(&entry.schema, &entry.name))
+                .collect(),
             schema_admin: user.schema_admin,
             cross_shard_disabled: user
                 .cross_shard_disabled
@@ -255,6 +290,7 @@ impl Cluster {
             multi_tenant,
             rw_strategy,
             rw_split,
+            write_functions,
             schema_admin,
             cross_shard_disabled,
             two_pc,
@@ -317,6 +353,7 @@ impl Cluster {
             multi_tenant: multi_tenant.clone(),
             rw_strategy,
             rw_split,
+            write_functions,
             schema_admin,
             stats: Arc::new(Mutex::new(MirrorStats::default())),
             cross_shard_disabled,
@@ -550,6 +587,7 @@ impl Cluster {
             tables: self.sharded_tables.clone(),
             schemas: self.sharded_schemas.clone(),
             rewrite: self.rewrite.clone(),
+            write_functions: self.write_functions.clone(),
             query_parser_engine: self.query_parser_engine,
             log_min_duration_parse: self.log_min_duration_parse,
             log_query_sample_length: self.log_query_sample_length,
@@ -665,7 +703,7 @@ impl Cluster {
 
 #[cfg(test)]
 mod test {
-    use std::{sync::Arc, time::Duration};
+    use std::{collections::HashSet, sync::Arc, time::Duration};
 
     use pgdog_config::{
         ConfigAndUsers, OmnishardedTable, PoolerMode, QueryParserLevel, ShardedSchema,
@@ -687,9 +725,19 @@ mod test {
         net::Query,
     };
 
-    use super::{Cluster, DatabaseUser};
+    use super::{Cluster, DatabaseUser, WriteFunction};
 
     impl Cluster {
+        fn test_write_functions(config: &ConfigAndUsers, database: &str) -> HashSet<WriteFunction> {
+            config
+                .config
+                .write_functions
+                .iter()
+                .filter(|entry| entry.database == database)
+                .map(|entry| WriteFunction::from_config(&entry.schema, &entry.name))
+                .collect()
+        }
+
         pub fn new_test(config: &ConfigAndUsers) -> Self {
             let identifier = Arc::new(DatabaseUser {
                 user: "pgdog".into(),
@@ -800,6 +848,7 @@ mod test {
                     config.config.general.query_parser,
                 ),
                 rewrite: config.config.rewrite.clone(),
+                write_functions: Self::test_write_functions(config, "pgdog"),
                 two_phase_commit: config.config.general.two_phase_commit,
                 two_phase_commit_auto: config.config.general.two_phase_commit_auto.unwrap_or(false),
                 ..Default::default()
@@ -881,6 +930,7 @@ mod test {
                     config.config.general.query_parser,
                 ),
                 rewrite: config.config.rewrite.clone(),
+                write_functions: Self::test_write_functions(config, "pgdog"),
                 two_phase_commit: config.config.general.two_phase_commit,
                 two_phase_commit_auto: config.config.general.two_phase_commit_auto.unwrap_or(false),
                 ..Default::default()
@@ -929,6 +979,21 @@ mod test {
         assert!(cluster.load_schema());
     }
 
+    #[test]
+    fn test_write_function_identifier_normalization() {
+        let wf = WriteFunction::from_config("PartMan", "Create_Partition");
+        assert_eq!(wf.schema, "partman");
+        assert_eq!(wf.name, "create_partition");
+
+        let wf = WriteFunction::from_config(r#""PartMan""#, r#""Create_Partition""#);
+        assert_eq!(wf.schema, "PartMan");
+        assert_eq!(wf.name, "Create_Partition");
+
+        let wf = WriteFunction::from_config(r#""my.schema""#, "my_fn");
+        assert_eq!(wf.schema, "my.schema");
+        assert_eq!(wf.name, "my_fn");
+    }
+
     #[test]
     fn test_load_schema_multiple_shards_with_schemas() {
         let config = ConfigAndUsers::default();