feat: SQL transformations related to GROUP BY

Includes a version bump to the latest `sqltk` version which vendors
`sqlparser` as `sqltk-parser` and includes an extra parameter in
`Transform::transform` to allow contextual transformations (instead of
solely by type).

The following refactorings were done as part of this PR:

1. Break sole `Transform` impl into modular rules using the
   `TransformationRule` trait.

2. Removed `EqlFunctionTracker` - this existed to support `Expr`
   transformations in `ORDER BY` clauses. With contextual transformation
   it was no longer required.

3. Removed `TypeCell` - it did not carry its weight and made debugging
   really hard. The `TypeRegistry` now associates nodes to Arc<Type>
   using an intermediary `TypeVar`.

An extra method was added to the post type checking step in the
`EqlMapper` to resolve all `Value` nodes which have no concrete type
after the inference step to `NativeValue`.

Author: freshtonic

Cargo.lock

@@ -35,7 +35,8 @@ strip = "none"
 debug = true
 
 [workspace.dependencies]
-sqlparser = { version = "^0.52", features = ["bigdecimal", "serde"] }
+sqltk = { version = "0.5.0" }
+sqltk-parser = { version = "0.52.0" }
 thiserror = "2.0.9"
 tokio = { version = "1.42.0", features = ["full"] }
 tracing = "0.1"

Cargo.toml

@@ -40,7 +40,8 @@ rustls-pki-types = "1.10.0"
 serde = "1.0"
 serde_json = "1.0"
 socket2 = "0.5.7"
