diff --git a/.schema/pgdog.schema.json b/.schema/pgdog.schema.json index 3900e97cd..1b2a05eb5 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.\nEach entry can be schema-qualified.", + "type": "array", + "default": [], + "items": { + "type": "string" + } + } + }, + "additionalProperties": false, + "required": [ + "database" + ] } } } \ No newline at end of file diff --git a/example.pgdog.toml b/example.pgdog.toml index cc648c890..1dfb830b4 100644 --- a/example.pgdog.toml +++ b/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. # 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..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" @@ -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) diff --git a/pgdog-config/src/core.rs b/pgdog-config/src/core.rs index 586e91e7b..5fb25a54b 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 4b5f83c1e..1f8f10114 100644 --- a/pgdog-config/src/sharding.rs +++ b/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, +} + #[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()); + } } diff --git a/pgdog/src/backend/pool/cluster.rs b/pgdog/src/backend/pool/cluster.rs index 78f39af2e..8f8498dd1 100644 --- a/pgdog/src/backend/pool/cluster.rs +++ b/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, 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, @@ -113,6 +116,54 @@ 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(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 }) + } +} + #[derive(Debug)] pub struct ClusterShardConfig { pub primary: Option, @@ -148,6 +199,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 +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 { + 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(); diff --git a/pgdog/src/frontend/router/parser/function.rs b/pgdog/src/frontend/router/parser/function.rs index 2a70a5f15..5bd77c678 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,33 @@ 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 = 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, + 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 +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")] @@ -193,4 +222,40 @@ mod test { }, ); } + + #[test] + #[cfg(feature = "new_parser")] + fn test_configured_write_function_pg_identifier_semantics() { + let mut configured = HashSet::new(); + 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/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..3b7db00b6 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,63 @@ 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(), + 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] +#[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(), + 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()); + + 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(), + 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(), + functions: vec!["partman.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()); +}