feat: add configurable write_functions

.schema/pgdog.schema.json

@@ -264,6 +264,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,
@@ -2327,6 +2335,28 @@
       "required": [
         "values"
       ]
+    },
+    "WriteFunctions": {
+      "description": "Per-database write function classification.",
+      "type": "object",
+      "properties": {
+        "database": {
+          "description": "Database name.",
+          "type": "string"
+        },
+        "functions": {
+          "description": "Functions that should be treated as write-only.\nEach entry can be schema-qualified.",
+          "type": "array",
+          "default": [],
+          "items": {
+            "type": "string"
+          }
+        }
+      },
+      "additionalProperties": false,
+      "required": [
+        "database"
+      ]
     }
   }
 }

example.pgdog.toml

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

integration/load_balancer/pgdog.toml

@@ -72,6 +72,10 @@ database_name = "postgres"
 role = "auto"
 port = 45002
 
+[[write_functions]]
+database = "postgres"
+functions = ["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 lb_write_fn($1)", int64(i))
+		assert.NoError(t, err)
+	}
+
+	primaryCalls := LoadStatsForPrimary("SELECT lb_write_fn")
+	assert.Equal(t, int64(20), primaryCalls.Calls)
+
+	replicaCalls := LoadStatsForReplicas("SELECT 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

@@ -565,6 +565,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,
+    /// Functions that should be treated as write-only.
+    /// Each entry can be schema-qualified.
+    #[serde(default)]
+    pub functions: Vec<String>,
+}
+
 #[cfg(test)]
 mod tests {
     use crate::Config;
@@ -588,4 +600,18 @@ database = "production"
             QueryParserEngine::PgQueryProtobuf
         );
     }
+
+    #[test]
+    fn write_functions_reads_default_values_from_config() {
+        let source = r#"
+[[write_functions]]
+database = "production"
+"#;
+
+        let config: Config = toml::from_str(source).unwrap();
+
+        assert_eq!(config.write_functions.len(), 1);
+        assert_eq!(config.write_functions[0].database, "production");
+        assert!(config.write_functions[0].functions.is_empty());
+    }
 }

pgdog/src/backend/pool/cluster.rs

@@ -7,13 +7,14 @@ use pgdog_config::{
     RewriteMode, users::PasswordKind,
 };
 use std::{
+    collections::HashSet,
     sync::{
         Arc,
         atomic::{AtomicBool, Ordering},
     },
     time::Duration,
 };
-use tracing::error;
+use tracing::{error, warn};
 
 use crate::frontend::router::sharding::ShardedTable;
 use crate::tasks;
@@ -62,6 +63,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,
@@ -101,6 +103,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>,
@@ -113,6 +116,54 @@ impl ShardingSchema {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct WriteFunction {
+    pub schema: Option<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(entry: &str) -> Option<Self> {
+        fn split_qualified(entry: &str) -> impl Iterator<Item = &str> {
+            let mut parts = Vec::new();
+            let mut start = 0;
+            let mut in_quotes = false;
+            for (i, c) in entry.char_indices() {
+                match c {
+                    '"' => in_quotes = !in_quotes,
+                    '.' if !in_quotes => {
+                        parts.push(&entry[start..i]);
+                        start = i + 1;
+                    }
+                    _ => (),
+                }
+            }
+            parts.push(&entry[start..]);
+            parts.into_iter().rev()
+        }
+        let mut parts = split_qualified(entry);
+        let name = Self::normalize_identifier(parts.next()?);
+        let schema = match parts.next() {
+            Some(schema) if parts.next().is_none() => Some(Self::normalize_identifier(schema)),
+            Some(_) => return None,
+            None => None,
+        };
+
+        Some(Self { schema, name })
+    }
+}
+
 #[derive(Debug)]
 pub struct ClusterShardConfig {
     pub primary: Option<PoolConfig>,
@@ -148,6 +199,7 @@ pub struct ClusterConfig<'a> {
     pub multi_tenant: &'a Option<MultiTenant>,
     pub rw_strategy: ReadWriteStrategy,
     pub rw_split: ReadWriteSplit,
+    pub write_functions: HashSet<WriteFunction>,
     pub schema_admin: bool,
     pub cross_shard_disabled: bool,
     pub two_pc: bool,
@@ -207,6 +259,19 @@ 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)
+                .flat_map(|entry| entry.functions.iter())
+                .filter_map(|func| {
+                    let parsed = WriteFunction::from_config(func);
+                    if parsed.is_none() {
+                        warn!("ignoring invalid write_functions entry: \"{}\"", func);
+                    }
+                    parsed
+                })
+                .collect(),
             schema_admin: user.schema_admin,
             cross_shard_disabled: user
                 .cross_shard_disabled
@@ -258,6 +323,7 @@ impl Cluster {
             multi_tenant,
             rw_strategy,
             rw_split,
+            write_functions,
             schema_admin,
             cross_shard_disabled,
             two_pc,
@@ -318,6 +384,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,
@@ -551,6 +618,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,
@@ -741,7 +809,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,
@@ -762,9 +830,20 @@ 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)
+                .flat_map(|entry| entry.functions.iter())
+                .filter_map(|func| WriteFunction::from_config(func))
+                .collect()
+        }
+
         pub fn new_test(config: &ConfigAndUsers) -> Self {
             let identifier = Arc::new(DatabaseUser {
                 user: "pgdog".into(),
@@ -874,6 +953,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()
@@ -953,6 +1033,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()
@@ -1000,6 +1081,29 @@ mod test {
         assert!(cluster.load_schema());
     }
 
+    #[test]
+    fn test_write_function_identifier_normalization() {
+        let wf = WriteFunction::from_config("Create_Partition").unwrap();
+        assert_eq!(wf.schema, None);
+        assert_eq!(wf.name, "create_partition");
+
+        let wf = WriteFunction::from_config("PartMan.Create_Partition").unwrap();
+        assert_eq!(wf.schema.as_deref(), Some("partman"));
+        assert_eq!(wf.name, "create_partition");
+
+        let wf = WriteFunction::from_config(r#""PartMan"."Create_Partition""#).unwrap();
+        assert_eq!(wf.schema.as_deref(), Some("PartMan"));
+        assert_eq!(wf.name, "Create_Partition");
+
+        // Dots inside quoted identifiers are not separators.
+        let wf = WriteFunction::from_config(r#""my.schema".my_fn"#).unwrap();
+        assert_eq!(wf.schema.as_deref(), Some("my.schema"));
+        assert_eq!(wf.name, "my_fn");
+
+        // More than two parts is invalid.
+        assert_eq!(WriteFunction::from_config("a.b.c"), None);
+    }
+
     #[test]
     fn test_load_schema_multiple_shards_with_schemas() {
         let config = ConfigAndUsers::default();