From a366b4bd9f7497a530c57d334502dd06047b01aa Mon Sep 17 00:00:00 2001
From: Ning Sun <sunning@greptime.com>
Date: Wed, 8 Oct 2025 03:09:39 +0800
Subject: [PATCH] fix: hook return for multi statement in simple query

---
 datafusion-postgres/src/handlers.rs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/datafusion-postgres/src/handlers.rs b/datafusion-postgres/src/handlers.rs
index a3a8374..8d4474e 100644
--- a/datafusion-postgres/src/handlers.rs
+++ b/datafusion-postgres/src/handlers.rs
@@ -471,7 +471,7 @@ impl SimpleQueryHandler for DfSessionService {
         }
 
         let mut results = vec![];
-        for statement in statements {
+        'stmt: for statement in statements {
             // TODO: improve statement check by using statement directly
             let query = statement.to_string();
             let query_lower = query.to_lowercase().trim().to_string();
@@ -495,7 +495,8 @@ impl SimpleQueryHandler for DfSessionService {
                     .handle_query(&statement, &self.session_context, client)
                     .await
                 {
-                    return result.map(|response| vec![response]);
+                    results.push(result?);
+                    break 'stmt;
                 }
             }