feat: TransformationRule dry-run support

The Proxy must be able to reliably detect whether a statement requires
transformation because if a statement does not require transformation
then the potentially expensive AST rebuilding step can be skipped.

The result of type-checking is insufficient in general to tell whether a
statement requires transformation unless the `TransformationRule` logic
is duplicated in the Proxy - which we don't want of course.

This commit extends the `TransformationRule` trait with a
`would_edit` method which answers the question "would this rule change
the AST if it was applied?".

Additionally, a new `TranformationRule` impl `DryRunnable` wraps another
rule in such a way that it can "pretend" to be performing a
`Transform` (as far as `sqltk` is concerned) when really its doing a
dry-run after which it will tell us if it *would* change the AST.

This is all wrapped up by `TypedStatement::transform`Vjjj which now returns
`Result<Cow<Statement>, _>`.

When the statement does not need to be modified it returns
`Cow::Borrowed(unmodified_statement)` and when the statement *has* been
modified it returns a `Cow::Owned(modified_statement)`

The Proxy `frontend` can now distinguish between the two cases via the
`Cow`, and manage its passthrough counter correctly without having to be
coupled to implementation details of the EQL Mapper.

Author: freshtonic

mise.local.example.toml

@@ -42,6 +42,6 @@ CS_LOG__AUTHENTICATION_LEVEL = "info"
 CS_LOG__CONTEXT_LEVEL = "info"
 CS_LOG__KEYSET_LEVEL = "info"
 CS_LOG__PROTOCOL_LEVEL = "info"
-CS_LOG__MAPPER_LEVEL = "info"
+CS_LOG__MAPPER_LEVEL = "debug"
 CS_LOG__SCHEMA_LEVEL = "info"
 CS_LOG__CONFIG_LEVEL = "info"

packages/cipherstash-proxy/src/error.rs

@@ -1,6 +1,7 @@
 use crate::{postgresql::Column, Identifier};
 use bytes::BytesMut;
 use cipherstash_client::encryption;
+use eql_mapper::EqlMapperError;
 use metrics_exporter_prometheus::BuildError;
 use std::{io, time::Duration};
 use thiserror::Error;
