♻️ refactor: extract column processing logic into dedicated ColumnMapper service

- Create ColumnMapper service to handle column mapping and encryption configuration
- Move get_projection_columns, get_param_columns, get_literal_columns from Frontend
- Remove ~150 lines of column processing logic from Frontend
- Add Context delegation methods for column processing operations
- Remove unused imports (Identifier, TableColumn, Type) from Frontend
- Fix test configuration with proper CRN format and required environment variables

This refactoring achieves clear separation of concerns:
- Frontend focuses solely on PostgreSQL message handling
- ColumnMapper handles column mapping and encryption configuration
- Context serves as coordinator with clean delegation methods
- Maintains all existing functionality while improving code organization

Author: tobyhede

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

@@ -0,0 +1,167 @@
+use crate::{
+    eql::Identifier,
+    error::{EncryptError, Error},
+    log::MAPPER,
+    postgresql::Column,
+    services::SchemaService,
+};
+use eql_mapper::{EqlTerm, TableColumn, TypeCheckedStatement};
+use postgres_types::Type;
+use std::sync::Arc;
+use tracing::{debug, warn};
+
+/// Service responsible for processing columns from type-checked SQL statements
+/// and mapping them to encryption configurations.
+#[derive(Clone)]
+pub struct ColumnMapper {
+    schema_service: Arc<dyn SchemaService>,
+}
+
+impl ColumnMapper {
+    /// Create a new ColumnProcessor with the given schema service and client ID
+    pub fn new(schema_service: Arc<dyn SchemaService>) -> Self {
+        Self { schema_service }
+    }
+
+    /// Maps typed statement projection columns to an Encrypt column configuration
+    ///
+    /// The returned `Vec` is of `Option<Column>` because the Projection columns are a mix of native and EQL types.
+    /// Only EQL colunms will have a configuration. Native types are always None.
+    ///
+    /// Preserves the ordering and semantics of the projection to reduce the complexity of positional encryption.
+    pub fn get_projection_columns(
+        &self,
+        typed_statement: &TypeCheckedStatement<'_>,
+    ) -> Result<Vec<Option<Column>>, Error> {
+        let mut projection_columns = vec![];
+
+        for col in typed_statement.projection.columns() {
+            let eql_mapper::ProjectionColumn { ty, .. } = col;
+            let configured_column = match &**ty {
+                eql_mapper::Type::Value(eql_mapper::Value::Eql(eql_term)) => {
+                    let TableColumn { table, column } = eql_term.table_column();
+                    let identifier: Identifier = Identifier::from((table, column));
+
+                    debug!(
+                        target: MAPPER,
+                        msg = "Configured column",
+                        column = ?identifier,
+                        ?eql_term,
+                    );
+                    self.get_column(identifier, eql_term)?
+                }
+                _ => None,
+            };
+            projection_columns.push(configured_column)
+        }
+
+        Ok(projection_columns)
+    }
+
+    /// Maps typed statement param columns to an Encrypt column configuration
+    ///
+    /// The returned `Vec` is of `Option<Column>` because the Param columns are a mix of native and EQL types.
+    /// Only EQL colunms will have a configuration. Native types are always None.
+    ///
+    /// Preserves the ordering and semantics of the projection to reduce the complexity of positional encryption.
+    pub fn get_param_columns(
+        &self,
+        typed_statement: &TypeCheckedStatement<'_>,
+    ) -> Result<Vec<Option<Column>>, Error> {
+        let mut param_columns = vec![];
+
+        for param in typed_statement.params.iter() {
+            let configured_column = match param {
+                (_, eql_mapper::Value::Eql(eql_term)) => {
+                    let TableColumn { table, column } = eql_term.table_column();
+                    let identifier = Identifier::from((table, column));
+
+                    debug!(
+                        target: MAPPER,
+                        msg = "Encrypted parameter",
+                        column = ?identifier,
+                        ?eql_term,
+                    );
+
+                    self.get_column(identifier, eql_term)?
+                }
+                _ => None,
+            };
+            param_columns.push(configured_column);
+        }
+
+        Ok(param_columns)
+    }
+
+    /// Maps typed statement literal columns to an Encrypt column configuration
+    pub fn get_literal_columns(
+        &self,
+        typed_statement: &TypeCheckedStatement<'_>,
+    ) -> Result<Vec<Option<Column>>, Error> {
+        let mut literal_columns = vec![];
+
+        for (eql_term, _) in typed_statement.literals.iter() {
+            let TableColumn { table, column } = eql_term.table_column();
+            let identifier = Identifier::from((table, column));
+
+            debug!(
+                target: MAPPER,
+                msg = "Encrypted literal",
+                column = ?identifier,
+                ?eql_term,
+            );
+            let col = self.get_column(identifier, eql_term)?;
+            if col.is_some() {
+                literal_columns.push(col);
+            }
+        }
+
+        Ok(literal_columns)
+    }
+
+    /// Get the column configuration for the Identifier
+    /// Returns `EncryptError::UnknownColumn` if configuration cannot be found for the Identified column
+    /// if mapping enabled, and None if mapping is disabled. It'll log a warning either way.
+    fn get_column(
+        &self,
+        identifier: Identifier,
+        eql_term: &EqlTerm,
+    ) -> Result<Option<Column>, Error> {
+        match self.schema_service.get_column_config(&identifier) {
+            Some(config) => {
+                debug!(
+                    target: MAPPER,
+                    msg = "Configured column",
+                    column = ?identifier
+                );
+
+                // IndexTerm::SteVecSelector
+                let postgres_type = if matches!(eql_term, EqlTerm::JsonPath(_)) {
+                    Some(Type::JSONPATH)
+                } else {
+                    None
+                };
+
+                let eql_term = eql_term.variant();
+                Ok(Some(Column::new(
+                    identifier,
+                    config,
+                    postgres_type,
+                    eql_term,
+                )))
+            }
+            None => {
+                warn!(
+                    target: MAPPER,
+                    msg = "Configured column not found. Encryption configuration may have been deleted.",
+                    ?identifier,
+                );
+                Err(EncryptError::UnknownColumn {
+                    table: identifier.table.to_owned(),
+                    column: identifier.column.to_owned(),
+                }
+                .into())
+            }
+        }
+    }
+}

