feat: pgdog.install_sharded_sequence() is cross-sharded (#881)

Make sure that installing sharded sequences is sent to all shards.

Author: levkk

pgdog/src/frontend/router/parser/function.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use once_cell::sync::Lazy;
 use pg_query::{protobuf, Node, NodeEnum};
@@ -20,6 +20,9 @@ static WRITE_ONLY: Lazy<HashMap<&'static str, LockingBehavior>> = Lazy::new(|| {
     ])
 });
 
+static CROSS_SHARD: Lazy<HashSet<(&'static str, &'static str)>> =
+    Lazy::new(|| HashSet::from([("pgdog", "install_sharded_sequence")]));
+
 #[derive(Debug, Copy, Clone, PartialEq, Default)]
 pub enum LockingBehavior {
     Lock,
@@ -32,6 +35,7 @@ pub enum LockingBehavior {
 pub struct FunctionBehavior {
     pub writes: bool,
     pub locking_behavior: LockingBehavior,
+    pub cross_shard: bool,
 }
 
 impl FunctionBehavior {
@@ -45,28 +49,49 @@ impl FunctionBehavior {
 
 pub struct Function<'a> {
     pub name: &'a str,
+    pub schema: Option<&'a str>,
 }
 
 impl<'a> Function<'a> {
-    fn from_string(node: &'a Option<NodeEnum>) -> Result<Self, ()> {
-        match node {
-            Some(NodeEnum::String(protobuf::String { sval })) => Ok(Self {
-                name: sval.as_str(),
+    /// Build a Function from a qualified name list (as found in `FuncCall.funcname`).
+    /// The last element is the function name; the preceding element (if any) is the
+    /// schema.
+    fn from_strings(parts: &'a [Node]) -> Result<Self, ()> {
+        let str_of = |node: &'a Node| match &node.node {
+            Some(NodeEnum::String(protobuf::String { sval })) => Ok(sval.as_str()),
+            _ => Err(()),
+        };
+        match parts {
+            [name] => Ok(Self {
+                name: str_of(name)?,
+                schema: None,
+            }),
+            [.., schema, name] => Ok(Self {
+                name: str_of(name)?,
+                schema: Some(str_of(schema)?),
             }),
-
             _ => Err(()),
         }
     }
 
     /// This function likely writes.
     pub fn behavior(&self) -> FunctionBehavior {
+        let cross_shard = self
+            .schema
+            .map(|schema| CROSS_SHARD.contains(&(schema, self.name)))
+            .unwrap_or(false);
+
         if let Some(locks) = WRITE_ONLY.get(&self.name) {
             FunctionBehavior {
                 writes: true,
                 locking_behavior: *locks,
+                cross_shard,
             }
         } else {
-            FunctionBehavior::default()
+            FunctionBehavior {
+                cross_shard,
+                ..FunctionBehavior::default()
+            }
         }
     }
 }
@@ -76,9 +101,7 @@ impl<'a> TryFrom<&'a Node> for Function<'a> {
     fn try_from(value: &'a Node) -> Result<Self, Self::Error> {
         match &value.node {
             Some(NodeEnum::FuncCall(func)) => {
-                if let Some(node) = func.funcname.last() {
-                    return Self::from_string(&node.node);
-                }
+                return Self::from_strings(&func.funcname);
             }
 
             Some(NodeEnum::TypeCast(cast)) => {
@@ -123,10 +146,52 @@ mod test {
                 for node in &stmt.target_list {
                     let func = Function::try_from(node).unwrap();
                     assert!(func.name.contains("advisory_lock"));
+                    assert!(func.schema.is_none());
+                    assert!(!func.behavior().cross_shard);
                 }
             }
 
             _ => panic!("not a select"),
         }
     }
+
+    fn first_func<R>(query: &str, check: impl FnOnce(Function<'_>) -> R) -> R {
+        let ast = parse(query).unwrap();
+        let root = ast.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap();
+        match root.node.as_ref() {
+            Some(NodeEnum::SelectStmt(stmt)) => {
+                let target = stmt.target_list.first().unwrap();
+                check(Function::try_from(target).unwrap())
+            }
+            _ => panic!("not a select"),
+        }
+    }
+
+    #[test]
+    fn test_cross_shard_function() {
+        first_func(
+            "SELECT pgdog.install_sharded_sequence('foo', 'id')",
+            |func| {
+                assert_eq!(func.name, "install_sharded_sequence");
+                assert_eq!(func.schema, Some("pgdog"));
+                assert!(func.behavior().cross_shard);
+            },
+        );
+
+        // Same function name without the schema should not be flagged.
+        first_func("SELECT install_sharded_sequence('foo', 'id')", |func| {
+            assert_eq!(func.name, "install_sharded_sequence");
+            assert!(func.schema.is_none());
+            assert!(!func.behavior().cross_shard);
+        });
+
+        // Different schema should not be flagged.
+        first_func(
+            "SELECT other.install_sharded_sequence('foo', 'id')",
+            |func| {
+                assert_eq!(func.schema, Some("other"));
+                assert!(!func.behavior().cross_shard);
+            },
+        );
+    }
 }

pgdog/src/frontend/router/parser/query/select.rs

@@ -21,21 +21,27 @@ impl QueryParser {
         context: &mut QueryParserContext,
     ) -> Result<Command, Error> {
         let cte_writes = Self::cte_writes(stmt);
-        let mut writes = Self::functions(stmt)?;
+        let mut overrides = Self::functions(stmt)?;
 
         // Write overwrite because of conservative read/write split.
         if self.write_override {
-            writes.writes = true;
+            overrides.writes = true;
         }
 
         if cte_writes {
-            writes.writes = true;
+            overrides.writes = true;
+        }
+
+        if overrides.cross_shard {
+            context
+                .shards_calculator
+                .push(ShardWithPriority::new_override_cross_shard_function());
         }
 
         // Early return for any direct-to-shard queries.
         if context.shards_calculator.shard().is_direct() {
             return Ok(Command::Query(
-                Route::read(context.shards_calculator.shard().clone()).with_write(writes),
+                Route::read(context.shards_calculator.shard().clone()).with_functions(overrides),
             ));
         }
 
@@ -83,7 +89,7 @@ impl QueryParser {
                 .push(ShardWithPriority::new_rr_no_table(shard));
 
             return Ok(Command::Query(
-                Route::read(context.shards_calculator.shard().clone()).with_write(writes),
+                Route::read(context.shards_calculator.shard().clone()).with_functions(overrides),
             ));
         }
 
@@ -201,7 +207,7 @@ impl QueryParser {
             query.with_aggregate_rewrite_plan_mut(cached_ast.rewrite_plan.aggregates.clone());
         }
 
-        Ok(Command::Query(query.with_write(writes)))
+        Ok(Command::Query(query.with_functions(overrides)))
     }
 
     /// Handle the `ORDER BY` clause of a `SELECT` statement.

pgdog/src/frontend/router/parser/query/test/test_functions.rs

@@ -31,3 +31,29 @@ fn test_write_function_nextval() {
     assert!(command.route().is_write());
     assert!(!command.route().is_lock_session());
 }
+
+#[test]
+fn test_cross_shard_install_sharded_sequence() {
+    let mut test = QueryParserTest::new();
+
+    let command = test.execute(vec![Query::new(
+        "SELECT pgdog.install_sharded_sequence('foo', 'id')",
+    )
+    .into()]);
+
+    assert!(command.route().is_cross_shard());
+}
+
+#[test]
+fn test_install_sharded_sequence_without_schema_not_cross_shard() {
+    // Without the `pgdog.` schema qualifier we should not flag the call
+    // as a cross-shard function — it could be any user-defined function.
+    let mut test = QueryParserTest::new();
+
+    let command = test.execute(vec![Query::new(
+        "SELECT install_sharded_sequence('foo', 'id')",
+    )
+    .into()]);
+
+    assert!(!command.route().is_cross_shard());
+}

pgdog/src/frontend/router/parser/route.rs

@@ -273,16 +273,17 @@ impl Route {
         self.rollback_savepoint
     }
 
-    pub fn with_write(mut self, write: FunctionBehavior) -> Self {
-        self.set_write(write);
+    pub fn with_functions(mut self, function: FunctionBehavior) -> Self {
+        self.set_functions(function);
         self
     }
 
-    pub fn set_write(&mut self, write: FunctionBehavior) {
+    pub fn set_functions(&mut self, function: FunctionBehavior) {
         let FunctionBehavior {
             writes,
             locking_behavior,
-        } = write;
+            ..
+        } = function;
         self.read = !writes;
         self.lock_session = match locking_behavior {
             LockingBehavior::Lock => Some(true),
@@ -364,6 +365,7 @@ pub enum OverrideReason {
     Transaction,
     OnlyOneShard,
     RewriteUpdate,
+    CrossShardFunction,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)]
@@ -424,6 +426,13 @@ impl ShardWithPriority {
         }
     }
 
+    pub fn new_override_cross_shard_function() -> Self {
+        Self {
+            shard: Shard::All,
+            source: ShardSource::Override(OverrideReason::CrossShardFunction),
+        }
+    }
+
     pub fn new_override_dry_run(shard: Shard) -> Self {
         Self {
             shard,