feat: improve support for current_user and privilege functions

Signed-off-by: Ning Sun <[email protected]>

Author: sunng87

datafusion-postgres/src/handlers.rs

@@ -3,9 +3,10 @@ use std::sync::Arc;
 
 use crate::auth::{AuthManager, Permission, ResourceType};
 use crate::sql::{
-    parse, rewrite, AliasDuplicatedProjectionRewrite, BlacklistSqlRewriter, FixArrayLiteral,
-    PrependUnqualifiedPgTableName, RemoveTableFunctionQualifier, RemoveUnsupportedTypes,
-    ResolveUnqualifiedIdentifer, RewriteArrayAnyAllOperation, SqlStatementRewriteRule,
+    parse, rewrite, AliasDuplicatedProjectionRewrite, BlacklistSqlRewriter,
+    CurrentUserVariableToSessionUserFunctionCall, FixArrayLiteral, PrependUnqualifiedPgTableName,
+    RemoveTableFunctionQualifier, RemoveUnsupportedTypes, ResolveUnqualifiedIdentifer,
+    RewriteArrayAnyAllOperation, SqlStatementRewriteRule,
 };
 use async_trait::async_trait;
 use datafusion::arrow::datatypes::DataType;
@@ -107,6 +108,7 @@ impl DfSessionService {
             Arc::new(PrependUnqualifiedPgTableName),
             Arc::new(FixArrayLiteral),
             Arc::new(RemoveTableFunctionQualifier),
+            Arc::new(CurrentUserVariableToSessionUserFunctionCall),
         ];
         let parser = Arc::new(Parser {
             session_context: session_context.clone(),

datafusion-postgres/src/pg_catalog.rs

@@ -836,7 +836,7 @@ pub fn create_pg_get_userbyid_udf() -> ScalarUDF {
     )
 }
 
-pub fn create_pg_table_is_visible() -> ScalarUDF {
+pub fn create_pg_table_is_visible(name: &str) -> ScalarUDF {
     // Define the function implementation
     let func = move |args: &[ColumnarValue]| {
         let args = ColumnarValue::values_to_arrays(args)?;
@@ -855,20 +855,20 @@ pub fn create_pg_table_is_visible() -> ScalarUDF {
 
     // Wrap the implementation in a scalar function
     create_udf(
-        "pg_catalog.pg_table_is_visible",
+        name,
         vec![DataType::Int32],
         DataType::Boolean,
         Volatility::Stable,
         Arc::new(func),
     )
 }
 
-pub fn create_has_table_privilege_3param_udf() -> ScalarUDF {
+pub fn create_has_privilege_3param_udf(name: &str) -> ScalarUDF {
     // Define the function implementation for 3-parameter version
     let func = move |args: &[ColumnarValue]| {
         let args = ColumnarValue::values_to_arrays(args)?;
         let user = &args[0]; // User (can be name or OID)
-        let _table = &args[1]; // Table (can be name or OID)
+        let _obj = &args[1]; // Table (can be name or OID)
         let _privilege = &args[2]; // Privilege type (SELECT, INSERT, etc.)
 
         // For now, always return true (full access)
@@ -884,24 +884,24 @@ pub fn create_has_table_privilege_3param_udf() -> ScalarUDF {
 
     // Wrap the implementation in a scalar function
     create_udf(
-        "has_table_privilege",
+        name,
         vec![DataType::Utf8, DataType::Utf8, DataType::Utf8],
         DataType::Boolean,
         Volatility::Stable,
         Arc::new(func),
     )
 }
 
-pub fn create_has_table_privilege_2param_udf() -> ScalarUDF {
+pub fn create_has_privilege_2param_udf(name: &str) -> ScalarUDF {
     // Define the function implementation for 2-parameter version (current user, table, privilege)
     let func = move |args: &[ColumnarValue]| {
         let args = ColumnarValue::values_to_arrays(args)?;
-        let table = &args[0]; // Table (can be name or OID)
+        let obj = &args[0]; // Table (can be name or OID)
         let _privilege = &args[1]; // Privilege type (SELECT, INSERT, etc.)
 
         // For now, always return true (full access for current user)
-        let mut builder = BooleanArray::builder(table.len());
-        for _ in 0..table.len() {
+        let mut builder = BooleanArray::builder(obj.len());
+        for _ in 0..obj.len() {
             builder.append_value(true);
         }
         let array: ArrayRef = Arc::new(builder.finish());
@@ -911,7 +911,7 @@ pub fn create_has_table_privilege_2param_udf() -> ScalarUDF {
 
     // Wrap the implementation in a scalar function
     create_udf(
-        "has_table_privilege",
+        name,
         vec![DataType::Utf8, DataType::Utf8],
         DataType::Boolean,
         Volatility::Stable,
@@ -970,7 +970,6 @@ pub fn create_pg_get_partkeydef_udf() -> ScalarUDF {
         let args = ColumnarValue::values_to_arrays(args)?;
         let oid = &args[0];
 
-        // For now, always return true (full access for current user)
         let mut builder = StringBuilder::new();
         for _ in 0..oid.len() {
             builder.append_value("");
@@ -1016,9 +1015,32 @@ pub fn setup_pg_catalog(
     session_context.register_udf(create_current_schemas_udf("pg_catalog.current_schemas"));
     session_context.register_udf(create_version_udf());
     session_context.register_udf(create_pg_get_userbyid_udf());
-    session_context.register_udf(create_has_table_privilege_2param_udf());
-    session_context.register_udf(create_has_table_privilege_3param_udf());
-    session_context.register_udf(create_pg_table_is_visible());
+    session_context.register_udf(create_has_privilege_2param_udf("has_table_privilege"));
+    session_context.register_udf(create_has_privilege_2param_udf(
+        "pg_catalog.has_table_privilege",
+    ));
+    session_context.register_udf(create_has_privilege_2param_udf("has_schema_privilege"));
+    session_context.register_udf(create_has_privilege_2param_udf(
+        "pg_catalog.has_schema_privilege",
+    ));
+    session_context.register_udf(create_has_privilege_2param_udf("has_any_column_privilege"));
+    session_context.register_udf(create_has_privilege_2param_udf(
+        "pg_catalog.has_any_column_privilege",
+    ));
+    session_context.register_udf(create_has_privilege_3param_udf("has_table_privilege"));
+    session_context.register_udf(create_has_privilege_3param_udf(
+        "pg_catalog.has_table_privilege",
+    ));
+    session_context.register_udf(create_has_privilege_3param_udf("has_schema_privilege"));
+    session_context.register_udf(create_has_privilege_3param_udf(
+        "pg_catalog.has_schema_privilege",
+    ));
+    session_context.register_udf(create_has_privilege_3param_udf("has_any_column_privilege"));
+    session_context.register_udf(create_has_privilege_3param_udf(
+        "pg_catalog.has_any_column_privilege",
+    ));
+    session_context.register_udf(create_pg_table_is_visible("pg_catalog"));
+    session_context.register_udf(create_pg_table_is_visible("pg_catalog.pg_table_is_visible"));
     session_context.register_udf(create_format_type_udf());
     session_context.register_udf(create_session_user_udf());
     session_context.register_udtf("pg_get_keywords", static_tables.pg_get_keywords.clone());

datafusion-postgres/src/sql.rs

@@ -613,6 +613,57 @@ impl SqlStatementRewriteRule for RemoveTableFunctionQualifier {
     }
 }
 
+/// Replace `current_user` with `session_user()`
+#[derive(Debug)]
+pub struct CurrentUserVariableToSessionUserFunctionCall;
+
+struct CurrentUserVariableToSessionUserFunctionCallVisitor;
+
+impl VisitorMut for CurrentUserVariableToSessionUserFunctionCallVisitor {
+    type Break = ();
+
+    fn pre_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow<Self::Break> {
+        if let Expr::Identifier(ident) = expr {
+            if ident.quote_style.is_none() && ident.value.to_lowercase() == "current_user" {
+                *expr = Expr::Function(Function {
+                    name: ObjectName::from(vec![Ident::new("session_user")]),
+                    args: FunctionArguments::None,
+                    uses_odbc_syntax: false,
+                    parameters: FunctionArguments::None,
+                    filter: None,
+                    null_treatment: None,
+                    over: None,
+                    within_group: vec![],
+                });
+            }
+        }
+
+        if let Expr::Function(func) = expr {
+            let fname = func
+                .name
+                .0
+                .iter()
+                .map(|ident| ident.to_string())
+                .collect::<Vec<String>>()
+                .join(".");
+            if fname.to_lowercase() == "current_user" {
+                func.name = ObjectName::from(vec![Ident::new("session_user")])
+            }
+        }
+
+        ControlFlow::Continue(())
+    }
+}
+
+impl SqlStatementRewriteRule for CurrentUserVariableToSessionUserFunctionCall {
+    fn rewrite(&self, mut s: Statement) -> Statement {
+        let mut visitor = CurrentUserVariableToSessionUserFunctionCallVisitor;
+
+        let _ = s.visit(&mut visitor);
+        s
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -802,4 +853,20 @@ mod tests {
             "SELECT * FROM pg_get_keywords()"
         );
     }
+
+    #[test]
+    fn test_current_user() {
+        let rules: Vec<Arc<dyn SqlStatementRewriteRule>> =
+            vec![Arc::new(CurrentUserVariableToSessionUserFunctionCall)];
+
+        assert_rewrite!(&rules, "SELECT current_user", "SELECT session_user");
+
+        assert_rewrite!(&rules, "SELECT CURRENT_USER", "SELECT session_user");
+
+        assert_rewrite!(
+            &rules,
+            "SELECT is_null(current_user)",
+            "SELECT is_null(session_user)"
+        );
+    }
 }