Merge branch 'master' into feat/pg-catalog-sql-dbeaver

sunng87 · sunng87 · commit 7c1aa144ca33 · 2025-08-30T22:18:50.000+08:00
diff --git a/arrow-pg/src/encoder.rs b/arrow-pg/src/encoder.rs
@@ -513,38 +513,31 @@ pub fn encode_value<T: Encoder>(
             if arr.is_null(idx) {
                 return encoder.encode_field_with_type_and_format(&None::<i8>, type_, format);
             }
-            // Get the dictionary values, ignoring keys
-            // We'll use Int32Type as a common key type, but we're only interested in values
-            macro_rules! get_dict_values {
+            // Get the dictionary values and the mapped row index
+            macro_rules! get_dict_values_and_index {
                 ($key_type:ty) => {
                     arr.as_any()
                         .downcast_ref::<DictionaryArray<$key_type>>()
-                        .map(|dict| dict.values())
+                        .map(|dict| (dict.values(), dict.keys().value(idx) as usize))
                 };
             }
 
             // Try to extract values using different key types
-            let values = get_dict_values!(Int8Type)
-                .or_else(|| get_dict_values!(Int16Type))
-                .or_else(|| get_dict_values!(Int32Type))
-                .or_else(|| get_dict_values!(Int64Type))
-                .or_else(|| get_dict_values!(UInt8Type))
-                .or_else(|| get_dict_values!(UInt16Type))
-                .or_else(|| get_dict_values!(UInt32Type))
-                .or_else(|| get_dict_values!(UInt64Type))
+            let (values, idx) = get_dict_values_and_index!(Int8Type)
+                .or_else(|| get_dict_values_and_index!(Int16Type))
+                .or_else(|| get_dict_values_and_index!(Int32Type))
+                .or_else(|| get_dict_values_and_index!(Int64Type))
+                .or_else(|| get_dict_values_and_index!(UInt8Type))
+                .or_else(|| get_dict_values_and_index!(UInt16Type))
+                .or_else(|| get_dict_values_and_index!(UInt32Type))
+                .or_else(|| get_dict_values_and_index!(UInt64Type))
                 .ok_or_else(|| {
                     ToSqlError::from(format!(
                         "Unsupported dictionary key type for value type {value_type}"
                     ))
                 })?;
 
-            // If the dictionary has only one value, treat it as a primitive
-            if values.len() == 1 {
-                encode_value(encoder, values, 0, type_, format)?
-            } else {
-                // Otherwise, use value directly indexed by values array
-                encode_value(encoder, values, idx, type_, format)?
-            }
+            encode_value(encoder, values, idx, type_, format)?
         }
         _ => {
             return Err(PgWireError::ApiError(ToSqlError::from(format!(
@@ -557,3 +550,48 @@ pub fn encode_value<T: Encoder>(
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn encodes_dictionary_array() {
+        #[derive(Default)]
+        struct MockEncoder {
+            encoded_value: String,
+        }
+
+        impl Encoder for MockEncoder {
+            fn encode_field_with_type_and_format<T>(
+                &mut self,
+                value: &T,
+                data_type: &Type,
+                _format: FieldFormat,
+            ) -> PgWireResult<()>
+            where
+                T: ToSql + ToSqlText + Sized,
+            {
+                let mut bytes = BytesMut::new();
+                let _sql_text = value.to_sql_text(data_type, &mut bytes);
+                let string = String::from_utf8((&bytes).to_vec());
+                self.encoded_value = string.unwrap();
+                Ok(())
+            }
+        }
+
+        let val = "~!@&$[]()@@!!";
+        let value = StringArray::from_iter_values([val]);
+        let keys = Int8Array::from_iter_values([0, 0, 0, 0]);
+        let dict_arr: Arc<dyn Array> =
+            Arc::new(DictionaryArray::<Int8Type>::try_new(keys, Arc::new(value)).unwrap());
+
+        let mut encoder = MockEncoder::default();
+
+        let result = encode_value(&mut encoder, &dict_arr, 2, &Type::TEXT, FieldFormat::Text);
+
+        assert!(result.is_ok());
+
+        assert!(encoder.encoded_value == val);
+    }
+}
diff --git a/datafusion-postgres/src/handlers.rs b/datafusion-postgres/src/handlers.rs
@@ -2,7 +2,10 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use crate::auth::{AuthManager, Permission, ResourceType};
-use crate::sql::{parse, rewrite, AliasDuplicatedProjectionRewrite, SqlStatementRewriteRule};
+use crate::sql::{
+    parse, rewrite, AliasDuplicatedProjectionRewrite, RemoveUnsupportedTypes,
+    ResolveUnqualifiedIdentifer, SqlStatementRewriteRule,
+};
 use async_trait::async_trait;
 use datafusion::arrow::datatypes::DataType;
 use datafusion::logical_expr::LogicalPlan;
@@ -74,8 +77,11 @@ impl DfSessionService {
         session_context: Arc<SessionContext>,
         auth_manager: Arc<AuthManager>,
     ) -> DfSessionService {
-        let sql_rewrite_rules: Vec<Arc<dyn SqlStatementRewriteRule>> =
-            vec![Arc::new(AliasDuplicatedProjectionRewrite)];
+        let sql_rewrite_rules: Vec<Arc<dyn SqlStatementRewriteRule>> = vec![
+            Arc::new(AliasDuplicatedProjectionRewrite),
+            Arc::new(ResolveUnqualifiedIdentifer),
+            Arc::new(RemoveUnsupportedTypes::new()),
+        ];
         let parser = Arc::new(Parser {
             session_context: session_context.clone(),
             sql_rewrite_rules: sql_rewrite_rules.clone(),
diff --git a/datafusion-postgres/src/sql.rs b/datafusion-postgres/src/sql.rs
@@ -1,12 +1,18 @@
+use std::collections::HashSet;
 use std::sync::Arc;
 
 use datafusion::sql::sqlparser::ast::Expr;
 use datafusion::sql::sqlparser::ast::Ident;
+use datafusion::sql::sqlparser::ast::OrderByKind;
+use datafusion::sql::sqlparser::ast::Query;
 use datafusion::sql::sqlparser::ast::Select;
 use datafusion::sql::sqlparser::ast::SelectItem;
 use datafusion::sql::sqlparser::ast::SelectItemQualifiedWildcardKind;
 use datafusion::sql::sqlparser::ast::SetExpr;
 use datafusion::sql::sqlparser::ast::Statement;
+use datafusion::sql::sqlparser::ast::TableFactor;
+use datafusion::sql::sqlparser::ast::TableWithJoins;
+use datafusion::sql::sqlparser::ast::Value;
 use datafusion::sql::sqlparser::dialect::PostgreSqlDialect;
 use datafusion::sql::sqlparser::parser::Parser;
 use datafusion::sql::sqlparser::parser::ParserError;
@@ -135,6 +141,188 @@ impl SqlStatementRewriteRule for AliasDuplicatedProjectionRewrite {
     }
 }
 
+/// Prepend qualifier for order by or filter when there is qualified wildcard
+///
+/// Postgres allows unqualified identifier in ORDER BY and FILTER but it's not
+/// accepted by datafusion.
+#[derive(Debug)]
+pub struct ResolveUnqualifiedIdentifer;
+
+impl ResolveUnqualifiedIdentifer {
+    fn rewrite_unqualified_identifiers(query: &mut Box<Query>) {
+        if let SetExpr::Select(select) = query.body.as_mut() {
+            // Step 1: Find all table aliases from FROM and JOIN clauses.
+            let table_aliases = Self::get_table_aliases(&select.from);
+
+            // Step 2: Check for a single qualified wildcard in the projection.
+            let qualified_wildcard_alias = Self::get_qualified_wildcard_alias(&select.projection);
+            if qualified_wildcard_alias.is_none() || table_aliases.is_empty() {
+                return; // Conditions not met.
+            }
+
+            let wildcard_alias = qualified_wildcard_alias.unwrap();
+
+            // Step 3: Rewrite expressions in the WHERE and ORDER BY clauses.
+            if let Some(selection) = &mut select.selection {
+                Self::rewrite_expr(selection, &wildcard_alias, &table_aliases);
+            }
+
+            if let Some(OrderByKind::Expressions(order_by_exprs)) =
+                query.order_by.as_mut().map(|o| &mut o.kind)
+            {
+                for order_by_expr in order_by_exprs {
+                    Self::rewrite_expr(&mut order_by_expr.expr, &wildcard_alias, &table_aliases);
+                }
+            }
+        }
+    }
+
+    fn get_table_aliases(tables: &[TableWithJoins]) -> HashSet<String> {
+        let mut aliases = HashSet::new();
+        for table_with_joins in tables {
+            if let TableFactor::Table {
+                alias: Some(alias), ..
+            } = &table_with_joins.relation
+            {
+                aliases.insert(alias.name.value.clone());
+            }
+            for join in &table_with_joins.joins {
+                if let TableFactor::Table {
+                    alias: Some(alias), ..
+                } = &join.relation
+                {
+                    aliases.insert(alias.name.value.clone());
+                }
+            }
+        }
+        aliases
+    }
+
+    fn get_qualified_wildcard_alias(projection: &[SelectItem]) -> Option<String> {
+        let mut qualified_wildcards = projection
+            .iter()
+            .filter_map(|item| {
+                if let SelectItem::QualifiedWildcard(
+                    SelectItemQualifiedWildcardKind::ObjectName(objname),
+                    _,
+                ) = item
+                {
+                    Some(
+                        objname
+                            .0
+                            .iter()
+                            .map(|v| v.as_ident().unwrap().value.clone())
+                            .collect::<Vec<_>>()
+                            .join("."),
+                    )
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>();
+
+        if qualified_wildcards.len() == 1 {
+            Some(qualified_wildcards.remove(0))
+        } else {
+            None
+        }
+    }
+
+    fn rewrite_expr(expr: &mut Expr, wildcard_alias: &str, table_aliases: &HashSet<String>) {
+        match expr {
+            Expr::Identifier(ident) => {
+                // If the identifier is not a table alias itself, rewrite it.
+                if !table_aliases.contains(&ident.value) {
+                    *expr = Expr::CompoundIdentifier(vec![
+                        Ident::new(wildcard_alias.to_string()),
+                        ident.clone(),
+                    ]);
+                }
+            }
+            Expr::BinaryOp { left, right, .. } => {
+                Self::rewrite_expr(left, wildcard_alias, table_aliases);
+                Self::rewrite_expr(right, wildcard_alias, table_aliases);
+            }
+            // Add more cases for other expression types as needed (e.g., `InList`, `Between`, etc.)
+            _ => {}
+        }
+    }
+}
+
+impl SqlStatementRewriteRule for ResolveUnqualifiedIdentifer {
+    fn rewrite(&self, mut statement: Statement) -> Statement {
+        if let Statement::Query(query) = &mut statement {
+            Self::rewrite_unqualified_identifiers(query);
+        }
+
+        statement
+    }
+}
+
+/// Remove datafusion unsupported type annotations
+#[derive(Debug)]
+pub struct RemoveUnsupportedTypes {
+    unsupported_types: HashSet<String>,
+}
+
+impl RemoveUnsupportedTypes {
+    pub fn new() -> Self {
+        let mut unsupported_types = HashSet::new();
+        unsupported_types.insert("regclass".to_owned());
+
+        Self { unsupported_types }
+    }
+
+    fn rewrite_expr_unsupported_types(&self, expr: &mut Expr) {
+        match expr {
+            // This is the key part: identify constants with type annotations.
+            Expr::TypedString { value, data_type } => {
+                if self
+                    .unsupported_types
+                    .contains(data_type.to_string().to_lowercase().as_str())
+                {
+                    *expr =
+                        Expr::Value(Value::SingleQuotedString(value.to_string()).with_empty_span());
+                }
+            }
+            Expr::Cast {
+                data_type,
+                expr: value,
+                ..
+            } => {
+                if self
+                    .unsupported_types
+                    .contains(data_type.to_string().to_lowercase().as_str())
+                {
+                    *expr = *value.clone();
+                }
+            }
+            // Handle binary operations by recursively rewriting both sides.
+            Expr::BinaryOp { left, right, .. } => {
+                self.rewrite_expr_unsupported_types(left);
+                self.rewrite_expr_unsupported_types(right);
+            }
+            // Add more match arms for other expression types (e.g., `Function`, `InList`) as needed.
+            _ => {}
+        }
+    }
+}
+
+impl SqlStatementRewriteRule for RemoveUnsupportedTypes {
+    fn rewrite(&self, mut s: Statement) -> Statement {
+        // Traverse the AST to find the WHERE clause and rewrite it.
+        if let Statement::Query(query) = &mut s {
+            if let SetExpr::Select(select) = query.body.as_mut() {
+                if let Some(expr) = &mut select.selection {
+                    self.rewrite_expr_unsupported_types(expr);
+                }
+            }
+        }
+
+        s
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -180,4 +368,61 @@ mod tests {
             "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"
         );
     }
+
+    #[test]
+    fn test_qualifier_prepend() {
+        let rules: Vec<Arc<dyn SqlStatementRewriteRule>> =
+            vec![Arc::new(ResolveUnqualifiedIdentifer)];
+
+        let sql = "SELECT n.* FROM pg_catalog.pg_namespace n WHERE nspname = 'pg_catalog' ORDER BY nspname";
+        let statement = parse(sql).expect("Failed to parse").remove(0);
+
+        let statement = rewrite(statement, &rules);
+        assert_eq!(
+            statement.to_string(),
+            "SELECT n.* FROM pg_catalog.pg_namespace AS n WHERE n.nspname = 'pg_catalog' ORDER BY n.nspname"
+        );
+
+        let sql = "SELECT * FROM pg_catalog.pg_namespace ORDER BY nspname";
+        let statement = parse(sql).expect("Failed to parse").remove(0);
+
+        let statement = rewrite(statement, &rules);
+        assert_eq!(
+            statement.to_string(),
+            "SELECT * FROM pg_catalog.pg_namespace ORDER BY nspname"
+        );
+
+        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, 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 n.nspsname"
+        );
+    }
+
+    #[test]
+    fn test_remove_unsupported_types() {
+        let rules: Vec<Arc<dyn SqlStatementRewriteRule>> =
+            vec![Arc::new(RemoveUnsupportedTypes::new())];
+
+        let sql = "SELECT n.* FROM pg_catalog.pg_namespace n WHERE n.nspname = 'pg_catalog'::regclass ORDER BY n.nspname";
+        let statement = parse(sql).expect("Failed to parse").remove(0);
+
+        let statement = rewrite(statement, &rules);
+        assert_eq!(
+            statement.to_string(),
+            "SELECT n.* FROM pg_catalog.pg_namespace AS n WHERE n.nspname = 'pg_catalog' ORDER BY n.nspname"
+        );
+
+        let sql = "SELECT n.* FROM pg_catalog.pg_namespace n WHERE n.nspname = 'pg_catalog' ORDER BY n.nspname";
+        let statement = parse(sql).expect("Failed to parse").remove(0);
+
+        let statement = rewrite(statement, &rules);
+        assert_eq!(
+            statement.to_string(),
+            "SELECT n.* FROM pg_catalog.pg_namespace AS n WHERE n.nspname = 'pg_catalog' ORDER BY n.nspname"
+        );
+    }
 }
