From f39f480082ca08e9861d6a64f347c25a5ddce303 Mon Sep 17 00:00:00 2001 From: Christopher Watford Date: Thu, 30 Oct 2025 13:47:25 -0400 Subject: [PATCH 1/2] Fix #227: handle multiple statements with hooks --- datafusion-postgres/src/handlers.rs | 30 ++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/datafusion-postgres/src/handlers.rs b/datafusion-postgres/src/handlers.rs index bc43b4f..3d1014e 100644 --- a/datafusion-postgres/src/handlers.rs +++ b/datafusion-postgres/src/handlers.rs @@ -309,7 +309,7 @@ impl SimpleQueryHandler for DfSessionService { .await { results.push(result?); - break 'stmt; + continue 'stmt; } } @@ -762,4 +762,32 @@ mod tests { let result = hook.handle_simple_query(stmt, &ctx, &mut client).await; assert!(result.is_none()); } + + #[tokio::test] + async fn test_multiple_statements_with_hook_continue() { + // Bug #227: when a hook returned a result, the code used `break 'stmt` + // which would exit the entire statement loop, preventing subsequent statements + // from being processed. + let session_context = Arc::new(SessionContext::new()); + let auth_manager = Arc::new(AuthManager::new()); + + let hooks: Vec> = vec![Arc::new(TestHook)]; + let service = DfSessionService::new_with_hooks(session_context, auth_manager, hooks); + + let mut client = MockClient::new(); + + // Mix of queries with hooks and those without + let query = "SELECT magic; SELECT 1; SELECT magic; SELECT 1"; + + let results = ::do_query(&service, &mut client, query) + .await + .unwrap(); + + assert_eq!(results.len(), 4, "Expected 4 responses"); + + assert!(matches!(results[0], Response::EmptyQuery)); + assert!(matches!(results[1], Response::Query(_))); + assert!(matches!(results[2], Response::EmptyQuery)); + assert!(matches!(results[3], Response::Query(_))); + } } From b5b4dc1cec2521d10c032eb4bad968d983dd6ae4 Mon Sep 17 00:00:00 2001 From: Christopher Watford Date: Fri, 31 Oct 2025 08:57:15 -0400 Subject: [PATCH 2/2] cargo fmt --- datafusion-postgres/src/handlers.rs | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/datafusion-postgres/src/handlers.rs b/datafusion-postgres/src/handlers.rs index 3d1014e..55fafd2 100644 --- a/datafusion-postgres/src/handlers.rs +++ b/datafusion-postgres/src/handlers.rs @@ -767,24 +767,25 @@ mod tests { async fn test_multiple_statements_with_hook_continue() { // Bug #227: when a hook returned a result, the code used `break 'stmt` // which would exit the entire statement loop, preventing subsequent statements - // from being processed. + // from being processed. let session_context = Arc::new(SessionContext::new()); let auth_manager = Arc::new(AuthManager::new()); - + let hooks: Vec> = vec![Arc::new(TestHook)]; let service = DfSessionService::new_with_hooks(session_context, auth_manager, hooks); - + let mut client = MockClient::new(); - + // Mix of queries with hooks and those without let query = "SELECT magic; SELECT 1; SELECT magic; SELECT 1"; - - let results = ::do_query(&service, &mut client, query) - .await - .unwrap(); - + + let results = + ::do_query(&service, &mut client, query) + .await + .unwrap(); + assert_eq!(results.len(), 4, "Expected 4 responses"); - + assert!(matches!(results[0], Response::EmptyQuery)); assert!(matches!(results[1], Response::Query(_))); assert!(matches!(results[2], Response::EmptyQuery));