-sqlparser = { workspace = true }
+sqltk = { workspace = true }
+sqltk-parser = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }
 tokio-postgres = { version = "0.7", features = [

packages/cipherstash-proxy/Cargo.toml

@@ -4,7 +4,7 @@ use crate::error::Error;
 use crate::{connect, log::SCHEMA};
 use arc_swap::ArcSwap;
 use eql_mapper::{Column, Schema, Table};
-use sqlparser::ast::Ident;
+use sqltk_parser::ast::Ident;
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::{task::JoinHandle, time};

packages/cipherstash-proxy/src/encrypt/schema/manager.rs

@@ -3,7 +3,7 @@ use cipherstash_client::{
     zerokms::{encrypted_record, EncryptedRecord},
 };
 use serde::{Deserialize, Serialize};
-use sqlparser::ast::Ident;
+use sqltk_parser::ast::Ident;
 
 #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
 pub struct Plaintext {

packages/cipherstash-proxy/src/eql/mod.rs

@@ -320,8 +320,8 @@ impl From<serde_json::Error> for Error {
     }
 }
 
-impl From<sqlparser::parser::ParserError> for Error {
-    fn from(e: sqlparser::parser::ParserError) -> Self {
+impl From<sqltk_parser::parser::ParserError> for Error {
+    fn from(e: sqltk_parser::parser::ParserError) -> Self {
         Error::Mapping(MappingError::InvalidSqlStatement(e.to_string()))
     }
 }

packages/cipherstash-proxy/src/error.rs

@@ -11,7 +11,7 @@ use cipherstash_config::ColumnType;
 use postgres_types::FromSql;
 use postgres_types::Type;
 use rust_decimal::Decimal;
-use sqlparser::ast::Value;
+use sqltk_parser::ast::Value;
 use std::str::FromStr;
 use tracing::debug;
 

packages/cipherstash-proxy/src/postgresql/data/from_sql.rs

@@ -26,13 +26,14 @@ use crate::prometheus::{
 use crate::Encrypted;
 use bytes::BytesMut;
 use cipherstash_client::encryption::Plaintext;
-use eql_mapper::{self, EqlMapperError, EqlValue, NodeKey, TableColumn, TypedStatement};
+use eql_mapper::{self, EqlMapperError, EqlValue, TableColumn, TypedStatement};
 use metrics::{counter, histogram};
 use pg_escape::quote_literal;
 use serde::Serialize;
-use sqlparser::ast::{self, Expr, Value};
-use sqlparser::dialect::PostgreSqlDialect;
-use sqlparser::parser::Parser;
+use sqltk::AsNodeKey;
+use sqltk_parser::ast::{self, Value};
+use sqltk_parser::dialect::PostgreSqlDialect;
+use sqltk_parser::parser::Parser;
 use std::collections::HashMap;
 use std::time::Instant;
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
@@ -285,7 +286,7 @@ where
 
             match self.to_encryptable_statement(&typed_statement, vec![])? {
                 Some(statement) => {
-                    if statement.has_literals() || typed_statement.has_nodes_to_wrap() {
+                    if statement.has_literals() {
                         let encrypted_literals = self
                             .encrypt_literals(&typed_statement, &statement.literal_columns)
                             .await?;
@@ -421,16 +422,15 @@ where
             .literals
             .iter()
             .zip(encrypted_expressions.into_iter())
-            .filter_map(|((_, original_node), en)| en.map(|en| (NodeKey::new(*original_node), en)))
+            .filter_map(|((_, original_node), en)| en.map(|en| (original_node.as_node_key(), en)))
             .collect::<HashMap<_, _>>();
 
         debug!(target: MAPPER,
             client_id = self.context.client_id,
-            nodes_to_wrap = typed_statement.nodes_to_wrap.len(),
             literals = encrypted_nodes.len(),
         );
 
-        if !typed_statement.has_nodes_to_wrap() && encrypted_nodes.is_empty() {
+        if encrypted_nodes.is_empty() {
             return Ok(None);
         }
 
@@ -500,7 +500,7 @@ where
 
         match self.to_encryptable_statement(&typed_statement, param_types)? {
             Some(statement) => {
-                if statement.has_literals() || typed_statement.has_nodes_to_wrap() {
+                if statement.has_literals() {
                     let encrypted_literals = self
                         .encrypt_literals(&typed_statement, &statement.literal_columns)
                         .await?;
@@ -620,7 +620,6 @@ where
         if (param_columns.is_empty() || no_encrypted_param_columns)
             && (projection_columns.is_empty() || no_encrypted_projection_columns)
             && literal_columns.is_empty()
-            && !typed_statement.has_nodes_to_wrap()
         {
             return Ok(None);
         }
@@ -781,7 +780,7 @@ where
         typed_statement: &eql_mapper::TypedStatement<'_>,
     ) -> Result<Vec<Option<Column>>, Error> {
         let mut projection_columns = vec![];
-        if let Some(eql_mapper::Projection::WithColumns(columns)) = &typed_statement.projection {
+        if let eql_mapper::Projection::WithColumns(columns) = &typed_statement.projection {
             for col in columns {
                 let eql_mapper::ProjectionColumn { ty, .. } = col;
                 let configured_column = match ty {
@@ -819,7 +818,7 @@ where
 
         for param in typed_statement.params.iter() {
             let configured_column = match param {
-                eql_mapper::Value::Eql(EqlValue(TableColumn { table, column })) => {
+                (_, eql_mapper::Value::Eql(EqlValue(TableColumn { table, column }))) => {
                     let identifier = Identifier::from((table, column));
 
                     debug!(
@@ -967,9 +966,9 @@ fn literals_to_plaintext(
     Ok(plaintexts)
 }
 
-fn to_json_literal_expr<T>(literal: &T) -> Result<Expr, Error>
+fn to_json_literal_expr<T>(literal: &T) -> Result<Value, Error>
 where
     T: ?Sized + Serialize,
 {
-    Ok(serde_json::to_string(literal).map(|json| Expr::Value(Value::SingleQuotedString(json)))?)
+    Ok(serde_json::to_string(literal).map(Value::SingleQuotedString)?)
 }

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

@@ -12,9 +12,10 @@ authors = [
 
 [dependencies]
 derive_more = { version = "^1.0", features = ["display", "constructor"] }
+impl-trait-for-tuples = "0.2.3"
 itertools = "^0.13"
-sqlparser = { workspace = true }
-sqltk = { git = "https://github.com/cipherstash/sqltk/", rev = "214f9b90e4f07d4414292813ffd6e45dec075fbb" }
+sqltk = { workspace = true }
+sqltk-parser = { workspace = true }
 thiserror = { workspace = true }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }

packages/eql-mapper/Cargo.toml

@@ -0,0 +1,130 @@
+use std::{
+    collections::HashMap,
+    fmt::{Debug, Display},
+};
+
+use sqltk::NodeKey;
+use sqltk_parser::ast::{
+    Delete, Expr, Function, Insert, Query, Select, SelectItem, SetExpr, Statement, Value, Values,
+};
+
+use crate::{EqlValue, Param, Type};
+
+#[derive(Debug, Clone, Copy, Hash, Eq, PartialEq)]
+pub struct Fmt<T>(pub(crate) T);
+
+#[derive(Debug, Clone, Copy, Hash, Eq, PartialEq)]
+pub struct FmtAst<T>(pub(crate) T);
+
+#[derive(Debug, Clone, Copy, Hash, Eq, PartialEq)]
+pub struct FmtAstVec<T>(pub(crate) T);
+
+impl Display for Fmt<NodeKey<'_>> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if let Some(node) = self.0.get_as::<Statement>() {
+            return Display::fmt(&FmtAst(node), f);
+        }
+        if let Some(node) = self.0.get_as::<Query>() {
+            return Display::fmt(&FmtAst(node), f);
+        }
+        if let Some(node) = self.0.get_as::<Insert>() {
+            return Display::fmt(&FmtAst(node), f);
+        }
+        if let Some(node) = self.0.get_as::<Delete>() {
+            return Display::fmt(&FmtAst(node), f);
+        }
+        if let Some(node) = self.0.get_as::<Expr>() {
+            return Display::fmt(&FmtAst(node), f);
+        }
+        if let Some(node) = self.0.get_as::<SetExpr>() {
+            return Display::fmt(&FmtAst(node), f);
+        }
+        if let Some(node) = self.0.get_as::<Select>() {
+            return Display::fmt(&FmtAst(node), f);
+        }
+        if let Some(node) = self.0.get_as::<SelectItem>() {
+            return Display::fmt(&FmtAst(node), f);
+        }
+        if let Some(node) = self.0.get_as::<Vec<SelectItem>>() {
+            return Display::fmt(&FmtAstVec(node), f);
+        }
+        if let Some(node) = self.0.get_as::<Function>() {
+            return Display::fmt(&FmtAst(node), f);
+        }
+        if let Some(node) = self.0.get_as::<Values>() {
+            return Display::fmt(&FmtAst(node), f);
+        }
+        if let Some(node) = self.0.get_as::<Value>() {
+            return Display::fmt(&FmtAst(node), f);
+        }
+
+        f.write_str("!! CANNOT RENDER SQL NODE !!!")?;
+
+        Ok(())
+    }
+}
+
+impl Display for Fmt<&HashMap<NodeKey<'_>, Type>> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut out: Vec<String> = Vec::new();
+        out.push("{ ".into());
+        for (k, v) in self.0.iter() {
+            out.push(format!("{}: {}", Fmt(*k), v));
+        }
+        out.push(" }".into());
+        f.write_str(&out.join(", "))
+    }
+}
+
+impl<T: Display> Display for FmtAstVec<&Vec<T>> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("![")?;
+        let children = self
+            .0
+            .iter()
+            .map(|n| n.to_string())
+            .collect::<Vec<_>>()
+            .join(", ");
+        f.write_str(&children)?;
+        f.write_str("]!")
+    }
+}
+
+impl<T: Display> Display for FmtAst<&T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        T::fmt(self.0, f)
+    }
+}
+
+impl Display for Fmt<&Vec<(Param, crate::Value)>> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let formatted = self
+            .0
+            .iter()
+            .map(|(p, v)| format!("{}: {}", p, v))
+            .collect::<Vec<_>>()
+            .join(", ");
+        f.write_str(&formatted)
+    }
+}
+
+impl Display for Fmt<&Vec<(EqlValue, &sqltk_parser::ast::Value)>> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let formatted = self
+            .0
+            .iter()
+            .map(|(e, n)| format!("{}: {}", n, e))
+            .collect::<Vec<_>>()
+            .join(", ");
+        f.write_str(&formatted)
+    }
+}
+
+impl<T: Display> Display for Fmt<Option<T>> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match &self.0 {
+            Some(t) => t.fmt(f),
+            None => Display::fmt("<not initialised>", f),
+        }
+    }
+}