From 5007cfc3c944eced4c552359657419af27c066b4 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Sun, 19 Jul 2026 11:13:16 +0530 Subject: [PATCH 01/10] feat: add configurable write_functions --- example.pgdog.toml | 10 +++++ integration/load_balancer/pgdog.toml | 4 ++ .../pgx/read_write_split_test.go | 29 ++++++++++++++ pgdog-config/src/core.rs | 6 ++- pgdog-config/src/sharding.rs | 25 ++++++++++++ pgdog/src/backend/pool/cluster.rs | 29 +++++++++++++- pgdog/src/frontend/router/parser/function.rs | 39 ++++++++++++++++++- .../frontend/router/parser/query/select.rs | 6 ++- .../parser/query/test/test_functions.rs | 31 +++++++++++++++ 9 files changed, 174 insertions(+), 5 deletions(-) diff --git a/example.pgdog.toml b/example.pgdog.toml index cc648c890..953fd2ee2 100644 --- a/example.pgdog.toml +++ b/example.pgdog.toml @@ -276,6 +276,16 @@ cross_shard_disabled = false # dns_ttl = 5_000 +# +# Per-database list of functions that should always route to the primary. +# +# Function matching is case-insensitive. +# +# [[write_functions]] +# database = "pgdog" +# functions = ["my_write_fn", "set_customer_balance"] +# + # # Admin database used for stats and system admin. # diff --git a/integration/load_balancer/pgdog.toml b/integration/load_balancer/pgdog.toml index bf1f9cb99..86fa17ab4 100644 --- a/integration/load_balancer/pgdog.toml +++ b/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" diff --git a/integration/load_balancer/pgx/read_write_split_test.go b/integration/load_balancer/pgx/read_write_split_test.go index e6a7e5245..3f02c38f8 100644 --- a/integration/load_balancer/pgx/read_write_split_test.go +++ b/integration/load_balancer/pgx/read_write_split_test.go @@ -151,6 +151,35 @@ func TestWriteFunctions(t *testing.T) { assert.Equal(t, int64(25), calls.Calls) } +func TestConfiguredWriteFunctionsRouteToPrimary(t *testing.T) { + 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) diff --git a/pgdog-config/src/core.rs b/pgdog-config/src/core.rs index a6da90096..6681e7cf1 100644 --- a/pgdog-config/src/core.rs +++ b/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, + + /// Per-database functions treated as writes by the query parser. + #[serde(default)] + pub write_functions: Vec, } impl Config { diff --git a/pgdog-config/src/sharding.rs b/pgdog-config/src/sharding.rs index 32b85d8b0..609f5e227 100644 --- a/pgdog-config/src/sharding.rs +++ b/pgdog-config/src/sharding.rs @@ -565,6 +565,17 @@ 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. + #[serde(default)] + pub functions: Vec, +} + #[cfg(test)] mod tests { use crate::Config; @@ -588,4 +599,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()); + } } diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 9f037de79..16cd5514a 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -7,6 +7,7 @@ use pgdog_config::{ RewriteMode, users::PasswordKind, }; use std::{ + collections::HashSet, sync::{ Arc, atomic::{AtomicBool, Ordering}, @@ -62,6 +63,7 @@ pub struct Cluster { multi_tenant: Option, rw_strategy: ReadWriteStrategy, rw_split: ReadWriteSplit, + write_functions: HashSet, schema_admin: bool, stats: Arc>, cross_shard_disabled: bool, @@ -101,6 +103,7 @@ pub struct ShardingSchema { pub schemas: ShardedSchemas, /// Rewrite config. pub rewrite: Rewrite, + pub write_functions: HashSet, /// Query parser engine. pub query_parser_engine: QueryParserEngine, pub log_min_duration_parse: Option, @@ -148,6 +151,7 @@ pub struct ClusterConfig<'a> { pub multi_tenant: &'a Option, pub rw_strategy: ReadWriteStrategy, pub rw_split: ReadWriteSplit, + pub write_functions: HashSet, pub schema_admin: bool, pub cross_shard_disabled: bool, pub two_pc: bool, @@ -207,6 +211,13 @@ 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()) + .map(|func| func.to_ascii_lowercase()) + .collect(), schema_admin: user.schema_admin, cross_shard_disabled: user .cross_shard_disabled @@ -258,6 +269,7 @@ impl Cluster { multi_tenant, rw_strategy, rw_split, + write_functions, schema_admin, cross_shard_disabled, two_pc, @@ -318,6 +330,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 +564,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 +755,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, @@ -765,6 +779,17 @@ mod test { use super::{Cluster, DatabaseUser}; impl Cluster { + fn test_write_functions(config: &ConfigAndUsers, database: &str) -> HashSet { + config + .config + .write_functions + .iter() + .filter(|entry| entry.database == database) + .flat_map(|entry| entry.functions.iter()) + .map(|func| func.to_ascii_lowercase()) + .collect() + } + pub fn new_test(config: &ConfigAndUsers) -> Self { let identifier = Arc::new(DatabaseUser { user: "pgdog".into(), @@ -874,6 +899,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 +979,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() diff --git a/pgdog/src/frontend/router/parser/function.rs b/pgdog/src/frontend/router/parser/function.rs index 2a70a5f15..b87ac8b47 100644 --- a/pgdog/src/frontend/router/parser/function.rs +++ b/pgdog/src/frontend/router/parser/function.rs @@ -2,6 +2,8 @@ use pg_query::{Node, NodeEnum, protobuf}; #[cfg(feature = "new_parser")] use pg_raw_parse::{Node, nodes}; +#[cfg(feature = "new_parser")] +use std::collections::HashSet; const WRITE_ONLY: &[&str] = &["nextval", "setval"]; @@ -34,11 +36,35 @@ impl<'a> Function<'a> { /// This function likely writes. pub(crate) fn behavior(&self) -> FunctionBehavior { FunctionBehavior { - writes: WRITE_ONLY.contains(&self.name), + writes: WRITE_ONLY + .iter() + .any(|write_only| self.name.eq_ignore_ascii_case(write_only)), cross_shard: CROSS_SHARD.contains(&(self.schema, self.name)), } } + #[cfg(feature = "new_parser")] + pub(crate) fn behavior_with_write_functions( + &self, + configured_write_functions: &HashSet, + ) -> FunctionBehavior { + let base = self.behavior(); + let configured_match = if configured_write_functions.is_empty() { + false + } else { + configured_write_functions.contains(&self.name.to_ascii_lowercase()) + || self.schema.is_some_and(|schema| { + configured_write_functions + .contains(&format!("{}.{}", schema, self.name).to_ascii_lowercase()) + }) + }; + + FunctionBehavior { + writes: configured_match || base.writes, + cross_shard: base.cross_shard, + } + } + #[cfg(feature = "new_parser")] pub(crate) fn extract_func_call(node: Node<'a>) -> Option<&'a nodes::FuncCall> { match node { @@ -193,4 +219,15 @@ mod test { }, ); } + + #[test] + #[cfg(feature = "new_parser")] + fn test_configured_write_function_case_insensitive() { + let mut configured = HashSet::new(); + configured.insert("my_write_fn".to_string()); + + first_func("SELECT My_Write_Fn(1)", |func| { + assert!(func.behavior_with_write_functions(&configured).writes); + }); + } } diff --git a/pgdog/src/frontend/router/parser/query/select.rs b/pgdog/src/frontend/router/parser/query/select.rs index 627770658..2f937dcdb 100644 --- a/pgdog/src/frontend/router/parser/query/select.rs +++ b/pgdog/src/frontend/router/parser/query/select.rs @@ -42,8 +42,10 @@ impl QueryParser { if let Some(f) = Function::from_strings(f.funcname().into_iter().filter_map(Node::as_str)) { - cross_shard = cross_shard || f.behavior().cross_shard; - writes = writes || f.behavior().writes; + let behavior = + f.behavior_with_write_functions(&context.sharding_schema.write_functions); + cross_shard = cross_shard || behavior.cross_shard; + writes = writes || behavior.writes; } } _ => (), diff --git a/pgdog/src/frontend/router/parser/query/test/test_functions.rs b/pgdog/src/frontend/router/parser/query/test/test_functions.rs index 4908992bc..b0d14b15d 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_functions.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_functions.rs @@ -1,6 +1,9 @@ use bytes::Bytes; +use pgdog_config::WriteFunctions; +use std::ops::Deref; use super::setup::*; +use crate::config::config; #[test] fn test_write_function_advisory_lock() { @@ -63,3 +66,31 @@ fn test_install_sharded_sequence_without_schema_not_cross_shard() { assert!(!command.route().is_cross_shard()); } + +#[test] +fn test_configured_write_function_routes_to_primary() { + let mut updated = config().deref().clone(); + updated.config.write_functions = vec![WriteFunctions { + database: "pgdog".into(), + functions: vec!["my_write_fn".into()], + }]; + + let mut test = QueryParserTest::new_with_config(&updated); + let command = test.execute(vec![Query::new("SELECT my_write_fn(1)").into()]); + + assert!(command.route().is_write()); +} + +#[test] +fn test_configured_write_function_case_insensitive_and_any_match() { + let mut updated = config().deref().clone(); + updated.config.write_functions = vec![WriteFunctions { + database: "pgdog".into(), + functions: vec!["my_write_fn".into()], + }]; + + let mut test = QueryParserTest::new_with_config(&updated); + let command = test.execute(vec![Query::new("SELECT now(), My_Write_Fn(1)").into()]); + + assert!(command.route().is_write()); +} From 6db6130994471820cb4b437525a351d501a9863f Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Sun, 19 Jul 2026 11:19:51 +0530 Subject: [PATCH 02/10] schema fix --- .schema/pgdog.schema.json | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index e5d0a4ff9..e7b495474 100644 --- a/.schema/pgdog.schema.json +++ b/.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.", + "type": "array", + "default": [], + "items": { + "type": "string" + } + } + }, + "additionalProperties": false, + "required": [ + "database" + ] } } } \ No newline at end of file From d2113abceef06cc2fa78a9a65669534848077d1b Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Sun, 19 Jul 2026 11:29:13 +0530 Subject: [PATCH 03/10] wire test to new parser --- integration/load_balancer/pgx/read_write_split_test.go | 6 ++++++ .../src/frontend/router/parser/query/test/test_functions.rs | 5 +++++ 2 files changed, 11 insertions(+) diff --git a/integration/load_balancer/pgx/read_write_split_test.go b/integration/load_balancer/pgx/read_write_split_test.go index 3f02c38f8..881334555 100644 --- a/integration/load_balancer/pgx/read_write_split_test.go +++ b/integration/load_balancer/pgx/read_write_split_test.go @@ -5,6 +5,8 @@ import ( "errors" "fmt" "math" + "os" + "strings" "testing" "time" @@ -152,6 +154,10 @@ func TestWriteFunctions(t *testing.T) { } 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() diff --git a/pgdog/src/frontend/router/parser/query/test/test_functions.rs b/pgdog/src/frontend/router/parser/query/test/test_functions.rs index b0d14b15d..8c5d1c650 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_functions.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_functions.rs @@ -1,8 +1,11 @@ use bytes::Bytes; +#[cfg(feature = "new_parser")] use pgdog_config::WriteFunctions; +#[cfg(feature = "new_parser")] use std::ops::Deref; use super::setup::*; +#[cfg(feature = "new_parser")] use crate::config::config; #[test] @@ -68,6 +71,7 @@ fn test_install_sharded_sequence_without_schema_not_cross_shard() { } #[test] +#[cfg(feature = "new_parser")] fn test_configured_write_function_routes_to_primary() { let mut updated = config().deref().clone(); updated.config.write_functions = vec![WriteFunctions { @@ -82,6 +86,7 @@ fn test_configured_write_function_routes_to_primary() { } #[test] +#[cfg(feature = "new_parser")] fn test_configured_write_function_case_insensitive_and_any_match() { let mut updated = config().deref().clone(); updated.config.write_functions = vec![WriteFunctions { From 9a26426b1baa397c79850917b41f76e2af05c22e Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 20 Jul 2026 17:18:36 +0530 Subject: [PATCH 04/10] schema --- pgdog-config/src/sharding.rs | 3 + pgdog/src/backend/pool/cluster.rs | 63 ++++++++++++++++--- pgdog/src/frontend/router/parser/function.rs | 52 +++++++++++---- .../parser/query/test/test_functions.rs | 36 ++++++++++- 4 files changed, 132 insertions(+), 22 deletions(-) diff --git a/pgdog-config/src/sharding.rs b/pgdog-config/src/sharding.rs index 609f5e227..eb0897108 100644 --- a/pgdog-config/src/sharding.rs +++ b/pgdog-config/src/sharding.rs @@ -571,6 +571,8 @@ pub struct QueryParser { pub struct WriteFunctions { /// Database name. pub database: String, + /// Optional schema qualifier for all functions in this entry. + pub schema: Option, /// Functions that should be treated as write-only. #[serde(default)] pub functions: Vec, @@ -611,6 +613,7 @@ database = "production" assert_eq!(config.write_functions.len(), 1); assert_eq!(config.write_functions[0].database, "production"); + assert_eq!(config.write_functions[0].schema, None); assert!(config.write_functions[0].functions.is_empty()); } } diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 16cd5514a..b6476c851 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -63,7 +63,7 @@ pub struct Cluster { multi_tenant: Option, rw_strategy: ReadWriteStrategy, rw_split: ReadWriteSplit, - write_functions: HashSet, + write_functions: HashSet, schema_admin: bool, stats: Arc>, cross_shard_disabled: bool, @@ -103,7 +103,7 @@ pub struct ShardingSchema { pub schemas: ShardedSchemas, /// Rewrite config. pub rewrite: Rewrite, - pub write_functions: HashSet, + pub write_functions: HashSet, /// Query parser engine. pub query_parser_engine: QueryParserEngine, pub log_min_duration_parse: Option, @@ -116,6 +116,32 @@ impl ShardingSchema { } } +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct WriteFunction { + pub schema: Option, + 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: Option<&str>, name: &str) -> Self { + Self { + schema: schema.map(Self::normalize_identifier), + name: Self::normalize_identifier(name), + } + } +} + #[derive(Debug)] pub struct ClusterShardConfig { pub primary: Option, @@ -151,7 +177,7 @@ pub struct ClusterConfig<'a> { pub multi_tenant: &'a Option, pub rw_strategy: ReadWriteStrategy, pub rw_split: ReadWriteSplit, - pub write_functions: HashSet, + pub write_functions: HashSet, pub schema_admin: bool, pub cross_shard_disabled: bool, pub two_pc: bool, @@ -215,8 +241,12 @@ impl<'a> ClusterConfig<'a> { .write_functions .iter() .filter(|entry| entry.database == user.database) - .flat_map(|entry| entry.functions.iter()) - .map(|func| func.to_ascii_lowercase()) + .flat_map(|entry| { + entry + .functions + .iter() + .map(move |func| WriteFunction::from_config(entry.schema.as_deref(), func)) + }) .collect(), schema_admin: user.schema_admin, cross_shard_disabled: user @@ -776,17 +806,21 @@ mod test { net::Query, }; - use super::{Cluster, DatabaseUser}; + use super::{Cluster, DatabaseUser, WriteFunction}; impl Cluster { - fn test_write_functions(config: &ConfigAndUsers, database: &str) -> HashSet { + fn test_write_functions(config: &ConfigAndUsers, database: &str) -> HashSet { config .config .write_functions .iter() .filter(|entry| entry.database == database) - .flat_map(|entry| entry.functions.iter()) - .map(|func| func.to_ascii_lowercase()) + .flat_map(|entry| { + entry + .functions + .iter() + .map(move |func| WriteFunction::from_config(entry.schema.as_deref(), func)) + }) .collect() } @@ -1027,6 +1061,17 @@ mod test { assert!(cluster.load_schema()); } + #[test] + fn test_write_function_identifier_normalization() { + let wf = WriteFunction::from_config(Some("PartMan"), "Create_Partition"); + assert_eq!(wf.schema.as_deref(), Some("partman")); + assert_eq!(wf.name, "create_partition"); + + let wf = WriteFunction::from_config(Some(r#""PartMan""#), r#""Create_Partition""#); + assert_eq!(wf.schema.as_deref(), Some("PartMan")); + assert_eq!(wf.name, "Create_Partition"); + } + #[test] fn test_load_schema_multiple_shards_with_schemas() { let config = ConfigAndUsers::default(); diff --git a/pgdog/src/frontend/router/parser/function.rs b/pgdog/src/frontend/router/parser/function.rs index b87ac8b47..5bd77c678 100644 --- a/pgdog/src/frontend/router/parser/function.rs +++ b/pgdog/src/frontend/router/parser/function.rs @@ -5,6 +5,9 @@ use pg_raw_parse::{Node, nodes}; #[cfg(feature = "new_parser")] use std::collections::HashSet; +#[cfg(feature = "new_parser")] +use crate::backend::pool::cluster::WriteFunction; + const WRITE_ONLY: &[&str] = &["nextval", "setval"]; const CROSS_SHARD: &[(Option<&str>, &str)] = &[(Some("pgdog"), "install_sharded_sequence")]; @@ -46,18 +49,16 @@ impl<'a> Function<'a> { #[cfg(feature = "new_parser")] pub(crate) fn behavior_with_write_functions( &self, - configured_write_functions: &HashSet, + configured_write_functions: &HashSet, ) -> FunctionBehavior { let base = self.behavior(); - let configured_match = if configured_write_functions.is_empty() { - false - } else { - configured_write_functions.contains(&self.name.to_ascii_lowercase()) - || self.schema.is_some_and(|schema| { - configured_write_functions - .contains(&format!("{}.{}", schema, self.name).to_ascii_lowercase()) - }) - }; + let configured_match = configured_write_functions.contains(&WriteFunction { + schema: self.schema.map(ToOwned::to_owned), + name: self.name.to_owned(), + }) || configured_write_functions.contains(&WriteFunction { + schema: None, + name: self.name.to_owned(), + }); FunctionBehavior { writes: configured_match || base.writes, @@ -119,6 +120,8 @@ impl<'a> TryFrom<&'a Node> for Function<'a> { #[cfg(test)] mod test { + #[cfg(feature = "new_parser")] + use crate::backend::pool::cluster::WriteFunction; #[cfg(not(feature = "new_parser"))] use pg_query::parse; #[cfg(feature = "new_parser")] @@ -222,12 +225,37 @@ mod test { #[test] #[cfg(feature = "new_parser")] - fn test_configured_write_function_case_insensitive() { + fn test_configured_write_function_pg_identifier_semantics() { let mut configured = HashSet::new(); - configured.insert("my_write_fn".to_string()); + configured.insert(WriteFunction { + schema: None, + name: "my_write_fn".to_string(), + }); first_func("SELECT My_Write_Fn(1)", |func| { assert!(func.behavior_with_write_functions(&configured).writes); }); + + first_func(r#"SELECT "My_Write_Fn"(1)"#, |func| { + assert!(!func.behavior_with_write_functions(&configured).writes); + }); + } + + #[test] + #[cfg(feature = "new_parser")] + fn test_configured_write_function_with_schema() { + let mut configured = HashSet::new(); + configured.insert(WriteFunction { + schema: Some("partman".to_string()), + name: "create_partition".to_string(), + }); + + first_func("SELECT partman.create_partition('foo')", |func| { + assert!(func.behavior_with_write_functions(&configured).writes); + }); + + first_func("SELECT other.create_partition('foo')", |func| { + assert!(!func.behavior_with_write_functions(&configured).writes); + }); } } diff --git a/pgdog/src/frontend/router/parser/query/test/test_functions.rs b/pgdog/src/frontend/router/parser/query/test/test_functions.rs index 8c5d1c650..ec052cf73 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_functions.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_functions.rs @@ -76,6 +76,7 @@ fn test_configured_write_function_routes_to_primary() { let mut updated = config().deref().clone(); updated.config.write_functions = vec![WriteFunctions { database: "pgdog".into(), + schema: None, functions: vec!["my_write_fn".into()], }]; @@ -87,10 +88,11 @@ fn test_configured_write_function_routes_to_primary() { #[test] #[cfg(feature = "new_parser")] -fn test_configured_write_function_case_insensitive_and_any_match() { +fn test_configured_write_function_pg_identifier_semantics() { let mut updated = config().deref().clone(); updated.config.write_functions = vec![WriteFunctions { database: "pgdog".into(), + schema: None, functions: vec!["my_write_fn".into()], }]; @@ -98,4 +100,36 @@ fn test_configured_write_function_case_insensitive_and_any_match() { let command = test.execute(vec![Query::new("SELECT now(), My_Write_Fn(1)").into()]); assert!(command.route().is_write()); + + let command = test.execute(vec![Query::new(r#"SELECT "My_Write_Fn"(1)"#).into()]); + assert!(command.route().is_read()); + + updated.config.write_functions = vec![WriteFunctions { + database: "pgdog".into(), + schema: None, + functions: vec![r#""My_Write_Fn""#.into()], + }]; + let mut test = QueryParserTest::new_with_config(&updated); + let command = test.execute(vec![Query::new(r#"SELECT "My_Write_Fn"(1)"#).into()]); + assert!(command.route().is_write()); +} + +#[test] +#[cfg(feature = "new_parser")] +fn test_configured_write_function_with_schema() { + let mut updated = config().deref().clone(); + updated.config.write_functions = vec![WriteFunctions { + database: "pgdog".into(), + schema: Some("partman".into()), + functions: vec!["create_partition".into()], + }]; + + let mut test = QueryParserTest::new_with_config(&updated); + let command = test.execute(vec![ + Query::new("SELECT partman.create_partition(1)").into(), + ]); + assert!(command.route().is_write()); + + let command = test.execute(vec![Query::new("SELECT other.create_partition(1)").into()]); + assert!(command.route().is_read()); } From a7fefc4179695295237deb0be0e37678ead132e9 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 20 Jul 2026 17:24:04 +0530 Subject: [PATCH 05/10] edit example --- .schema/pgdog.schema.json | 7 +++++++ example.pgdog.toml | 5 ++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index e7b495474..c5644ce8c 100644 --- a/.schema/pgdog.schema.json +++ b/.schema/pgdog.schema.json @@ -2351,6 +2351,13 @@ "items": { "type": "string" } + }, + "schema": { + "description": "Optional schema qualifier for all functions in this entry.", + "type": [ + "string", + "null" + ] } }, "additionalProperties": false, diff --git a/example.pgdog.toml b/example.pgdog.toml index 953fd2ee2..71a960ea3 100644 --- a/example.pgdog.toml +++ b/example.pgdog.toml @@ -279,10 +279,13 @@ dns_ttl = 5_000 # # Per-database list of functions that should always route to the primary. # -# Function matching is case-insensitive. +# Function matching follows PostgreSQL identifier semantics: +# - unquoted identifiers are folded to lowercase +# - quoted identifiers preserve case # # [[write_functions]] # database = "pgdog" +# schema = "partman" # optional # functions = ["my_write_fn", "set_customer_balance"] # From 27e776652c3338173515abc4acc62042d0be0777 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Mon, 20 Jul 2026 17:27:00 +0530 Subject: [PATCH 06/10] extend example --- example.pgdog.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/example.pgdog.toml b/example.pgdog.toml index 71a960ea3..beb1619f0 100644 --- a/example.pgdog.toml +++ b/example.pgdog.toml @@ -286,7 +286,7 @@ dns_ttl = 5_000 # [[write_functions]] # database = "pgdog" # schema = "partman" # optional -# functions = ["my_write_fn", "set_customer_balance"] +# functions = ["my_write_fn", "\"Set_Customer_Balance\""] # # From 6d6f69d75768f0b372d7312f86dddf773348e2f0 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 21 Jul 2026 21:14:12 +0530 Subject: [PATCH 07/10] fix --- .schema/pgdog.schema.json | 9 +-- example.pgdog.toml | 6 +- pgdog-config/src/sharding.rs | 4 +- pgdog/src/backend/pool/cluster.rs | 68 ++++++++++++++----- .../parser/query/test/test_functions.rs | 6 +- 5 files changed, 56 insertions(+), 37 deletions(-) diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index c5644ce8c..3bf77e46d 100644 --- a/.schema/pgdog.schema.json +++ b/.schema/pgdog.schema.json @@ -2345,19 +2345,12 @@ "type": "string" }, "functions": { - "description": "Functions that should be treated as write-only.", + "description": "Functions that should be treated as write-only.\nEach entry can be schema-qualified.", "type": "array", "default": [], "items": { "type": "string" } - }, - "schema": { - "description": "Optional schema qualifier for all functions in this entry.", - "type": [ - "string", - "null" - ] } }, "additionalProperties": false, diff --git a/example.pgdog.toml b/example.pgdog.toml index beb1619f0..1dfb830b4 100644 --- a/example.pgdog.toml +++ b/example.pgdog.toml @@ -279,14 +279,14 @@ dns_ttl = 5_000 # # Per-database list of functions that should always route to the primary. # -# Function matching follows PostgreSQL identifier semantics: +# 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" -# schema = "partman" # optional -# functions = ["my_write_fn", "\"Set_Customer_Balance\""] +# functions = ["my_write_fn", "partman.create_partition", '"Set_Customer_Balance"'] # # diff --git a/pgdog-config/src/sharding.rs b/pgdog-config/src/sharding.rs index eb0897108..3231868f2 100644 --- a/pgdog-config/src/sharding.rs +++ b/pgdog-config/src/sharding.rs @@ -571,9 +571,8 @@ pub struct QueryParser { pub struct WriteFunctions { /// Database name. pub database: String, - /// Optional schema qualifier for all functions in this entry. - pub schema: Option, /// Functions that should be treated as write-only. + /// Each entry can be schema-qualified. #[serde(default)] pub functions: Vec, } @@ -613,7 +612,6 @@ database = "production" assert_eq!(config.write_functions.len(), 1); assert_eq!(config.write_functions[0].database, "production"); - assert_eq!(config.write_functions[0].schema, None); assert!(config.write_functions[0].functions.is_empty()); } } diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index b6476c851..6990e0391 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -15,7 +15,7 @@ use std::{ time::Duration, }; use tokio::spawn; -use tracing::error; +use tracing::{error, warn}; use crate::frontend::router::sharding::ShardedTable; use crate::{ @@ -134,11 +134,33 @@ impl WriteFunction { } } - fn from_config(schema: Option<&str>, name: &str) -> Self { - Self { - schema: schema.map(Self::normalize_identifier), - name: Self::normalize_identifier(name), + fn from_config(entry: &str) -> Option { + fn split_qualified(entry: &str) -> impl Iterator { + 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 }) } } @@ -241,11 +263,13 @@ impl<'a> ClusterConfig<'a> { .write_functions .iter() .filter(|entry| entry.database == user.database) - .flat_map(|entry| { - entry - .functions - .iter() - .map(move |func| WriteFunction::from_config(entry.schema.as_deref(), func)) + .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, @@ -815,12 +839,8 @@ mod test { .write_functions .iter() .filter(|entry| entry.database == database) - .flat_map(|entry| { - entry - .functions - .iter() - .map(move |func| WriteFunction::from_config(entry.schema.as_deref(), func)) - }) + .flat_map(|entry| entry.functions.iter()) + .filter_map(|func| WriteFunction::from_config(func)) .collect() } @@ -1063,13 +1083,25 @@ mod test { #[test] fn test_write_function_identifier_normalization() { - let wf = WriteFunction::from_config(Some("PartMan"), "Create_Partition"); + 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(Some(r#""PartMan""#), r#""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] diff --git a/pgdog/src/frontend/router/parser/query/test/test_functions.rs b/pgdog/src/frontend/router/parser/query/test/test_functions.rs index ec052cf73..3b7db00b6 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_functions.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_functions.rs @@ -76,7 +76,6 @@ fn test_configured_write_function_routes_to_primary() { let mut updated = config().deref().clone(); updated.config.write_functions = vec![WriteFunctions { database: "pgdog".into(), - schema: None, functions: vec!["my_write_fn".into()], }]; @@ -92,7 +91,6 @@ fn test_configured_write_function_pg_identifier_semantics() { let mut updated = config().deref().clone(); updated.config.write_functions = vec![WriteFunctions { database: "pgdog".into(), - schema: None, functions: vec!["my_write_fn".into()], }]; @@ -106,7 +104,6 @@ fn test_configured_write_function_pg_identifier_semantics() { updated.config.write_functions = vec![WriteFunctions { database: "pgdog".into(), - schema: None, functions: vec![r#""My_Write_Fn""#.into()], }]; let mut test = QueryParserTest::new_with_config(&updated); @@ -120,8 +117,7 @@ fn test_configured_write_function_with_schema() { let mut updated = config().deref().clone(); updated.config.write_functions = vec![WriteFunctions { database: "pgdog".into(), - schema: Some("partman".into()), - functions: vec!["create_partition".into()], + functions: vec!["partman.create_partition".into()], }]; let mut test = QueryParserTest::new_with_config(&updated); From a41b9e6e3d3b1514262138d0fa150cf8216a46c6 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 21 Jul 2026 21:19:44 +0530 Subject: [PATCH 08/10] clippy fix --- pgdog/src/backend/pool/cluster.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 8d5ff35bb..8f8498dd1 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -14,7 +14,6 @@ use std::{ }, time::Duration, }; -use tokio::spawn; use tracing::{error, warn}; use crate::frontend::router::sharding::ShardedTable; From fd2e0ef3eb6bd4db3af2e690bacb5e3181731eb8 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 28 Jul 2026 21:44:26 +0530 Subject: [PATCH 09/10] enforce schema --- .schema/pgdog.schema.json | 18 ++--- example.pgdog.toml | 8 +-- integration/load_balancer/pgdog.toml | 3 +- .../pgx/read_write_split_test.go | 6 +- pgdog-config/src/sharding.rs | 32 +++++++-- pgdog/src/backend/pool/cluster.rs | 66 ++++--------------- pgdog/src/frontend/router/parser/function.rs | 23 ++++--- .../parser/query/test/test_functions.rs | 32 ++++++--- 8 files changed, 96 insertions(+), 92 deletions(-) diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index 1b2a05eb5..8aa1004f2 100644 --- a/.schema/pgdog.schema.json +++ b/.schema/pgdog.schema.json @@ -2344,18 +2344,20 @@ "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" - } + "name": { + "description": "Function that should be treated as write-only.", + "type": "string" + }, + "schema": { + "description": "Schema containing the function.", + "type": "string" } }, "additionalProperties": false, "required": [ - "database" + "database", + "schema", + "name" ] } } diff --git a/example.pgdog.toml b/example.pgdog.toml index 1dfb830b4..3028737c6 100644 --- a/example.pgdog.toml +++ b/example.pgdog.toml @@ -277,16 +277,16 @@ cross_shard_disabled = false dns_ttl = 5_000 # -# Per-database list of functions that should always route to the primary. +# Per-database function that should always route to the primary. # -# Each function can be schema-qualified. Matching follows PostgreSQL -# identifier semantics: +# 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"'] +# schema = "partman" +# name = "create_partition" # # diff --git a/integration/load_balancer/pgdog.toml b/integration/load_balancer/pgdog.toml index 86fa17ab4..6001749d0 100644 --- a/integration/load_balancer/pgdog.toml +++ b/integration/load_balancer/pgdog.toml @@ -74,7 +74,8 @@ port = 45002 [[write_functions]] database = "postgres" -functions = ["Lb_Write_Fn"] +schema = "public" +name = "Lb_Write_Fn" [[plugins]] diff --git a/integration/load_balancer/pgx/read_write_split_test.go b/integration/load_balancer/pgx/read_write_split_test.go index 881334555..f0eac920e 100644 --- a/integration/load_balancer/pgx/read_write_split_test.go +++ b/integration/load_balancer/pgx/read_write_split_test.go @@ -173,14 +173,14 @@ func TestConfiguredWriteFunctionsRouteToPrimary(t *testing.T) { ResetStats() for i := range 20 { - _, err = pool.Exec(context.Background(), "SELECT lb_write_fn($1)", int64(i)) + _, err = pool.Exec(context.Background(), "SELECT public.lb_write_fn($1)", int64(i)) assert.NoError(t, err) } - primaryCalls := LoadStatsForPrimary("SELECT lb_write_fn") + primaryCalls := LoadStatsForPrimary("SELECT public.lb_write_fn") assert.Equal(t, int64(20), primaryCalls.Calls) - replicaCalls := LoadStatsForReplicas("SELECT lb_write_fn") + replicaCalls := LoadStatsForReplicas("SELECT public.lb_write_fn") for _, call := range replicaCalls { assert.Equal(t, int64(0), call.Calls) } diff --git a/pgdog-config/src/sharding.rs b/pgdog-config/src/sharding.rs index 1f8f10114..d6994d145 100644 --- a/pgdog-config/src/sharding.rs +++ b/pgdog-config/src/sharding.rs @@ -571,10 +571,10 @@ pub struct QueryParser { 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, + /// Schema containing the function. + pub schema: String, + /// Function that should be treated as write-only. + pub name: String, } #[cfg(test)] @@ -602,16 +602,36 @@ database = "production" } #[test] - fn write_functions_reads_default_values_from_config() { + 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!(config.write_functions[0].functions.is_empty()); + 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::(missing_schema).is_err()); + + let missing_name = r#" +[[write_functions]] +database = "production" +schema = "partman" +"#; + assert!(toml::from_str::(missing_name).is_err()); } } diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 8f8498dd1..a90a4c347 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -14,7 +14,7 @@ use std::{ }, time::Duration, }; -use tracing::{error, warn}; +use tracing::error; use crate::frontend::router::sharding::ShardedTable; use crate::tasks; @@ -118,7 +118,7 @@ impl ShardingSchema { #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct WriteFunction { - pub schema: Option, + pub schema: String, pub name: String, } @@ -134,33 +134,11 @@ impl WriteFunction { } } - fn from_config(entry: &str) -> Option { - fn split_qualified(entry: &str) -> impl Iterator { - 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() + fn from_config(schema: &str, name: &str) -> Self { + Self { + schema: Self::normalize_identifier(schema), + name: Self::normalize_identifier(name), } - 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 }) } } @@ -263,14 +241,7 @@ impl<'a> ClusterConfig<'a> { .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 - }) + .map(|entry| WriteFunction::from_config(&entry.schema, &entry.name)) .collect(), schema_admin: user.schema_admin, cross_shard_disabled: user @@ -839,8 +810,7 @@ mod test { .write_functions .iter() .filter(|entry| entry.database == database) - .flat_map(|entry| entry.functions.iter()) - .filter_map(|func| WriteFunction::from_config(func)) + .map(|entry| WriteFunction::from_config(&entry.schema, &entry.name)) .collect() } @@ -1083,25 +1053,17 @@ mod test { #[test] fn test_write_function_identifier_normalization() { - let wf = WriteFunction::from_config("Create_Partition").unwrap(); - assert_eq!(wf.schema, None); + let wf = WriteFunction::from_config("PartMan", "Create_Partition"); + assert_eq!(wf.schema, "partman"); 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")); + let wf = WriteFunction::from_config(r#""PartMan""#, r#""Create_Partition""#); + assert_eq!(wf.schema, "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")); + let wf = WriteFunction::from_config(r#""my.schema""#, "my_fn"); + assert_eq!(wf.schema, "my.schema"); assert_eq!(wf.name, "my_fn"); - - // More than two parts is invalid. - assert_eq!(WriteFunction::from_config("a.b.c"), None); } #[test] diff --git a/pgdog/src/frontend/router/parser/function.rs b/pgdog/src/frontend/router/parser/function.rs index 5bd77c678..58ca8d788 100644 --- a/pgdog/src/frontend/router/parser/function.rs +++ b/pgdog/src/frontend/router/parser/function.rs @@ -52,12 +52,11 @@ impl<'a> Function<'a> { configured_write_functions: &HashSet, ) -> FunctionBehavior { let base = self.behavior(); - let configured_match = configured_write_functions.contains(&WriteFunction { - schema: self.schema.map(ToOwned::to_owned), - name: self.name.to_owned(), - }) || configured_write_functions.contains(&WriteFunction { - schema: None, - name: self.name.to_owned(), + let configured_match = self.schema.is_some_and(|schema| { + configured_write_functions.contains(&WriteFunction { + schema: schema.to_owned(), + name: self.name.to_owned(), + }) }); FunctionBehavior { @@ -228,15 +227,15 @@ mod test { fn test_configured_write_function_pg_identifier_semantics() { let mut configured = HashSet::new(); configured.insert(WriteFunction { - schema: None, + schema: "public".to_string(), name: "my_write_fn".to_string(), }); - first_func("SELECT My_Write_Fn(1)", |func| { + first_func("SELECT PUBLIC.My_Write_Fn(1)", |func| { assert!(func.behavior_with_write_functions(&configured).writes); }); - first_func(r#"SELECT "My_Write_Fn"(1)"#, |func| { + first_func(r#"SELECT public."My_Write_Fn"(1)"#, |func| { assert!(!func.behavior_with_write_functions(&configured).writes); }); } @@ -246,7 +245,7 @@ mod test { fn test_configured_write_function_with_schema() { let mut configured = HashSet::new(); configured.insert(WriteFunction { - schema: Some("partman".to_string()), + schema: "partman".to_string(), name: "create_partition".to_string(), }); @@ -257,5 +256,9 @@ mod test { first_func("SELECT other.create_partition('foo')", |func| { assert!(!func.behavior_with_write_functions(&configured).writes); }); + + first_func("SELECT create_partition('foo')", |func| { + assert!(!func.behavior_with_write_functions(&configured).writes); + }); } } diff --git a/pgdog/src/frontend/router/parser/query/test/test_functions.rs b/pgdog/src/frontend/router/parser/query/test/test_functions.rs index 3b7db00b6..464f97dc7 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_functions.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_functions.rs @@ -76,11 +76,12 @@ fn test_configured_write_function_routes_to_primary() { let mut updated = config().deref().clone(); updated.config.write_functions = vec![WriteFunctions { database: "pgdog".into(), - functions: vec!["my_write_fn".into()], + schema: "public".into(), + name: "my_write_fn".into(), }]; let mut test = QueryParserTest::new_with_config(&updated); - let command = test.execute(vec![Query::new("SELECT my_write_fn(1)").into()]); + let command = test.execute(vec![Query::new("SELECT public.my_write_fn(1)").into()]); assert!(command.route().is_write()); } @@ -91,24 +92,35 @@ fn test_configured_write_function_pg_identifier_semantics() { let mut updated = config().deref().clone(); updated.config.write_functions = vec![WriteFunctions { database: "pgdog".into(), - functions: vec!["my_write_fn".into()], + schema: "public".into(), + name: "my_write_fn".into(), }]; let mut test = QueryParserTest::new_with_config(&updated); - let command = test.execute(vec![Query::new("SELECT now(), My_Write_Fn(1)").into()]); + let command = test.execute(vec![ + Query::new("SELECT now(), PUBLIC.My_Write_Fn(1)").into(), + ]); assert!(command.route().is_write()); - let command = test.execute(vec![Query::new(r#"SELECT "My_Write_Fn"(1)"#).into()]); + let command = test.execute(vec![Query::new(r#"SELECT public."My_Write_Fn"(1)"#).into()]); assert!(command.route().is_read()); updated.config.write_functions = vec![WriteFunctions { database: "pgdog".into(), - functions: vec![r#""My_Write_Fn""#.into()], + schema: r#""AppSchema""#.into(), + name: r#""My_Write_Fn""#.into(), }]; let mut test = QueryParserTest::new_with_config(&updated); - let command = test.execute(vec![Query::new(r#"SELECT "My_Write_Fn"(1)"#).into()]); + let command = test.execute(vec![ + Query::new(r#"SELECT "AppSchema"."My_Write_Fn"(1)"#).into(), + ]); assert!(command.route().is_write()); + + let command = test.execute(vec![ + Query::new(r#"SELECT "appschema"."My_Write_Fn"(1)"#).into(), + ]); + assert!(command.route().is_read()); } #[test] @@ -117,7 +129,8 @@ fn test_configured_write_function_with_schema() { let mut updated = config().deref().clone(); updated.config.write_functions = vec![WriteFunctions { database: "pgdog".into(), - functions: vec!["partman.create_partition".into()], + schema: "partman".into(), + name: "create_partition".into(), }]; let mut test = QueryParserTest::new_with_config(&updated); @@ -128,4 +141,7 @@ fn test_configured_write_function_with_schema() { let command = test.execute(vec![Query::new("SELECT other.create_partition(1)").into()]); assert!(command.route().is_read()); + + let command = test.execute(vec![Query::new("SELECT create_partition(1)").into()]); + assert!(command.route().is_read()); } From 6d9a1d11c968bc7597e71855d72da2de020a54f6 Mon Sep 17 00:00:00 2001 From: Nupur Agrawal Date: Tue, 28 Jul 2026 22:05:15 +0530 Subject: [PATCH 10/10] fmt fixes --- pgdog/src/backend/pool/cluster.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index c58e4ea63..967faedd9 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/pgdog/src/backend/pool/cluster.rs @@ -6,11 +6,7 @@ use pgdog_config::{ LoadSchema, PreparedStatements, QueryParser, QueryParserEngine, QueryParserLevel, Rewrite, RewriteMode, users::PasswordKind, }; -use std::{ - collections::HashSet, - sync::Arc, - time::Duration, -}; +use std::{collections::HashSet, sync::Arc, time::Duration}; use crate::backend::schema::SchemaCache; use crate::frontend::router::sharding::ShardedTable;