feat: add more rewrite rules

Author: sunng87

datafusion-postgres/src/handlers.rs

@@ -3,8 +3,9 @@ use std::sync::Arc;
 
 use crate::auth::{AuthManager, Permission, ResourceType};
 use crate::sql::{
-    parse, rewrite, AliasDuplicatedProjectionRewrite, RemoveUnsupportedTypes,
-    ResolveUnqualifiedIdentifer, RewriteArrayAnyOperation, SqlStatementRewriteRule,
+    parse, rewrite, AliasDuplicatedProjectionRewrite, PrependUnqualifiedTableName,
+    RemoveUnsupportedTypes, ResolveUnqualifiedIdentifer, RewriteArrayAnyOperation,
+    SqlStatementRewriteRule,
 };
 use async_trait::async_trait;
 use datafusion::arrow::datatypes::DataType;
@@ -82,6 +83,7 @@ impl DfSessionService {
             Arc::new(ResolveUnqualifiedIdentifer),
             Arc::new(RemoveUnsupportedTypes::new()),
             Arc::new(RewriteArrayAnyOperation),
+            Arc::new(PrependUnqualifiedTableName::new()),
         ];
         let parser = Arc::new(Parser {
             session_context: session_context.clone(),
@@ -297,8 +299,8 @@ impl DfSessionService {
                     Ok(Some(Response::Query(resp)))
                 }
                 "show search_path" => {
-                    let default_catalog = "datafusion";
-                    let resp = Self::mock_show_response("search_path", default_catalog)?;
+                    let default_schema = "public";
+                    let resp = Self::mock_show_response("search_path", default_schema)?;
                     Ok(Some(Response::Query(resp)))
                 }
                 _ => Err(PgWireError::UserError(Box::new(

datafusion-postgres/src/sql.rs

@@ -11,6 +11,7 @@ use datafusion::sql::sqlparser::ast::FunctionArgumentList;
 use datafusion::sql::sqlparser::ast::FunctionArguments;
 use datafusion::sql::sqlparser::ast::Ident;
 use datafusion::sql::sqlparser::ast::ObjectName;
+use datafusion::sql::sqlparser::ast::ObjectNamePart;
 use datafusion::sql::sqlparser::ast::OrderByKind;
 use datafusion::sql::sqlparser::ast::Query;
 use datafusion::sql::sqlparser::ast::Select;
@@ -401,6 +402,63 @@ impl SqlStatementRewriteRule for RewriteArrayAnyOperation {
     }
 }
 
+/// Prepend qualifier to table_name
+///
+/// Postgres has pg_catalog in search_path by default so it allow access to
+/// `pg_namespace` without `pg_catalog.` qualifier
+#[derive(Debug)]
+pub struct PrependUnqualifiedTableName {
+    table_names: HashSet<String>,
+}
+
+impl PrependUnqualifiedTableName {
+    pub fn new() -> Self {
+        let mut table_names = HashSet::new();
+
+        table_names.insert("pg_namespace".to_owned());
+
+        Self { table_names }
+    }
+}
+
+struct PrependUnqualifiedTableNameVisitor<'a> {
+    table_names: &'a HashSet<String>,
+}
+
+impl<'a> VisitorMut for PrependUnqualifiedTableNameVisitor<'a> {
+    type Break = ();
+
+    fn pre_visit_table_factor(
+        &mut self,
+        table_factor: &mut TableFactor,
+    ) -> ControlFlow<Self::Break> {
+        if let TableFactor::Table { name, .. } = table_factor {
+            if name.0.len() == 1 {
+                let ObjectNamePart::Identifier(ident) = &name.0[0];
+                if self.table_names.contains(&ident.to_string()) {
+                    *name = ObjectName(vec![
+                        ObjectNamePart::Identifier(Ident::new("pg_catalog")),
+                        name.0[0].clone(),
+                    ]);
+                }
+            }
+        }
+
+        ControlFlow::Continue(())
+    }
+}
+
+impl SqlStatementRewriteRule for PrependUnqualifiedTableName {
+    fn rewrite(&self, mut s: Statement) -> Statement {
+        let mut visitor = PrependUnqualifiedTableNameVisitor {
+            table_names: &self.table_names,
+        };
+
+        let _ = s.visit(&mut visitor);
+        s
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -524,4 +582,28 @@ mod tests {
             "SELECT a FROM tbl WHERE array_contains(current_schemas(true), a)"
         );
     }
+
+    #[test]
+    fn test_prepend_unqualified_table_name() {
+        let rules: Vec<Arc<dyn SqlStatementRewriteRule>> =
+            vec![Arc::new(PrependUnqualifiedTableName::new())];
+
+        assert_rewrite!(
+            &rules,
+            "SELECT * FROM pg_catalog.pg_namespace",
+            "SELECT * FROM pg_catalog.pg_namespace"
+        );
+
+        assert_rewrite!(
+            &rules,
+            "SELECT * FROM pg_namespace",
+            "SELECT * FROM pg_catalog.pg_namespace"
+        );
+
+        assert_rewrite!(
+            &rules,
+            "SELECT typtype, typname, pg_type.oid FROM pg_catalog.pg_type LEFT JOIN pg_namespace as ns ON ns.oid = oid",
+            "SELECT typtype, typname, pg_type.oid FROM pg_catalog.pg_type LEFT JOIN pg_catalog.pg_namespace AS ns ON ns.oid = oid"
+        );
+    }
 }

datafusion-postgres/tests/dbeaver.rs

@@ -13,7 +13,7 @@ const DBEAVER_QUERIES: &[&str] = &[
     "SELECT typinput='pg_catalog.array_in'::regproc as is_array, typtype, typname, pg_type.oid   FROM pg_catalog.pg_type   LEFT JOIN (select ns.oid as nspoid, ns.nspname, r.r           from pg_namespace as ns           join ( select s.r, (current_schemas(false))[s.r] as nspname                    from generate_series(1, array_upper(current_schemas(false), 1)) as s(r) ) as r          using ( nspname )        ) as sp     ON sp.nspoid = typnamespace  WHERE pg_type.oid = 1034  ORDER BY sp.r, pg_type.oid DESC",
     "SHOW search_path",
     "SELECT db.oid,db.* FROM pg_catalog.pg_database db WHERE datname='postgres'",
-    "SELECT * FROM pg_catalog.pg_settinngs where name='standard_conforming_strings'",
+    "SELECT * FROM pg_catalog.pg_settings where name='standard_conforming_strings'",
     "SELECT string_agg(word, ',' ) from pg_catalog.pg_get_keywords() where word <> ALL ('{a,abs,absolute,action,ada,add,admin,after,all,allocate,alter,aIways,and,any,are,array,as,asc,asenstitive,assertion,assignment,asymmetric,at,atomic,attribute,attributes,authorization,avg,before,begin,bernoulli,between,bigint,binary,blob,boolean,both,breaadth,by,c,call,called,cardinaliity,cascade,cascaded,case,cast,catalog,catalog_name,ceil,ceiling,chain,char,char_length,character,character_length,character_set_catalog,character_set_name,character_set_schema,characteristics,characters,check,checkeed,class_origin,clob,close,coalesce,coboI,code_units,collate,collation,collaition_catalog,collaition_name,collaition_schema,collect,colum,column_name,command_function,command_function_code,commit,committed,condiition,condiition_number,connect,connection_name,constraint,constraint_catalog,constraint_name,constraint_schema,constraints,constructors,contains,continue,convert,corr,correspondiing,count,covar_pop,covar_samp,create,cross,cube,cume_dist,current,current_collation,current_date,current_default_transfom_group,current_path,current_role,current_time,current_timestamp,current_transfom_group_for_type,current_user,cursor,cursor_name,cycle,data,date,datetime_interval_code,datetime_interval_precision,day,deallocate,dec,decimaI,declare,default,defaults,not,null,nullable,nullif,nulls,number,numeric,object,octeet_length,octets,of,old,on,only,open,option,options,or,order,ordering,ordinaliity,others,out,outer,output,over,overlaps,overlay,overriding,pad,parameter,parameter_mode,parameter_name,parameter_ordinal_position,parameter_speciific_catalog,parameter_speciific_name,parameter_speciific_schema,partiaI,partitioon,pascal,path,percent_rank,percentile_cont,percentile_disc,placing,pli,position,power,preceding,precision,prepare,preseerv,primary,prior,privileges,procedure,public,range,rank,read,reads,real,recursivve,ref,references,referencing,regr_avgx,regr_avgy,regr_count,regr_intercept,regr_r2,regr_slope,regr_sxx,regr_sxy,regr_sy y,relative,release,repeatable,restart,result,retun,returned_cardinality,returned_length,returned_octeet_length,returned_sqlstate,returns,revoe,right,role,rollback,rollup,routine,routine_catalog,routine_name,routine_schema,row,row_count,row_number,rows,savepoint,scale,schema,schema_name,scope_catalog,scope_name,scope_schema,scroll,search,second,section,security,select,self,sensitive,sequence,seriializeable,server_name,session,session_user,set,sets,similar,simple,size,smalIint,some,source,space,specifiic,speciific_name,speciifictype,sql,sqlexception,sqlstate,sqlwarning,sqrt,start,state,statement,static,stddev_pop,stddev_samp,structure,style,subclass_origin,submultiset,substring,sum,symmetric,system,system_user,table,table_name,tablesample,temporary,then,ties,time,timesamp,timezone_hour,timezone_minute,to,top_level_count,trailing,transaction,transaction_active,transactions_committed,transactions_rolled_back,transfor,transforms,translate,translation,treat,trigger,trigger_catalog,trigger_name,trigger_schema,trim,true,type,unbounde,undefined,uncommitted,under,union,unique,unknown,unnaamed,unnest,update,upper,usage,user,user_defined_type_catalog,user_defined_type_code,user_defined_type_name,user_defined_type_schema,using,value,values,var_pop,var_samp,varchar,varying,view,when,whenever,where,width_bucket,window,with,within,without,work,write,year,zone",
     "SELECT version()",
     "SELECT * FROM pg_catalog.pg_enum WHERE 1<>1 LIMIT 1",