@@ -92,6 +93,9 @@ pub enum MappingError {
 
     #[error("Statement encountered an internal error. This may be a bug in the statement mapping module of CipherStash Proxy. Please visit {}#mapping-internal-error for more information.", ERROR_DOC_BASE_URL)]
     Internal(String),
+
+    #[error(transparent)]
+    EqlMapper(#[from] EqlMapperError),
 }
 
 #[derive(Error, Debug)]

packages/cipherstash-proxy/src/postgresql/frontend.rs

@@ -26,11 +26,11 @@ use crate::prometheus::{
 use crate::Encrypted;
 use bytes::BytesMut;
 use cipherstash_client::encryption::Plaintext;
-use eql_mapper::{self, EqlMapperError, EqlValue, TableColumn, TypedStatement};
+use eql_mapper::{self, EqlMapperError, EqlValue, TableColumn, TypeCheckedStatement};
 use metrics::{counter, histogram};
 use pg_escape::quote_literal;
 use serde::Serialize;
-use sqltk::AsNodeKey;
+use sqltk::NodeKey;
 use sqltk_parser::ast::{self, Value};
 use sqltk_parser::dialect::PostgreSqlDialect;
 use sqltk_parser::parser::Parser;
@@ -286,22 +286,23 @@ where
 
             match self.to_encryptable_statement(&typed_statement, vec![])? {
                 Some(statement) => {
-                    let encrypted_literals = self
-                        .encrypt_literals(&typed_statement, &statement.literal_columns)
-                        .await?;
-
-                    if let Some(transformed_statement) = self
-                        .transform_statement(&typed_statement, &encrypted_literals)
-                        .await?
-                    {
-                        debug!(target: MAPPER,
-                            client_id = self.context.client_id,
-                            transformed_statement = ?transformed_statement,
-                            transformed_statement_text = %transformed_statement,
-                        );
+                    if typed_statement.requires_transform() {
+                        let encrypted_literals = self
+                            .encrypt_literals(&typed_statement, &statement.literal_columns)
+                            .await?;
+
+                        if let Some(transformed_statement) = self
+                            .transform_statement(&typed_statement, &encrypted_literals)
+                            .await?
+                        {
+                            debug!(target: MAPPER,
+                                client_id = self.context.client_id,
+                                transformed_statement = ?transformed_statement,
+                            );
 
-                        transformed_statements.push(transformed_statement);
-                        encrypted = true;
+                            transformed_statements.push(transformed_statement);
+                            encrypted = true;
+                        }
                     }
                     debug!(target: MAPPER,
                         client_id = self.context.client_id,
@@ -356,7 +357,7 @@ where
     ///
     async fn encrypt_literals(
         &mut self,
-        typed_statement: &TypedStatement<'_>,
+        typed_statement: &TypeCheckedStatement<'_>,
         literal_columns: &Vec<Option<Column>>,
     ) -> Result<Vec<Option<Encrypted>>, Error> {
         let literal_values = typed_statement.literal_values();
@@ -402,14 +403,14 @@ where
     ///
     async fn transform_statement(
         &mut self,
-        typed_statement: &TypedStatement<'_>,
+        typed_statement: &TypeCheckedStatement<'_>,
         encrypted_literals: &Vec<Option<Encrypted>>,
     ) -> Result<Option<ast::Statement>, Error> {
         // Convert literals to ast Expr
         let mut encrypted_expressions = vec![];
         for encrypted in encrypted_literals {
             let e = match encrypted {
-                Some(en) => Some(to_json_literal_expr(&en)?),
+                Some(en) => Some(to_json_literal_value(&en)?),
                 None => None,
             };
             encrypted_expressions.push(e);
@@ -421,14 +422,18 @@ where
             .literals
             .iter()
             .zip(encrypted_expressions.into_iter())
-            .filter_map(|((_, original_node), en)| en.map(|en| (original_node.as_node_key(), en)))
+            .filter_map(|((_, original_node), en)| en.map(|en| (NodeKey::new(*original_node), en)))
             .collect::<HashMap<_, _>>();
 
         debug!(target: MAPPER,
             client_id = self.context.client_id,
             literals = encrypted_nodes.len(),
         );
 
+        if !typed_statement.requires_transform() {
+            return Ok(None);
+        }
+
         let transformed_statement = typed_statement
             .transform(encrypted_nodes)
             .map_err(|e| MappingError::StatementCouldNotBeTransformed(e.to_string()))?;
@@ -495,21 +500,22 @@ where
 
         match self.to_encryptable_statement(&typed_statement, param_types)? {
             Some(statement) => {
-                let encrypted_literals = self
-                    .encrypt_literals(&typed_statement, &statement.literal_columns)
-                    .await?;
-
-                if let Some(transformed_statement) = self
-                    .transform_statement(&typed_statement, &encrypted_literals)
-                    .await?
-                {
-                    debug!(target: MAPPER,
-                        client_id = self.context.client_id,
-                        transformed_statement = ?transformed_statement,
-                        transformed_statement_text = %transformed_statement,
-                    );
+                if typed_statement.requires_transform() {
+                    let encrypted_literals = self
+                        .encrypt_literals(&typed_statement, &statement.literal_columns)
+                        .await?;
+
+                    if let Some(transformed_statement) = self
+                        .transform_statement(&typed_statement, &encrypted_literals)
+                        .await?
+                    {
+                        debug!(target: MAPPER,
+                            client_id = self.context.client_id,
+                            transformed_statement = ?transformed_statement,
+                        );
 
-                    message.rewrite_statement(transformed_statement.to_string());
+                        message.rewrite_statement(transformed_statement.to_string());
+                    }
                 }
 
                 counter!(STATEMENTS_ENCRYPTED_TOTAL).increment(1);
@@ -526,14 +532,19 @@ where
                 counter!(STATEMENTS_PASSTHROUGH_TOTAL).increment(1);
             }
         }
-        let bytes = BytesMut::try_from(message)?;
 
-        debug!(target: MAPPER,
+        if message.requires_rewrite() {
+            let bytes = BytesMut::try_from(message)?;
+
+            debug!(target: MAPPER,
                 client_id = self.context.client_id,
                 msg = "Rewrite Parse",
                 bytes = ?bytes);
 
-        Ok(Some(bytes))
+            Ok(Some(bytes))
+        } else {
+            Ok(None)
+        }
     }
 
     ///
@@ -596,13 +607,23 @@ where
     ///
     fn to_encryptable_statement(
         &self,
-        typed_statement: &TypedStatement<'_>,
+        typed_statement: &TypeCheckedStatement<'_>,
         param_types: Vec<i32>,
     ) -> Result<Option<Statement>, Error> {
         let param_columns = self.get_param_columns(typed_statement)?;
         let projection_columns = self.get_projection_columns(typed_statement)?;
         let literal_columns = self.get_literal_columns(typed_statement)?;
 
+        let no_encrypted_param_columns = param_columns.iter().all(|c| c.is_none());
+        let no_encrypted_projection_columns = projection_columns.iter().all(|c| c.is_none());
+
+        if (param_columns.is_empty() || no_encrypted_param_columns)
+            && (projection_columns.is_empty() || no_encrypted_projection_columns)
+            && !typed_statement.requires_transform()
+        {
+            return Ok(None);
+        }
+
         debug!(target: MAPPER,
             client_id = self.context.client_id,
             msg = "Encryptable Statement",
@@ -713,7 +734,10 @@ where
         Ok(encrypted)
     }
 
-    fn type_check<'a>(&self, statement: &'a ast::Statement) -> Result<TypedStatement<'a>, Error> {
+    fn type_check<'a>(
+        &self,
+        statement: &'a ast::Statement,
+    ) -> Result<TypeCheckedStatement<'a>, Error> {
         match eql_mapper::type_check(self.context.get_table_resolver(), statement) {
             Ok(typed_statement) => {
                 debug!(target: MAPPER,
@@ -756,7 +780,7 @@ where
     ///
     fn get_projection_columns(
         &self,
-        typed_statement: &eql_mapper::TypedStatement<'_>,
+        typed_statement: &eql_mapper::TypeCheckedStatement<'_>,
     ) -> Result<Vec<Option<Column>>, Error> {
         let mut projection_columns = vec![];
         if let eql_mapper::Projection::WithColumns(columns) = &typed_statement.projection {
@@ -791,7 +815,7 @@ where
     ///
     fn get_param_columns(
         &self,
-        typed_statement: &eql_mapper::TypedStatement<'_>,
+        typed_statement: &eql_mapper::TypeCheckedStatement<'_>,
     ) -> Result<Vec<Option<Column>>, Error> {
         let mut param_columns = vec![];
 
@@ -819,7 +843,7 @@ where
 
     fn get_literal_columns(
         &self,
-        typed_statement: &eql_mapper::TypedStatement<'_>,
+        typed_statement: &eql_mapper::TypeCheckedStatement<'_>,
     ) -> Result<Vec<Option<Column>>, Error> {
         let mut literal_columns = vec![];
 
@@ -945,7 +969,7 @@ fn literals_to_plaintext(
     Ok(plaintexts)
 }
 
-fn to_json_literal_expr<T>(literal: &T) -> Result<Value, Error>
+fn to_json_literal_value<T>(literal: &T) -> Result<Value, Error>
 where
     T: ?Sized + Serialize,
 {

packages/eql-mapper/src/encrypted_statement.rs

@@ -2,8 +2,8 @@ use super::importer::{ImportError, Importer};
 use crate::{
     inference::{TypeError, TypeInferencer},
     unifier::{EqlValue, Unifier},
-    DepMut, Fmt, Param, ParamError, ScopeError, ScopeTracker, TableResolver, Type, TypeRegistry,
-    TypedStatement, Value,
+    DepMut, Fmt, Param, ParamError, ScopeError, ScopeTracker, TableResolver, Type,
+    TypeCheckedStatement, TypeRegistry, Value,
 };
 use sqltk::{Break, NodeKey, Visitable, Visitor};
 use sqltk_parser::ast::{self as ast, Statement};
@@ -22,19 +22,19 @@ use tracing::{event, span, Level};
 /// - all operators and functions used with literals destined to be transformed to EQL types are semantically valid for
 ///   that EQL type
 ///
-/// A successful type check will return a [`TypedStatement`] which can be interrogated to discover the required params
+/// A successful type check will return a [`TypeCheckedStatement`] which can be interrogated to discover the required params
 /// and their types, the types and plaintext values of all literals, and an optional projection type (the optionality
 /// depending on the specific statement).
 ///
-/// Invoking [`TypedStatement::transform`] will return an updated [`Statement`] where all plaintext literals have been
+/// Invoking [`TypeCheckedStatement::transform`] will return an updated [`Statement`] where all plaintext literals have been
 /// replaced with their EQL (encrypted) equivalent and specific SQL operators and functions will have been rewritten to
 /// invoke the EQL equivalents.
 ///
 /// An [`EqlMapperError`] is returned if type checking fails.
 pub fn type_check<'ast>(
     resolver: Arc<TableResolver>,
     statement: &'ast Statement,
-) -> Result<TypedStatement<'ast>, EqlMapperError> {
+) -> Result<TypeCheckedStatement<'ast>, EqlMapperError> {
     let mut mapper = EqlMapper::<'ast>::new_with_resolver(resolver);
     match statement.accept(&mut mapper) {
         ControlFlow::Continue(()) => mapper.resolve(statement),
@@ -137,7 +137,7 @@ impl<'ast> EqlMapper<'ast> {
     pub fn resolve(
         self,
         statement: &'ast Statement,
-    ) -> Result<TypedStatement<'ast>, EqlMapperError> {
+    ) -> Result<TypeCheckedStatement<'ast>, EqlMapperError> {
         let span_begin = span!(
             target: "eqlmapper::spans",
             Level::TRACE,
@@ -169,7 +169,7 @@ impl<'ast> EqlMapper<'ast> {
                     node_types = %Fmt(&node_types)
                 );
 
-                Ok(TypedStatement {
+                Ok(TypeCheckedStatement {
                     statement,
                     projection,
                     params,