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` 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/config/tandem.rs

@@ -426,36 +426,40 @@ mod tests {
 
     #[test]
     fn prometheus_config() {
-        let config = TandemConfig::build("tests/config/cipherstash-proxy-test.toml").unwrap();
-        assert!(!config.prometheus_enabled());
-
-        temp_env::with_vars([("CS_PROMETHEUS__ENABLED", Some("true"))], || {
-            let config = TandemConfig::build("tests/config/cipherstash-proxy-test.toml").unwrap();
-            assert!(config.prometheus_enabled());
-            assert!(config.prometheus.enabled);
-            assert_eq!(config.prometheus.port, 9930);
-        });
-
-        temp_env::with_vars([("CS_PROMETHEUS__PORT", Some("7777"))], || {
+        temp_env::with_vars_unset(["CS_PROMETHEUS__ENABLED"], || {
             let config = TandemConfig::build("tests/config/cipherstash-proxy-test.toml").unwrap();
             assert!(!config.prometheus_enabled());
-            assert!(!config.prometheus.enabled);
-            assert_eq!(config.prometheus.port, 7777);
-        });
 
-        temp_env::with_vars(
-            [
-                ("CS_PROMETHEUS__ENABLED", Some("true")),
-                ("CS_PROMETHEUS__PORT", Some("7777")),
-            ],
-            || {
+            temp_env::with_vars([("CS_PROMETHEUS__ENABLED", Some("true"))], || {
                 let config =
                     TandemConfig::build("tests/config/cipherstash-proxy-test.toml").unwrap();
                 assert!(config.prometheus_enabled());
                 assert!(config.prometheus.enabled);
+                assert_eq!(config.prometheus.port, 9930);
+            });
+
+            temp_env::with_vars([("CS_PROMETHEUS__PORT", Some("7777"))], || {
+                let config =
+                    TandemConfig::build("tests/config/cipherstash-proxy-test.toml").unwrap();
+                assert!(!config.prometheus_enabled());
+                assert!(!config.prometheus.enabled);
                 assert_eq!(config.prometheus.port, 7777);
-            },
-        );
+            });
+
+            temp_env::with_vars(
+                [
+                    ("CS_PROMETHEUS__ENABLED", Some("true")),
+                    ("CS_PROMETHEUS__PORT", Some("7777")),
+                ],
+                || {
+                    let config =
+                        TandemConfig::build("tests/config/cipherstash-proxy-test.toml").unwrap();
+                    assert!(config.prometheus_enabled());
+                    assert!(config.prometheus.enabled);
+                    assert_eq!(config.prometheus.port, 7777);
+                },
+            );
+        });
     }
 
     #[test]

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

@@ -34,6 +34,7 @@ use sqltk::AsNodeKey;
 use sqltk_parser::ast::{self, Value};
 use sqltk_parser::dialect::PostgreSqlDialect;
 use sqltk_parser::parser::Parser;
+use std::borrow::Cow;
 use std::collections::HashMap;
 use std::time::Instant;
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
@@ -255,7 +256,7 @@ where
 
         // Simple Query may contain many statements
         let parsed_statements = self.parse_statements(&query.statement)?;
-        let mut transformed_statements = vec![];
+        let mut transformed_statements: Vec<Cow<'_, ast::Statement>> = vec![];
 
         debug!(target: MAPPER,
             client_id = self.context.client_id,
@@ -265,15 +266,15 @@ where
         let mut portal = Portal::passthrough();
         let mut encrypted = false;
 
-        for statement in parsed_statements {
-            self.check_for_schema_change(&statement);
+        for statement in parsed_statements.iter() {
+            self.check_for_schema_change(statement);
 
-            if !eql_mapper::requires_type_check(&statement) {
+            if !eql_mapper::requires_type_check(statement) {
                 counter!(STATEMENTS_PASSTHROUGH_TOTAL).increment(1);
                 continue;
             }
 
-            let typed_statement = match self.type_check(&statement) {
+            let typed_statement = match self.type_check(statement) {
                 Ok(ts) => ts,
                 Err(err) => {
                     if self.encrypt.config.mapping_errors_enabled() {
@@ -284,25 +285,25 @@ 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,
-                        );
+            let statement = self.to_encryptable_statement(&typed_statement, vec![])?;
+            let encrypted_literals = self
+                .encrypt_literals(&typed_statement, &statement.literal_columns)
+                .await?;
 
-                        transformed_statements.push(transformed_statement);
-                        encrypted = true;
-                    }
+            let 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,
+                transform_was_noop = matches!(transformed_statement, Cow::Borrowed(_)),
+            );
+
+            match &transformed_statement {
+                Cow::Owned(_) => {
+                    encrypted = true;
                     debug!(target: MAPPER,
                         client_id = self.context.client_id,
                         msg = "Encrypted Statement"
@@ -312,15 +313,16 @@ where
                     // Set Encrypted portal
                     portal = Portal::encrypted_text();
                 }
-                None => {
+                Cow::Borrowed(_) => {
                     debug!(target: MAPPER,
                         client_id = self.context.client_id,
                         msg = "Passthrough Statement"
                     );
                     counter!(STATEMENTS_PASSTHROUGH_TOTAL).increment(1);
-                    transformed_statements.push(statement);
                 }
-            };
+            }
+
+            transformed_statements.push(transformed_statement);
         }
 
         self.context.add_portal(Name::unnamed(), portal);
@@ -400,11 +402,11 @@ where
     ///  - rewrites any encrypted literal values
     ///  - wraps any nodes in appropriate EQL function
     ///
-    async fn transform_statement(
+    async fn transform_statement<'ast>(
         &mut self,
-        typed_statement: &TypedStatement<'_>,
+        typed_statement: &TypedStatement<'ast>,
         encrypted_literals: &Vec<Option<Encrypted>>,
-    ) -> Result<Option<ast::Statement>, Error> {
+    ) -> Result<Cow<'ast, ast::Statement>, Error> {
         // Convert literals to ast Expr
         let mut encrypted_expressions = vec![];
         for encrypted in encrypted_literals {
@@ -429,11 +431,9 @@ where
             literals = encrypted_nodes.len(),
         );
 
-        let transformed_statement = typed_statement
+        Ok(typed_statement
             .transform(encrypted_nodes)
-            .map_err(|e| MappingError::StatementCouldNotBeTransformed(e.to_string()))?;
-
-        Ok(Some(transformed_statement))
+            .map_err(|e| MappingError::StatementCouldNotBeTransformed(e.to_string()))?)
     }
 
     ///
@@ -493,39 +493,37 @@ where
         // These override the underlying column type
         let param_types = message.param_types.clone();
 
-        match self.to_encryptable_statement(&typed_statement, param_types)? {
-            Some(statement) => {
-                let encrypted_literals = self
-                    .encrypt_literals(&typed_statement, &statement.literal_columns)
-                    .await?;
+        let statement = self.to_encryptable_statement(&typed_statement, param_types)?;
+        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,
-                    );
+        let transformed_statement = self
+            .transform_statement(&typed_statement, &encrypted_literals)
+            .await?;
 
-                    message.rewrite_statement(transformed_statement.to_string());
-                }
+        debug!(target: MAPPER,
+            client_id = self.context.client_id,
+            transformed_statement = ?transformed_statement,
+            transformed_statement_text = %transformed_statement,
+            transform_was_noop = matches!(transformed_statement, Cow::Borrowed(_)),
+        );
 
-                counter!(STATEMENTS_ENCRYPTED_TOTAL).increment(1);
+        if let Cow::Owned(transformed_statement) = transformed_statement {
+            message.rewrite_statement(transformed_statement.to_string());
+            counter!(STATEMENTS_ENCRYPTED_TOTAL).increment(1);
 
-                message.rewrite_param_types(&statement.param_columns);
-                self.context
-                    .add_statement(message.name.to_owned(), statement);
-            }
-            _ => {
-                debug!(target: MAPPER,
-                    client_id = self.context.client_id,
-                    msg = "Passthrough Parse"
-                );
-                counter!(STATEMENTS_PASSTHROUGH_TOTAL).increment(1);
-            }
+            message.rewrite_param_types(&statement.param_columns);
+            self.context
+                .add_statement(message.name.to_owned(), statement);
+        } else {
+            debug!(target: MAPPER,
+                client_id = self.context.client_id,
+                msg = "Passthrough Parse"
+            );
+            counter!(STATEMENTS_PASSTHROUGH_TOTAL).increment(1);
         }
+
         let bytes = BytesMut::try_from(message)?;
 
         debug!(target: MAPPER,
@@ -598,7 +596,7 @@ where
         &self,
         typed_statement: &TypedStatement<'_>,
         param_types: Vec<i32>,
-    ) -> Result<Option<Statement>, Error> {
+    ) -> Result<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)?;
@@ -618,7 +616,7 @@ where
             param_types,
         );
 
-        Ok(Some(statement))
+        Ok(statement)
     }
 
     ///

packages/eql-mapper/src/encrypted_statement.rs

@@ -3,22 +3,25 @@ use std::{collections::HashMap, sync::Arc};
 use sqltk::{NodeKey, NodePath, Transform, Visitable};
 
 use crate::{
-    EqlMapperError, FailOnPlaceholderChange, GroupByEqlCol, PreserveEffectiveAliases,
-    ReplacePlaintextEqlLiterals, TransformationRule, Type, UseEquivalentSqlFuncForEqlTypes,
-    WrapEqlColsInOrderByWithOreFn, WrapGroupedEqlColInAggregateFn,
+    DryRunnable, EqlMapperError, FailOnPlaceholderChange, GroupByEqlCol, Mode,
+    PreserveEffectiveAliases, ReplacePlaintextEqlLiterals, TransformationRule, Type,
+    UseEquivalentSqlFuncForEqlTypes, WrapEqlColsInOrderByWithOreFn, WrapGroupedEqlColInAggregateFn,
 };
 
 #[derive(Debug)]
 pub(crate) struct EncryptedStatement<'ast> {
-    transformation_rules: (
-        WrapGroupedEqlColInAggregateFn<'ast>,
-        GroupByEqlCol<'ast>,
-        WrapEqlColsInOrderByWithOreFn<'ast>,
-        PreserveEffectiveAliases,
-        ReplacePlaintextEqlLiterals<'ast>,
-        UseEquivalentSqlFuncForEqlTypes<'ast>,
-        FailOnPlaceholderChange,
-    ),
+    transformation_rules: DryRunnable<
+        'ast,
+        (
+            WrapGroupedEqlColInAggregateFn<'ast>,
+            GroupByEqlCol<'ast>,
+            WrapEqlColsInOrderByWithOreFn<'ast>,
+            PreserveEffectiveAliases,
+            ReplacePlaintextEqlLiterals<'ast>,
+            UseEquivalentSqlFuncForEqlTypes<'ast>,
+            FailOnPlaceholderChange<'ast>,
+        ),
+    >,
 }
 
 impl<'ast> EncryptedStatement<'ast> {
@@ -27,17 +30,25 @@ impl<'ast> EncryptedStatement<'ast> {
         node_types: Arc<HashMap<NodeKey<'ast>, Type>>,
     ) -> Self {
         Self {
-            transformation_rules: (
+            transformation_rules: DryRunnable::new((
                 WrapGroupedEqlColInAggregateFn::new(Arc::clone(&node_types)),
                 GroupByEqlCol::new(Arc::clone(&node_types)),
                 WrapEqlColsInOrderByWithOreFn::new(Arc::clone(&node_types)),
                 PreserveEffectiveAliases,
                 ReplacePlaintextEqlLiterals::new(encrypted_literals),
                 UseEquivalentSqlFuncForEqlTypes::new(Arc::clone(&node_types)),
-                FailOnPlaceholderChange,
-            ),
+                FailOnPlaceholderChange::new(),
+            )),
         }
     }
+
+    pub(crate) fn set_mode(&mut self, new_mode: Mode) {
+        self.transformation_rules.set_mode(new_mode);
+    }
+
+    pub(crate) fn did_edit(&self) -> bool {
+        self.transformation_rules.did_edit()
+    }
 }
 
 /// Applies all of the transormation rules from the `EncryptedStatement`.