From 5744b726f3a75b9ab77642f6091024bc2bea2ec1 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Mon, 5 Oct 2026 22:43:56 -0700 Subject: [PATCH] fix: panic in multi-statement queries containing RESET ALL --- integration/rust/tests/integration/reset.rs | 25 ++++++++++++++++- .../frontend/client/query_engine/discard.rs | 16 +++++------ pgdog/src/frontend/router/parser/error.rs | 3 ++ pgdog/src/frontend/router/parser/query/set.rs | 21 +++++++++----- .../src/frontend/router/parser/query/split.rs | 1 + .../router/parser/query/test/test_set.rs | 28 +++++++++++++++++++ 6 files changed, 77 insertions(+), 17 deletions(-) diff --git a/integration/rust/tests/integration/reset.rs b/integration/rust/tests/integration/reset.rs index 0a334eed3..817e97b7f 100644 --- a/integration/rust/tests/integration/reset.rs +++ b/integration/rust/tests/integration/reset.rs @@ -1,5 +1,28 @@ use crate::setup::{admin_sqlx, connections_sqlx}; -use sqlx::Executor; +use sqlx::{Executor, Row}; + +#[tokio::test] +async fn test_npgsql_reset_batch() -> Result<(), sqlx::Error> { + let pools = connections_sqlx().await; + let mut conn = pools[1].acquire().await?; + + conn.execute("SET statement_timeout TO '5s'; SET lock_timeout TO '3s'") + .await?; + let row = conn.fetch_one("SHOW statement_timeout").await?; + assert_eq!(row.get::(0), "5s"); + let row = conn.fetch_one("SHOW lock_timeout").await?; + assert_eq!(row.get::(0), "3s"); + + conn.execute( + "SET SESSION AUTHORIZATION DEFAULT;RESET ALL;CLOSE ALL;UNLISTEN *;SELECT pg_advisory_unlock_all();DISCARD SEQUENCES;DISCARD TEMP", + ).await?; + + let row = conn.fetch_one("SHOW statement_timeout").await?; + assert_eq!(row.get::(0), "0"); + let row = conn.fetch_one("SHOW lock_timeout").await?; + assert_eq!(row.get::(0), "0"); + Ok(()) +} async fn run_reset_single_param() { let pools = connections_sqlx().await; diff --git a/pgdog/src/frontend/client/query_engine/discard.rs b/pgdog/src/frontend/client/query_engine/discard.rs index 08cb84980..c8a98d4e2 100644 --- a/pgdog/src/frontend/client/query_engine/discard.rs +++ b/pgdog/src/frontend/client/query_engine/discard.rs @@ -1,7 +1,7 @@ +use super::fake::FakeResponse; use crate::frontend::{ client::Transaction, client::TransactionType, router::parameter_hints::PGDOG_PIN, }; -use crate::net::{CommandComplete, Protocol, ReadyForQuery}; use super::*; @@ -65,14 +65,12 @@ impl QueryEngine { DiscardTarget::Plans | DiscardTarget::Sequences | DiscardTarget::Temp => {} } - let bytes_sent = context - .stream - .send_many(&[ - CommandComplete::new("DISCARD").message(), - ReadyForQuery::in_transaction(context.in_transaction()).message(), - ]) - .await?; - self.stats.sent(bytes_sent); + self.fake_command_response( + context, + &client_request.messages, + &FakeResponse::command("DISCARD"), + ) + .await?; Ok(()) } diff --git a/pgdog/src/frontend/router/parser/error.rs b/pgdog/src/frontend/router/parser/error.rs index ed884393a..9dd6e757b 100644 --- a/pgdog/src/frontend/router/parser/error.rs +++ b/pgdog/src/frontend/router/parser/error.rs @@ -94,6 +94,9 @@ pub(crate) enum Error { #[error("multi-statement queries cannot mix SET with other commands")] MultiStatementMixedSet, + #[error("unsupported SET statement kind: {0}")] + UnsupportedSetKind(u32), + #[error("multi-query statement cannot be safely executed")] MultiStatementSafety, diff --git a/pgdog/src/frontend/router/parser/query/set.rs b/pgdog/src/frontend/router/parser/query/set.rs index aadf8d240..1178ab2ed 100644 --- a/pgdog/src/frontend/router/parser/query/set.rs +++ b/pgdog/src/frontend/router/parser/query/set.rs @@ -39,7 +39,7 @@ impl QueryParser { } else if stmt.kind == VAR_RESET || stmt.kind == VAR_SET_DEFAULT { None } else { - panic!("parse_set_param called on invalid kind {}", stmt.kind); + return Err(Error::UnsupportedSetKind(stmt.kind)); }; match value { @@ -60,23 +60,28 @@ impl QueryParser { /// /// - All SETs → returns `Ok(Some(Command::Set { .. }))` /// - No SETs → returns `Ok(None)`, caller falls through to default parsing - /// - Mix of SET + non-SET → returns `Err(MultiStatementMixedSet)` - /// - /// In session mode, returns `Ok(Some(Command::Query(..)))` immediately so that - /// all multi-statement queries are forwarded to the server verbatim. + /// - Mix of SET + non-SET, or RESET ALL → returns `Err(MultiStatementMixedSet)` + /// so the caller can split the statements and preserve their order. pub(super) fn try_multi_set<'a>( &self, stmts: impl IntoIterator, context: &QueryParserContext, ) -> Result, Error> { let mut has_other = false; + let mut has_reset_all = false; let params = stmts .into_iter() .filter_map(|stmt| match stmt.stmt() { - Node::VariableSetStmt(stmt) if stmt.kind != VAR_SET_MULTI => { + Node::VariableSetStmt(stmt) + if matches!(stmt.kind, VAR_SET_VALUE | VAR_SET_DEFAULT | VAR_RESET) => + { Some(Self::parse_set_param(stmt)) } + Node::VariableSetStmt(stmt) if stmt.kind == VAR_RESET_ALL => { + has_reset_all = true; + None + } _ => { has_other = true; None @@ -84,7 +89,9 @@ impl QueryParser { }) .collect::, _>>()?; - if params.is_empty() { + if has_reset_all { + Err(Error::MultiStatementMixedSet) + } else if params.is_empty() { Ok(None) } else if has_other { Err(Error::MultiStatementMixedSet) diff --git a/pgdog/src/frontend/router/parser/query/split.rs b/pgdog/src/frontend/router/parser/query/split.rs index 8a444f89e..f45d50fbf 100644 --- a/pgdog/src/frontend/router/parser/query/split.rs +++ b/pgdog/src/frontend/router/parser/query/split.rs @@ -114,6 +114,7 @@ impl QueryParser { Node::VariableSetStmt(_) => (), Node::VariableShowStmt(_) => (), Node::DeallocateStmt(_) => (), + Node::ClosePortalStmt(_) | Node::UnlistenStmt(_) | Node::DiscardStmt(_) => (), Node::VacuumRelation(_) | Node::VacuumStmt(_) => (), Node::PrepareStmt(_) => (), // We intercept prepared statements and handle them ourselves. Node::SelectStmt(_) diff --git a/pgdog/src/frontend/router/parser/query/test/test_set.rs b/pgdog/src/frontend/router/parser/query/test/test_set.rs index 8917f2993..94f2f44fa 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_set.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_set.rs @@ -13,6 +13,34 @@ use crate::{ use super::Error; use super::setup::*; +#[test] +fn test_npgsql_reset_is_split() { + for mut test in [ + QueryParserTest::new(), + QueryParserTest::new_single_shard(&config()), + ] { + let command = test.execute(vec![Query::new( + "SET SESSION AUTHORIZATION DEFAULT;RESET ALL;CLOSE ALL;UNLISTEN *;SELECT pg_advisory_unlock_all();DISCARD SEQUENCES;DISCARD TEMP", + ).into()]); + + let Command::Split(queries) = command else { + panic!("expected Npgsql reset batch to split, got {command:?}"); + }; + assert_eq!( + queries, + [ + "SET session_authorization TO DEFAULT", + "RESET ALL", + "CLOSE ALL", + "UNLISTEN *", + "SELECT pg_advisory_unlock_all()", + "DISCARD SEQUENCES", + "DISCARD TEMP", + ], + ); + } +} + #[test] fn test_mixed_set_passthrough_in_session_mode() { let mut test = QueryParserTest::new_session_mode(&config());