feat: add rewrite to query handlers

Author: sunng87

datafusion-postgres/src/handlers.rs

@@ -2,10 +2,12 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use crate::auth::{AuthManager, Permission, ResourceType};
+use crate::sql::{parse, rewrite, AliasDuplicatedProjectionRewrite, SqlStatementRewriteRule};
 use async_trait::async_trait;
 use datafusion::arrow::datatypes::DataType;
 use datafusion::logical_expr::LogicalPlan;
 use datafusion::prelude::*;
+use datafusion::sql::parser::Statement;
 use pgwire::api::auth::noop::NoopStartupHandler;
 use pgwire::api::auth::StartupHandler;
 use pgwire::api::portal::{Format, Portal};
@@ -63,21 +65,26 @@ pub struct DfSessionService {
     parser: Arc<Parser>,
     timezone: Arc<Mutex<String>>,
     auth_manager: Arc<AuthManager>,
+    sql_rewrite_rules: Vec<Arc<dyn SqlStatementRewriteRule>>,
 }
 
 impl DfSessionService {
     pub fn new(
         session_context: Arc<SessionContext>,
         auth_manager: Arc<AuthManager>,
     ) -> DfSessionService {
+        let sql_rewrite_rules: Vec<Arc<dyn SqlStatementRewriteRule>> =
+            vec![Arc::new(AliasDuplicatedProjectionRewrite)];
         let parser = Arc::new(Parser {
             session_context: session_context.clone(),
+            sql_rewrite_rules: sql_rewrite_rules.clone(),
         });
         DfSessionService {
             session_context,
             parser,
             timezone: Arc::new(Mutex::new("UTC".to_string())),
             auth_manager,
+            sql_rewrite_rules,
         }
     }
 
@@ -308,8 +315,17 @@ impl SimpleQueryHandler for DfSessionService {
     where
         C: ClientInfo + Unpin + Send + Sync,
     {
-        let query_lower = query.to_lowercase().trim().to_string();
         log::debug!("Received query: {query}"); // Log the query for debugging
+        let mut statements = parse(query).map_err(|e| PgWireError::ApiError(Box::new(e)))?;
+
+        // TODO: deal with multiple statements
+        let mut statement = statements.remove(0);
+
+        // Attempt to rewrite
+        statement = rewrite(statement, &self.sql_rewrite_rules);
+
+        // TODO: improve statement check by using statement directly
+        let query_lower = statement.to_string().to_lowercase().trim().to_string();
 
         // Check permissions for the query (skip for SET, transaction, and SHOW statements)
         if !query_lower.starts_with("set")
@@ -526,6 +542,7 @@ impl ExtendedQueryHandler for DfSessionService {
 
 pub struct Parser {
     session_context: Arc<SessionContext>,
+    sql_rewrite_rules: Vec<Arc<dyn SqlStatementRewriteRule>>,
 }
 
 #[async_trait]
@@ -538,14 +555,23 @@ impl QueryParser for Parser {
         sql: &str,
         _types: &[Type],
     ) -> PgWireResult<Self::Statement> {
-        log::debug!("Received parse extended query: {sql}"); // Log for debugging
+        log::debug!("Received parse extended query: {sql}"); // Log for
+                                                             // debugging
+        let mut statements = parse(sql).map_err(|e| PgWireError::ApiError(Box::new(e)))?;
+        let mut statement = statements.remove(0);
+
+        // Attempt to rewrite
+        statement = rewrite(statement, &self.sql_rewrite_rules);
+
+        let query = statement.to_string();
+
         let context = &self.session_context;
         let state = context.state();
         let logical_plan = state
-            .create_logical_plan(sql)
+            .statement_to_plan(Statement::Statement(Box::new(statement)))
             .await
             .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
-        Ok((sql.to_string(), logical_plan))
+        Ok((query, logical_plan))
     }
 }
 

datafusion-postgres/src/sql.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use datafusion::sql::sqlparser::ast::Expr;
 use datafusion::sql::sqlparser::ast::Ident;
 use datafusion::sql::sqlparser::ast::Select;
@@ -15,15 +17,15 @@ pub fn parse(sql: &str) -> Result<Vec<Statement>, ParserError> {
     Parser::parse_sql(&dialect, sql)
 }
 
-pub fn rewrite(mut s: Statement, rules: &[Box<dyn SqlStatementRewriteRule>]) -> Statement {
+pub fn rewrite(mut s: Statement, rules: &[Arc<dyn SqlStatementRewriteRule>]) -> Statement {
     for rule in rules {
         s = rule.rewrite(s);
     }
 
     s
 }
 
-pub trait SqlStatementRewriteRule {
+pub trait SqlStatementRewriteRule: Send + Sync {
     fn rewrite(&self, s: Statement) -> Statement;
 }
 
@@ -35,7 +37,8 @@ pub trait SqlStatementRewriteRule {
 ///
 /// This rule will add alias to column, when there is a wildcard found in
 /// projection.
-struct AliasDuplicatedProjectionRewrite;
+#[derive(Debug)]
+pub struct AliasDuplicatedProjectionRewrite;
 
 impl AliasDuplicatedProjectionRewrite {
     // Rewrites a SELECT statement to alias explicit columns from the same table as a qualified wildcard.
@@ -103,7 +106,7 @@ impl AliasDuplicatedProjectionRewrite {
                     };
 
                     if let Some(name) = alias_partial {
-                        let alias = format!("__alias_{}", name);
+                        let alias = format!("__alias_{name}");
                         new_projection.push(SelectItem::ExprWithAlias {
                             expr,
                             alias: Ident::new(alias),
@@ -138,8 +141,8 @@ mod tests {
 
     #[test]
     fn test_alias_rewrite() {
-        let rules: Vec<Box<dyn SqlStatementRewriteRule>> =
-            vec![Box::new(AliasDuplicatedProjectionRewrite)];
+        let rules: Vec<Arc<dyn SqlStatementRewriteRule>> =
+            vec![Arc::new(AliasDuplicatedProjectionRewrite)];
 
         let sql = "SELECT n.oid, n.* FROM pg_catalog.pg_namespace n";
         let statement = parse(sql).expect("Failed to parse").remove(0);