Merge pull request #317 from cipherstash/refactor-frontend-parser

♻️ refactor: extract SQL parsing logic into dedicated Parser module

Author: tobyhede

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

@@ -6,6 +6,7 @@ use super::messages::execute::Execute;
 use super::messages::parse::Parse;
 use super::messages::query::Query;
 use super::messages::FrontendCode as Code;
+use super::parser::SqlParser;
 use super::protocol::{self};
 use crate::connect::Sender;
 use crate::eql::Identifier;
@@ -22,7 +23,7 @@ use crate::prometheus::{
     CLIENTS_BYTES_RECEIVED_TOTAL, ENCRYPTED_VALUES_TOTAL, ENCRYPTION_DURATION_SECONDS,
     ENCRYPTION_ERROR_TOTAL, ENCRYPTION_REQUESTS_TOTAL, SERVER_BYTES_SENT_TOTAL,
     STATEMENTS_ENCRYPTED_TOTAL, STATEMENTS_PASSTHROUGH_MAPPING_DISABLED_TOTAL,
-    STATEMENTS_PASSTHROUGH_TOTAL, STATEMENTS_TOTAL, STATEMENTS_UNMAPPABLE_TOTAL,
+    STATEMENTS_PASSTHROUGH_TOTAL, STATEMENTS_UNMAPPABLE_TOTAL,
 };
 use crate::proxy::Proxy;
 use crate::EqlEncrypted;
@@ -34,17 +35,13 @@ use pg_escape::quote_literal;
 use postgres_types::Type;
 use serde::Serialize;
 use sqltk::parser::ast::{self, Value};
-use sqltk::parser::dialect::PostgreSqlDialect;
-use sqltk::parser::parser::Parser;
 use sqltk::NodeKey;
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Instant;
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 use tracing::{debug, error, info, warn};
 
-const DIALECT: PostgreSqlDialect = PostgreSqlDialect {};
-
 /// The PostgreSQL proxy frontend that handles client-to-server message processing.
 ///
 /// The Frontend intercepts messages from PostgreSQL clients, analyzes SQL statements for
@@ -384,7 +381,7 @@ where
         let mut query = Query::try_from(bytes)?;
 
         // Simple Query may contain many statements
-        let parsed_statements = self.parse_statements(&query.statement)?;
+        let parsed_statements = SqlParser::parse_statements(&query.statement)?;
         let mut transformed_statements = vec![];
 
         debug!(target: MAPPER,
@@ -659,7 +656,7 @@ where
             parse = ?message
         );
 
-        let statement = self.parse_statement(&message.statement)?;
+        let statement = SqlParser::parse_statement(&message.statement)?;
 
         if let Some(mapping_disabled) = self.context.maybe_set_unsafe_disable_mapping(&statement) {
             warn!(
@@ -748,42 +745,6 @@ where
         }
     }
 
-    ///
-    /// Parse a SQL statement string into an SqlParser AST
-    ///
-    fn parse_statement(&mut self, statement: &str) -> Result<ast::Statement, Error> {
-        let statement = Parser::new(&DIALECT)
-            .try_with_sql(statement)?
-            .parse_statement()?;
-
-        debug!(target: MAPPER,
-            client_id = self.context.client_id,
-            statement = %statement
-        );
-
-        counter!(STATEMENTS_TOTAL).increment(1);
-
-        Ok(statement)
-    }
-
-    ///
-    /// Parse a SQL String potentially containing multiple statements into parsed SqlParser AST
-    ///
-    fn parse_statements(&mut self, statement: &str) -> Result<Vec<ast::Statement>, Error> {
-        let statement = Parser::new(&DIALECT)
-            .try_with_sql(statement)?
-            .parse_statements()?;
-
-        debug!(target: MAPPER,
-            client_id = self.context.client_id,
-            statement = ?statement
-        );
-
-        counter!(STATEMENTS_TOTAL).increment(statement.len() as u64);
-
-        Ok(statement)
-    }
-
     ///
     /// Check the Statement AST for DDL
     /// Sets a schema changed flag in the Context

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

@@ -7,6 +7,7 @@ mod frontend;
 mod handler;
 mod message_buffer;
 mod messages;
+mod parser;
 mod protocol;
 mod startup;
 

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

@@ -0,0 +1,34 @@
+use crate::error::Error;
+use crate::prometheus::STATEMENTS_TOTAL;
+use metrics::counter;
+use sqltk::parser::ast;
+use sqltk::parser::dialect::PostgreSqlDialect;
+use sqltk::parser::parser::Parser;
+
+const DIALECT: PostgreSqlDialect = PostgreSqlDialect {};
+
+pub struct SqlParser;
+
+impl SqlParser {
+    /// Parse a SQL statement string into an SqlParser AST
+    pub fn parse_statement(statement: &str) -> Result<ast::Statement, Error> {
+        let statement = Parser::new(&DIALECT)
+            .try_with_sql(statement)?
+            .parse_statement()?;
+
+        counter!(STATEMENTS_TOTAL).increment(1);
+
+        Ok(statement)
+    }
+
+    /// Parse a SQL String potentially containing multiple statements into parsed SqlParser AST
+    pub fn parse_statements(statement: &str) -> Result<Vec<ast::Statement>, Error> {
+        let statement = Parser::new(&DIALECT)
+            .try_with_sql(statement)?
+            .parse_statements()?;
+
+        counter!(STATEMENTS_TOTAL).increment(statement.len() as u64);
+
+        Ok(statement)
+    }
+}