Use parsed AST instead of double-parsing SELECT stmt (#527)

Author: levkk

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

@@ -3,14 +3,15 @@ use super::*;
 impl QueryParser {
     pub(super) fn explain(
         &mut self,
+        ast: &pg_query::ParseResult,
         stmt: &ExplainStmt,
         context: &mut QueryParserContext,
     ) -> Result<Command, Error> {
         let query = stmt.query.as_ref().ok_or(Error::EmptyQuery)?;
         let node = query.node.as_ref().ok_or(Error::EmptyQuery)?;
 
         match node {
-            NodeEnum::SelectStmt(ref stmt) => self.select(stmt, context),
+            NodeEnum::SelectStmt(ref stmt) => self.select(ast, stmt, context),
             NodeEnum::InsertStmt(ref stmt) => Self::insert(stmt, context),
             NodeEnum::UpdateStmt(ref stmt) => Self::update(stmt, context),
             NodeEnum::DeleteStmt(ref stmt) => Self::delete(stmt, context),

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

@@ -212,7 +212,7 @@ impl QueryParser {
                 return Ok(Command::Deallocate);
             }
             // SELECT statements.
-            Some(NodeEnum::SelectStmt(ref stmt)) => self.select(stmt, context),
+            Some(NodeEnum::SelectStmt(ref stmt)) => self.select(statement.ast(), stmt, context),
             // COPY statements.
             Some(NodeEnum::CopyStmt(ref stmt)) => Self::copy(stmt, context),
             // INSERT statements.
@@ -258,7 +258,7 @@ impl QueryParser {
                 return Ok(Command::Unlisten(stmt.conditionname.clone()));
             }
 
-            Some(NodeEnum::ExplainStmt(ref stmt)) => self.explain(stmt, context),
+            Some(NodeEnum::ExplainStmt(ref stmt)) => self.explain(statement.ast(), stmt, context),
 
             // VACUUM.
             Some(NodeEnum::VacuumRelation(_)) | Some(NodeEnum::VacuumStmt(_)) => {

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

@@ -12,6 +12,7 @@ impl QueryParser {
     ///
     pub(super) fn select(
         &mut self,
+        ast: &pg_query::ParseResult,
         stmt: &SelectStmt,
         context: &mut QueryParserContext,
     ) -> Result<Command, Error> {
@@ -93,8 +94,11 @@ impl QueryParser {
         // Only rewrite if query is cross-shard.
         if query.is_cross_shard() && context.shards > 1 {
             if let Some(buffered_query) = context.router_context.query.as_ref() {
-                let rewrite =
-                    RewriteEngine::new().rewrite_select(buffered_query.query(), query.aggregate());
+                let rewrite = RewriteEngine::new().rewrite_select(
+                    ast,
+                    buffered_query.query(),
+                    query.aggregate(),
+                );
                 if !rewrite.plan.is_noop() {
                     if let BufferedQuery::Prepared(parse) = buffered_query {
                         let name = parse.name().to_owned();

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

@@ -1,7 +1,6 @@
 use super::{Aggregate, AggregateFunction, HelperMapping, RewriteOutput, RewritePlan};
 use pg_query::protobuf::{FuncCall, Node, ResTarget, String as PgString};
 use pg_query::{NodeEnum, ParseResult};
-use tracing::debug;
 
 /// Query rewrite engine. Currently supports injecting helper aggregates for AVG.
 #[derive(Default)]
@@ -13,19 +12,18 @@ impl RewriteEngine {
     }
 
     /// Rewrite a SELECT query, adding helper aggregates when necessary.
-    pub fn rewrite_select(&self, sql: &str, aggregate: &Aggregate) -> RewriteOutput {
-        match pg_query::parse(sql) {
-            Ok(parsed) => self.rewrite_parsed(parsed, aggregate, sql),
-            Err(err) => {
-                debug!("rewrite failed to parse SELECT: {}", err);
-                RewriteOutput::new(sql.to_string(), RewritePlan::new())
-            }
-        }
+    pub fn rewrite_select(
+        &self,
+        ast: &ParseResult,
+        sql: &str,
+        aggregate: &Aggregate,
+    ) -> RewriteOutput {
+        self.rewrite_parsed(ast, aggregate, sql)
     }
 
     fn rewrite_parsed(
         &self,
-        parsed: ParseResult,
+        parsed: &ParseResult,
         aggregate: &Aggregate,
         original_sql: &str,
     ) -> RewriteOutput {
@@ -171,7 +169,7 @@ mod tests {
             None => panic!("empty"),
         };
         let aggregate = Aggregate::parse(stmt).unwrap();
-        RewriteEngine::new().rewrite_select(sql, &aggregate)
+        RewriteEngine::new().rewrite_select(&ast, sql, &aggregate)
     }
 
     #[test]