From 7572b7892e976236c8a6ed0aa07da7dc290c3b46 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 3 Sep 2026 17:47:00 +0800 Subject: [PATCH 1/5] fix(server-api): handle empty queries in extended query protocol like PostgreSQL An empty query string (no statement: "", ";", ";;", whitespace-only) sent with Parse used to be dispatched to the user's QueryParser, which typically fails with a syntax error. Real PostgreSQL accepts it: Parse succeeds, the statement name is remembered as empty, and later Bind/Describe/Execute on that name follow empty-query semantics. Do we need to store the empty statement server-side? Yes, but only as a name marker, not a statement: Parse of an empty query must replace any statement previously stored under the same name, and Bind on that name must succeed and shadow portals of the same name. There is no parsed value to store, so instead of changing PortalStore/StoredStatement/Portal to carry an empty variant (a breaking change for every implementor), the default extended-query handlers track empty statement and portal names in a private per-connection registry kept in SessionExtensions: - on_parse: empty queries never reach QueryParser; the name is marked empty (replacing any stored statement) - on_bind: binding an empty statement succeeds with zero parameters (binding parameters is rejected with 08P01, like PostgreSQL) and shadows portals of the same name - on_execute: an empty portal answers EmptyQueryResponse and never calls do_query; the portal stays valid across repeated Execute - on_describe: an empty statement describes as ParameterDescription (no parameters) + NoData; an empty portal as NoData - on_close/on_sync: drop empty statements/portals like real ones Behavior verified message-for-message against PostgreSQL 18.4 driven over the wire (raw socket probe), and covered by unit tests with a mock client asserting the exact message sequences. --- CHANGELOG.md | 10 + src/api/query.rs | 818 ++++++++++++++++++++++++++++++++++++++++++++++- src/api/stmt.rs | 5 + 3 files changed, 820 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0a1db5f0..08228609 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,16 @@ Versioning](https://semver.org/spec/v2.0.0.html). ### Fixed +- Extended query protocol: empty queries (a query string without any + statement, such as `""` or `";;"`) are now handled like PostgreSQL instead + of being dispatched to the query parser: `Parse` succeeds without calling + `QueryParser`, `Describe` answers `ParameterDescription` (no parameters) + + `NoData`, `Bind` succeeds (rejecting bound parameters with `08P01`), and + `Execute` returns `EmptyQueryResponse` without reaching `do_query`. An + empty `Parse` replaces any statement previously stored under the same name, + and `Close`/`Sync` drop empty statements and the unnamed empty portal like + real ones. Empty queries are tracked internally per connection, so no + changes to `PortalStore`, `StoredStatement`, or `Portal` were required. - Client API: backend messages are now decoded with the rules of the protocol version the client actually advertised, instead of always 3.2. Previously a 4-byte protocol 3.0 cancel key was decoded as `SecretKey::Bytes` instead of diff --git a/src/api/query.rs b/src/api/query.rs index 67583b42..a3276f1b 100644 --- a/src/api/query.rs +++ b/src/api/query.rs @@ -1,7 +1,8 @@ use std::cmp::max; +use std::collections::HashSet; use std::fmt::Debug; use std::ops::Deref; -use std::sync::Arc; +use std::sync::{Arc, RwLock}; use async_trait::async_trait; use futures::channel::oneshot; @@ -48,6 +49,96 @@ where Some(handle.start_query().await) } +/// Tracks statement and portal names that refer to empty queries. +/// +/// Empty queries (a query string without any statement, such as `""` or +/// `";;"`) have no parsed representation, so there is nothing to store in +/// the [`PortalStore`] as a `StoredStatement`. Instead, the default extended +/// query handlers record their names in this per-connection registry, +/// mirroring PostgreSQL's extended-query behavior for empty queries: +/// +/// - `Parse` of an empty query succeeds and replaces any statement previously +/// stored under the same name; +/// - `Bind` on that name succeeds (binding zero parameters), replacing any +/// portal previously stored under the portal name; +/// - `Describe` answers `ParameterDescription` (no parameters) + `NoData`; +/// - `Execute` returns `EmptyQueryResponse` without dispatching to +/// [`ExtendedQueryHandler::do_query`]; +/// - `Close` removes the marker, and `Sync` drops the unnamed empty portal. +/// +/// The registry lives in [`SessionExtensions`](super::SessionExtensions), so +/// this is purely internal bookkeeping: neither the [`PortalStore`] API nor +/// the `StoredStatement`/`Portal` types need to represent an empty variant. +#[derive(Debug, Default)] +struct EmptyStatementRegistry { + statements: RwLock>, + portals: RwLock>, +} + +impl EmptyStatementRegistry { + fn for_client(client: &C) -> Arc { + client + .session_extensions() + .get_or_insert_with(Self::default) + } + + /// Record that `name` was last `Parse`d from an empty query. + fn mark_statement(client: &C, name: &str) { + Self::for_client(client) + .statements + .write() + .unwrap() + .insert(name.to_owned()); + } + + /// Forget an empty-statement marker, e.g. after re-`Parse` of a real + /// statement or a `Close`. + fn unmark_statement(client: &C, name: &str) { + Self::for_client(client) + .statements + .write() + .unwrap() + .remove(name); + } + + /// Test whether `name` is an empty statement. + fn is_empty_statement(client: &C, name: &str) -> bool { + Self::for_client(client) + .statements + .read() + .unwrap() + .contains(name) + } + + /// Record that `name` was `Bind`ed from an empty statement. + fn mark_portal(client: &C, name: &str) { + Self::for_client(client) + .portals + .write() + .unwrap() + .insert(name.to_owned()); + } + + /// Forget an empty-portal marker, e.g. after `Bind` of a real portal or a + /// `Close`. + fn unmark_portal(client: &C, name: &str) { + Self::for_client(client) + .portals + .write() + .unwrap() + .remove(name); + } + + /// Test whether `name` is an empty portal. + fn is_empty_portal(client: &C, name: &str) -> bool { + Self::for_client(client) + .portals + .read() + .unwrap() + .contains(name) + } +} + /// handler for processing simple query. #[async_trait] pub trait SimpleQueryHandler: Send + Sync { @@ -186,8 +277,15 @@ pub trait ExtendedQueryHandler: Send + Sync { /// Called when client sends `parse` command. /// - /// The default implementation parsed query with `Self::QueryParser` and - /// stores it in `Self::PortalStore`. + /// The default implementation parses the query with + /// `Self::QueryParser` and stores it in `Self::PortalStore`. + /// + /// If the query is an empty query (a string without any statement, such + /// as `""` or `";;"`), the parser is not called at all: like PostgreSQL, + /// the statement name is remembered as empty (replacing any statement + /// previously stored under the same name), `ParseComplete` is returned, + /// and later `Bind`/`Describe`/`Execute` on that name follow PostgreSQL's + /// empty-query behavior. async fn on_parse(&self, client: &mut C, message: Parse) -> PgWireResult<()> where C: ClientInfo + ClientPortalStore + Sink + Unpin + Send + Sync, @@ -195,9 +293,23 @@ pub trait ExtendedQueryHandler: Send + Sync { C::Error: Debug, PgWireError: From<>::Error>, { - let parser = self.query_parser(); - let stmt = StoredStatement::parse(client, &message, parser).await?; - client.portal_store().put_statement(Arc::new(stmt)); + let name = message + .name + .clone() + .unwrap_or_else(|| DEFAULT_NAME.to_owned()); + + if is_empty_query(&message.query) { + // An empty query has no parsed statement to store. Remember the + // name instead, and drop any previously stored statement of the + // same name: `Parse` always replaces its target. + client.portal_store().rm_statement(&name); + EmptyStatementRegistry::mark_statement(client, &name); + } else { + let parser = self.query_parser(); + let stmt = StoredStatement::parse(client, &message, parser).await?; + EmptyStatementRegistry::unmark_statement(client, &name); + client.portal_store().put_statement(Arc::new(stmt)); + } client .send(PgWireBackendMessage::ParseComplete(ParseComplete::new())) .await?; @@ -207,8 +319,13 @@ pub trait ExtendedQueryHandler: Send + Sync { /// Called when client sends `bind` command. /// - /// The default implementation associate parameters with previous parsed - /// statement and stores in `Self::PortalStore` as well. + /// The default implementation associates parameters with a previously + /// parsed statement and stores the result in `Self::PortalStore` as well. + /// + /// Binding to a statement parsed from an empty query also succeeds + /// (with zero parameters, like PostgreSQL): the portal name is remembered + /// as an empty portal, replacing any portal previously stored under the + /// same name. async fn on_bind(&self, client: &mut C, message: Bind) -> PgWireResult<()> where C: ClientInfo + ClientPortalStore + Sink + Unpin + Send + Sync, @@ -217,10 +334,34 @@ pub trait ExtendedQueryHandler: Send + Sync { PgWireError: From<>::Error>, { let statement_name = message.statement_name.as_deref().unwrap_or(DEFAULT_NAME); + let portal_name = message.portal_name.as_deref().unwrap_or(DEFAULT_NAME); if let Some(statement) = client.portal_store().get_statement(statement_name) { let portal = Portal::try_new(&message, statement)?; client.portal_store().put_portal(Arc::new(portal)); + // a real portal replaces a possibly existing empty-portal marker + EmptyStatementRegistry::unmark_portal(client, portal_name); + client + .send(PgWireBackendMessage::BindComplete(BindComplete::new())) + .await?; + Ok(()) + } else if EmptyStatementRegistry::is_empty_statement(client, statement_name) { + if !message.parameters.is_empty() { + return Err(PgWireError::UserError(Box::new(ErrorInfo::new( + "ERROR".to_owned(), + "08P01".to_owned(), + format!( + "bind message supplies {} parameters, but prepared statement {:?} requires 0", + message.parameters.len(), + statement_name + ), + )))); + } + // an empty statement binds successfully: remember the portal name + // as an empty portal, and drop any previously stored portal of the + // same name: `Bind` always replaces its target. + client.portal_store().rm_portal(portal_name); + EmptyStatementRegistry::mark_portal(client, portal_name); client .send(PgWireBackendMessage::BindComplete(BindComplete::new())) .await?; @@ -235,9 +376,8 @@ pub trait ExtendedQueryHandler: Send + Sync { /// The default implementation delegates the query to `self::do_query` and /// sends response messages according to `Response` from `self::do_query`. /// - /// Note that, different from `SimpleQueryHandler`, this implementation - /// won't check empty query because it cannot understand parsed - /// `Self::Statement`. + /// Portals bound from empty statements are handled here like PostgreSQL: + /// they respond with `EmptyQueryResponse` and never reach `do_query`. async fn on_execute(&self, client: &mut C, message: Execute) -> PgWireResult<()> where C: ClientInfo + ClientPortalStore + Sink + Unpin + Send + Sync, @@ -271,6 +411,17 @@ pub trait ExtendedQueryHandler: Send + Sync { let max_rows = message.max_rows as usize; let Some(portal) = client.portal_store().get_portal(portal_name) else { + if EmptyStatementRegistry::is_empty_portal(client, portal_name) { + // An empty portal never reaches `do_query`: it directly + // answers `EmptyQueryResponse`, like PostgreSQL. The portal + // stays valid until it is replaced, closed, or (for the + // unnamed portal) dropped by `Sync`. + client + .feed(PgWireBackendMessage::EmptyQueryResponse(EmptyQueryResponse)) + .await?; + client.set_state(super::PgWireConnectionState::ReadyForQuery); + return Ok(()); + } return Err(PgWireError::PortalNotFound(portal_name.to_owned())); }; // Execute query if the portal hasn't been started yet @@ -314,6 +465,7 @@ pub trait ExtendedQueryHandler: Send + Sync { // remove unnamed portal when transaction ends client.portal_store().rm_portal(DEFAULT_NAME); + EmptyStatementRegistry::unmark_portal(client, DEFAULT_NAME); false } @@ -404,6 +556,10 @@ pub trait ExtendedQueryHandler: Send + Sync { if let Some(stmt) = client.portal_store().get_statement(name) { let describe_response = self.do_describe_statement(client, &stmt).await?; send_describe_response(client, &describe_response).await?; + } else if EmptyStatementRegistry::is_empty_statement(client, name) { + // an empty statement has no parameters and no result data + let describe_response = DescribeStatementResponse::no_data(); + send_describe_response(client, &describe_response).await?; } else { return Err(PgWireError::StatementNotFound(name.to_owned())); } @@ -412,6 +568,9 @@ pub trait ExtendedQueryHandler: Send + Sync { if let Some(portal) = client.portal_store().get_portal(name) { let describe_response = self.do_describe_portal(client, &portal).await?; send_describe_response(client, &describe_response).await?; + } else if EmptyStatementRegistry::is_empty_portal(client, name) { + let describe_response = DescribePortalResponse::no_data(); + send_describe_response(client, &describe_response).await?; } else { return Err(PgWireError::PortalNotFound(name.to_owned())); } @@ -438,7 +597,8 @@ pub trait ExtendedQueryHandler: Send + Sync { /// Called when client sends `sync` command. /// /// The default implementation flushes client buffer and sends - /// `READY_FOR_QUERY` response to client + /// `READY_FOR_QUERY` response to client. The unnamed portal, including an + /// empty one, is removed, like PostgreSQL. async fn on_sync(&self, client: &mut C, _message: PgSync) -> PgWireResult<()> where C: ClientInfo + ClientPortalStore + Sink + Unpin + Send + Sync, @@ -447,6 +607,7 @@ pub trait ExtendedQueryHandler: Send + Sync { PgWireError: From<>::Error>, { client.portal_store().rm_portal(DEFAULT_NAME); + EmptyStatementRegistry::unmark_portal(client, DEFAULT_NAME); client .send(PgWireBackendMessage::ReadyForQuery(ReadyForQuery::new( @@ -459,7 +620,8 @@ pub trait ExtendedQueryHandler: Send + Sync { /// Called when client sends `close` command. /// - /// The default implementation closes certain statement or portal. + /// The default implementation closes certain statement or portal, + /// including empty statements and empty portals. async fn on_close(&self, client: &mut C, message: Close) -> PgWireResult<()> where C: ClientInfo + ClientPortalStore + Sink + Unpin + Send + Sync, @@ -471,9 +633,11 @@ pub trait ExtendedQueryHandler: Send + Sync { match message.target_type { TARGET_TYPE_BYTE_STATEMENT => { client.portal_store().rm_statement(name); + EmptyStatementRegistry::unmark_statement(client, name); } TARGET_TYPE_BYTE_PORTAL => { client.portal_store().rm_portal(name); + EmptyStatementRegistry::unmark_portal(client, name); } _ => {} } @@ -820,3 +984,631 @@ mod tests { assert!(!is_empty_query("';'")); } } + +/// Unit tests for extended-query empty statement handling. +/// +/// The expected message sequences mirror the behavior of a real PostgreSQL +/// server (verified against PostgreSQL 18) driven over the wire with +/// Parse/Bind/Describe/Execute/Sync of empty and semicolon-only queries. +#[cfg(test)] +mod extended_empty_query_tests { + use std::net::SocketAddr; + use std::pin::Pin; + use std::sync::Mutex; + use std::task::{Context, Poll}; + + use async_trait::async_trait; + use bytes::Bytes; + use futures::Sink; + + use super::*; + use crate::api::results::Tag; + use crate::api::{DefaultClient, PgWireConnectionState}; + use crate::messages::response::TransactionStatus; + + /// A client test-double implementing everything the query handlers + /// require. Backend messages are recorded instead of being encoded. + struct TestClient { + inner: DefaultClient, + sent: Mutex>, + } + + impl TestClient { + fn new() -> Self { + let mut inner = DefaultClient::new(SocketAddr::from(([127, 0, 0, 1], 5432)), false); + inner.set_state(PgWireConnectionState::ReadyForQuery); + TestClient { + inner, + sent: Mutex::new(Vec::new()), + } + } + + /// Short names of all backend messages sent so far. + fn sent(&self) -> Vec<&'static str> { + self.sent + .lock() + .unwrap() + .iter() + .map(|m| match m { + PgWireBackendMessage::ParseComplete(_) => "ParseComplete", + PgWireBackendMessage::BindComplete(_) => "BindComplete", + PgWireBackendMessage::CloseComplete(_) => "CloseComplete", + PgWireBackendMessage::EmptyQueryResponse(_) => "EmptyQueryResponse", + PgWireBackendMessage::ParameterDescription(_) => "ParameterDescription", + PgWireBackendMessage::NoData(_) => "NoData", + PgWireBackendMessage::CommandComplete(_) => "CommandComplete", + PgWireBackendMessage::ReadyForQuery(_) => "ReadyForQuery", + _ => "other", + }) + .collect() + } + + /// Number of parameter types in the `ParameterDescription` at + /// `idx` among the sent messages. + fn parameter_description_len(&self, idx: usize) -> usize { + self.sent + .lock() + .unwrap() + .iter() + .filter_map(|m| match m { + PgWireBackendMessage::ParameterDescription(p) => Some(p.types.len()), + _ => None, + }) + .nth(idx) + .unwrap() + } + } + + impl ClientInfo for TestClient { + fn socket_addr(&self) -> SocketAddr { + self.inner.socket_addr() + } + + fn is_secure(&self) -> bool { + self.inner.is_secure() + } + + fn protocol_version(&self) -> crate::messages::ProtocolVersion { + self.inner.protocol_version() + } + + fn set_protocol_version(&mut self, version: crate::messages::ProtocolVersion) { + self.inner.set_protocol_version(version) + } + + fn pid_and_secret_key(&self) -> (i32, crate::messages::startup::SecretKey) { + self.inner.pid_and_secret_key() + } + + fn set_pid_and_secret_key( + &mut self, + pid: i32, + secret_key: crate::messages::startup::SecretKey, + ) { + self.inner.set_pid_and_secret_key(pid, secret_key) + } + + fn state(&self) -> PgWireConnectionState { + self.inner.state() + } + + fn set_state(&mut self, new_state: PgWireConnectionState) { + self.inner.set_state(new_state) + } + + fn transaction_status(&self) -> TransactionStatus { + self.inner.transaction_status() + } + + fn set_transaction_status(&mut self, new_status: TransactionStatus) { + self.inner.set_transaction_status(new_status) + } + + fn metadata(&self) -> &std::collections::HashMap { + self.inner.metadata() + } + + fn metadata_mut(&mut self) -> &mut std::collections::HashMap { + self.inner.metadata_mut() + } + + fn session_extensions(&self) -> &crate::api::SessionExtensions { + self.inner.session_extensions() + } + + #[cfg(any(feature = "_ring", feature = "_aws-lc-rs"))] + fn sni_server_name(&self) -> Option<&str> { + self.inner.sni_server_name() + } + + #[cfg(any(feature = "_ring", feature = "_aws-lc-rs"))] + fn client_certificates<'a>(&self) -> Option<&[rustls_pki_types::CertificateDer<'a>]> { + self.inner.client_certificates() + } + } + + impl ClientPortalStore for TestClient { + type PortalStore = crate::api::store::MemPortalStore; + + fn portal_store(&self) -> &Self::PortalStore { + self.inner.portal_store() + } + } + + impl Sink for TestClient { + type Error = PgWireError; + + fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn start_send(self: Pin<&mut Self>, item: PgWireBackendMessage) -> Result<(), Self::Error> { + self.sent.lock().unwrap().push(item); + Ok(()) + } + + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + /// A parser that records every query it receives and refuses empty ones: + /// after this change, empty queries must never reach a parser. + #[derive(Default)] + struct RecordingParser { + calls: Mutex>, + } + + #[async_trait] + impl QueryParser for RecordingParser { + type Statement = String; + + async fn parse_sql( + &self, + _client: &C, + sql: &str, + _types: &[Option], + ) -> PgWireResult + where + C: ClientInfo + Unpin + Send + Sync, + { + assert!( + !is_empty_query(sql), + "parser must never be called for an empty query, got {sql:?}" + ); + self.calls.lock().unwrap().push(sql.to_owned()); + Ok(sql.to_owned()) + } + + fn get_parameter_types(&self, _stmt: &Self::Statement) -> PgWireResult> { + Ok(vec![]) + } + + fn get_result_schema( + &self, + _stmt: &Self::Statement, + _column_format: Option<&crate::api::portal::Format>, + ) -> PgWireResult> { + Ok(vec![]) + } + } + + struct TestHandler { + parser: Arc, + } + + impl TestHandler { + fn new() -> Self { + TestHandler { + parser: Arc::new(RecordingParser::default()), + } + } + } + + #[async_trait] + impl ExtendedQueryHandler for TestHandler { + type Statement = String; + type QueryParser = RecordingParser; + + fn query_parser(&self) -> Arc { + self.parser.clone() + } + + async fn do_query( + &self, + _client: &mut C, + _portal: &Portal, + _max_rows: usize, + ) -> PgWireResult + where + C: ClientInfo + ClientPortalStore + Sink + Unpin + Send + Sync, + C::PortalStore: PortalStore, + C::Error: Debug, + PgWireError: From<>::Error>, + { + Ok(Response::Execution(Tag::new("OK"))) + } + } + + fn parse(name: Option<&str>, query: &str) -> Parse { + Parse { + name: name.map(str::to_owned), + query: query.to_owned(), + type_oids: vec![], + } + } + + fn bind(portal: Option<&str>, statement: Option<&str>) -> Bind { + Bind { + portal_name: portal.map(str::to_owned), + statement_name: statement.map(str::to_owned), + parameter_format_codes: vec![], + parameters: vec![], + result_column_format_codes: vec![], + } + } + + fn describe(target_type: u8, name: Option<&str>) -> Describe { + Describe { + target_type, + name: name.map(str::to_owned), + } + } + + fn close(target_type: u8, name: Option<&str>) -> Close { + Close { + target_type, + name: name.map(str::to_owned), + } + } + + /// Parse/Bind/Describe/Execute/Sync of an empty query behaves exactly + /// like PostgreSQL: ParseComplete, ParameterDescription(0)+NoData, + /// BindComplete, NoData, EmptyQueryResponse (repeatable), ReadyForQuery, + /// and the parser is never invoked. + #[tokio::test] + async fn empty_query_extended_protocol_sequence() { + let handler = TestHandler::new(); + let mut client = TestClient::new(); + + handler + .on_parse(&mut client, parse(None, "")) + .await + .unwrap(); + assert_eq!(client.sent(), ["ParseComplete"]); + + handler + ._on_describe(&mut client, describe(TARGET_TYPE_BYTE_STATEMENT, None)) + .await + .unwrap(); + assert_eq!(client.sent()[1..], ["ParameterDescription", "NoData"]); + assert_eq!(client.parameter_description_len(0), 0); + + handler + .on_bind(&mut client, bind(None, None)) + .await + .unwrap(); + assert_eq!(client.sent()[3..], ["BindComplete"]); + + handler + ._on_describe(&mut client, describe(TARGET_TYPE_BYTE_PORTAL, None)) + .await + .unwrap(); + assert_eq!(client.sent()[4..], ["NoData"]); + + // empty portals respond EmptyQueryResponse and stay valid across + // repeated Execute, without ever reaching do_query + for _ in 0..2 { + handler + ._on_execute( + &mut client, + Execute { + name: None, + max_rows: 0, + }, + ) + .await + .unwrap(); + } + assert_eq!( + client.sent()[5..], + ["EmptyQueryResponse", "EmptyQueryResponse"] + ); + assert!(matches!( + client.state(), + PgWireConnectionState::ReadyForQuery + )); + + handler.on_sync(&mut client, PgSync).await.unwrap(); + assert_eq!(client.sent()[7..], ["ReadyForQuery"]); + + // Sync removed the unnamed empty portal + assert!(matches!( + handler + ._on_execute( + &mut client, + Execute { + name: None, + max_rows: 0 + }, + ) + .await, + Err(PgWireError::PortalNotFound(_)) + )); + + assert!(handler.parser.calls.lock().unwrap().is_empty()); + } + + /// Semicolon-only and whitespace-only queries are empty in the extended + /// protocol as well. + #[tokio::test] + async fn semicolon_only_queries_are_empty() { + let handler = TestHandler::new(); + let mut client = TestClient::new(); + + handler + .on_parse(&mut client, parse(Some("s1"), ";")) + .await + .unwrap(); + handler + .on_parse(&mut client, parse(Some("s2"), " \n;\t")) + .await + .unwrap(); + assert_eq!(client.sent(), ["ParseComplete", "ParseComplete"]); + + handler + .on_bind(&mut client, bind(Some("p1"), Some("s1"))) + .await + .unwrap(); + handler + ._on_execute( + &mut client, + Execute { + name: Some("p1".to_owned()), + max_rows: 0, + }, + ) + .await + .unwrap(); + assert_eq!(client.sent()[2..], ["BindComplete", "EmptyQueryResponse"]); + + assert!(handler.parser.calls.lock().unwrap().is_empty()); + } + + /// A string literal containing only a semicolon is a real query and must + /// reach the parser. + #[tokio::test] + async fn string_literal_semicolon_reaches_parser() { + let handler = TestHandler::new(); + let mut client = TestClient::new(); + + handler + .on_parse(&mut client, parse(None, "';'")) + .await + .unwrap(); + assert_eq!(client.sent(), ["ParseComplete"]); + assert_eq!(handler.parser.calls.lock().unwrap().as_slice(), ["';'"]); + + handler + .on_bind(&mut client, bind(None, None)) + .await + .unwrap(); + handler + ._on_execute( + &mut client, + Execute { + name: None, + max_rows: 0, + }, + ) + .await + .unwrap(); + // do_query ran and returned an execution response + assert_eq!(client.sent()[1..], ["BindComplete", "CommandComplete"]); + } + + /// An empty Parse replaces a previously stored statement of the same + /// name, like PostgreSQL. + #[tokio::test] + async fn empty_parse_replaces_stored_statement() { + let handler = TestHandler::new(); + let mut client = TestClient::new(); + + handler + .on_parse(&mut client, parse(None, "select 1")) + .await + .unwrap(); + handler + .on_parse(&mut client, parse(None, "")) + .await + .unwrap(); + assert_eq!(client.sent(), ["ParseComplete", "ParseComplete"]); + + // describing the unnamed statement now reports the empty statement, + // not the previously parsed `select 1` + handler + ._on_describe(&mut client, describe(TARGET_TYPE_BYTE_STATEMENT, None)) + .await + .unwrap(); + assert_eq!(client.sent()[2..], ["ParameterDescription", "NoData"]); + + handler + .on_bind(&mut client, bind(None, None)) + .await + .unwrap(); + handler + ._on_execute( + &mut client, + Execute { + name: None, + max_rows: 0, + }, + ) + .await + .unwrap(); + assert_eq!(client.sent()[4..], ["BindComplete", "EmptyQueryResponse"]); + } + + /// Binding an empty statement to a portal name replaces a previously + /// bound real portal of the same name. + #[tokio::test] + async fn empty_bind_replaces_stored_portal() { + let handler = TestHandler::new(); + let mut client = TestClient::new(); + + handler + .on_parse(&mut client, parse(None, "select 1")) + .await + .unwrap(); + handler + .on_bind(&mut client, bind(Some("p"), None)) + .await + .unwrap(); + handler + ._on_execute( + &mut client, + Execute { + name: Some("p".to_owned()), + max_rows: 0, + }, + ) + .await + .unwrap(); + assert_eq!(client.sent()[1..], ["BindComplete", "CommandComplete"]); + + // re-parse the unnamed statement as empty, re-bind the same portal + handler + .on_parse(&mut client, parse(None, "")) + .await + .unwrap(); + handler + .on_bind(&mut client, bind(Some("p"), None)) + .await + .unwrap(); + handler + ._on_execute( + &mut client, + Execute { + name: Some("p".to_owned()), + max_rows: 0, + }, + ) + .await + .unwrap(); + assert_eq!(client.sent()[4..], ["BindComplete", "EmptyQueryResponse"]); + } + + /// Bind on a statement that was never parsed is an error, and binding + /// parameters to an empty statement is a protocol violation (PostgreSQL + /// answers 08P01). + #[tokio::test] + async fn bind_errors() { + let handler = TestHandler::new(); + let mut client = TestClient::new(); + + assert!(matches!( + handler + .on_bind(&mut client, bind(Some("p"), Some("missing"))) + .await, + Err(PgWireError::StatementNotFound(_)) + )); + + handler + .on_parse(&mut client, parse(Some("e"), "")) + .await + .unwrap(); + let mut with_param = bind(Some("p"), Some("e")); + with_param.parameters.push(Some(Bytes::from_static(b"1"))); + match handler.on_bind(&mut client, with_param).await { + Err(PgWireError::UserError(info)) => { + assert_eq!(info.code, "08P01"); + assert!(info.message.contains("requires 0")); + } + other => panic!("expected 08P01 user error, got {other:?}"), + } + } + + /// `Close` removes empty statements and empty portals like real ones. + #[tokio::test] + async fn close_removes_empty_statement_and_portal() { + let handler = TestHandler::new(); + let mut client = TestClient::new(); + + handler + .on_parse(&mut client, parse(Some("e"), "")) + .await + .unwrap(); + handler + .on_close(&mut client, close(TARGET_TYPE_BYTE_STATEMENT, Some("e"))) + .await + .unwrap(); + assert_eq!(client.sent(), ["ParseComplete", "CloseComplete"]); + assert!(matches!( + handler + .on_bind(&mut client, bind(Some("p"), Some("e"))) + .await, + Err(PgWireError::StatementNotFound(_)) + )); + + handler + .on_parse(&mut client, parse(Some("e2"), "")) + .await + .unwrap(); + handler + .on_bind(&mut client, bind(Some("p2"), Some("e2"))) + .await + .unwrap(); + handler + .on_close(&mut client, close(TARGET_TYPE_BYTE_PORTAL, Some("p2"))) + .await + .unwrap(); + assert!(matches!( + handler + ._on_execute( + &mut client, + Execute { + name: Some("p2".to_owned()), + max_rows: 0, + }, + ) + .await, + Err(PgWireError::PortalNotFound(_)) + )); + } + + /// After a failed Bind the empty-statement marker survives (Close is what + /// clears it), matching PostgreSQL where the statement stays prepared. + #[tokio::test] + async fn marker_survives_failed_bind() { + let handler = TestHandler::new(); + let mut client = TestClient::new(); + + handler + .on_parse(&mut client, parse(Some("e"), "")) + .await + .unwrap(); + let mut bad = bind(Some("p"), Some("e")); + bad.parameters.push(Some(Bytes::from_static(b"1"))); + assert!(handler.on_bind(&mut client, bad).await.is_err()); + + handler + .on_bind(&mut client, bind(Some("p"), Some("e"))) + .await + .unwrap(); + handler + ._on_execute( + &mut client, + Execute { + name: Some("p".to_owned()), + max_rows: 0, + }, + ) + .await + .unwrap(); + assert_eq!(client.sent()[1..], ["BindComplete", "EmptyQueryResponse"]); + } +} diff --git a/src/api/stmt.rs b/src/api/stmt.rs index 3cefa650..9e6a8c76 100644 --- a/src/api/stmt.rs +++ b/src/api/stmt.rs @@ -63,6 +63,11 @@ pub trait QueryParser { /// /// The client may or may not provide type information with any parameters /// from the sql. + /// + /// Note that the default `on_parse` implementation never calls this + /// method with an empty query (a query string without any statement, such + /// as `""` or `";;"`): those are handled by the extended query protocol + /// itself and never reach a query parser or executor. async fn parse_sql( &self, client: &C, From 60eeb39042b1ca726756f13aa25250dda9a9b82d Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 3 Sep 2026 17:57:25 +0800 Subject: [PATCH 2/5] refactor: move empty statement tracking into PortalStore MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit After review: there are almost no direct PortalStore implementors or callers outside the crate (only examples/cursor.rs), so representing the empty variant in the store costs far less breakage than any change to StoredStatement/Portal (which every do_query implementation touches through portal.statement.statement) — and it is the better home for the state anyway. PortalStore now models three states per name (missing / empty / real): - StatementEntry { Empty, Statement(Arc>) } returned by get_statement, with as_statement()/is_empty() helpers - PortalEntry { Empty, Portal(Arc>) } returned by get_portal - new put_empty_statement/put_empty_portal store empty markers; like every put_*, they replace whatever was stored under the name, so rm_*/clear_portals remove empty entries along with regular ones MemPortalStore stores the entries in its maps directly. The private EmptyStatementRegistry (SessionExtensions) from the previous commit is removed: replacement/removal/clearing now fall out of normal store semantics instead of a second bookkeeping structure that could drift (e.g. clear_portals not clearing markers). Wire behavior is unchanged: verified message-for-message against PostgreSQL 18 again with the raw-socket probe (Parse/Bind/Describe/ Execute/Sync of empty queries, statement and portal shadowing, 08P01 on bound parameters, Close/Sync cleanup). --- CHANGELOG.md | 27 ++++-- examples/cursor.rs | 4 +- src/api/query.rs | 199 +++++++++++---------------------------------- src/api/store.rs | 158 +++++++++++++++++++++++++++++++++-- 4 files changed, 221 insertions(+), 167 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 08228609..d73e0f67 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,18 +23,33 @@ Versioning](https://semver.org/spec/v2.0.0.html). by PostgreSQL 18+ and the minor-only form used by older servers) and adopts the negotiated version for the rest of the connection. +### Changed + +- Breaking: `PortalStore` now represents empty statements and portals. + `get_statement` returns `Option>` and `get_portal` + returns `Option>`: the new `Empty` variant marks a name + under which an empty prepared statement or portal is stored, alongside the + new `put_empty_statement`/`put_empty_portal` methods. Like every `put_*`, + storing an empty entry replaces whatever was previously stored under that + name, and `rm_*`/`clear_portals` remove empty entries along with regular + ones. `StoredStatement` and `Portal` themselves are unchanged — the + impact is limited to `PortalStore` implementors and code calling + `get_statement`/`get_portal` directly (`StatementEntry::as_statement`/ + `PortalEntry::as_portal` help with the migration). + ### Fixed - Extended query protocol: empty queries (a query string without any statement, such as `""` or `";;"`) are now handled like PostgreSQL instead of being dispatched to the query parser: `Parse` succeeds without calling - `QueryParser`, `Describe` answers `ParameterDescription` (no parameters) + - `NoData`, `Bind` succeeds (rejecting bound parameters with `08P01`), and + `QueryParser` and stores an empty statement, `Describe` answers + `ParameterDescription` (no parameters) + `NoData`, `Bind` succeeds + (rejecting bound parameters with `08P01`) and stores an empty portal, and `Execute` returns `EmptyQueryResponse` without reaching `do_query`. An - empty `Parse` replaces any statement previously stored under the same name, - and `Close`/`Sync` drop empty statements and the unnamed empty portal like - real ones. Empty queries are tracked internally per connection, so no - changes to `PortalStore`, `StoredStatement`, or `Portal` were required. + empty `Parse` replaces any statement previously stored under the same + name, and `Close`/`Sync` drop empty statements and the unnamed empty + portal like real ones. Behavior verified message-for-message against + PostgreSQL 18. - Client API: backend messages are now decoded with the rules of the protocol version the client actually advertised, instead of always 3.2. Previously a 4-byte protocol 3.0 cancel key was decoded as `SecretKey::Bytes` instead of diff --git a/examples/cursor.rs b/examples/cursor.rs index 3f0332e2..65994a2a 100644 --- a/examples/cursor.rs +++ b/examples/cursor.rs @@ -11,7 +11,7 @@ use pgwire::api::portal::Portal; use pgwire::api::query::SimpleQueryHandler; use pgwire::api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse, Response, Tag}; use pgwire::api::stmt::StoredStatement; -use pgwire::api::store::{MemPortalStore, PortalStore}; +use pgwire::api::store::{MemPortalStore, PortalEntry, PortalStore}; use pgwire::api::{ClientInfo, ClientPortalStore, PgWireServerHandlers, Type}; use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; use pgwire::messages::response::NoticeResponse; @@ -210,7 +210,7 @@ async fn handle_fetch( ) -> PgWireResult> { println!("FETCH {} FROM {}", count, cursor_name); - let Some(portal) = portal_store.get_portal(cursor_name) else { + let Some(PortalEntry::Portal(portal)) = portal_store.get_portal(cursor_name) else { return Err(PgWireError::UserError(Box::new(ErrorInfo::new( "ERROR".to_owned(), "34000".to_owned(), diff --git a/src/api/query.rs b/src/api/query.rs index a3276f1b..9d24bf41 100644 --- a/src/api/query.rs +++ b/src/api/query.rs @@ -1,8 +1,7 @@ use std::cmp::max; -use std::collections::HashSet; use std::fmt::Debug; use std::ops::Deref; -use std::sync::{Arc, RwLock}; +use std::sync::Arc; use async_trait::async_trait; use futures::channel::oneshot; @@ -13,7 +12,7 @@ use futures::stream::StreamExt; use super::portal::Portal; use super::results::{Tag, into_row_description}; use super::stmt::{NoopQueryParser, QueryParser, StoredStatement}; -use super::store::PortalStore; +use super::store::{PortalEntry, PortalStore, StatementEntry}; use super::{ClientInfo, ClientPortalStore, ConnectionHandle, DEFAULT_NAME, copy}; use crate::api::PgWireConnectionState; use crate::api::Type; @@ -49,96 +48,6 @@ where Some(handle.start_query().await) } -/// Tracks statement and portal names that refer to empty queries. -/// -/// Empty queries (a query string without any statement, such as `""` or -/// `";;"`) have no parsed representation, so there is nothing to store in -/// the [`PortalStore`] as a `StoredStatement`. Instead, the default extended -/// query handlers record their names in this per-connection registry, -/// mirroring PostgreSQL's extended-query behavior for empty queries: -/// -/// - `Parse` of an empty query succeeds and replaces any statement previously -/// stored under the same name; -/// - `Bind` on that name succeeds (binding zero parameters), replacing any -/// portal previously stored under the portal name; -/// - `Describe` answers `ParameterDescription` (no parameters) + `NoData`; -/// - `Execute` returns `EmptyQueryResponse` without dispatching to -/// [`ExtendedQueryHandler::do_query`]; -/// - `Close` removes the marker, and `Sync` drops the unnamed empty portal. -/// -/// The registry lives in [`SessionExtensions`](super::SessionExtensions), so -/// this is purely internal bookkeeping: neither the [`PortalStore`] API nor -/// the `StoredStatement`/`Portal` types need to represent an empty variant. -#[derive(Debug, Default)] -struct EmptyStatementRegistry { - statements: RwLock>, - portals: RwLock>, -} - -impl EmptyStatementRegistry { - fn for_client(client: &C) -> Arc { - client - .session_extensions() - .get_or_insert_with(Self::default) - } - - /// Record that `name` was last `Parse`d from an empty query. - fn mark_statement(client: &C, name: &str) { - Self::for_client(client) - .statements - .write() - .unwrap() - .insert(name.to_owned()); - } - - /// Forget an empty-statement marker, e.g. after re-`Parse` of a real - /// statement or a `Close`. - fn unmark_statement(client: &C, name: &str) { - Self::for_client(client) - .statements - .write() - .unwrap() - .remove(name); - } - - /// Test whether `name` is an empty statement. - fn is_empty_statement(client: &C, name: &str) -> bool { - Self::for_client(client) - .statements - .read() - .unwrap() - .contains(name) - } - - /// Record that `name` was `Bind`ed from an empty statement. - fn mark_portal(client: &C, name: &str) { - Self::for_client(client) - .portals - .write() - .unwrap() - .insert(name.to_owned()); - } - - /// Forget an empty-portal marker, e.g. after `Bind` of a real portal or a - /// `Close`. - fn unmark_portal(client: &C, name: &str) { - Self::for_client(client) - .portals - .write() - .unwrap() - .remove(name); - } - - /// Test whether `name` is an empty portal. - fn is_empty_portal(client: &C, name: &str) -> bool { - Self::for_client(client) - .portals - .read() - .unwrap() - .contains(name) - } -} - /// handler for processing simple query. #[async_trait] pub trait SimpleQueryHandler: Send + Sync { @@ -299,15 +208,13 @@ pub trait ExtendedQueryHandler: Send + Sync { .unwrap_or_else(|| DEFAULT_NAME.to_owned()); if is_empty_query(&message.query) { - // An empty query has no parsed statement to store. Remember the - // name instead, and drop any previously stored statement of the - // same name: `Parse` always replaces its target. - client.portal_store().rm_statement(&name); - EmptyStatementRegistry::mark_statement(client, &name); + // An empty query has no parsed statement to store. Like + // PostgreSQL, remember the name as an empty statement; it + // replaces any statement previously stored under that name. + client.portal_store().put_empty_statement(&name); } else { let parser = self.query_parser(); let stmt = StoredStatement::parse(client, &message, parser).await?; - EmptyStatementRegistry::unmark_statement(client, &name); client.portal_store().put_statement(Arc::new(stmt)); } client @@ -336,39 +243,34 @@ pub trait ExtendedQueryHandler: Send + Sync { let statement_name = message.statement_name.as_deref().unwrap_or(DEFAULT_NAME); let portal_name = message.portal_name.as_deref().unwrap_or(DEFAULT_NAME); - if let Some(statement) = client.portal_store().get_statement(statement_name) { - let portal = Portal::try_new(&message, statement)?; - client.portal_store().put_portal(Arc::new(portal)); - // a real portal replaces a possibly existing empty-portal marker - EmptyStatementRegistry::unmark_portal(client, portal_name); - client - .send(PgWireBackendMessage::BindComplete(BindComplete::new())) - .await?; - Ok(()) - } else if EmptyStatementRegistry::is_empty_statement(client, statement_name) { - if !message.parameters.is_empty() { - return Err(PgWireError::UserError(Box::new(ErrorInfo::new( - "ERROR".to_owned(), - "08P01".to_owned(), - format!( - "bind message supplies {} parameters, but prepared statement {:?} requires 0", - message.parameters.len(), - statement_name - ), - )))); + match client.portal_store().get_statement(statement_name) { + Some(StatementEntry::Statement(statement)) => { + let portal = Portal::try_new(&message, statement)?; + client.portal_store().put_portal(Arc::new(portal)); } - // an empty statement binds successfully: remember the portal name - // as an empty portal, and drop any previously stored portal of the - // same name: `Bind` always replaces its target. - client.portal_store().rm_portal(portal_name); - EmptyStatementRegistry::mark_portal(client, portal_name); - client - .send(PgWireBackendMessage::BindComplete(BindComplete::new())) - .await?; - Ok(()) - } else { - Err(PgWireError::StatementNotFound(statement_name.to_owned())) + Some(StatementEntry::Empty) => { + if !message.parameters.is_empty() { + return Err(PgWireError::UserError(Box::new(ErrorInfo::new( + "ERROR".to_owned(), + "08P01".to_owned(), + format!( + "bind message supplies {} parameters, but prepared statement {:?} requires 0", + message.parameters.len(), + statement_name + ), + )))); + } + // an empty statement binds to an empty portal, which + // replaces any portal previously stored under that name + client.portal_store().put_empty_portal(portal_name); + } + None => return Err(PgWireError::StatementNotFound(statement_name.to_owned())), } + + client + .send(PgWireBackendMessage::BindComplete(BindComplete::new())) + .await?; + Ok(()) } /// Called when client sends `execute` command. @@ -410,8 +312,9 @@ pub trait ExtendedQueryHandler: Send + Sync { let portal_name = message.name.as_deref().unwrap_or(DEFAULT_NAME); let max_rows = message.max_rows as usize; - let Some(portal) = client.portal_store().get_portal(portal_name) else { - if EmptyStatementRegistry::is_empty_portal(client, portal_name) { + let portal = match client.portal_store().get_portal(portal_name) { + Some(PortalEntry::Portal(portal)) => portal, + Some(PortalEntry::Empty) => { // An empty portal never reaches `do_query`: it directly // answers `EmptyQueryResponse`, like PostgreSQL. The portal // stays valid until it is replaced, closed, or (for the @@ -422,7 +325,7 @@ pub trait ExtendedQueryHandler: Send + Sync { client.set_state(super::PgWireConnectionState::ReadyForQuery); return Ok(()); } - return Err(PgWireError::PortalNotFound(portal_name.to_owned())); + None => return Err(PgWireError::PortalNotFound(portal_name.to_owned())), }; // Execute query if the portal hasn't been started yet let needs_fetch = if matches!( @@ -465,7 +368,6 @@ pub trait ExtendedQueryHandler: Send + Sync { // remove unnamed portal when transaction ends client.portal_store().rm_portal(DEFAULT_NAME); - EmptyStatementRegistry::unmark_portal(client, DEFAULT_NAME); false } @@ -552,29 +454,29 @@ pub trait ExtendedQueryHandler: Send + Sync { { let name = message.name.as_deref().unwrap_or(DEFAULT_NAME); match message.target_type { - TARGET_TYPE_BYTE_STATEMENT => { - if let Some(stmt) = client.portal_store().get_statement(name) { + TARGET_TYPE_BYTE_STATEMENT => match client.portal_store().get_statement(name) { + Some(StatementEntry::Statement(stmt)) => { let describe_response = self.do_describe_statement(client, &stmt).await?; send_describe_response(client, &describe_response).await?; - } else if EmptyStatementRegistry::is_empty_statement(client, name) { - // an empty statement has no parameters and no result data + } + // an empty statement has no parameters and no result data + Some(StatementEntry::Empty) => { let describe_response = DescribeStatementResponse::no_data(); send_describe_response(client, &describe_response).await?; - } else { - return Err(PgWireError::StatementNotFound(name.to_owned())); } - } - TARGET_TYPE_BYTE_PORTAL => { - if let Some(portal) = client.portal_store().get_portal(name) { + None => return Err(PgWireError::StatementNotFound(name.to_owned())), + }, + TARGET_TYPE_BYTE_PORTAL => match client.portal_store().get_portal(name) { + Some(PortalEntry::Portal(portal)) => { let describe_response = self.do_describe_portal(client, &portal).await?; send_describe_response(client, &describe_response).await?; - } else if EmptyStatementRegistry::is_empty_portal(client, name) { + } + Some(PortalEntry::Empty) => { let describe_response = DescribePortalResponse::no_data(); send_describe_response(client, &describe_response).await?; - } else { - return Err(PgWireError::PortalNotFound(name.to_owned())); } - } + None => return Err(PgWireError::PortalNotFound(name.to_owned())), + }, _ => return Err(PgWireError::InvalidTargetType(message.target_type)), } @@ -607,7 +509,6 @@ pub trait ExtendedQueryHandler: Send + Sync { PgWireError: From<>::Error>, { client.portal_store().rm_portal(DEFAULT_NAME); - EmptyStatementRegistry::unmark_portal(client, DEFAULT_NAME); client .send(PgWireBackendMessage::ReadyForQuery(ReadyForQuery::new( @@ -633,11 +534,9 @@ pub trait ExtendedQueryHandler: Send + Sync { match message.target_type { TARGET_TYPE_BYTE_STATEMENT => { client.portal_store().rm_statement(name); - EmptyStatementRegistry::unmark_statement(client, name); } TARGET_TYPE_BYTE_PORTAL => { client.portal_store().rm_portal(name); - EmptyStatementRegistry::unmark_portal(client, name); } _ => {} } diff --git a/src/api/store.rs b/src/api/store.rs index f1582625..44b480b9 100644 --- a/src/api/store.rs +++ b/src/api/store.rs @@ -5,7 +5,72 @@ use std::sync::{Arc, RwLock}; use super::portal::Portal; use super::stmt::StoredStatement; +/// The stored state of a prepared statement name. +/// +/// [`StatementEntry::Empty`] represents a statement that was `Parse`d from an +/// empty query (a query string without any statement, such as `""` or +/// `";;"`). Following PostgreSQL's extended-query protocol, an empty +/// statement has no parsed representation: it binds with zero parameters, +/// describes as `ParameterDescription` (no parameters) + `NoData`, and +/// executes to `EmptyQueryResponse` without reaching the query handler. +#[derive(Debug, Clone)] +pub enum StatementEntry { + /// The name holds a statement parsed from an empty query. + Empty, + /// The name holds a parsed statement. + Statement(Arc>), +} + +impl StatementEntry { + /// Get the stored statement, if this entry is not empty. + pub fn as_statement(&self) -> Option<&Arc>> { + match self { + StatementEntry::Empty => None, + StatementEntry::Statement(stmt) => Some(stmt), + } + } + + /// Test whether this entry is an empty statement. + pub fn is_empty(&self) -> bool { + matches!(self, StatementEntry::Empty) + } +} + +/// The stored state of a portal name. +/// +/// [`PortalEntry::Empty`] represents a portal bound from an empty statement: +/// it describes as `NoData` and executes to `EmptyQueryResponse`. +#[derive(Debug, Clone)] +pub enum PortalEntry { + /// The name holds a portal bound from an empty statement. + Empty, + /// The name holds a bound portal. + Portal(Arc>), +} + +impl PortalEntry { + /// Get the bound portal, if this entry is not empty. + pub fn as_portal(&self) -> Option<&Arc>> { + match self { + PortalEntry::Empty => None, + PortalEntry::Portal(portal) => Some(portal), + } + } + + /// Test whether this entry is an empty portal. + pub fn is_empty(&self) -> bool { + matches!(self, PortalEntry::Empty) + } +} + /// Storage trait for prepared statements and portals. +/// +/// Both statements and portals can also be *empty*: a `Parse` of an empty +/// query or a `Bind` on an empty statement stores an empty marker under the +/// target name (replacing whatever was stored under that name before), like +/// PostgreSQL. Empty entries are returned by `get_statement`/`get_portal` as +/// [`StatementEntry::Empty`]/[`PortalEntry::Empty`]; removing or clearing +/// removes them along with regular entries. pub trait PortalStore: Any + Send + Sync + 'static { type Statement; @@ -15,32 +80,40 @@ pub trait PortalStore: Any + Send + Sync + 'static { /// Store a prepared statement by name. fn put_statement(&self, statement: Arc>); + /// Store an empty prepared statement by name, replacing any statement + /// previously stored under the same name. + fn put_empty_statement(&self, name: &str); + /// Remove a prepared statement by name. fn rm_statement(&self, name: &str); /// Retrieve a prepared statement by name. - fn get_statement(&self, name: &str) -> Option>>; + fn get_statement(&self, name: &str) -> Option>; /// Store a portal by name. fn put_portal(&self, portal: Arc>); + /// Store an empty portal by name, replacing any portal previously stored + /// under the same name. + fn put_empty_portal(&self, name: &str); + /// Remove a portal by name. fn rm_portal(&self, name: &str); - /// Remove all portals. + /// Remove all portals, including empty ones. fn clear_portals(&self); /// Retrieve a portal by name. - fn get_portal(&self, name: &str) -> Option>>; + fn get_portal(&self, name: &str) -> Option>; } /// In-memory implementation of `PortalStore` backed by `BTreeMap`. #[derive(Debug, Default, new)] pub struct MemPortalStore { #[new(default)] - statements: RwLock>>>, + statements: RwLock>>, #[new(default)] - portals: RwLock>>>, + portals: RwLock>>, } impl PortalStore for MemPortalStore { @@ -51,8 +124,14 @@ impl PortalStore for MemPortalStore { } fn put_statement(&self, statement: Arc>) { + let name = statement.id.to_owned(); let mut guard = self.statements.write().unwrap(); - guard.insert(statement.id.to_owned(), statement); + guard.insert(name, StatementEntry::Statement(statement)); + } + + fn put_empty_statement(&self, name: &str) { + let mut guard = self.statements.write().unwrap(); + guard.insert(name.to_owned(), StatementEntry::Empty); } fn rm_statement(&self, name: &str) { @@ -60,14 +139,19 @@ impl PortalStore for MemPortalStore { guard.remove(name); } - fn get_statement(&self, name: &str) -> Option>> { + fn get_statement(&self, name: &str) -> Option> { let guard = self.statements.read().unwrap(); guard.get(name).cloned() } fn put_portal(&self, portal: Arc>) { let mut guard = self.portals.write().unwrap(); - guard.insert(portal.name.to_owned(), portal); + guard.insert(portal.name.to_owned(), PortalEntry::Portal(portal)); + } + + fn put_empty_portal(&self, name: &str) { + let mut guard = self.portals.write().unwrap(); + guard.insert(name.to_owned(), PortalEntry::Empty); } fn rm_portal(&self, name: &str) { @@ -80,8 +164,64 @@ impl PortalStore for MemPortalStore { guard.clear(); } - fn get_portal(&self, name: &str) -> Option>> { + fn get_portal(&self, name: &str) -> Option> { let guard = self.portals.read().unwrap(); guard.get(name).cloned() } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn statement_entries_replace_each_other() { + let store: MemPortalStore = MemPortalStore::new(); + assert!(store.get_statement("s").is_none()); + + store.put_empty_statement("s"); + assert!(store.get_statement("s").unwrap().is_empty()); + + // a real statement replaces the empty marker + store.put_statement(Arc::new(StoredStatement::new( + "s".to_owned(), + "select 1".to_owned(), + vec![], + ))); + assert_eq!( + store + .get_statement("s") + .and_then(|e| e.as_statement().map(|s| s.statement.clone())), + Some("select 1".to_owned()) + ); + + // and an empty marker replaces the real statement + store.put_empty_statement("s"); + assert!(store.get_statement("s").unwrap().is_empty()); + + store.rm_statement("s"); + assert!(store.get_statement("s").is_none()); + } + + #[test] + fn portal_entries_replace_each_other_and_clear() { + let store: MemPortalStore = MemPortalStore::new(); + let statement = Arc::new(StoredStatement::new( + "s".to_owned(), + "select 1".to_owned(), + vec![], + )); + let portal = Portal::new_cursor("p".to_owned(), statement); + + store.put_portal(Arc::new(portal)); + assert!(store.get_portal("p").unwrap().as_portal().is_some()); + + store.put_empty_portal("p"); + assert!(store.get_portal("p").unwrap().is_empty()); + + store.put_empty_portal("p2"); + store.clear_portals(); + assert!(store.get_portal("p").is_none()); + assert!(store.get_portal("p2").is_none()); + } +} From 44aacb3ba1cc9f50301db9896c4fe53bc0c261d9 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 3 Sep 2026 18:29:00 +0800 Subject: [PATCH 3/5] chore: trim verbose comments Keep only brief comments where they add non-obvious information; drop narration of self-explanatory code. --- src/api/query.rs | 68 ++++++++++++++---------------------------------- src/api/stmt.rs | 6 ++--- src/api/store.rs | 47 ++++++++++----------------------- 3 files changed, 35 insertions(+), 86 deletions(-) diff --git a/src/api/query.rs b/src/api/query.rs index 9d24bf41..36c494a5 100644 --- a/src/api/query.rs +++ b/src/api/query.rs @@ -187,14 +187,9 @@ pub trait ExtendedQueryHandler: Send + Sync { /// Called when client sends `parse` command. /// /// The default implementation parses the query with - /// `Self::QueryParser` and stores it in `Self::PortalStore`. - /// - /// If the query is an empty query (a string without any statement, such - /// as `""` or `";;"`), the parser is not called at all: like PostgreSQL, - /// the statement name is remembered as empty (replacing any statement - /// previously stored under the same name), `ParseComplete` is returned, - /// and later `Bind`/`Describe`/`Execute` on that name follow PostgreSQL's - /// empty-query behavior. + /// `Self::QueryParser` and stores it in `Self::PortalStore`. Empty + /// queries are not parsed: like PostgreSQL, an empty statement is stored + /// instead, which binds, describes and executes as an empty query. async fn on_parse(&self, client: &mut C, message: Parse) -> PgWireResult<()> where C: ClientInfo + ClientPortalStore + Sink + Unpin + Send + Sync, @@ -208,9 +203,6 @@ pub trait ExtendedQueryHandler: Send + Sync { .unwrap_or_else(|| DEFAULT_NAME.to_owned()); if is_empty_query(&message.query) { - // An empty query has no parsed statement to store. Like - // PostgreSQL, remember the name as an empty statement; it - // replaces any statement previously stored under that name. client.portal_store().put_empty_statement(&name); } else { let parser = self.query_parser(); @@ -228,11 +220,8 @@ pub trait ExtendedQueryHandler: Send + Sync { /// /// The default implementation associates parameters with a previously /// parsed statement and stores the result in `Self::PortalStore` as well. - /// - /// Binding to a statement parsed from an empty query also succeeds - /// (with zero parameters, like PostgreSQL): the portal name is remembered - /// as an empty portal, replacing any portal previously stored under the - /// same name. + /// Binding an empty statement stores an empty portal, with zero + /// parameters, like PostgreSQL. async fn on_bind(&self, client: &mut C, message: Bind) -> PgWireResult<()> where C: ClientInfo + ClientPortalStore + Sink + Unpin + Send + Sync, @@ -260,8 +249,6 @@ pub trait ExtendedQueryHandler: Send + Sync { ), )))); } - // an empty statement binds to an empty portal, which - // replaces any portal previously stored under that name client.portal_store().put_empty_portal(portal_name); } None => return Err(PgWireError::StatementNotFound(statement_name.to_owned())), @@ -277,9 +264,7 @@ pub trait ExtendedQueryHandler: Send + Sync { /// /// The default implementation delegates the query to `self::do_query` and /// sends response messages according to `Response` from `self::do_query`. - /// - /// Portals bound from empty statements are handled here like PostgreSQL: - /// they respond with `EmptyQueryResponse` and never reach `do_query`. + /// Empty portals answer `EmptyQueryResponse` and never reach `do_query`. async fn on_execute(&self, client: &mut C, message: Execute) -> PgWireResult<()> where C: ClientInfo + ClientPortalStore + Sink + Unpin + Send + Sync, @@ -315,10 +300,7 @@ pub trait ExtendedQueryHandler: Send + Sync { let portal = match client.portal_store().get_portal(portal_name) { Some(PortalEntry::Portal(portal)) => portal, Some(PortalEntry::Empty) => { - // An empty portal never reaches `do_query`: it directly - // answers `EmptyQueryResponse`, like PostgreSQL. The portal - // stays valid until it is replaced, closed, or (for the - // unnamed portal) dropped by `Sync`. + // never reaches do_query; stays valid for repeated Execute client .feed(PgWireBackendMessage::EmptyQueryResponse(EmptyQueryResponse)) .await?; @@ -459,7 +441,6 @@ pub trait ExtendedQueryHandler: Send + Sync { let describe_response = self.do_describe_statement(client, &stmt).await?; send_describe_response(client, &describe_response).await?; } - // an empty statement has no parameters and no result data Some(StatementEntry::Empty) => { let describe_response = DescribeStatementResponse::no_data(); send_describe_response(client, &describe_response).await?; @@ -499,8 +480,7 @@ pub trait ExtendedQueryHandler: Send + Sync { /// Called when client sends `sync` command. /// /// The default implementation flushes client buffer and sends - /// `READY_FOR_QUERY` response to client. The unnamed portal, including an - /// empty one, is removed, like PostgreSQL. + /// `READY_FOR_QUERY` response to client async fn on_sync(&self, client: &mut C, _message: PgSync) -> PgWireResult<()> where C: ClientInfo + ClientPortalStore + Sink + Unpin + Send + Sync, @@ -521,8 +501,7 @@ pub trait ExtendedQueryHandler: Send + Sync { /// Called when client sends `close` command. /// - /// The default implementation closes certain statement or portal, - /// including empty statements and empty portals. + /// The default implementation closes certain statement or portal. async fn on_close(&self, client: &mut C, message: Close) -> PgWireResult<()> where C: ClientInfo + ClientPortalStore + Sink + Unpin + Send + Sync, @@ -884,11 +863,8 @@ mod tests { } } -/// Unit tests for extended-query empty statement handling. -/// -/// The expected message sequences mirror the behavior of a real PostgreSQL -/// server (verified against PostgreSQL 18) driven over the wire with -/// Parse/Bind/Describe/Execute/Sync of empty and semicolon-only queries. +/// Extended-query empty statement handling, mirroring PostgreSQL 18 +/// message sequences. #[cfg(test)] mod extended_empty_query_tests { use std::net::SocketAddr; @@ -905,8 +881,8 @@ mod extended_empty_query_tests { use crate::api::{DefaultClient, PgWireConnectionState}; use crate::messages::response::TransactionStatus; - /// A client test-double implementing everything the query handlers - /// require. Backend messages are recorded instead of being encoded. + /// A client test-double recording backend messages instead of encoding + /// them. struct TestClient { inner: DefaultClient, sent: Mutex>, @@ -1055,8 +1031,7 @@ mod extended_empty_query_tests { } } - /// A parser that records every query it receives and refuses empty ones: - /// after this change, empty queries must never reach a parser. + /// A parser that fails on empty queries: they must never reach a parser. #[derive(Default)] struct RecordingParser { calls: Mutex>, @@ -1166,9 +1141,7 @@ mod extended_empty_query_tests { } /// Parse/Bind/Describe/Execute/Sync of an empty query behaves exactly - /// like PostgreSQL: ParseComplete, ParameterDescription(0)+NoData, - /// BindComplete, NoData, EmptyQueryResponse (repeatable), ReadyForQuery, - /// and the parser is never invoked. + /// like PostgreSQL. #[tokio::test] async fn empty_query_extended_protocol_sequence() { let handler = TestHandler::new(); @@ -1199,8 +1172,7 @@ mod extended_empty_query_tests { .unwrap(); assert_eq!(client.sent()[4..], ["NoData"]); - // empty portals respond EmptyQueryResponse and stay valid across - // repeated Execute, without ever reaching do_query + // empty portals stay valid across repeated Execute for _ in 0..2 { handler ._on_execute( @@ -1306,7 +1278,6 @@ mod extended_empty_query_tests { ) .await .unwrap(); - // do_query ran and returned an execution response assert_eq!(client.sent()[1..], ["BindComplete", "CommandComplete"]); } @@ -1327,8 +1298,7 @@ mod extended_empty_query_tests { .unwrap(); assert_eq!(client.sent(), ["ParseComplete", "ParseComplete"]); - // describing the unnamed statement now reports the empty statement, - // not the previously parsed `select 1` + // describes the empty statement, not the previously parsed `select 1` handler ._on_describe(&mut client, describe(TARGET_TYPE_BYTE_STATEMENT, None)) .await @@ -1479,8 +1449,8 @@ mod extended_empty_query_tests { )); } - /// After a failed Bind the empty-statement marker survives (Close is what - /// clears it), matching PostgreSQL where the statement stays prepared. + /// After a failed Bind the empty statement stays prepared, matching + /// PostgreSQL. #[tokio::test] async fn marker_survives_failed_bind() { let handler = TestHandler::new(); diff --git a/src/api/stmt.rs b/src/api/stmt.rs index 9e6a8c76..50caf6c4 100644 --- a/src/api/stmt.rs +++ b/src/api/stmt.rs @@ -64,10 +64,8 @@ pub trait QueryParser { /// The client may or may not provide type information with any parameters /// from the sql. /// - /// Note that the default `on_parse` implementation never calls this - /// method with an empty query (a query string without any statement, such - /// as `""` or `";;"`): those are handled by the extended query protocol - /// itself and never reach a query parser or executor. + /// Empty queries are never passed to this method; they are handled by + /// the extended query protocol itself. async fn parse_sql( &self, client: &C, diff --git a/src/api/store.rs b/src/api/store.rs index 44b480b9..2d0d619b 100644 --- a/src/api/store.rs +++ b/src/api/store.rs @@ -5,24 +5,16 @@ use std::sync::{Arc, RwLock}; use super::portal::Portal; use super::stmt::StoredStatement; -/// The stored state of a prepared statement name. -/// -/// [`StatementEntry::Empty`] represents a statement that was `Parse`d from an -/// empty query (a query string without any statement, such as `""` or -/// `";;"`). Following PostgreSQL's extended-query protocol, an empty -/// statement has no parsed representation: it binds with zero parameters, -/// describes as `ParameterDescription` (no parameters) + `NoData`, and -/// executes to `EmptyQueryResponse` without reaching the query handler. +/// A stored prepared statement: parsed from a query, or empty when `Parse`d +/// from an empty query (no statement to parse). #[derive(Debug, Clone)] pub enum StatementEntry { - /// The name holds a statement parsed from an empty query. Empty, - /// The name holds a parsed statement. Statement(Arc>), } impl StatementEntry { - /// Get the stored statement, if this entry is not empty. + /// The stored statement, if any. pub fn as_statement(&self) -> Option<&Arc>> { match self { StatementEntry::Empty => None, @@ -30,26 +22,21 @@ impl StatementEntry { } } - /// Test whether this entry is an empty statement. + /// Whether this is an empty statement. pub fn is_empty(&self) -> bool { matches!(self, StatementEntry::Empty) } } -/// The stored state of a portal name. -/// -/// [`PortalEntry::Empty`] represents a portal bound from an empty statement: -/// it describes as `NoData` and executes to `EmptyQueryResponse`. +/// A stored portal: bound from a parsed statement, or from an empty one. #[derive(Debug, Clone)] pub enum PortalEntry { - /// The name holds a portal bound from an empty statement. Empty, - /// The name holds a bound portal. Portal(Arc>), } impl PortalEntry { - /// Get the bound portal, if this entry is not empty. + /// The bound portal, if any. pub fn as_portal(&self) -> Option<&Arc>> { match self { PortalEntry::Empty => None, @@ -57,7 +44,7 @@ impl PortalEntry { } } - /// Test whether this entry is an empty portal. + /// Whether this is an empty portal. pub fn is_empty(&self) -> bool { matches!(self, PortalEntry::Empty) } @@ -65,12 +52,10 @@ impl PortalEntry { /// Storage trait for prepared statements and portals. /// -/// Both statements and portals can also be *empty*: a `Parse` of an empty -/// query or a `Bind` on an empty statement stores an empty marker under the -/// target name (replacing whatever was stored under that name before), like -/// PostgreSQL. Empty entries are returned by `get_statement`/`get_portal` as -/// [`StatementEntry::Empty`]/[`PortalEntry::Empty`]; removing or clearing -/// removes them along with regular entries. +/// Statements `Parse`d from empty queries and portals bound from them are +/// stored as empty entries, like PostgreSQL: every `put_*` replaces whatever +/// was previously stored under the name, and `rm_*`/`clear_portals` remove +/// empty entries along with regular ones. pub trait PortalStore: Any + Send + Sync + 'static { type Statement; @@ -80,8 +65,7 @@ pub trait PortalStore: Any + Send + Sync + 'static { /// Store a prepared statement by name. fn put_statement(&self, statement: Arc>); - /// Store an empty prepared statement by name, replacing any statement - /// previously stored under the same name. + /// Store an empty prepared statement by name. fn put_empty_statement(&self, name: &str); /// Remove a prepared statement by name. @@ -93,14 +77,13 @@ pub trait PortalStore: Any + Send + Sync + 'static { /// Store a portal by name. fn put_portal(&self, portal: Arc>); - /// Store an empty portal by name, replacing any portal previously stored - /// under the same name. + /// Store an empty portal by name. fn put_empty_portal(&self, name: &str); /// Remove a portal by name. fn rm_portal(&self, name: &str); - /// Remove all portals, including empty ones. + /// Remove all portals. fn clear_portals(&self); /// Retrieve a portal by name. @@ -182,7 +165,6 @@ mod tests { store.put_empty_statement("s"); assert!(store.get_statement("s").unwrap().is_empty()); - // a real statement replaces the empty marker store.put_statement(Arc::new(StoredStatement::new( "s".to_owned(), "select 1".to_owned(), @@ -195,7 +177,6 @@ mod tests { Some("select 1".to_owned()) ); - // and an empty marker replaces the real statement store.put_empty_statement("s"); assert!(store.get_statement("s").unwrap().is_empty()); From c470b396a16308fcf00c3bb15839462140bd1f04 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 3 Sep 2026 23:47:28 +0800 Subject: [PATCH 4/5] feat: allow parsers to report empty queries QueryParser::parse_sql now returns Option and StoredStatement::parse returns Option>: None denotes an empty query, stored as an empty statement and executed to EmptyQueryResponse. The syntactic empty-query check (semicolons and whitespace only) moved into StoredStatement::parse, so on_parse has a single branch and custom on_parse overrides calling it get empty-query handling for free. --- CHANGELOG.md | 7 +++++ src/api/query.rs | 66 +++++++++++++++++++++++++++++++++++++++--------- src/api/stmt.rs | 45 ++++++++++++++++++++------------- 3 files changed, 89 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d73e0f67..18893f26 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,13 @@ Versioning](https://semver.org/spec/v2.0.0.html). ### Changed +- Breaking: `QueryParser::parse_sql` now returns + `PgWireResult>`, and `StoredStatement::parse` + correspondingly returns `Option>`. `None` denotes an + empty query: it is stored as an empty statement and executes to + `EmptyQueryResponse`. This allows a query parser to report its own notion + of empty query; syntactically empty queries (semicolons and whitespace + only) are still never passed to the parser. - Breaking: `PortalStore` now represents empty statements and portals. `get_statement` returns `Option>` and `get_portal` returns `Option>`: the new `Empty` variant marks a name diff --git a/src/api/query.rs b/src/api/query.rs index 36c494a5..4e460552 100644 --- a/src/api/query.rs +++ b/src/api/query.rs @@ -30,7 +30,7 @@ use crate::messages::extendedquery::{ use crate::messages::response::{EmptyQueryResponse, ReadyForQuery, TransactionStatus}; use crate::messages::simplequery::Query; -fn is_empty_query(q: &str) -> bool { +pub(crate) fn is_empty_query(q: &str) -> bool { // A query string that contains only semicolons and whitespace parses to no // statements, which PostgreSQL treats as an empty query and answers with // `EmptyQueryResponse` instead of dispatching to the executor. This covers @@ -188,8 +188,8 @@ pub trait ExtendedQueryHandler: Send + Sync { /// /// The default implementation parses the query with /// `Self::QueryParser` and stores it in `Self::PortalStore`. Empty - /// queries are not parsed: like PostgreSQL, an empty statement is stored - /// instead, which binds, describes and executes as an empty query. + /// queries are stored as empty statements instead, like PostgreSQL: + /// they bind, describe and execute as empty queries. async fn on_parse(&self, client: &mut C, message: Parse) -> PgWireResult<()> where C: ClientInfo + ClientPortalStore + Sink + Unpin + Send + Sync, @@ -202,12 +202,10 @@ pub trait ExtendedQueryHandler: Send + Sync { .clone() .unwrap_or_else(|| DEFAULT_NAME.to_owned()); - if is_empty_query(&message.query) { - client.portal_store().put_empty_statement(&name); - } else { - let parser = self.query_parser(); - let stmt = StoredStatement::parse(client, &message, parser).await?; - client.portal_store().put_statement(Arc::new(stmt)); + let parser = self.query_parser(); + match StoredStatement::parse(client, &message, parser).await? { + Some(stmt) => client.portal_store().put_statement(Arc::new(stmt)), + None => client.portal_store().put_empty_statement(&name), } client .send(PgWireBackendMessage::ParseComplete(ParseComplete::new())) @@ -1031,7 +1029,8 @@ mod extended_empty_query_tests { } } - /// A parser that fails on empty queries: they must never reach a parser. + /// A parser that fails on syntactically empty queries (they must never + /// reach a parser) and reports comment-only queries as empty. #[derive(Default)] struct RecordingParser { calls: Mutex>, @@ -1046,7 +1045,7 @@ mod extended_empty_query_tests { _client: &C, sql: &str, _types: &[Option], - ) -> PgWireResult + ) -> PgWireResult> where C: ClientInfo + Unpin + Send + Sync, { @@ -1055,7 +1054,11 @@ mod extended_empty_query_tests { "parser must never be called for an empty query, got {sql:?}" ); self.calls.lock().unwrap().push(sql.to_owned()); - Ok(sql.to_owned()) + if sql.starts_with("--") { + Ok(None) + } else { + Ok(Some(sql.to_owned())) + } } fn get_parameter_types(&self, _stmt: &Self::Statement) -> PgWireResult> { @@ -1214,6 +1217,45 @@ mod extended_empty_query_tests { assert!(handler.parser.calls.lock().unwrap().is_empty()); } + /// A query the parser reports as empty (`None`) is stored and executed + /// as an empty query. + #[tokio::test] + async fn parser_reported_empty_query_executes_as_empty() { + let handler = TestHandler::new(); + let mut client = TestClient::new(); + + handler + .on_parse(&mut client, parse(Some("c"), "-- comment only")) + .await + .unwrap(); + assert_eq!( + handler.parser.calls.lock().unwrap().as_slice(), + ["-- comment only"] + ); + + handler + ._on_describe(&mut client, describe(TARGET_TYPE_BYTE_STATEMENT, Some("c"))) + .await + .unwrap(); + assert_eq!(client.sent()[1..], ["ParameterDescription", "NoData"]); + + handler + .on_bind(&mut client, bind(Some("p"), Some("c"))) + .await + .unwrap(); + handler + ._on_execute( + &mut client, + Execute { + name: Some("p".to_owned()), + max_rows: 0, + }, + ) + .await + .unwrap(); + assert_eq!(client.sent()[3..], ["BindComplete", "EmptyQueryResponse"]); + } + /// Semicolon-only and whitespace-only queries are empty in the extended /// protocol as well. #[tokio::test] diff --git a/src/api/stmt.rs b/src/api/stmt.rs index 50caf6c4..1bec9155 100644 --- a/src/api/stmt.rs +++ b/src/api/stmt.rs @@ -9,6 +9,7 @@ use crate::messages::PgWireBackendMessage; use crate::messages::extendedquery::Parse; use super::portal::Format; +use super::query::is_empty_query; use super::results::FieldInfo; use super::{ClientInfo, DEFAULT_NAME}; @@ -26,12 +27,15 @@ pub struct StoredStatement { } impl StoredStatement { - /// Parse a `Parse` message into a stored statement using the given query parser. + /// Parse a `Parse` message into a stored statement using the given query + /// parser. + /// + /// Returns `None` for an empty query: there is no statement to store. pub async fn parse( client: &C, parse: &Parse, parser: Q, - ) -> PgWireResult> + ) -> PgWireResult>> where C: ClientInfo + Sink + Unpin + Send + Sync, Q: QueryParser, @@ -41,15 +45,20 @@ impl StoredStatement { .iter() .map(|oid| Type::from_oid(*oid)) .collect::>(); - let statement = parser.parse_sql(client, &parse.query, &types).await?; - Ok(StoredStatement { - id: parse - .name - .clone() - .unwrap_or_else(|| DEFAULT_NAME.to_owned()), - statement, - parameter_types: types, - }) + if is_empty_query(&parse.query) { + return Ok(None); + } + Ok(parser + .parse_sql(client, &parse.query, &types) + .await? + .map(|statement| StoredStatement { + id: parse + .name + .clone() + .unwrap_or_else(|| DEFAULT_NAME.to_owned()), + statement, + parameter_types: types, + })) } } @@ -64,14 +73,16 @@ pub trait QueryParser { /// The client may or may not provide type information with any parameters /// from the sql. /// - /// Empty queries are never passed to this method; they are handled by - /// the extended query protocol itself. + /// Return `Ok(None)` for an empty query; it is stored as an empty + /// statement and executes to `EmptyQueryResponse`, like PostgreSQL. + /// Syntactically empty queries (only semicolons and whitespace) never + /// reach this method. async fn parse_sql( &self, client: &C, sql: &str, types: &[Option], - ) -> PgWireResult + ) -> PgWireResult> where C: ClientInfo + Unpin + Send + Sync; @@ -106,7 +117,7 @@ where client: &C, sql: &str, types: &[Option], - ) -> PgWireResult + ) -> PgWireResult> where C: ClientInfo + Unpin + Send + Sync, { @@ -139,11 +150,11 @@ impl QueryParser for NoopQueryParser { _client: &C, sql: &str, _types: &[Option], - ) -> PgWireResult + ) -> PgWireResult> where C: ClientInfo + Unpin + Send + Sync, { - Ok(sql.to_owned()) + Ok(Some(sql.to_owned())) } fn get_parameter_types(&self, _stmt: &Self::Statement) -> PgWireResult> { From 7b16316cb6e39566008549cf0a1813af3af07e96 Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Thu, 3 Sep 2026 23:51:58 +0800 Subject: [PATCH 5/5] refactor: generic Entry for stored statements and portals Replace StatementEntry/PortalEntry with a single Entry: Empty marker or Value(Arc), where T is StoredStatement or Portal. Clone is implemented manually so the payload type does not need to be Clone. --- CHANGELOG.md | 20 ++++++------- examples/cursor.rs | 4 +-- src/api/query.rs | 18 ++++++------ src/api/store.rs | 70 +++++++++++++++++++--------------------------- 4 files changed, 50 insertions(+), 62 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 18893f26..86997e85 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,16 +33,16 @@ Versioning](https://semver.org/spec/v2.0.0.html). of empty query; syntactically empty queries (semicolons and whitespace only) are still never passed to the parser. - Breaking: `PortalStore` now represents empty statements and portals. - `get_statement` returns `Option>` and `get_portal` - returns `Option>`: the new `Empty` variant marks a name - under which an empty prepared statement or portal is stored, alongside the - new `put_empty_statement`/`put_empty_portal` methods. Like every `put_*`, - storing an empty entry replaces whatever was previously stored under that - name, and `rm_*`/`clear_portals` remove empty entries along with regular - ones. `StoredStatement` and `Portal` themselves are unchanged — the - impact is limited to `PortalStore` implementors and code calling - `get_statement`/`get_portal` directly (`StatementEntry::as_statement`/ - `PortalEntry::as_portal` help with the migration). + `get_statement` returns `Option>>` and + `get_portal` returns `Option>>`: the new `Entry::Empty` + variant marks a name under which an empty prepared statement or portal is + stored, alongside the new `put_empty_statement`/`put_empty_portal` + methods. Like every `put_*`, storing an empty entry replaces whatever was + previously stored under that name, and `rm_*`/`clear_portals` remove + empty entries along with regular ones. `StoredStatement` and `Portal` + themselves are unchanged — the impact is limited to `PortalStore` + implementors and code calling `get_statement`/`get_portal` directly + (`Entry::value` helps with the migration). ### Fixed diff --git a/examples/cursor.rs b/examples/cursor.rs index 65994a2a..1b5bca9e 100644 --- a/examples/cursor.rs +++ b/examples/cursor.rs @@ -11,7 +11,7 @@ use pgwire::api::portal::Portal; use pgwire::api::query::SimpleQueryHandler; use pgwire::api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse, Response, Tag}; use pgwire::api::stmt::StoredStatement; -use pgwire::api::store::{MemPortalStore, PortalEntry, PortalStore}; +use pgwire::api::store::{Entry, MemPortalStore, PortalStore}; use pgwire::api::{ClientInfo, ClientPortalStore, PgWireServerHandlers, Type}; use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; use pgwire::messages::response::NoticeResponse; @@ -210,7 +210,7 @@ async fn handle_fetch( ) -> PgWireResult> { println!("FETCH {} FROM {}", count, cursor_name); - let Some(PortalEntry::Portal(portal)) = portal_store.get_portal(cursor_name) else { + let Some(Entry::Value(portal)) = portal_store.get_portal(cursor_name) else { return Err(PgWireError::UserError(Box::new(ErrorInfo::new( "ERROR".to_owned(), "34000".to_owned(), diff --git a/src/api/query.rs b/src/api/query.rs index 4e460552..e0ad9421 100644 --- a/src/api/query.rs +++ b/src/api/query.rs @@ -12,7 +12,7 @@ use futures::stream::StreamExt; use super::portal::Portal; use super::results::{Tag, into_row_description}; use super::stmt::{NoopQueryParser, QueryParser, StoredStatement}; -use super::store::{PortalEntry, PortalStore, StatementEntry}; +use super::store::{Entry, PortalStore}; use super::{ClientInfo, ClientPortalStore, ConnectionHandle, DEFAULT_NAME, copy}; use crate::api::PgWireConnectionState; use crate::api::Type; @@ -231,11 +231,11 @@ pub trait ExtendedQueryHandler: Send + Sync { let portal_name = message.portal_name.as_deref().unwrap_or(DEFAULT_NAME); match client.portal_store().get_statement(statement_name) { - Some(StatementEntry::Statement(statement)) => { + Some(Entry::Value(statement)) => { let portal = Portal::try_new(&message, statement)?; client.portal_store().put_portal(Arc::new(portal)); } - Some(StatementEntry::Empty) => { + Some(Entry::Empty) => { if !message.parameters.is_empty() { return Err(PgWireError::UserError(Box::new(ErrorInfo::new( "ERROR".to_owned(), @@ -296,8 +296,8 @@ pub trait ExtendedQueryHandler: Send + Sync { let max_rows = message.max_rows as usize; let portal = match client.portal_store().get_portal(portal_name) { - Some(PortalEntry::Portal(portal)) => portal, - Some(PortalEntry::Empty) => { + Some(Entry::Value(portal)) => portal, + Some(Entry::Empty) => { // never reaches do_query; stays valid for repeated Execute client .feed(PgWireBackendMessage::EmptyQueryResponse(EmptyQueryResponse)) @@ -435,22 +435,22 @@ pub trait ExtendedQueryHandler: Send + Sync { let name = message.name.as_deref().unwrap_or(DEFAULT_NAME); match message.target_type { TARGET_TYPE_BYTE_STATEMENT => match client.portal_store().get_statement(name) { - Some(StatementEntry::Statement(stmt)) => { + Some(Entry::Value(stmt)) => { let describe_response = self.do_describe_statement(client, &stmt).await?; send_describe_response(client, &describe_response).await?; } - Some(StatementEntry::Empty) => { + Some(Entry::Empty) => { let describe_response = DescribeStatementResponse::no_data(); send_describe_response(client, &describe_response).await?; } None => return Err(PgWireError::StatementNotFound(name.to_owned())), }, TARGET_TYPE_BYTE_PORTAL => match client.portal_store().get_portal(name) { - Some(PortalEntry::Portal(portal)) => { + Some(Entry::Value(portal)) => { let describe_response = self.do_describe_portal(client, &portal).await?; send_describe_response(client, &describe_response).await?; } - Some(PortalEntry::Empty) => { + Some(Entry::Empty) => { let describe_response = DescribePortalResponse::no_data(); send_describe_response(client, &describe_response).await?; } diff --git a/src/api/store.rs b/src/api/store.rs index 2d0d619b..e51439c2 100644 --- a/src/api/store.rs +++ b/src/api/store.rs @@ -5,48 +5,36 @@ use std::sync::{Arc, RwLock}; use super::portal::Portal; use super::stmt::StoredStatement; -/// A stored prepared statement: parsed from a query, or empty when `Parse`d -/// from an empty query (no statement to parse). -#[derive(Debug, Clone)] -pub enum StatementEntry { +/// An entry stored in a [`PortalStore`] under a statement or portal name: +/// either the stored value, or an empty marker for a query that parsed to +/// no statement. +#[derive(Debug)] +pub enum Entry { Empty, - Statement(Arc>), + Value(Arc), } -impl StatementEntry { - /// The stored statement, if any. - pub fn as_statement(&self) -> Option<&Arc>> { +impl Clone for Entry { + fn clone(&self) -> Self { match self { - StatementEntry::Empty => None, - StatementEntry::Statement(stmt) => Some(stmt), + Entry::Empty => Entry::Empty, + Entry::Value(value) => Entry::Value(Arc::clone(value)), } } - - /// Whether this is an empty statement. - pub fn is_empty(&self) -> bool { - matches!(self, StatementEntry::Empty) - } -} - -/// A stored portal: bound from a parsed statement, or from an empty one. -#[derive(Debug, Clone)] -pub enum PortalEntry { - Empty, - Portal(Arc>), } -impl PortalEntry { - /// The bound portal, if any. - pub fn as_portal(&self) -> Option<&Arc>> { +impl Entry { + /// The stored value, if any. + pub fn value(&self) -> Option<&Arc> { match self { - PortalEntry::Empty => None, - PortalEntry::Portal(portal) => Some(portal), + Entry::Empty => None, + Entry::Value(value) => Some(value), } } - /// Whether this is an empty portal. + /// Whether this is an empty entry. pub fn is_empty(&self) -> bool { - matches!(self, PortalEntry::Empty) + matches!(self, Entry::Empty) } } @@ -72,7 +60,7 @@ pub trait PortalStore: Any + Send + Sync + 'static { fn rm_statement(&self, name: &str); /// Retrieve a prepared statement by name. - fn get_statement(&self, name: &str) -> Option>; + fn get_statement(&self, name: &str) -> Option>>; /// Store a portal by name. fn put_portal(&self, portal: Arc>); @@ -87,16 +75,16 @@ pub trait PortalStore: Any + Send + Sync + 'static { fn clear_portals(&self); /// Retrieve a portal by name. - fn get_portal(&self, name: &str) -> Option>; + fn get_portal(&self, name: &str) -> Option>>; } /// In-memory implementation of `PortalStore` backed by `BTreeMap`. #[derive(Debug, Default, new)] pub struct MemPortalStore { #[new(default)] - statements: RwLock>>, + statements: RwLock>>>, #[new(default)] - portals: RwLock>>, + portals: RwLock>>>, } impl PortalStore for MemPortalStore { @@ -109,12 +97,12 @@ impl PortalStore for MemPortalStore { fn put_statement(&self, statement: Arc>) { let name = statement.id.to_owned(); let mut guard = self.statements.write().unwrap(); - guard.insert(name, StatementEntry::Statement(statement)); + guard.insert(name, Entry::Value(statement)); } fn put_empty_statement(&self, name: &str) { let mut guard = self.statements.write().unwrap(); - guard.insert(name.to_owned(), StatementEntry::Empty); + guard.insert(name.to_owned(), Entry::Empty); } fn rm_statement(&self, name: &str) { @@ -122,19 +110,19 @@ impl PortalStore for MemPortalStore { guard.remove(name); } - fn get_statement(&self, name: &str) -> Option> { + fn get_statement(&self, name: &str) -> Option>> { let guard = self.statements.read().unwrap(); guard.get(name).cloned() } fn put_portal(&self, portal: Arc>) { let mut guard = self.portals.write().unwrap(); - guard.insert(portal.name.to_owned(), PortalEntry::Portal(portal)); + guard.insert(portal.name.to_owned(), Entry::Value(portal)); } fn put_empty_portal(&self, name: &str) { let mut guard = self.portals.write().unwrap(); - guard.insert(name.to_owned(), PortalEntry::Empty); + guard.insert(name.to_owned(), Entry::Empty); } fn rm_portal(&self, name: &str) { @@ -147,7 +135,7 @@ impl PortalStore for MemPortalStore { guard.clear(); } - fn get_portal(&self, name: &str) -> Option> { + fn get_portal(&self, name: &str) -> Option>> { let guard = self.portals.read().unwrap(); guard.get(name).cloned() } @@ -173,7 +161,7 @@ mod tests { assert_eq!( store .get_statement("s") - .and_then(|e| e.as_statement().map(|s| s.statement.clone())), + .and_then(|e| e.value().map(|s| s.statement.clone())), Some("select 1".to_owned()) ); @@ -195,7 +183,7 @@ mod tests { let portal = Portal::new_cursor("p".to_owned(), statement); store.put_portal(Arc::new(portal)); - assert!(store.get_portal("p").unwrap().as_portal().is_some()); + assert!(store.get_portal("p").unwrap().value().is_some()); store.put_empty_portal("p"); assert!(store.get_portal("p").unwrap().is_empty());