diff --git a/pgdog/src/frontend/router/parser/query/ddl.rs b/pgdog/src/frontend/router/parser/query/ddl.rs index 079d7c472..694b3d137 100644 --- a/pgdog/src/frontend/router/parser/query/ddl.rs +++ b/pgdog/src/frontend/router/parser/query/ddl.rs @@ -51,7 +51,7 @@ impl QueryParser { } ObjectType::ObjectSchema => { - if let Some(Node { + if let Some(PgNode { node: Some(NodeEnum::String(string)), }) = stmt.objects.first() && let Some(schema) = schema.schemas.get(Some(string.sval.as_str().into())) diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index 63f12ff8c..8978f77f9 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -37,6 +37,9 @@ mod transaction; mod update; use multi_tenant::MultiTenantCheck; +use pg_query::protobuf::Node as PgNode; +#[cfg(feature = "new_parser")] +use pg_raw_parse::Node; use pgdog_plugin::pg_query::{ NodeEnum, protobuf::{a_const::Val, *}, @@ -71,23 +74,39 @@ impl QueryParser { self.explain_recorder.as_mut() } - fn ensure_explain_recorder( - &mut self, - ast: &pg_query::ParseResult, - context: &QueryParserContext, - ) { + #[cfg(feature = "new_parser")] + fn ensure_explain_recorder(&mut self, node: Node<'_>, context: &QueryParserContext) { if self.explain_recorder.is_some() || !context.expanded_explain() { return; } - if let Some(root) = ast.protobuf.stmts.first() - && let Some(node) = root.stmt.as_ref().and_then(|stmt| stmt.node.as_ref()) - && matches!(node, NodeEnum::ExplainStmt(_)) - { + if matches!(node, Node::ExplainStmt(_)) { self.explain_recorder = Some(ExplainRecorder::new()); } } + cfg_select! { + not(feature = "new_parser") => { + fn ensure_explain_recorder( + &mut self, + ast: &pg_query::ParseResult, + context: &QueryParserContext, + ) { + if self.explain_recorder.is_some() || !context.expanded_explain() { + return; + } + + if let Some(root) = ast.protobuf.stmts.first() + && let Some(node) = root.stmt.as_ref().and_then(|stmt| stmt.node.as_ref()) + && matches!(node, NodeEnum::ExplainStmt(_)) + { + self.explain_recorder = Some(ExplainRecorder::new()); + } + } + } + _ => {} + } + fn attach_explain(&mut self, command: &mut Command) { if let (Some(recorder), Command::Query(route)) = (self.explain_recorder.take(), command) { let summary = ExplainSummary { @@ -213,6 +232,11 @@ impl QueryParser { .clone() .ok_or(Error::EmptyQuery)?; + #[cfg(feature = "new_parser")] + if let Some(stmt) = statement.new_ast.stmts().next() { + self.ensure_explain_recorder(stmt, context); + } + #[cfg(not(feature = "new_parser"))] self.ensure_explain_recorder(statement.parse_result(), context); // Parse hardcoded shard from a query comment. @@ -624,12 +648,12 @@ fn extract_set_config(stmt: &SelectStmt) -> Option<&FuncCall> { ]; // FIXME(sage): Dear god we need some pattern macros for this if let [ - Node { + PgNode { node: Some(NodeEnum::ResTarget(r)), }, ] = &*stmt.target_list && let ResTarget { val: Some(n), .. } = &**r - && let Node { + && let PgNode { node: Some(NodeEnum::FuncCall(f)), } = &**n && SET_CONFIG.iter().any(|&n| n == f.funcname) diff --git a/pgdog/src/frontend/router/parser/query/set_config.rs b/pgdog/src/frontend/router/parser/query/set_config.rs index 63c907f52..387ae27c3 100644 --- a/pgdog/src/frontend/router/parser/query/set_config.rs +++ b/pgdog/src/frontend/router/parser/query/set_config.rs @@ -30,7 +30,7 @@ fn parse_args(fcall: &FuncCall) -> Option { } /// Returns None if the name could not be parsed -fn parse_config_name(arg: &Node) -> Option { +fn parse_config_name(arg: &PgNode) -> Option { match &arg.node { Some(NodeEnum::AConst(AConst { val: Some(Val::Sval(PgString { sval })), @@ -43,7 +43,7 @@ fn parse_config_name(arg: &Node) -> Option { /// Returns None if the value could not be parsed, Some(None) if the value /// is NULL, and Some if the value was successfully parsed -fn parse_config_value(arg: &Node) -> Option> { +fn parse_config_value(arg: &PgNode) -> Option> { match &arg.node { Some(NodeEnum::AConst(AConst { val: Some(Val::Sval(PgString { sval })), @@ -57,7 +57,7 @@ fn parse_config_value(arg: &Node) -> Option> { } /// Returns None if the node was not a constant boolean -fn parse_is_local(arg: &Node) -> Option { +fn parse_is_local(arg: &PgNode) -> Option { match &arg.node { Some(NodeEnum::AConst(AConst { val: Some(Val::Boolval(Boolean { boolval })), diff --git a/pgdog/src/frontend/router/parser/query/transaction.rs b/pgdog/src/frontend/router/parser/query/transaction.rs index ee4ff01a6..8dd4ec861 100644 --- a/pgdog/src/frontend/router/parser/query/transaction.rs +++ b/pgdog/src/frontend/router/parser/query/transaction.rs @@ -61,7 +61,7 @@ impl QueryParser { } #[inline] - fn transaction_type(options: &[Node]) -> Option { + fn transaction_type(options: &[PgNode]) -> Option { for option_node in options { let node_enum = option_node.node.as_ref()?; if let NodeEnum::DefElem(def_elem) = node_enum