refactor: consolidate multi-step planning, assembly and execution into QueryPlanner type

pgdog/src/frontend/client/query_engine/context.rs

@@ -18,7 +18,7 @@ pub(crate) struct QueryEngineContext<'a> {
     /// Client session parameters.
     pub(super) params: &'a mut Parameters,
     /// Request.
-    pub(super) client_request: &'a mut ClientRequest,
+    pub(crate) client_request: &'a mut ClientRequest,
     /// How many requests are left to execute in an extended pipeline.
     pub(super) pipeline: Pipeline,
     /// Client's socket to send responses to.
@@ -106,4 +106,15 @@ impl<'a> QueryEngineContext<'a> {
     pub(crate) fn in_error(&self) -> bool {
         self.transaction.map(|t| t.error()).unwrap_or_default()
     }
+
+    /// Mark the transaction failed
+    pub(crate) fn set_transaction_error(&mut self) {
+        self.transaction = match self.transaction {
+            Some(TransactionType::ReadOnly) => Some(TransactionType::ErrorReadOnly),
+            Some(TransactionType::ReadWrite | TransactionType::Implicit) => {
+                Some(TransactionType::ErrorReadWrite)
+            }
+            _ => None,
+        };
+    }
 }

pgdog/src/frontend/client/query_engine/mod.rs

@@ -146,8 +146,8 @@ impl QueryEngine {
         }
 
         // Rewrite statement if necessary.
-        let rewrite_result = match self.parse_and_rewrite(context) {
-            Ok(rewrite_result) => rewrite_result,
+        let (query_planner, offset_plan) = match self.parse_and_rewrite(context).await {
+            Ok(result) => result,
             Err(e) => {
                 self.error_response(context, ErrorResponse::syntax(e.to_string()))
                     .await?;
@@ -162,7 +162,7 @@ impl QueryEngine {
         }
 
         // Route transaction to the right servers.
-        if !self.route_query(context, rewrite_result.as_ref()).await? {
+        if !self.route_query(context, offset_plan.as_ref()).await? {
             self.update_stats(context);
             debug!("query has nowhere to go");
             return Ok(QueryEngineResult::Done(context.transaction()));
@@ -240,11 +240,11 @@ impl QueryEngine {
 
                 context.params.rollback();
             }
-            Command::Query(_) => self.execute(context, rewrite_result).await?,
+            Command::Query(_) => self.execute(context, query_planner).await?,
             Command::Listen { .. } | Command::Notify { .. } | Command::Unlisten(_)
                 if self.backend.session_mode() =>
             {
-                self.execute(context, rewrite_result).await?
+                self.execute(context, query_planner).await?
             }
             Command::Listen { channel, shard } => {
                 self.listen(context, &channel.clone(), shard.clone())
@@ -268,7 +268,7 @@ impl QueryEngine {
             Command::ResetAll => {
                 self.reset_all(context).await?;
             }
-            Command::Copy(_) => self.execute(context, rewrite_result).await?,
+            Command::Copy(_) => self.execute(context, query_planner).await?,
             Command::Deallocate => self.deallocate(context).await?,
             Command::Discard { extended } => self.discard(context, *extended).await?,
             Command::Split(queries) => return Ok(Self::build_simple_split(queries)),

pgdog/src/frontend/client/query_engine/multi_step/error.rs

@@ -7,9 +7,18 @@ pub(crate) enum Error {
     #[error("{0}")]
     Update(#[from] UpdateError),
 
+    #[error("{0}")]
+    Insert(#[from] InsertError),
+
     #[error("frontend: {0}")]
     Frontend(Box<crate::frontend::Error>),
 
+    #[error("parser: {0}")]
+    Parser(#[from] crate::frontend::router::parser::Error),
+
+    #[error("deparse: {0}")]
+    Deparse(#[from] pg_raw_parse::Error),
+
     #[error("backend: {0}")]
     Backend(#[from] crate::backend::Error),
 
@@ -26,11 +35,30 @@ pub(crate) enum Error {
     Net(#[from] crate::net::Error),
 }
 
+impl Error {
+    /// Errors the client should see as an `ErrorResponse`.
+    /// Otherwise, it's an internal failure and propagates up which closes the connection.
+    pub(crate) fn into_client_error(self) -> Result<ErrorResponse, Self> {
+        match self {
+            Self::Execution(error) => Ok(*error),
+            err @ (Self::Update(_) | Self::Insert(_) | Self::Rewrite(_)) => {
+                Ok(ErrorResponse::from_err(&err))
+            }
+            err => Err(err),
+        }
+    }
+}
+
 #[derive(Debug, Error)]
 pub(crate) enum UpdateError {
     #[error("sharding key updates are forbidden")]
     Disabled,
 
+    /// Parser flagged a sharding key update but the planner can't continue.
+    /// If we let it continue, this could cause unintended side effects.
+    #[error("sharding key update plan doesn't match the parsed statement")]
+    PlanMismatch,
+
     #[error("sharding key update must be executed inside a transaction")]
     TransactionRequired,
 
@@ -42,6 +70,26 @@ pub(crate) enum UpdateError {
 
     #[error("sharding key update would move a row referenced by an ON DELETE foreign key")]
     ForeignKeyOnDelete,
+
+    #[error("sharding key update expected an UPDATE statement")]
+    NotAnUpdate,
+
+    #[error("sharding key update step \"{0}\" response is missing or incomplete")]
+    MissingStepResponse(&'static str),
+}
+
+#[derive(Debug, Error)]
+pub(crate) enum InsertError {
+    #[error("multi-tuple insert requires multi-shard binding")]
+    MultiShardRequired,
+
+    /// Parser flagged a multi insert but the planner can't continue.
+    /// If we let it continue, this could cause unintended side effects.
+    #[error("multi-tuple insert plan doesn't match the parsed statement")]
+    PlanMismatch,
+
+    #[error("cache: {0}")]
+    Cache(String),
 }
 
 impl From<crate::frontend::Error> for Error {

pgdog/src/frontend/client/query_engine/multi_step/mod.rs

@@ -1,14 +1,10 @@
-pub(crate) mod error;
+pub mod error;
 pub(crate) mod forward_check;
-pub(crate) mod insert;
-pub(crate) mod state;
-pub(crate) mod update;
+pub(crate) mod shared;
 
-pub(crate) use error::{Error, UpdateError};
-pub(crate) use forward_check::*;
-pub(crate) use insert::InsertMulti;
-pub(crate) use state::{CommandType, MultiServerState};
-pub(crate) use update::UpdateMulti;
+pub(crate) mod types;
+
+pub(crate) mod ops;
 
 #[cfg(test)]
 mod test;