diff --git a/pgdog/src/frontend/router/parser/rewrite/ee/nextval.rs b/pgdog/src/frontend/router/parser/rewrite/ee/nextval.rs index a1631108c..f18ac17fb 100644 --- a/pgdog/src/frontend/router/parser/rewrite/ee/nextval.rs +++ b/pgdog/src/frontend/router/parser/rewrite/ee/nextval.rs @@ -3,3 +3,11 @@ use super::Error; pub(crate) async fn nextval(_name: &str) -> Result { Err(Error::EERequired) } + +pub(crate) async fn currval(_name: &str) -> Result { + Err(Error::EERequired) +} + +pub(crate) async fn setval(_name: &str, _value: i64, _is_called: bool) -> Result { + Err(Error::EERequired) +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs b/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs index 5c6b88902..ba0f9223a 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs @@ -227,6 +227,7 @@ fn is_bigint_type(data_type: &str) -> bool { #[cfg(test)] mod tests { + use super::super::nextval::SequenceCall; use super::super::plan::GeneratedId; use crate::frontend::router::sharding::ShardedTable; use indexmap::IndexMap; @@ -691,8 +692,14 @@ mod tests { assert_eq!( plan.generated_ids, vec![ - (1, GeneratedId::Sequence(sequence.to_owned())), - (2, GeneratedId::Sequence(sequence.to_owned())), + ( + 1, + GeneratedId::Sequence(SequenceCall::Nextval(sequence.to_owned())) + ), + ( + 2, + GeneratedId::Sequence(SequenceCall::Nextval(sequence.to_owned())) + ), ] ); } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index 59b05c95e..87c600ca9 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -166,7 +166,7 @@ impl<'a> StatementRewrite<'a> { } Ok(None) => { if let Some(replacement) = - self.rewrite_nextval(node.as_ref(), mem, &mut next_param, &mut plan) + self.rewrite_sequence(node.as_ref(), mem, &mut next_param, &mut plan) { node.replace(replacement); None diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/nextval.rs b/pgdog/src/frontend/router/parser/rewrite/statement/nextval.rs index 75b30f327..210584e1b 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/nextval.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/nextval.rs @@ -7,17 +7,43 @@ use crate::frontend::router::parser::rewrite::ee; use super::plan::GeneratedId; use super::{Error, RewritePlan, StatementRewrite}; +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum SequenceCall { + Nextval(String), + Currval(String), + Setval { + name: String, + value: i64, + is_called: bool, + }, +} + +impl SequenceCall { + pub(super) async fn execute(&self) -> Result { + match self { + Self::Nextval(name) => ee::nextval(name).await, + Self::Currval(name) => ee::currval(name).await, + Self::Setval { + name, + value, + is_called, + } => ee::setval(name, *value, *is_called).await, + } + } +} + impl RewritePlan { /// Replace simple-protocol sequence calls in the current statement with /// freshly fetched bigint literals. The cached plan remains unchanged. /// Returns `None` when the plan has no statement. - pub(crate) async fn rewrite_nextval_simple(&self) -> Result, Error> { - self.rewrite_nextval_simple_with(ee::nextval).await + pub(crate) async fn rewrite_sequence_simple(&self) -> Result, Error> { + self.rewrite_sequence_simple_with(SequenceCall::execute) + .await } - async fn rewrite_nextval_simple_with( + async fn rewrite_sequence_simple_with( &self, - mut nextval: impl AsyncFnMut(&str) -> Result, + mut execute: impl AsyncFnMut(&SequenceCall) -> Result, ) -> Result, Error> { let Some(stmt) = &self.stmt else { return Ok(None); @@ -29,9 +55,9 @@ impl RewritePlan { for stmt in ast.stmts() { walk::walk(stmt, |node| { if let Node::FuncCall(func) = node - && let Some(name) = sequence_name(node) + && let Some(call) = sequence_call(node) { - calls.push((func.location, name)); + calls.push((func.location, call)); } }); } @@ -43,8 +69,8 @@ impl RewritePlan { // AST mutation callbacks are synchronous; resolve values before // entering the parser's memory context to replace the calls. let mut values = HashMap::with_capacity(calls.len()); - for (location, name) in calls { - values.insert(location, nextval(&name).await?); + for (location, call) in calls { + values.insert(location, execute(&call).await?); } let rewritten = make::owned(|mem| { let mut copy = mem.make_unique(&**ast); @@ -82,19 +108,19 @@ impl RewritePlan { impl StatementRewrite<'_> { /// Record a sequence call and return its replacement in extended protocol. /// Simple protocol retains the call for asynchronous rewriting. - pub(super) fn rewrite_nextval<'mem>( + pub(super) fn rewrite_sequence<'mem>( &mut self, node: Node<'_>, mem: make::MemoryToken<'mem>, next_param: &mut i32, plan: &mut RewritePlan, ) -> Option>> { - let sequence = sequence_name(node)?; + let sequence = sequence_call(node)?; let param = *next_param; *next_param += 1; plan.generated_ids .push((param as u16, GeneratedId::Sequence(sequence))); - // Retain simple-protocol SQL even when nextval is the only rewrite. + // Retain simple-protocol SQL even when a sequence call is the only rewrite. self.rewritten = true; self.extended.then(|| { mem.make_type_cast( @@ -109,26 +135,50 @@ impl StatementRewrite<'_> { } } -/// Extract the literal sequence name from a pgdog.nextval() call. -fn sequence_name(node: Node<'_>) -> Option { +/// Recognize pgdog sequence calls with literal arguments. +fn sequence_call(node: Node<'_>) -> Option { let Node::FuncCall(func) = node else { return None; }; - if !func - .funcname() - .iter() - .map(Node::as_str) - .eq([Some("pgdog"), Some("nextval")]) - { + let mut names = func.funcname().iter(); + if names.next()?.as_str()? != "pgdog" { + return None; + } + let function = names.next()?.as_str()?; + if names.next().is_some() { return None; } let mut args = func.args().iter(); - let mut arg = args.next()?; + let name = sequence_name(args.next()?)?; + let call = match function { + "nextval" => SequenceCall::Nextval(name), + "currval" => SequenceCall::Currval(name), + "setval" => { + let Node::A_Const(value) = args.next()? else { + return None; + }; + let is_called = match args.next() { + None => true, + Some(Node::A_Const(is_called)) => is_called.val()?.bool_value()?, + _ => return None, + }; + SequenceCall::Setval { + name, + value: value.val()?.numeric_value::()?, + is_called, + } + } + _ => return None, + }; if args.next().is_some() { return None; } + Some(call) +} +/// Extract a literal sequence name, optionally cast to regclass. +fn sequence_name(mut arg: Node<'_>) -> Option { if let Node::TypeCast(cast) = arg { let type_name = cast.type_name()?; let names = type_name.names(); @@ -192,10 +242,10 @@ mod tests { } #[test] - fn test_nextval_extended_parameter_indexes() { + fn test_sequence_extended_parameter_indexes() { let (sql, plan) = rewrite( "SELECT pgdog.nextval('sequence.name'), $1, pgdog.unique_id(), \ - pgdog.nextval('other.seq'::regclass), pgdog.nextval('sequence.name')", + pgdog.currval('other.seq'::regclass), pgdog.setval('sequence.name', 42, false)", true, ); assert_eq!( @@ -207,10 +257,23 @@ mod tests { assert_eq!( plan.generated_ids, vec![ - (2, GeneratedId::Sequence("sequence.name".to_owned())), + ( + 2, + GeneratedId::Sequence(SequenceCall::Nextval("sequence.name".to_owned())) + ), (3, GeneratedId::UniqueId), - (4, GeneratedId::Sequence("other.seq".to_owned())), - (5, GeneratedId::Sequence("sequence.name".to_owned())), + ( + 4, + GeneratedId::Sequence(SequenceCall::Currval("other.seq".to_owned())) + ), + ( + 5, + GeneratedId::Sequence(SequenceCall::Setval { + name: "sequence.name".to_owned(), + value: 42, + is_called: false, + }) + ), ] ); assert_eq!(plan.stmt.as_deref(), Some(sql.as_str())); @@ -225,8 +288,14 @@ mod tests { assert_eq!( plan.generated_ids, vec![ - (1, GeneratedId::Sequence("sequence.name".to_owned())), - (2, GeneratedId::Sequence("sequence.name".to_owned())), + ( + 1, + GeneratedId::Sequence(SequenceCall::Nextval("sequence.name".to_owned())) + ), + ( + 2, + GeneratedId::Sequence(SequenceCall::Nextval("sequence.name".to_owned())) + ), ] ); assert_eq!(plan.unique_ids, 0); @@ -248,9 +317,14 @@ mod tests { vec![ ( 1, - GeneratedId::Sequence("\"My Schema\".\"My Sequence\"".to_owned()) + GeneratedId::Sequence(SequenceCall::Nextval( + "\"My Schema\".\"My Sequence\"".to_owned() + )) + ), + ( + 2, + GeneratedId::Sequence(SequenceCall::Nextval("other.seq".to_owned())) ), - (2, GeneratedId::Sequence("other.seq".to_owned())), ] ); } @@ -270,7 +344,10 @@ mod tests { let mut names = Vec::new(); let mut values = [i64::MIN, 0, i64::MAX].into_iter(); let sql = plan - .rewrite_nextval_simple_with(async |name: &str| { + .rewrite_sequence_simple_with(async |call: &SequenceCall| { + let SequenceCall::Nextval(name) = call else { + panic!("expected nextval"); + }; names.push(name.to_owned()); Ok(values.next().expect("one value per call")) }) @@ -295,13 +372,13 @@ mod tests { false, ); let mut value = 0i64; - let mut nextval = async |_: &str| { + let mut nextval = async |_: &SequenceCall| { value += 1; Ok(value) }; for (first, second) in [(1, 2), (3, 4)] { let sql = plan - .rewrite_nextval_simple_with(&mut nextval) + .rewrite_sequence_simple_with(&mut nextval) .await .expect("rewrite succeeds") .expect("statement"); @@ -318,7 +395,7 @@ mod tests { let before = plan.stmt.clone(); let mut calls = 0; let error = plan - .rewrite_nextval_simple_with(async |_: &str| { + .rewrite_sequence_simple_with(async |_: &SequenceCall| { calls += 1; if calls == 1 { Ok(42) @@ -331,7 +408,7 @@ mod tests { assert!(matches!(error, Error::Enterprise(ee::Error::EERequired))); assert_eq!(plan.stmt, before); assert!(matches!( - plan.rewrite_nextval_simple().await, + plan.rewrite_sequence_simple().await, Err(Error::Enterprise(ee::Error::EERequired)) )); } @@ -340,7 +417,7 @@ mod tests { async fn test_nextval_simple_without_calls() { assert_eq!( RewritePlan::default() - .rewrite_nextval_simple() + .rewrite_sequence_simple() .await .expect("no SQL"), None @@ -351,7 +428,7 @@ mod tests { ..Default::default() }; assert_eq!( - plan.rewrite_nextval_simple() + plan.rewrite_sequence_simple() .await .expect("no global calls"), Some(sql.to_owned()) @@ -361,7 +438,7 @@ mod tests { ..Default::default() }; assert!(matches!( - invalid.rewrite_nextval_simple().await, + invalid.rewrite_sequence_simple().await, Err(Error::Parser(_)) )); } @@ -386,7 +463,10 @@ mod tests { let mut bind = Bind::new_params_codes("stmt", &original_params, &codes); let mut calls = Vec::new(); let mut value = -2i64; - plan.apply_generated_ids(&mut bind, async |name: &str| { + plan.apply_generated_ids(&mut bind, async |call: &SequenceCall| { + let SequenceCall::Nextval(name) = call else { + panic!("expected nextval"); + }; calls.push(name.to_owned()); value += 1; Ok(value) @@ -446,7 +526,7 @@ mod tests { async fn test_nextval_bind_reexecution_fetches_new_values() { let (_, plan) = rewrite("SELECT pgdog.nextval('a'), pgdog.nextval('a')", true); let mut value = 0i64; - let mut nextval = async |_: &str| { + let mut nextval = async |_: &SequenceCall| { value += 1; Ok(value) }; @@ -466,38 +546,247 @@ mod tests { } #[tokio::test] - async fn test_nextval_apply_propagates_enterprise_error() { - let (_, plan) = rewrite("SELECT pgdog.nextval('a')", true); - let mut request = ClientRequest::from(vec![ProtocolMessage::Bind(Bind::default())]); - let error = plan - .apply(&mut request) - .await - .expect_err("EE hook rejects sequence"); - assert!(matches!(error, Error::Enterprise(ee::Error::EERequired))); + async fn test_sequence_apply_propagates_enterprise_error() { + for call in [ + "nextval('a')", + "currval('a')", + "setval('a', 42)", + "setval('a', 42, true)", + ] { + let (_, plan) = rewrite(&format!("SELECT pgdog.{call}"), true); + let mut request = ClientRequest::from(vec![ProtocolMessage::Bind(Bind::default())]); + let error = plan + .apply(&mut request) + .await + .expect_err("EE hook rejects sequence"); + assert!(matches!(error, Error::Enterprise(ee::Error::EERequired))); + } } #[tokio::test] - async fn test_nextval_apply_fetches_for_query_but_not_parse() { - let original = "SELECT pgdog.nextval('a')"; - let (_, extended_plan) = rewrite(original, true); - let mut request = - ClientRequest::from(vec![ProtocolMessage::Parse(Parse::new_anonymous(original))]); - extended_plan - .apply(&mut request) - .await - .expect("prepare does not fetch"); + async fn test_sequence_apply_fetches_for_query_but_not_parse() { + for call in [ + "nextval('a')", + "currval('a')", + "setval('a', 42)", + "setval('a', 42, true)", + ] { + let original = format!("SELECT pgdog.{call}"); + let (_, extended_plan) = rewrite(&original, true); + let mut request = ClientRequest::from(vec![ProtocolMessage::Parse( + Parse::new_anonymous(&original), + )]); + extended_plan + .apply(&mut request) + .await + .expect("prepare does not fetch"); - let (_, simple_plan) = rewrite(original, false); - let mut request = ClientRequest::from(vec![ProtocolMessage::Query(Query::new(original))]); - let error = simple_plan - .apply(&mut request) - .await - .expect_err("simple query calls the EE hook"); - assert!(matches!(error, Error::Enterprise(ee::Error::EERequired))); - let ProtocolMessage::Query(query) = &request.messages[0] else { - panic!("expected Query"); - }; - assert_eq!(query.query(), original); + let (_, simple_plan) = rewrite(&original, false); + let mut request = + ClientRequest::from(vec![ProtocolMessage::Query(Query::new(&original))]); + let error = simple_plan + .apply(&mut request) + .await + .expect_err("simple query calls the EE hook"); + assert!(matches!(error, Error::Enterprise(ee::Error::EERequired))); + let ProtocolMessage::Query(query) = &request.messages[0] else { + panic!("expected Query"); + }; + assert_eq!(query.query(), original); + } + } + + #[test] + fn test_sequence_literal_arguments() { + let name = "\"My Schema\".\"My Sequence\""; + for argument in [ + format!("'{name}'"), + format!("'{name}'::regclass"), + format!("'{name}'::pg_catalog.regclass"), + format!("CAST('{name}' AS regclass)"), + ] { + let mut cases = vec![( + format!("pgdog.currval({argument})"), + SequenceCall::Currval(name.to_owned()), + )]; + for value in [i64::MIN, -2147483649, -1, 0, 2147483648, i64::MAX] { + for is_called in [None, Some(false), Some(true)] { + cases.push(( + match is_called { + Some(is_called) => { + format!("pgdog.setval({argument}, {value}, {is_called})") + } + None => format!("pgdog.setval({argument}, {value})"), + }, + SequenceCall::Setval { + name: name.to_owned(), + value, + is_called: is_called.unwrap_or(true), + }, + )); + } + } + for (call, expected) in cases { + for extended in [false, true] { + let original = format!("SELECT {call}"); + let (sql, plan) = rewrite(&original, extended); + assert_eq!( + plan.generated_ids, + [(1, GeneratedId::Sequence(expected.clone()))], + "{call}" + ); + let canonical = original.replace( + &format!("CAST('{name}' AS regclass)"), + &format!("'{name}'::regclass"), + ); + assert_eq!( + sql, + if extended { + "SELECT $1::bigint" + } else { + &canonical + }, + "{call}" + ); + assert_eq!(plan.stmt.as_deref(), Some(sql.as_str())); + } + } + } + } + + #[tokio::test] + async fn test_sequence_mixed_calls_reexecute_in_order() { + let original = "SELECT pgdog.nextval('a'), (SELECT pgdog.currval('a')), \ + pgdog.setval('a', -42, false), pgdog.nextval('a'), \ + pgdog.setval('a', 42), pgdog.nextval('a')"; + let expected_calls = [ + SequenceCall::Nextval("a".to_owned()), + SequenceCall::Currval("a".to_owned()), + SequenceCall::Setval { + name: "a".to_owned(), + value: -42, + is_called: false, + }, + SequenceCall::Nextval("a".to_owned()), + SequenceCall::Setval { + name: "a".to_owned(), + value: 42, + is_called: true, + }, + SequenceCall::Nextval("a".to_owned()), + ]; + for extended in [false, true] { + let (sql, plan) = rewrite(original, extended); + if extended { + assert_eq!( + sql, + "SELECT $1::bigint, (SELECT $2::bigint), $3::bigint, \ + $4::bigint, $5::bigint, $6::bigint" + ); + } + let before = plan.stmt.clone(); + let mut calls = Vec::new(); + let mut current = 0i64; + let mut called = true; + let mut execute = async |call: &SequenceCall| { + calls.push(call.clone()); + match call { + SequenceCall::Nextval(_) => { + if called { + current += 1; + } + called = true; + Ok(current) + } + SequenceCall::Currval(_) => Ok(current), + SequenceCall::Setval { + value, is_called, .. + } => { + current = *value; + called = *is_called; + Ok(*value) + } + } + }; + for first in [1, 44] { + if extended { + let mut bind = Bind::default(); + plan.apply_generated_ids(&mut bind, &mut execute) + .await + .expect("values appended"); + assert_eq!(bind.params_raw().len(), 6); + for (index, value) in [first, first, -42, -42, 42, 43].into_iter().enumerate() { + assert_eq!( + bind.parameter(index) + .expect("format") + .expect("parameter") + .bigint(), + Some(value) + ); + } + } else { + let sql = plan + .rewrite_sequence_simple_with(&mut execute) + .await + .expect("rewrite succeeds") + .expect("statement"); + assert_eq!( + sql, + format!( + "SELECT ({first})::bigint, (SELECT ({first})::bigint), \ + (-42)::bigint, (-42)::bigint, (42)::bigint, (43)::bigint" + ) + ); + } + assert_eq!(plan.stmt, before); + } + assert_eq!(calls.len(), expected_calls.len() * 2); + for execution in calls.chunks(expected_calls.len()) { + assert_eq!(execution, expected_calls); + } + } + } + + #[test] + fn test_sequence_unsupported_arguments() { + for call in [ + "currval('seq')", + "setval('seq', 42, true)", + "other.currval('seq')", + "other.setval('seq', 42, true)", + "pgdog.currval()", + "pgdog.currval('seq', 42)", + "pgdog.currval($1)", + "pgdog.currval(123)", + "pgdog.currval(NULL)", + "pgdog.currval('seq'::text)", + "pgdog.currval('seq'::other.regclass)", + "pgdog.setval()", + "pgdog.setval('seq')", + "pgdog.setval('seq', 42, true, false)", + "pgdog.setval($1, 42, true)", + "pgdog.setval(NULL, 42, true)", + "pgdog.setval('seq'::text, 42, true)", + "pgdog.setval('seq', $1, true)", + "pgdog.setval('seq', 1 + 2, true)", + "pgdog.setval('seq', NULL, true)", + "pgdog.setval('seq', '42', true)", + "pgdog.setval('seq', 1.5, true)", + "pgdog.setval('seq', 1e2, true)", + "pgdog.setval('seq', 9223372036854775808, true)", + "pgdog.setval('seq', -9223372036854775809, true)", + "pgdog.setval('seq', 42, $1)", + "pgdog.setval('seq', 42, NULL)", + "pgdog.setval('seq', 42, 'true')", + "pgdog.setval('seq', 42, 1)", + "pgdog.setval('seq', 42, 1 = 1)", + ] { + for extended in [false, true] { + let (_, plan) = rewrite(&format!("SELECT {call}"), extended); + assert!(plan.generated_ids.is_empty(), "{call}"); + assert!(plan.is_empty(), "{call}"); + } + } } #[test] @@ -514,8 +803,8 @@ mod tests { panic!("expected SELECT"); }; assert_eq!( - sequence_name(select.target_list().first().expect("target").val()), - Some("sequence.name".to_owned()), + sequence_call(select.target_list().first().expect("target").val()), + Some(SequenceCall::Nextval("sequence.name".to_owned())), "{argument}" ); } @@ -540,7 +829,7 @@ mod tests { panic!("expected SELECT"); }; assert_eq!( - sequence_name(select.target_list().first().expect("target").val()), + sequence_call(select.target_list().first().expect("target").val()), None, "{call}" ); diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs index 55695be8e..664fbf00d 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs @@ -5,6 +5,7 @@ use crate::unique_id::UniqueId; use super::super::ee; use super::insert::build_split_requests; +use super::nextval::SequenceCall; use super::offset::OffsetPlan; use super::{ Error, InsertSplit, PrepareExecute, ShardingKeyUpdate, aggregate::AggregateRewritePlan, @@ -13,7 +14,7 @@ use super::{ #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) enum GeneratedId { UniqueId, - Sequence(String), + Sequence(SequenceCall), } /// Statement rewrite plan. @@ -97,20 +98,20 @@ impl RewritePlan { /// Append generated unique IDs and sequence values to a Bind message. async fn apply_bind(&self, bind: &mut Bind) -> Result<(), Error> { - self.apply_generated_ids(bind, ee::nextval).await + self.apply_generated_ids(bind, SequenceCall::execute).await } /// Append values in the same order their placeholders were allocated. pub(super) async fn apply_generated_ids( &self, bind: &mut Bind, - mut nextval: impl AsyncFnMut(&str) -> Result, + mut execute: impl AsyncFnMut(&SequenceCall) -> Result, ) -> Result<(), Error> { let format = bind.default_param_format(); for (_, source) in &self.generated_ids { let id = match source { GeneratedId::UniqueId => UniqueId::generator()?.next_id(), - GeneratedId::Sequence(name) => nextval(name).await?, + GeneratedId::Sequence(call) => execute(call).await?, }; let param = match format { Format::Binary => Parameter::new(&id.to_be_bytes()), @@ -139,7 +140,7 @@ impl RewritePlan { .iter() .any(|(_, source)| matches!(source, GeneratedId::Sequence(_))) { - if let Some(stmt) = self.rewrite_nextval_simple().await? { + if let Some(stmt) = self.rewrite_sequence_simple().await? { query.set_query(&stmt); } } else if let Some(ref stmt) = self.stmt {