rewrite plan

Author: levkk

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

@@ -134,11 +134,7 @@ impl QueryEngine {
  context.prepared_statements,
  &mut self.rewrite_state,
  );
- match rewrite.execute() {
- Ok(ast) => context.ast = Some(ast),
- Err(rewrite::Error::EmptyQuery) => (),
- Err(err) => return Err(err.into()),
- }
+ context.ast = rewrite.execute()?;
  }
  }
 

pgdog/src/frontend/client_request.rs

@@ -140,6 +140,17 @@ impl ClientRequest {
  Ok(None)
  }
 
+ /// Get mutable reference to parameters, if any.
+ pub fn parameters_mut(&mut self) -> Result<Option<&mut Bind>, Error> {
+ for message in self.messages.iter_mut() {
+ if let ProtocolMessage::Bind(bind) = message {
+ return Ok(Some(bind));
+ }
+ }
+
+ Ok(None)
+ }
+
  /// Get all CopyData messages.
  pub fn copy_data(&self) -> Result<Vec<CopyData>, Error> {
  let mut rows = vec![];

pgdog/src/frontend/router/rewrite/context.rs

@@ -2,7 +2,9 @@
 
 use pg_query::protobuf::{ParseResult, RawStmt};
 
-use super::{output::RewriteActionKind, stats::RewriteStats, Error, RewriteAction, StepOutput};
+use super::{
+ output::RewriteActionKind, stats::RewriteStats, Error, RewriteAction, RewritePlan, StepOutput,
+};
 use crate::net::{Bind, Parse, ProtocolMessage, Query};
 
 #[derive(Debug, Clone)]
@@ -12,30 +14,23 @@ pub struct Context<'a> {
  original: &'a ParseResult,
  // If an in-place rewrite was done, the statement is saved here.
  rewrite: Option<ParseResult>,
- /// Original bind message, if any.
- bind: Option<&'a Bind>,
- /// Bind rewritten.
- rewrite_bind: Option<Bind>,
  /// Additional messages to add to the request.
  result: Vec<RewriteAction>,
  /// Extended protocol.
  parse: Option<&'a Parse>,
+ /// Rewrite plan.
+ plan: RewritePlan,
 }
 
 impl<'a> Context<'a> {
  /// Create new input.
- pub(super) fn new(
- original: &'a ParseResult,
- bind: Option<&'a Bind>,
- parse: Option<&'a Parse>,
- ) -> Self {
+ pub(super) fn new(original: &'a ParseResult, parse: Option<&'a Parse>) -> Self {
  Self {
  original,
- bind,
  rewrite: None,
- rewrite_bind: None,
  result: vec![],
  parse,
+ plan: RewritePlan::default(),
  }
  }
 
@@ -46,32 +41,12 @@ impl<'a> Context<'a> {
 
  /// We are rewriting an extended protocol request.
  pub fn extended(&self) -> bool {
- self.parse().is_some() || self.bind().is_some()
+ self.parse().is_some()
  }
 
- /// Get the Bind message, if set.
- pub fn bind(&'a self) -> Option<&'a Bind> {
- if let Some(ref rewrite_bind) = self.rewrite_bind {
- Some(rewrite_bind)
- } else {
- self.bind
- }
- }
-
- /// Take the Bind message for modification.
- /// Don't forget to return it.
- #[must_use]
- pub fn bind_take(&mut self) -> Option<Bind> {
- if self.rewrite_bind.is_none() {
- self.rewrite_bind = self.bind.cloned();
- }
-
- self.rewrite_bind.take()
- }
-
- /// Put the bind message back.
- pub fn bind_put(&mut self, bind: Option<Bind>) {
- self.rewrite_bind = bind;
+ /// Get reference to rewrite plan for modification.
+ pub fn plan(&mut self) -> &mut RewritePlan {
+ &mut self.plan
  }
 
  /// Get the original (or modified) statement.
@@ -121,31 +96,19 @@ impl<'a> Context<'a> {
  Ok(StepOutput::NoOp)
  } else {
  let mut stats = RewriteStats::default();
- let bind = self.rewrite_bind.take();
  let ast = self.rewrite.take().ok_or(Error::NoRewrite)?;
  let stmt = ast.deparse()?;
- let extended = self.extended();
  let mut parse = self.parse().cloned();
 
  let mut actions = self.result;
 
- if extended {
- if let Some(mut parse) = parse.take() {
- parse.set_query(&stmt);
- actions.push(RewriteAction {
- message: parse.into(),
- action: RewriteActionKind::Replace,
- });
- stats.parse += 1;
- }
-
- if let Some(bind) = bind {
- actions.push(RewriteAction {
- message: bind.into(),
- action: RewriteActionKind::Replace,
- });
- stats.bind += 1;
- }
+ if let Some(mut parse) = parse.take() {
+ parse.set_query(&stmt);
+ actions.push(RewriteAction {
+ message: parse.into(),
+ action: RewriteActionKind::Replace,
+ });
+ stats.parse += 1;
  } else {
  actions.push(RewriteAction {
  message: Query::new(stmt.clone()).into(),
@@ -159,6 +122,7 @@ impl<'a> Context<'a> {
  ast,
  actions,
  stats,
+ plan: self.plan.freeze(),
  })
  }
  }

pgdog/src/frontend/router/rewrite/error.rs

@@ -31,4 +31,7 @@ pub enum Error {
 
  #[error("parser: {0}")]
  Parser(#[from] crate::frontend::router::parser::Error),
+
+ #[error("no active rewrite plan set")]
+ NoActiveRewritePlan,
 }

pgdog/src/frontend/router/rewrite/insert_split/mod.rs

@@ -3,8 +3,6 @@ use pg_query::{
  Node, NodeEnum,
 };
 
-use crate::net::Bind;
-
 use super::*;
 
 #[derive(Default)]
@@ -39,20 +37,19 @@ impl RewriteModule for InsertSplitRewrite {
  let mut new_insert = proto_insert.clone();
  let mut new_select = proto_select.clone();
  let mut new_values = values.clone();
- let mut new_bind = Bind::default();
 
  // Rewrite the parameter references
  // and create new Bind message for each INSERT statement.
  if let Some(NodeEnum::List(list)) = new_values.node.as_mut() {
  for value in list.items.iter_mut() {
- if let Some(NodeEnum::ParamRef(param)) = value.node.as_mut() {
- let parameter = input
- .bind()
- .and_then(|bind| bind.parameter(param.number as usize - 1).ok())
- .flatten();
- if let Some(parameter) = parameter {
- param.number = new_bind.add_existing(parameter)?;
- }
+ if let Some(NodeEnum::ParamRef(_)) = value.node.as_mut() {
+ // let parameter = input
+ //  .bind()
+ //  .and_then(|bind| bind.parameter(param.number as usize - 1).ok())
+ //  .flatten();
+ // if let Some(parameter) = parameter {
+ //  param.number = new_bind.add_existing(parameter)?;
+ // }
  }
  }
  }
@@ -69,7 +66,7 @@ impl RewriteModule for InsertSplitRewrite {
  ..Default::default()
  }],
  };
- inserts.push((result, new_bind));
+ inserts.push(result);
  }
  }
  }

pgdog/src/frontend/router/rewrite/mod.rs

@@ -11,6 +11,7 @@ pub mod error;
 pub mod insert_split;
 pub mod interface;
 pub mod output;
+pub mod plan;
 pub mod prepared;
 pub mod request;
 pub mod state;
@@ -21,6 +22,7 @@ pub use context::Context;
 pub use error::Error;
 pub use interface::RewriteModule;
 pub use output::{RewriteAction, StepOutput};
+pub use plan::{ImmutableRewritePlan, RewritePlan, UniqueIdPlan};
 pub use request::RewriteRequest;
 pub use state::RewriteState;
 

pgdog/src/frontend/router/rewrite/output.rs

@@ -1,6 +1,7 @@
 use pg_query::protobuf::ParseResult;
 
 use super::stats::RewriteStats;
+use super::ImmutableRewritePlan;
 use crate::{frontend::ClientRequest, net::ProtocolMessage};
 
 use std::mem::discriminant;
@@ -53,6 +54,7 @@ pub enum StepOutput {
  ast: ParseResult,
  stmt: String,
  stats: RewriteStats,
+ plan: ImmutableRewritePlan,
  },
 }
 

pgdog/src/frontend/router/rewrite/plan.rs

@@ -0,0 +1,60 @@
+use std::{ops::Deref, sync::Arc};
+
+use super::Error;
+use crate::{
+ net::{Bind, Datum},
+ unique_id::UniqueId,
+};
+
+#[derive(Debug, Clone, Default)]
+pub struct UniqueIdPlan {
+ /// Parameter number.
+ pub(super) param_ref: i32,
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct RewritePlan {
+ /// How many unique IDs to add to the Bind message.
+ pub(super) unique_ids: Vec<UniqueIdPlan>,
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct ImmutableRewritePlan {
+ /// Compiled rewrite plan, that cannot be modified further.
+ pub(super) plan: Arc<RewritePlan>,
+}
+
+impl Deref for ImmutableRewritePlan {
+ type Target = RewritePlan;
+
+ fn deref(&self) -> &Self::Target {
+ &self.plan
+ }
+}
+
+impl RewritePlan {
+ /// Apply rewrite plan to Bind message.
+ ///
+ /// N.B. this isn't idempotent, run this only once.
+ ///
+ pub fn apply_bind(&self, bind: &mut Bind) -> Result<(), Error> {
+ for unique_id in &self.unique_ids {
+ let id = UniqueId::generator()?.next_id();
+ let counter = bind.add_parameter(Datum::Bigint(id))?;
+ // Params should be added to plan in order.
+ // This validates it.
+ if counter != unique_id.param_ref {
+ return Err(Error::ParameterCountMismatch);
+ }
+ }
+
+ Ok(())
+ }
+
+ /// Freeze rewrite plan, without any more modifications allowed.
+ pub fn freeze(self) -> ImmutableRewritePlan {
+ ImmutableRewritePlan {
+ plan: Arc::new(self),
+ }
+ }
+}