Port QueryParser::ensure_explain_recorder to the new parser

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()))

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)

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

@@ -30,7 +30,7 @@ fn parse_args(fcall: &FuncCall) -> Option<SetParam> {
 }
 
 /// Returns None if the name could not be parsed
-fn parse_config_name(arg: &Node) -> Option<String> {
+fn parse_config_name(arg: &PgNode) -> Option<String> {
     match &arg.node {
         Some(NodeEnum::AConst(AConst {
             val: Some(Val::Sval(PgString { sval })),
@@ -43,7 +43,7 @@ fn parse_config_name(arg: &Node) -> Option<String> {
 
 /// 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<Option<ParameterValue>> {
+fn parse_config_value(arg: &PgNode) -> Option<Option<ParameterValue>> {
     match &arg.node {
         Some(NodeEnum::AConst(AConst {
             val: Some(Val::Sval(PgString { sval })),
@@ -57,7 +57,7 @@ fn parse_config_value(arg: &Node) -> Option<Option<ParameterValue>> {
 }
 
 /// Returns None if the node was not a constant boolean
-fn parse_is_local(arg: &Node) -> Option<bool> {
+fn parse_is_local(arg: &PgNode) -> Option<bool> {
     match &arg.node {
         Some(NodeEnum::AConst(AConst {
             val: Some(Val::Boolval(Boolean { boolval })),

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

@@ -61,7 +61,7 @@ impl QueryParser {
     }
 
     #[inline]
-    fn transaction_type(options: &[Node]) -> Option<TransactionType> {
+    fn transaction_type(options: &[PgNode]) -> Option<TransactionType> {
         for option_node in options {
             let node_enum = option_node.node.as_ref()?;
             if let NodeEnum::DefElem(def_elem) = node_enum