feat: Support ORE column usage in GROUP BY

Cargo.toml

@@ -36,6 +36,8 @@ debug = true
 
 [workspace.dependencies]
 sqlparser = { version = "^0.52", features = ["bigdecimal", "serde"] }
+# sqltk = { git = "https://github.com/cipherstash/sqltk/", rev = "e16406f" }
+sqltk = { path = "../sqltk/packages/sqltk" }
 thiserror = "2.0.9"
 tokio = { version = "1.42.0", features = ["full"] }
 tracing = "0.1"

packages/cipherstash-proxy/Cargo.toml

@@ -41,6 +41,7 @@ serde = "1.0"
 serde_json = "1.0"
 socket2 = "0.5.7"
 sqlparser = { workspace = true }
+sqltk = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }
 tokio-postgres = { version = "0.7", features = [

packages/cipherstash-proxy/src/postgresql/frontend.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::ast::{self, Value};
 use sqlparser::dialect::PostgreSqlDialect;
 use sqlparser::parser::Parser;
+use sqltk::AsNodeKey;
 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/eql-mapper-macros/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "eql-mapper-macros"
+version.workspace = true
+edition.workspace = true
+
+[lib]
+proc-macro = true
+
+[dependencies]
+syn = { version = "2.0", features = ["full"] }
+quote = "1.0"
+proc-macro2 = "1.0"

packages/eql-mapper-macros/src/lib.rs

@@ -0,0 +1,115 @@
+use proc_macro::TokenStream;
+use quote::{quote, ToTokens};
+use syn::{
+    parse::Parse, parse_macro_input, parse_quote, Attribute, FnArg, Ident, ImplItem, ImplItemFn,
+    ItemImpl, Pat, PatType, Signature, Type, TypePath, TypeReference,
+};
+
+/// This macro generates consistently defined `#[tracing::instrument]` attributes for `InferType::infer_enter` &
+/// `InferType::infer_enter` implementations on `TypeInferencer`.
+///
+/// This attribute MUST be defined on the trait `impl` itself (not the trait method impls).
+#[proc_macro_attribute]
+pub fn trace_infer(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let mut input = parse_macro_input!(item as ItemImpl);
+
+    for item in &mut input.items {
+        if let ImplItem::Fn(ImplItemFn {
+            attrs,
+            sig:
+                Signature {
+                    ident: method,
+                    inputs,
+                    ..
+                },
+            ..
+        }) = item
+        {
+            let node_ident_and_type: Option<(&Ident, &Type)> =
+                if let Some(FnArg::Typed(PatType {
+                    ty: node_ty, pat, ..
+                })) = inputs.get(1)
+                {
+                    if let Pat::Ident(pat_ident) = &**pat {
+                        Some((&pat_ident.ident, node_ty))
+                    } else {
+                        None
+                    }
+                } else {
+                    None
+                };
+
+            let vec_ident: Ident = parse_quote!(Vec);
+
+            match node_ident_and_type {
+                Some((node_ident, node_ty)) => {
+                    let (formatter, node_ty_abbrev) = match node_ty {
+                        Type::Reference(TypeReference { elem, .. }) => match &**elem {
+                            Type::Path(TypePath { path, .. }) => {
+                                let last_segment = path.segments.last().unwrap();
+                                let last_segment_ident = &last_segment.ident;
+                                let last_segment_arguments = if last_segment.arguments.is_empty() {
+                                    None
+                                }  else {
+                                    let args = &last_segment.arguments;
+                                    Some(quote!(<#args>))
+                                };
+                                match last_segment_ident {
+                                    ident if vec_ident == *ident => {
+                                        (quote!(crate::FmtAstVec), quote!(#last_segment_ident #last_segment_arguments))
+                                    }
+                                    _ => (quote!(crate::FmtAst), quote!(#last_segment_ident #last_segment_arguments))
+                                }
+                            },
+                            _ => unreachable!("Infer::infer_enter/infer_exit has sig: infer_..(&mut self, delete: &'ast N) -> Result<(), TypeError>")
+                        },
+                            _ => unreachable!("Infer::infer_enter/infer_exit has sig: infer_..(&mut self, delete: &'ast N) -> Result<(), TypeError>")
+                    };
+
+                    let node_ty_abbrev = node_ty_abbrev
+                        .to_token_stream()
+                        .to_string()
+                        .replace(" ", "");
+
+                    let target = format!("eql-mapper::{}", method.to_string().to_uppercase());
+
+                    let attr: TracingInstrumentAttr = syn::parse2(quote! {
+                        #[tracing::instrument(
+                            target = #target,
+                            level = "trace",
+                            skip(self, #node_ident),
+                            fields(
+                                ast_ty = #node_ty_abbrev,
+                                ast = %#formatter(#node_ident),
+                            ),
+                            ret(Debug)
+                        )]
+                    })
+                    .unwrap();
+                    attrs.push(attr.attr);
+                }
+                None => {
+                    return quote!(compile_error!(
+                        "could not determine name of node argumemt in Infer impl"
+                    ))
+                    .to_token_stream()
+                    .into();
+                }
+            }
+        }
+    }
+
+    input.to_token_stream().into()
+}
+
+struct TracingInstrumentAttr {
+    attr: Attribute,
+}
+
+impl Parse for TracingInstrumentAttr {
+    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+        Ok(Self {
+            attr: Attribute::parse_outer(input)?.first().unwrap().clone(),
+        })
+    }
+}

packages/eql-mapper/Cargo.toml

@@ -11,10 +11,14 @@ authors = [
 ]
 
 [dependencies]
+eql-mapper-macros = { path = "../eql-mapper-macros" }
 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 = { git = "https://github.com/cipherstash/sqltk/", rev = "214f9b90e4f07d4414292813ffd6e45dec075fbb" }
+# sqltk = { git = "https://github.com/cipherstash/sqltk/", rev = "cdf2e7c" }
+sqltk = { workspace = true }
 thiserror = { workspace = true }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }