test: include queries for pgcli startup (#149)

* test: include queries for pgcli startup

* test: add compatibility tests for pgcli startup queries

* fix: fallback to stringarray or string if no type information provided

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

* chore: add pgcli on readme

* chore: udpate comments

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

* chore: remove dbg

---------

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

Author: sunng87

README.md

@@ -23,6 +23,7 @@ project.
   - Built-in `pg_catalog` tables
   - Built-in postgres functions for common meta queries
     - [x] DBeaver compatibility
+    - [x] pgcli compatibility
 - `datafusion-postgres-cli`: A cli tool starts a postgres compatible server for
   datafusion supported file formats, just like python's `SimpleHTTPServer`.
 - `arrow-pg`: A data type mapping, encoding/decoding library for arrow and

arrow-pg/src/datatypes/df.rs

@@ -66,9 +66,7 @@ where
         } else if let Some(infer_type) = inferenced_type {
             into_pg_type(infer_type)
         } else {
-            // Default to TEXT for untyped parameters in extended queries
-            // This allows arithmetic operations to work with implicit casting
-            Ok(Type::TEXT)
+            Ok(Type::UNKNOWN)
         }
     }
 
@@ -262,10 +260,28 @@ where
             }
             // TODO: add more advanced types (composite types, ranges, etc.)
             _ => {
-                // Default to string/text for unsupported parameter types
-                // This allows graceful degradation instead of fatal errors
+                // the client didn't provide type information and we are also
+                // unable to inference the type, or it's a type that we haven't
+                // supported:
+                //
+                // In this case we retry to resolve it as String or StringArray
                 let value = portal.parameter::<String>(i, &pg_type)?;
-                deserialized_params.push(ScalarValue::Utf8(value));
+                if let Some(value) = value {
+                    if value.starts_with('{') && value.ends_with('}') {
+                        // Looks like an array
+                        let items = value.trim_matches(|c| c == '{' || c == '}' || c == ' ');
+                        let items = items.split(',').map(|s| s.trim());
+                        let scalar_values: Vec<ScalarValue> = items
+                            .map(|s| ScalarValue::Utf8(Some(s.to_string())))
+                            .collect();
+
+                        deserialized_params.push(ScalarValue::List(
+                            ScalarValue::new_list_nullable(&scalar_values, &DataType::Utf8),
+                        ));
+                    } else {
+                        deserialized_params.push(ScalarValue::Utf8(Some(value)));
+                    }
+                }
             }
         }
     }

datafusion-postgres/src/handlers.rs

@@ -3,16 +3,16 @@ use std::sync::Arc;
 
 use crate::auth::{AuthManager, Permission, ResourceType};
 use crate::sql::{
-    parse, rewrite, AliasDuplicatedProjectionRewrite, FixArrayLiteral, PrependUnqualifiedTableName,
-    RemoveTableFunctionQualifier, RemoveUnsupportedTypes, ResolveUnqualifiedIdentifer,
-    RewriteArrayAnyAllOperation, SqlStatementRewriteRule,
+    parse, rewrite, AliasDuplicatedProjectionRewrite, BlacklistSqlRewriter, FixArrayLiteral,
+    PrependUnqualifiedPgTableName, RemoveTableFunctionQualifier, RemoveUnsupportedTypes,
+    ResolveUnqualifiedIdentifer, RewriteArrayAnyAllOperation, 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 log::warn;
+use log::{info, warn};
 use pgwire::api::auth::noop::NoopStartupHandler;
 use pgwire::api::auth::StartupHandler;
 use pgwire::api::portal::{Format, Portal};
@@ -23,7 +23,7 @@ use pgwire::api::results::{
 };
 use pgwire::api::stmt::QueryParser;
 use pgwire::api::stmt::StoredStatement;
-use pgwire::api::{ClientInfo, PgWireServerHandlers, Type};
+use pgwire::api::{ClientInfo, ErrorHandler, PgWireServerHandlers, Type};
 use pgwire::error::{PgWireError, PgWireResult};
 use pgwire::messages::response::TransactionStatus;
 use tokio::sync::Mutex;
@@ -65,6 +65,21 @@ impl PgWireServerHandlers for HandlerFactory {
     fn startup_handler(&self) -> Arc<impl StartupHandler> {
         Arc::new(SimpleStartupHandler)
     }
+
+    fn error_handler(&self) -> Arc<impl ErrorHandler> {
+        Arc::new(LoggingErrorHandler)
+    }
+}
+
+struct LoggingErrorHandler;
+
+impl ErrorHandler for LoggingErrorHandler {
+    fn on_error<C>(&self, _client: &C, error: &mut PgWireError)
+    where
+        C: ClientInfo,
+    {
+        info!("Sending error: {error}")
+    }
 }
 
 /// The pgwire handler backed by a datafusion `SessionContext`
@@ -82,11 +97,14 @@ impl DfSessionService {
         auth_manager: Arc<AuthManager>,
     ) -> DfSessionService {
         let sql_rewrite_rules: Vec<Arc<dyn SqlStatementRewriteRule>> = vec![
+            // make sure blacklist based rewriter it on the top to prevent sql
+            // being rewritten from other rewriters
+            Arc::new(BlacklistSqlRewriter::new()),
             Arc::new(AliasDuplicatedProjectionRewrite),
             Arc::new(ResolveUnqualifiedIdentifer),
             Arc::new(RemoveUnsupportedTypes::new()),
             Arc::new(RewriteArrayAnyAllOperation),
-            Arc::new(PrependUnqualifiedTableName::new()),
+            Arc::new(PrependUnqualifiedPgTableName),
             Arc::new(FixArrayLiteral),
             Arc::new(RemoveTableFunctionQualifier),
         ];
@@ -649,7 +667,9 @@ impl ExtendedQueryHandler for DfSessionService {
         let param_types = plan
             .get_parameter_types()
             .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
+
         let param_values = df::deserialize_parameters(portal, &ordered_param_types(&param_types))?; // Fixed: Use &param_types
+
         let plan = plan
             .clone()
             .replace_params_with_values(&param_values)
@@ -678,12 +698,10 @@ impl ExtendedQueryHandler for DfSessionService {
                 })?
                 .map_err(|e| PgWireError::ApiError(Box::new(e)))?
             } else {
-                match self.session_context.execute_logical_plan(optimised).await {
-                    Ok(df) => df,
-                    Err(e) => {
-                        return Err(PgWireError::ApiError(Box::new(e)));
-                    }
-                }
+                self.session_context
+                    .execute_logical_plan(optimised)
+                    .await
+                    .map_err(|e| PgWireError::ApiError(Box::new(e)))?
             }
         };
         let resp = df::encode_dataframe(dataframe, &portal.result_column_format).await?;

datafusion-postgres/src/pg_catalog.rs

@@ -23,6 +23,7 @@ use tokio::sync::RwLock;
 mod pg_attribute;
 mod pg_class;
 mod pg_database;
+mod pg_get_expr_udf;
 mod pg_namespace;
 mod pg_settings;
 
@@ -917,32 +918,6 @@ pub fn create_session_user_udf() -> ScalarUDF {
     )
 }
 
-pub fn create_pg_get_expr_udf() -> ScalarUDF {
-    let func = move |args: &[ColumnarValue]| {
-        let args = ColumnarValue::values_to_arrays(args)?;
-        let expr = &args[0];
-        let _oid = &args[1];
-
-        // For now, always return true (full access for current user)
-        let mut builder = StringBuilder::new();
-        for _ in 0..expr.len() {
-            builder.append_value("");
-        }
-
-        let array: ArrayRef = Arc::new(builder.finish());
-
-        Ok(ColumnarValue::Array(array))
-    };
-
-    create_udf(
-        "pg_catalog.pg_get_expr",
-        vec![DataType::Utf8, DataType::Int32],
-        DataType::Utf8,
-        Volatility::Stable,
-        Arc::new(func),
-    )
-}
-
 pub fn create_pg_get_partkeydef_udf() -> ScalarUDF {
     let func = move |args: &[ColumnarValue]| {
         let args = ColumnarValue::values_to_arrays(args)?;
@@ -996,7 +971,7 @@ pub fn setup_pg_catalog(
     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(create_pg_get_expr_udf());
+    session_context.register_udf(pg_get_expr_udf::PgGetExprUDF::new().into_scalar_udf());
     session_context.register_udf(create_pg_get_partkeydef_udf());
 
     Ok(())

datafusion-postgres/src/pg_catalog/pg_get_expr_udf.rs

@@ -0,0 +1,74 @@
+use std::sync::Arc;
+
+use datafusion::arrow::array::{ArrayRef, StringBuilder};
+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 PgGetExprUDF {
+    signature: Signature,
+    name: &'static str,
+}
+
+impl PgGetExprUDF {
+    pub(crate) fn new() -> PgGetExprUDF {
+        Self {
+            signature: Signature::one_of(
+                vec![
+                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Int32]),
+                    TypeSignature::Exact(vec![DataType::Utf8, DataType::UInt32]),
+                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Int64]),
+                    TypeSignature::Exact(vec![DataType::Utf8, DataType::UInt64]),
+                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Int32, DataType::Boolean]),
+                    TypeSignature::Exact(vec![DataType::Utf8, DataType::UInt32, DataType::Boolean]),
+                    TypeSignature::Exact(vec![DataType::Utf8, DataType::Int64, DataType::Boolean]),
+                    TypeSignature::Exact(vec![DataType::Utf8, DataType::UInt64, DataType::Boolean]),
+                ],
+                Volatility::Stable,
+            ),
+            name: "pg_catalog.pg_get_expr",
+        }
+    }
+
+    pub fn into_scalar_udf(self) -> ScalarUDF {
+        ScalarUDF::new_from_impl(self).with_aliases(vec!["pg_get_expr"])
+    }
+}
+
+impl ScalarUDFImpl for PgGetExprUDF {
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(DataType::Utf8)
+    }
+
+    fn name(&self) -> &str {
+        self.name
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        let args = ColumnarValue::values_to_arrays(&args.args)?;
+        let expr = &args[0];
+        let _oid = &args[1];
+
+        // For now, always return true (full access for current user)
+        let mut builder = StringBuilder::new();
+        for _ in 0..expr.len() {
+            builder.append_value("");
+        }
+
+        let array: ArrayRef = Arc::new(builder.finish());
+
+        Ok(ColumnarValue::Array(array))
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}