From 79226b73fec04aa7f04a6e474d3d4a444ea6d04f Mon Sep 17 00:00:00 2001 From: psteinroe Date: Thu, 14 Dec 2023 09:41:28 +0100 Subject: [PATCH 1/2] feat: add support for create aggregate with order by --- crates/codegen/src/get_node_properties.rs | 42 +++++++++++++++++++ crates/parser/src/parse/libpg_query_node.rs | 1 + .../tests/data/statements/valid/0039.sql | 10 +++++ 3 files changed, 53 insertions(+) create mode 100644 crates/parser/tests/data/statements/valid/0039.sql diff --git a/crates/codegen/src/get_node_properties.rs b/crates/codegen/src/get_node_properties.rs index 0caecb9d..51e32cec 100644 --- a/crates/codegen/src/get_node_properties.rs +++ b/crates/codegen/src/get_node_properties.rs @@ -481,6 +481,48 @@ fn custom_handlers(node: &Node) -> TokenStream { tokens.push(TokenProperty::from(Token::As)); } }, + "DefineStmt" => quote! { + tokens.push(TokenProperty::from(Token::Create)); + if n.replace { + tokens.push(TokenProperty::from(Token::Or)); + tokens.push(TokenProperty::from(Token::Replace)); + } + match n.kind() { + protobuf::ObjectType::ObjectAggregate => { + tokens.push(TokenProperty::from(Token::Aggregate)); + + // n.args is always an array with two nodes + assert_eq!(n.args.len(), 2, "DefineStmt of type ObjectAggregate does not have exactly 2 args"); + // the first is either a List or a Node { node: None } + + if let Some(node) = &n.args.first() { + if node.node.is_none() { + // if first element is a Node { node: None }, then it's "*" + tokens.push(TokenProperty::from(Token::Ascii42)); + } else if let Some(node) = &node.node { + if let NodeEnum::List(_) = node { + // there *seems* to be an integer node in the last position of args that + // defines whether the list contains an order by statement + let integer = n.args.last() + .and_then(|node| node.node.as_ref()) + .and_then(|node| if let NodeEnum::Integer(n) = node { Some(n.ival) } else { None }); + if integer.is_none() { + panic!("DefineStmt of type ObjectAggregate has no integer node in last position of args"); + } + // if the integer is 1, then there is an order by statement + // BUT: the order by tokens should be part of the List or maybe + // even the last FunctionParameter node in the list + if integer.unwrap() == 1 { + tokens.push(TokenProperty::from(Token::Order)); + tokens.push(TokenProperty::from(Token::By)); + } + } + } + } + }, + _ => panic!("Unknown DefineStmt {:#?}", n.kind()), + } + }, "CreateSchemaStmt" => quote! { tokens.push(TokenProperty::from(Token::Create)); tokens.push(TokenProperty::from(Token::Schema)); diff --git a/crates/parser/src/parse/libpg_query_node.rs b/crates/parser/src/parse/libpg_query_node.rs index b2c58830..bad223fc 100644 --- a/crates/parser/src/parse/libpg_query_node.rs +++ b/crates/parser/src/parse/libpg_query_node.rs @@ -52,6 +52,7 @@ impl<'p> LibpgQueryNodeParser<'p> { ) -> LibpgQueryNodeParser<'p> { let current_depth = parser.depth.clone(); debug!("Parsing node {:#?}", node); + println!("Parsing node {:#?}", node); Self { parser, token_range, diff --git a/crates/parser/tests/data/statements/valid/0039.sql b/crates/parser/tests/data/statements/valid/0039.sql new file mode 100644 index 00000000..91ad4381 --- /dev/null +++ b/crates/parser/tests/data/statements/valid/0039.sql @@ -0,0 +1,10 @@ +CREATE AGGREGATE aggregate1 (int4) (sfunc = sfunc1, stype = stype1); +CREATE AGGREGATE aggregate1 (int4, bool) (sfunc = sfunc1, stype = stype1); +CREATE AGGREGATE aggregate1 (*) (sfunc = sfunc1, stype = stype1); +CREATE AGGREGATE aggregate1 (int4) (sfunc = sfunc1, stype = stype1, finalfunc_extra, mfinalfuncextra); +CREATE AGGREGATE aggregate1 (int4) (sfunc = sfunc1, stype = stype1, finalfunc_modify = read_only, parallel = restricted); +CREATE AGGREGATE percentile_disc (float8 ORDER BY anyelement) (sfunc = ordered_set_transition, stype = internal, finalfunc = percentile_disc_final, finalfunc_extra); +CREATE AGGREGATE custom_aggregate (float8 ORDER BY column1, column2) (sfunc = sfunc1, stype = stype1); + + + From 07539124f5402b9600ed6cf147bd0f75bdb4ff2f Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sun, 17 Dec 2023 13:58:21 +0100 Subject: [PATCH 2/2] feat: pass parent to and finally fix define aggregate --- crates/codegen/src/get_node_properties.rs | 50 +++++++++++++---------- crates/codegen/src/get_nodes.rs | 6 +-- 2 files changed, 32 insertions(+), 24 deletions(-) diff --git a/crates/codegen/src/get_node_properties.rs b/crates/codegen/src/get_node_properties.rs index 51e32cec..a63401be 100644 --- a/crates/codegen/src/get_node_properties.rs +++ b/crates/codegen/src/get_node_properties.rs @@ -127,7 +127,7 @@ pub fn get_node_properties_mod(proto_file: &ProtoFile) -> proc_macro2::TokenStre } } - pub fn get_node_properties(node: &NodeEnum) -> Vec { + pub fn get_node_properties(node: &NodeEnum, parent: Option<&NodeEnum>) -> Vec { let mut tokens: Vec = Vec::new(); match node { @@ -481,6 +481,32 @@ fn custom_handlers(node: &Node) -> TokenStream { tokens.push(TokenProperty::from(Token::As)); } }, + "List" => quote! { + if parent.is_some() { + // if parent is `DefineStmt`, we need to check whether an ORDER BY needs to be added + if let NodeEnum::DefineStmt(define_stmt) = parent.unwrap() { + // there *seems* to be an integer node in the last position of the DefineStmt args that + // defines whether the list contains an order by statement + let integer = define_stmt.args.last() + .and_then(|node| node.node.as_ref()) + .and_then(|node| if let NodeEnum::Integer(n) = node { Some(n.ival) } else { None }); + if integer.is_none() { + panic!("DefineStmt of type ObjectAggregate has no integer node in last position of args"); + } + // if the integer is 1, then there is an order by statement + // we add it to the `List` node because that seems to make most sense based off the grammar definition + // ref: https://github.com/postgres/postgres/blob/REL_15_STABLE/src/backend/parser/gram.y#L8355 + // ``` + // aggr_args: + // | '(' aggr_args_list ORDER BY aggr_args_list ')' + // ``` + if integer.unwrap() == 1 { + tokens.push(TokenProperty::from(Token::Order)); + tokens.push(TokenProperty::from(Token::By)); + } + } + } + }, "DefineStmt" => quote! { tokens.push(TokenProperty::from(Token::Create)); if n.replace { @@ -499,26 +525,8 @@ fn custom_handlers(node: &Node) -> TokenStream { if node.node.is_none() { // if first element is a Node { node: None }, then it's "*" tokens.push(TokenProperty::from(Token::Ascii42)); - } else if let Some(node) = &node.node { - if let NodeEnum::List(_) = node { - // there *seems* to be an integer node in the last position of args that - // defines whether the list contains an order by statement - let integer = n.args.last() - .and_then(|node| node.node.as_ref()) - .and_then(|node| if let NodeEnum::Integer(n) = node { Some(n.ival) } else { None }); - if integer.is_none() { - panic!("DefineStmt of type ObjectAggregate has no integer node in last position of args"); - } - // if the integer is 1, then there is an order by statement - // BUT: the order by tokens should be part of the List or maybe - // even the last FunctionParameter node in the list - if integer.unwrap() == 1 { - tokens.push(TokenProperty::from(Token::Order)); - tokens.push(TokenProperty::from(Token::By)); - } - } - } - } + } } + // if its a list, we handle it in the handler for `List` }, _ => panic!("Unknown DefineStmt {:#?}", n.kind()), } diff --git a/crates/codegen/src/get_nodes.rs b/crates/codegen/src/get_nodes.rs index 69296501..73574f44 100644 --- a/crates/codegen/src/get_nodes.rs +++ b/crates/codegen/src/get_nodes.rs @@ -25,7 +25,7 @@ pub fn get_nodes_mod(proto_file: &ProtoFile) -> proc_macro2::TokenStream { let root_node_idx = g.add_node(Node { kind: SyntaxKind::from(node), depth: at_depth, - properties: get_node_properties(node), + properties: get_node_properties(node, None), location: get_location(node), }); @@ -45,12 +45,12 @@ pub fn get_nodes_mod(proto_file: &ProtoFile) -> proc_macro2::TokenStream { NodeEnum::BitString(n) => true, _ => false } { - g[parent_idx].properties.extend(get_node_properties(&c)); + g[parent_idx].properties.extend(get_node_properties(&c, Some(&node))); } else { let node_idx = g.add_node(Node { kind: SyntaxKind::from(&c), depth: current_depth, - properties: get_node_properties(&c), + properties: get_node_properties(&c, Some(&node)), location: get_location(&c), }); g.add_edge(parent_idx, node_idx, ());