refactor: remove custom information_schema and use datafusion built-in

Author: sunng87

datafusion-postgres-cli/src/main.rs

@@ -5,7 +5,7 @@ use std::sync::Arc;
 use datafusion::execution::options::{
     ArrowReadOptions, AvroReadOptions, CsvReadOptions, NdJsonReadOptions, ParquetReadOptions,
 };
-use datafusion::prelude::SessionContext;
+use datafusion::prelude::{SessionConfig, SessionContext};
 use datafusion_postgres::{serve, ServerOptions}; // Assuming the crate name is `datafusion_postgres`
 use structopt::StructOpt;
 
@@ -179,7 +179,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let mut opts = Opt::from_args();
     opts.include_directory_files()?;
 
-    let session_context = SessionContext::new();
+    let session_config = SessionConfig::new().with_information_schema(true);
+    let session_context = SessionContext::new_with_config(session_config);
 
     setup_session_context(&session_context, &opts).await?;
 

datafusion-postgres/src/handlers.rs

@@ -19,7 +19,6 @@ use pgwire::api::{ClientInfo, NoopErrorHandler, PgWireServerHandlers, Type};
 use tokio::sync::Mutex;
 
 use crate::datatypes;
-use crate::information_schema::{columns_df, schemata_df, tables_df};
 use pgwire::error::{PgWireError, PgWireResult};
 
 pub struct HandlerFactory(pub Arc<DfSessionService>);
@@ -189,39 +188,6 @@ impl DfSessionService {
             Ok(None)
         }
     }
-
-    async fn try_respond_information_schema<'a>(
-        &self,
-        query_lower: &str,
-    ) -> PgWireResult<Option<Response<'a>>> {
-        if query_lower.contains("information_schema.schemata") {
-            let df = schemata_df(&self.session_context)
-                .await
-                .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
-            let resp = datatypes::encode_dataframe(df, &Format::UnifiedText).await?;
-            return Ok(Some(Response::Query(resp)));
-        } else if query_lower.contains("information_schema.tables") {
-            let df = tables_df(&self.session_context)
-                .await
-                .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
-            let resp = datatypes::encode_dataframe(df, &Format::UnifiedText).await?;
-            return Ok(Some(Response::Query(resp)));
-        } else if query_lower.contains("information_schema.columns") {
-            let df = columns_df(&self.session_context)
-                .await
-                .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
-            let resp = datatypes::encode_dataframe(df, &Format::UnifiedText).await?;
-            return Ok(Some(Response::Query(resp)));
-        }
-
-        // Handle pg_catalog.pg_namespace for pgcli compatibility
-        if query_lower.contains("pg_catalog.pg_namespace") {
-            let resp = self.mock_pg_namespace().await?;
-            return Ok(Some(Response::Query(resp)));
-        }
-
-        Ok(None)
-    }
 }
 
 #[async_trait]
@@ -241,10 +207,6 @@ impl SimpleQueryHandler for DfSessionService {
             return Ok(vec![resp]);
         }
 
-        if let Some(resp) = self.try_respond_information_schema(&query_lower).await? {
-            return Ok(vec![resp]);
-        }
-
         let df = self
             .session_context
             .sql(query)
@@ -361,11 +323,8 @@ impl ExtendedQueryHandler for DfSessionService {
             return Ok(resp);
         }
 
-        if let Some(resp) = self.try_respond_information_schema(&query).await? {
-            return Ok(resp);
-        }
-
         let (_, plan) = &portal.statement.statement;
+
         let param_types = plan
             .get_parameter_types()
             .map_err(|e| PgWireError::ApiError(Box::new(e)))?;

datafusion-postgres/src/information_schema.rs

@@ -1,7 +1,6 @@
 mod datatypes;
 mod encoder;
 mod handlers;
-mod information_schema;
 
 use std::sync::Arc;