packages/cipherstash-proxy/src/postgresql/context/mod.rs

@@ -1,6 +1,7 @@
 pub mod column;
 
 use super::{
+    column_mapper::ColumnMapper,
     format_code::FormatCode,
     messages::{describe::Describe, Name, Target},
     Column,
@@ -44,6 +45,7 @@ pub struct Context {
     config: Arc<TandemConfig>,
     encryption: Arc<dyn EncryptionService>,
     schema: Arc<dyn SchemaService>,
+    column_processor: ColumnMapper,
     statements: Arc<RwLock<HashMap<Name, Arc<Statement>>>>,
     portals: Arc<RwLock<HashMap<Name, PortalQueue>>>,
     describe: Arc<RwLock<DescribeQueue>>,
@@ -124,6 +126,8 @@ impl Context {
         encryption: Arc<dyn EncryptionService>,
         schema_service: Arc<dyn SchemaService>,
     ) -> Context {
+        let column_processor = ColumnMapper::new(schema_service.clone());
+
         Context {
             statements: Arc::new(RwLock::new(HashMap::new())),
             portals: Arc::new(RwLock::new(HashMap::new())),
@@ -136,6 +140,7 @@ impl Context {
             config,
             encryption,
             schema: schema_service,
+            column_processor,
             unsafe_disable_mapping: false,
             keyset_id: Arc::new(RwLock::new(None)),
         }
@@ -532,6 +537,29 @@ impl Context {
         self.schema.is_empty_config()
     }
 
+    // Column processing delegation methods
+    pub fn get_projection_columns(
+        &self,
+        typed_statement: &eql_mapper::TypeCheckedStatement<'_>,
+    ) -> Result<Vec<Option<Column>>, Error> {
+        self.column_processor
+            .get_projection_columns(typed_statement)
+    }
+
+    pub fn get_param_columns(
+        &self,
+        typed_statement: &eql_mapper::TypeCheckedStatement<'_>,
+    ) -> Result<Vec<Option<Column>>, Error> {
+        self.column_processor.get_param_columns(typed_statement)
+    }
+
+    pub fn get_literal_columns(
+        &self,
+        typed_statement: &eql_mapper::TypeCheckedStatement<'_>,
+    ) -> Result<Vec<Option<Column>>, Error> {
+        self.column_processor.get_literal_columns(typed_statement)
+    }
+
     // Direct config access methods
     pub fn connection_timeout(&self) -> std::time::Duration {
         self.config
@@ -629,10 +657,20 @@ impl Context {
         std::env::set_var("CS_DATABASE__NAME", "test");
         std::env::set_var("CS_DATABASE__HOST", "localhost");
         std::env::set_var("CS_DATABASE__PORT", "5432");
-        std::env::set_var("CS_AUTH__WORKSPACE_CRN", "crn:ap-southeast-2.aws:test");
+        std::env::set_var(
+            "CS_AUTH__WORKSPACE_CRN",
+            "crn:ap-southeast-2.aws:3KISDURL3ZCWYZ2O",
+        );
         std::env::set_var("CS_AUTH__CLIENT_ACCESS_KEY", "test");
-        std::env::set_var("CS_ENCRYPT__CLIENT_ID", "test");
-        std::env::set_var("CS_ENCRYPT__CLIENT_KEY", "test");
+        std::env::set_var(
+            "CS_ENCRYPT__CLIENT_ID",
+            "e40f1692-6bb7-4bbd-a552-4c0f155be073",
+        );
+        std::env::set_var("CS_ENCRYPT__CLIENT_KEY", "a4627031a16b7065726d75746174696f6e90090e0805000b0d0c0106040f0a0302076770325f66726f6da16b7065726d75746174696f6e9007060a0b02090d080c00040f0305010e6570325f746fa16b7065726d75746174696f6e900a0206090b04050c070f0e010d030800627033a16b7065726d75746174696f6e98210514181d0818200a18190b1112181809130f15181a0717181e000e0103181f0d181c1602040c181b1006");
+        std::env::set_var(
+            "CS_ENCRYPT__KEYSET_ID",
+            "c50d8463-60e9-41a5-86cd-5782e03a503c",
+        );
 
         let config = Arc::new(
             crate::config::TandemConfig::build("tests/config/unknown.toml")