Skip to content

feat: add configurable write_functions #1236

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
murex971 wants to merge 9 commits into
base: main
Choose a base branch
from
Open

feat: add configurable write_functions #1236

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
5007cfc
feat: add configurable write_functions
murex971 Jul 19, 2026
6db6130
schema fix
murex971 Jul 19, 2026
d2113ab
wire test to new parser
murex971 Jul 19, 2026
9a26426
schema
murex971 Jul 20, 2026
a7fefc4
edit example
murex971 Jul 20, 2026
27e7766
extend example
murex971 Jul 20, 2026
6d6f69d
fix
murex971 Jul 21, 2026
773067f
Merge remote-tracking branch 'forkie/main' into murex-write-fn
murex971 Jul 21, 2026
a41b9e6
clippy fix
murex971 Jul 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions .schema/pgdog.schema.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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"
]
}
}
}
13 changes: 13 additions & 0 deletions example.pgdog.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
#
Expand Down
4 changes: 4 additions & 0 deletions integration/load_balancer/pgdog.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
35 changes: 35 additions & 0 deletions integration/load_balancer/pgx/read_write_split_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ import (
"errors"
"fmt"
"math"
"os"
"strings"
"testing"
"time"

Expand Down Expand Up @@ -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)
Expand Down
6 changes: 5 additions & 1 deletion pgdog-config/src/core.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 {
Expand Down
26 changes: 26 additions & 0 deletions pgdog-config/src/sharding.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
Comment thread
murex971 marked this conversation as resolved.
}

#[cfg(test)]
mod tests {
use crate::Config;
Expand All @@ -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());
}
}
110 changes: 107 additions & 3 deletions pgdog/src/backend/pool/cluster.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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>,
Expand All @@ -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>,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -258,6 +323,7 @@ impl Cluster {
multi_tenant,
rw_strategy,
rw_split,
write_functions,
schema_admin,
cross_shard_disabled,
two_pc,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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(),
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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();
Expand Down
Loading
Loading