fix: correct query feed into handler

Author: sunng87

datafusion-postgres/src/handlers.rs

@@ -324,7 +324,8 @@ impl SimpleQueryHandler for DfSessionService {
         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();
+        let query = statement.to_string();
+        let query_lower = query.to_lowercase().trim().to_string();
 
         // Check permissions for the query (skip for SET, transaction, and SHOW statements)
         if !query_lower.starts_with("set")
@@ -336,7 +337,7 @@ impl SimpleQueryHandler for DfSessionService {
             && !query_lower.starts_with("abort")
             && !query_lower.starts_with("show")
         {
-            self.check_query_permission(client, query).await?;
+            self.check_query_permission(client, &query).await?;
         }
 
         if let Some(resp) = self.try_respond_set_statements(&query_lower).await? {
@@ -366,7 +367,7 @@ impl SimpleQueryHandler for DfSessionService {
             )));
         }
 
-        let df_result = self.session_context.sql(query).await;
+        let df_result = self.session_context.sql(&query).await;
 
         // Handle query execution errors and transaction state
         let df = match df_result {

datafusion-postgres/src/sql.rs

@@ -170,5 +170,14 @@ mod tests {
             statement.to_string(),
             "SELECT t1.oid, t2.* FROM tbl1 AS t1 JOIN tbl2 AS t2 ON t1.id = t2.id"
         );
+
+        let sql = "SELECT n.oid,n.*,d.description FROM pg_catalog.pg_namespace n LEFT OUTER JOIN pg_catalog.pg_description d ON d.objoid=n.oid AND d.objsubid=0 AND d.classoid='pg_namespace' ORDER BY nspsname";
+        let statement = parse(sql).expect("Failed to parse").remove(0);
+
+        let statement = rewrite(statement, &rules);
+        assert_eq!(
+            statement.to_string(),
+            "SELECT n.oid AS __alias_oid, n.*, d.description FROM pg_catalog.pg_namespace AS n LEFT OUTER JOIN pg_catalog.pg_description AS d ON d.objoid = n.oid AND d.objsubid = 0 AND d.classoid = 'pg_namespace' ORDER BY nspsname"
+        );
     }
 }