fix: refactor has_privilege udfs

Author: sunng87

datafusion-postgres/src/pg_catalog.rs

@@ -4,8 +4,7 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 use datafusion::arrow::array::{
-    as_boolean_array, ArrayRef, BooleanArray, BooleanBuilder, RecordBatch, StringArray,
-    StringBuilder,
+    as_boolean_array, ArrayRef, BooleanBuilder, RecordBatch, StringArray, StringBuilder,
 };
 use datafusion::arrow::datatypes::{DataType, Field, SchemaRef};
 use datafusion::arrow::ipc::reader::FileReader;
@@ -21,6 +20,7 @@ use postgres_types::Oid;
 use tokio::sync::RwLock;
 
 mod empty_table;
+mod has_privilege_udf;
 mod pg_attribute;
 mod pg_class;
 mod pg_database;
@@ -863,62 +863,6 @@ pub fn create_pg_table_is_visible(name: &str) -> 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 _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)
-        let mut builder = BooleanArray::builder(user.len());
-        for _ in 0..user.len() {
-            builder.append_value(true);
-        }
-
-        let array: ArrayRef = Arc::new(builder.finish());
-
-        Ok(ColumnarValue::Array(array))
-    };
-
-    // Wrap the implementation in a scalar function
-    create_udf(
-        name,
-        vec![DataType::Utf8, DataType::Utf8, DataType::Utf8],
-        DataType::Boolean,
-        Volatility::Stable,
-        Arc::new(func),
-    )
-}
-
-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 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(obj.len());
-        for _ in 0..obj.len() {
-            builder.append_value(true);
-        }
-        let array: ArrayRef = Arc::new(builder.finish());
-
-        Ok(ColumnarValue::Array(array))
-    };
-
-    // Wrap the implementation in a scalar function
-    create_udf(
-        name,
-        vec![DataType::Utf8, DataType::Utf8],
-        DataType::Boolean,
-        Volatility::Stable,
-        Arc::new(func),
-    )
-}
-
 pub fn create_format_type_udf() -> ScalarUDF {
     let func = move |args: &[ColumnarValue]| {
         let args = ColumnarValue::values_to_arrays(args)?;
@@ -1015,36 +959,30 @@ 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_privilege_2param_udf("has_table_privilege"));
-    session_context.register_udf(create_has_privilege_2param_udf(
-        "pg_catalog.has_table_privilege",
+    session_context.register_udf(has_privilege_udf::create_has_privilege_udf(
+        "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(
+    session_context.register_udf(has_privilege_udf::create_has_privilege_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(
+    session_context.register_udf(has_privilege_udf::create_has_privilege_udf(
+        "has_schema_privilege",
+    ));
+    session_context.register_udf(has_privilege_udf::create_has_privilege_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(
+    session_context.register_udf(has_privilege_udf::create_has_privilege_udf(
+        "has_any_column_privilege",
+    ));
+    session_context.register_udf(has_privilege_udf::create_has_privilege_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());
-    session_context.register_udf(pg_get_expr_udf::PgGetExprUDF::new().into_scalar_udf());
+    session_context.register_udf(pg_get_expr_udf::create_pg_get_expr_udf());
     session_context.register_udf(create_pg_get_partkeydef_udf());
 
     Ok(())

datafusion-postgres/src/pg_catalog/has_privilege_udf.rs

@@ -0,0 +1,71 @@
+use std::sync::Arc;
+
+use datafusion::arrow::array::{ArrayRef, BooleanArray};
+use datafusion::error::Result;
+use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF};
+use datafusion::{
+    arrow::datatypes::DataType,
+    logical_expr::{ScalarUDFImpl, Signature, TypeSignature, Volatility},
+};
+
+#[derive(Debug)]
+pub struct PgHasPrivilegeUDF {
+    signature: Signature,
+    name: String,
+}
+
+impl PgHasPrivilegeUDF {
+    pub(crate) fn new(name: &str) -> PgHasPrivilegeUDF {
+        Self {
+            signature: Signature::one_of(
+                vec![
+                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8, DataType::Utf8]),
+                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Utf8]),
+                ],
+                Volatility::Stable,
+            ),
+            name: name.to_owned(),
+        }
+    }
+
+    pub fn into_scalar_udf(self) -> ScalarUDF {
+        ScalarUDF::new_from_impl(self)
+    }
+}
+
+impl ScalarUDFImpl for PgHasPrivilegeUDF {
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(DataType::Boolean)
+    }
+
+    fn name(&self) -> &str {
+        self.name.as_ref()
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        let args = ColumnarValue::values_to_arrays(&args.args)?;
+
+        let len = args[0].len();
+
+        // For now, always return true (full access for current user)
+        let mut builder = BooleanArray::builder(len);
+        for _ in 0..len {
+            builder.append_value(true);
+        }
+        let array: ArrayRef = Arc::new(builder.finish());
+
+        Ok(ColumnarValue::Array(array))
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+pub fn create_has_privilege_udf(name: &str) -> ScalarUDF {
+    PgHasPrivilegeUDF::new(name).into_scalar_udf()
+}

datafusion-postgres/src/pg_catalog/pg_get_expr_udf.rs

@@ -72,3 +72,7 @@ impl ScalarUDFImpl for PgGetExprUDF {
         self
     }
 }
+
+pub fn create_pg_get_expr_udf() -> ScalarUDF {
+    PgGetExprUDF::new().into_scalar_udf()
+}