diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index c8bb7a87e..af4a2c8fe 100644 --- a/.schema/pgdog.schema.json +++ b/.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" + ] } } } \ No newline at end of file diff --git a/example.pgdog.toml b/example.pgdog.toml index d110dd659..ef22543fd 100644 --- a/example.pgdog.toml +++ b/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. # diff --git a/integration/load_balancer/pgdog.toml b/integration/load_balancer/pgdog.toml index bf1f9cb99..6001749d0 100644 --- a/integration/load_balancer/pgdog.toml +++ b/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" diff --git a/integration/load_balancer/pgx/read_write_split_test.go b/integration/load_balancer/pgx/read_write_split_test.go index e6a7e5245..f0eac920e 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" @@ -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) diff --git a/pgdog-config/src/core.rs b/pgdog-config/src/core.rs index bc0b6f929..52462823f 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 1a26640ab..7ee113e54 100644 --- a/pgdog-config/src/sharding.rs +++ b/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::(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 9899b1072..967faedd9 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/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, rw_strategy: ReadWriteStrategy, rw_split: ReadWriteSplit, + write_functions: HashSet, schema_admin: bool, stats: Arc>, cross_shard_disabled: bool, @@ -91,6 +92,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, @@ -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, @@ -144,6 +172,7 @@ pub struct ClusterConfig<'a> { two_pc_auto: bool, sharded_schemas: ShardedSchemas, rewrite: &'a Rewrite, + write_functions: HashSet, 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 { + 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(); diff --git a/pgdog/src/frontend/router/parser/function.rs b/pgdog/src/frontend/router/parser/function.rs index 2a70a5f15..58ca8d788 100644 --- a/pgdog/src/frontend/router/parser/function.rs +++ b/pgdog/src/frontend/router/parser/function.rs @@ -2,6 +2,11 @@ use pg_query::{Node, NodeEnum, protobuf}; #[cfg(feature = "new_parser")] 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"]; @@ -34,11 +39,32 @@ 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 = self.schema.is_some_and(|schema| { + configured_write_functions.contains(&WriteFunction { + schema: schema.to_owned(), + name: self.name.to_owned(), + }) + }); + + 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 { @@ -93,6 +119,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")] @@ -193,4 +221,44 @@ mod test { }, ); } + + #[test] + #[cfg(feature = "new_parser")] + fn test_configured_write_function_pg_identifier_semantics() { + let mut configured = HashSet::new(); + configured.insert(WriteFunction { + schema: "public".to_string(), + name: "my_write_fn".to_string(), + }); + + first_func("SELECT PUBLIC.My_Write_Fn(1)", |func| { + assert!(func.behavior_with_write_functions(&configured).writes); + }); + + first_func(r#"SELECT public."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: "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); + }); + + first_func("SELECT create_partition('foo')", |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..464f97dc7 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,12 @@ 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] fn test_write_function_advisory_lock() { @@ -63,3 +69,79 @@ fn test_install_sharded_sequence_without_schema_not_cross_shard() { assert!(!command.route().is_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 { + database: "pgdog".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 public.my_write_fn(1)").into()]); + + assert!(command.route().is_write()); +} + +#[test] +#[cfg(feature = "new_parser")] +fn test_configured_write_function_pg_identifier_semantics() { + let mut updated = config().deref().clone(); + updated.config.write_functions = vec![WriteFunctions { + database: "pgdog".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(), PUBLIC.My_Write_Fn(1)").into(), + ]); + + assert!(command.route().is_write()); + + 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(), + 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 "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] +#[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: "partman".into(), + name: "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()); + + let command = test.execute(vec![Query::new("SELECT create_partition(1)").into()]); + assert!(command.route().is_read()); +}