feat: implment extended query part

sunng87 · sunng87 · commit 0b8d952c4d07 · 2025-10-12T11:36:11.000+08:00
diff --git a/datafusion-postgres/src/handlers.rs b/datafusion-postgres/src/handlers.rs
@@ -25,6 +25,7 @@ use pgwire::messages::response::TransactionStatus;
 
 use crate::auth::AuthManager;
 use crate::client;
+use crate::hooks::set_show::SetShowHook;
 use crate::hooks::QueryHook;
 use arrow_pg::datatypes::df;
 use arrow_pg::datatypes::{arrow_schema_to_pg_fields, into_pg_type};
@@ -43,12 +44,18 @@ pub struct HandlerFactory {
 }
 
 impl HandlerFactory {
-    pub fn new(
+    pub fn new(session_context: Arc<SessionContext>, auth_manager: Arc<AuthManager>) -> Self {
+        let session_service =
+            Arc::new(DfSessionService::new(session_context, auth_manager.clone()));
+        HandlerFactory { session_service }
+    }
+
+    pub fn new_with_hooks(
         session_context: Arc<SessionContext>,
         auth_manager: Arc<AuthManager>,
         query_hooks: Vec<Arc<dyn QueryHook>>,
     ) -> Self {
-        let session_service = Arc::new(DfSessionService::new(
+        let session_service = Arc::new(DfSessionService::new_with_hooks(
             session_context,
             auth_manager.clone(),
             query_hooks,
@@ -98,6 +105,14 @@ impl DfSessionService {
     pub fn new(
         session_context: Arc<SessionContext>,
         auth_manager: Arc<AuthManager>,
+    ) -> DfSessionService {
+        let hooks: Vec<Arc<dyn QueryHook>> = vec![Arc::new(SetShowHook)];
+        Self::new_with_hooks(session_context, auth_manager, hooks)
+    }
+
+    pub fn new_with_hooks(
+        session_context: Arc<SessionContext>,
+        auth_manager: Arc<AuthManager>,
         query_hooks: Vec<Arc<dyn QueryHook>>,
     ) -> DfSessionService {
         let parser = Arc::new(Parser {
@@ -601,7 +616,6 @@ impl Parser {
 
         // show statement may not be supported by datafusion
         if sql_trimmed.starts_with("show") {
-            // Return a dummy plan for transaction commands - they'll be handled by transaction handler
             let show_schema =
                 Arc::new(Schema::new(vec![Field::new("show", DataType::Utf8, false)]));
             let df_schema = show_schema.to_dfschema()?;
diff --git a/datafusion-postgres/src/hooks/mod.rs b/datafusion-postgres/src/hooks/mod.rs
@@ -5,7 +5,7 @@ use async_trait::async_trait;
 use datafusion::common::ParamValues;
 use datafusion::logical_expr::LogicalPlan;
 use datafusion::prelude::SessionContext;
-use datafusion::sql::sqlparser;
+use datafusion::sql::sqlparser::ast::Statement;
 use pgwire::api::results::Response;
 use pgwire::api::ClientInfo;
 use pgwire::error::PgWireResult;
@@ -15,23 +15,23 @@ pub trait QueryHook: Send + Sync {
     /// called in simple query handler to return response directly
     async fn handle_simple_query(
         &self,
-        statement: &sqlparser::ast::Statement,
+        statement: &Statement,
         session_context: &SessionContext,
         client: &mut (dyn ClientInfo + Send + Sync),
     ) -> Option<PgWireResult<Response>>;
 
     /// called at extended query parse phase, for generating `LogicalPlan`from statement
     async fn handle_extended_parse_query(
         &self,
-        statement: &sqlparser::ast::Statement,
+        sql: &Statement,
         session_context: &SessionContext,
         client: &(dyn ClientInfo + Send + Sync),
     ) -> Option<PgWireResult<LogicalPlan>>;
 
     /// called at extended query execute phase, for query execution
     async fn handle_extended_query(
         &self,
-        statement: &sqlparser::ast::Statement,
+        statement: &Statement,
         logical_plan: &LogicalPlan,
         params: &ParamValues,
         session_context: &SessionContext,
diff --git a/datafusion-postgres/src/hooks/set_show.rs b/datafusion-postgres/src/hooks/set_show.rs
@@ -1,15 +1,13 @@
 use std::sync::Arc;
 
 use async_trait::async_trait;
-use datafusion::common::ParamValues;
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::common::{ParamValues, ToDFSchema};
 use datafusion::logical_expr::LogicalPlan;
 use datafusion::prelude::SessionContext;
 use datafusion::sql::sqlparser::ast::Statement;
 use log::{info, warn};
-use pgwire::api::results::{
-    DataRowEncoder, DescribePortalResponse, DescribeResponse, DescribeStatementResponse,
-    FieldFormat, FieldInfo, QueryResponse, Response, Tag,
-};
+use pgwire::api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse, Response, Tag};
 use pgwire::api::ClientInfo;
 use pgwire::error::{PgWireError, PgWireResult};
 use postgres_types::Type;
@@ -48,22 +46,66 @@ impl QueryHook for SetShowHook {
 
     async fn handle_extended_parse_query(
         &self,
-        statement: &Statement,
-        session_context: &SessionContext,
-        client: &(dyn ClientInfo + Send + Sync),
+        stmt: &Statement,
+        _session_context: &SessionContext,
+        _client: &(dyn ClientInfo + Send + Sync),
     ) -> Option<PgWireResult<LogicalPlan>> {
-        None
+        let sql_lower = stmt.to_string().to_lowercase();
+        let sql_trimmed = sql_lower.trim();
+
+        if sql_trimmed.starts_with("show") {
+            let show_schema =
+                Arc::new(Schema::new(vec![Field::new("show", DataType::Utf8, false)]));
+            let result = show_schema
+                .to_dfschema()
+                .map(|df_schema| {
+                    LogicalPlan::EmptyRelation(datafusion::logical_expr::EmptyRelation {
+                        produce_one_row: true,
+                        schema: Arc::new(df_schema),
+                    })
+                })
+                .map_err(|e| PgWireError::ApiError(Box::new(e)));
+            Some(result)
+        } else if sql_trimmed.starts_with("set") {
+            let show_schema = Arc::new(Schema::new(Vec::<Field>::new()));
+            let result = show_schema
+                .to_dfschema()
+                .map(|df_schema| {
+                    LogicalPlan::EmptyRelation(datafusion::logical_expr::EmptyRelation {
+                        produce_one_row: true,
+                        schema: Arc::new(df_schema),
+                    })
+                })
+                .map_err(|e| PgWireError::ApiError(Box::new(e)));
+            Some(result)
+        } else {
+            None
+        }
     }
 
     async fn handle_extended_query(
         &self,
         statement: &Statement,
-        logical_plan: &LogicalPlan,
-        params: &ParamValues,
+        _logical_plan: &LogicalPlan,
+        _params: &ParamValues,
         session_context: &SessionContext,
         client: &mut (dyn ClientInfo + Send + Sync),
     ) -> Option<PgWireResult<Response>> {
-        None
+        match statement {
+            Statement::Set { .. } => {
+                let query = statement.to_string();
+                let query_lower = query.to_lowercase();
+
+                try_respond_set_statements(client, &query_lower, session_context).await
+            }
+            Statement::ShowVariable { .. } | Statement::ShowStatus { .. } => {
+                let query = statement.to_string();
+                let query_lower = query.to_lowercase();
+
+                try_respond_show_statements(client, &query_lower, session_context).await
+            }
+            _ => None,
+        }
     }
 }
 
diff --git a/datafusion-postgres/src/lib.rs b/datafusion-postgres/src/lib.rs
@@ -89,7 +89,25 @@ pub async fn serve(
     auth_manager: Arc<AuthManager>,
 ) -> Result<(), std::io::Error> {
     // Create the handler factory with authentication
-    let factory = Arc::new(HandlerFactory::new(session_context, auth_manager, vec![]));
+    let factory = Arc::new(HandlerFactory::new(session_context, auth_manager));
+
+    serve_with_handlers(factory, opts).await
+}
+
+/// Serve the Datafusion `SessionContext` with Postgres protocol, using custom
+/// query processing hooks.
+pub async fn serve_with_hooks(
+    session_context: Arc<SessionContext>,
+    opts: &ServerOptions,
+    auth_manager: Arc<AuthManager>,
+    hooks: Vec<Arc<dyn QueryHook>>,
+) -> Result<(), std::io::Error> {
+    // Create the handler factory with authentication
+    let factory = Arc::new(HandlerFactory::new_with_hooks(
+        session_context,
+        auth_manager,
+        hooks,
+    ));
 
     serve_with_handlers(factory, opts).await
 }
diff --git a/datafusion-postgres/src/testing.rs b/datafusion-postgres/src/testing.rs
@@ -21,11 +21,7 @@ pub fn setup_handlers() -> DfSessionService {
     )
     .expect("Failed to setup sesession context");
 
-    DfSessionService::new(
-        Arc::new(session_context),
-        Arc::new(AuthManager::new()),
-        vec![],
-    )
+    DfSessionService::new(Arc::new(session_context), Arc::new(AuthManager::new()))
 }
 
 #[derive(Debug, Default)]
diff --git a/datafusion-postgres/tests/common/mod.rs b/datafusion-postgres/tests/common/mod.rs
@@ -21,11 +21,7 @@ pub fn setup_handlers() -> DfSessionService {
     )
     .expect("Failed to setup sesession context");
 
-    DfSessionService::new(
-        Arc::new(session_context),
-        Arc::new(AuthManager::new()),
-        vec![],
-    )
+    DfSessionService::new(Arc::new(session_context), Arc::new(AuthManager::new()))
 }
 
 #[derive(Debug, Default)]

Original file line number	Diff line number	Diff line change
`@@ -21,11 +21,7 @@ pub fn setup_handlers() -> DfSessionService {`
`21`	`21`	`)`
`22`	`22`	`.expect("Failed to setup sesession context");`
`23`	`23`
`24`		`- DfSessionService::new(`
`25`		`- Arc::new(session_context),`
`26`		`- Arc::new(AuthManager::new()),`
`27`		`- vec![],`
`28`		`- )`
	`24`	`+ DfSessionService::new(Arc::new(session_context), Arc::new(AuthManager::new()))`
`29`	`25`	`}`
`30`	`26`
`31`	`27`	`#[derive(Debug, Default)]`