diff --git a/examples/cursor.rs b/examples/cursor.rs index 3f0332e2..c259a748 100644 --- a/examples/cursor.rs +++ b/examples/cursor.rs @@ -196,7 +196,11 @@ fn handle_declare( )))); } - let statement = StoredStatement::new(cursor_name.to_string(), inner_query.to_string(), vec![]); + let statement = StoredStatement::new( + cursor_name.to_string(), + Some(inner_query.to_string()), + vec![], + ); let portal = Portal::new_cursor(cursor_name.to_string(), Arc::new(statement)); portal_store.put_portal(Arc::new(portal)); @@ -224,7 +228,7 @@ async fn handle_fetch( portal.state().lock().await.deref(), pgwire::api::portal::PortalExecutionState::Initial ) { - let inner_query = &portal.statement.statement; + let inner_query = portal.statement.statement.as_ref().unwrap(); println!(" -> Lazy execution of: {}", inner_query); let response = execute_inner_query(inner_query)?; portal.start(response).await; diff --git a/examples/sqlite.rs b/examples/sqlite.rs index 049df48d..ddf263ea 100644 --- a/examples/sqlite.rs +++ b/examples/sqlite.rs @@ -210,7 +210,7 @@ impl ExtendedQueryHandler for SqliteBackend { C: ClientInfo + Unpin + Send + Sync, { let conn = self.conn.lock().unwrap(); - let query = &portal.statement.statement; + let query = portal.statement.statement.as_ref().unwrap(); let mut stmt = conn .prepare_cached(query) .map_err(|e| PgWireError::ApiError(Box::new(e)))?; @@ -250,7 +250,7 @@ impl ExtendedQueryHandler for SqliteBackend { .map(|t| t.clone().unwrap_or(Type::UNKNOWN)) .collect(); let stmt = conn - .prepare_cached(&stmt.statement) + .prepare_cached(stmt.statement.as_ref().unwrap()) .map_err(|e| PgWireError::ApiError(Box::new(e)))?; row_desc_from_stmt(&stmt, &Format::UnifiedBinary) .map(|fields| DescribeStatementResponse::new(param_types, fields)) @@ -266,7 +266,7 @@ impl ExtendedQueryHandler for SqliteBackend { { let conn = self.conn.lock().unwrap(); let stmt = conn - .prepare_cached(&portal.statement.statement) + .prepare_cached(portal.statement.statement.as_ref().unwrap()) .map_err(|e| PgWireError::ApiError(Box::new(e)))?; row_desc_from_stmt(&stmt, &portal.result_column_format).map(DescribePortalResponse::new) } diff --git a/src/api/query.rs b/src/api/query.rs index 621b339c..4ce24f12 100644 --- a/src/api/query.rs +++ b/src/api/query.rs @@ -11,7 +11,7 @@ use futures::stream::StreamExt; use super::portal::Portal; use super::results::{Tag, into_row_description}; -use super::stmt::{NoopQueryParser, QueryParser, StoredStatement}; +use super::stmt::{NoopQueryParser, QueryParser, StoredStatement, is_empty_query}; use super::store::PortalStore; use super::{ClientInfo, ClientPortalStore, ConnectionHandle, DEFAULT_NAME, copy}; use crate::api::PgWireConnectionState; @@ -30,11 +30,6 @@ use crate::messages::extendedquery::{ use crate::messages::response::{EmptyQueryResponse, ReadyForQuery, TransactionStatus}; use crate::messages::simplequery::Query; -fn is_empty_query(q: &str) -> bool { - let trimmed_query = q.trim(); - trimmed_query == ";" || trimmed_query.is_empty() -} - async fn get_cancel_receiver(client: &mut C) -> Option> where C: ClientInfo + ClientPortalStore + Unpin + Send + Sync, @@ -268,7 +263,12 @@ pub trait ExtendedQueryHandler: Send + Sync { return Err(PgWireError::PortalNotFound(portal_name.to_owned())); }; // Execute query if the portal hasn't been started yet - let needs_fetch = if matches!( + let needs_fetch = if portal.statement.statement.is_none() { + client + .feed(PgWireBackendMessage::EmptyQueryResponse(EmptyQueryResponse)) + .await?; + false + } else if matches!( portal.state().lock().await.deref(), PortalExecutionState::Initial ) { @@ -396,7 +396,11 @@ pub trait ExtendedQueryHandler: Send + Sync { match message.target_type { TARGET_TYPE_BYTE_STATEMENT => { if let Some(stmt) = client.portal_store().get_statement(name) { - let describe_response = self.do_describe_statement(client, &stmt).await?; + let describe_response = if stmt.statement.is_none() { + DescribeStatementResponse::no_data() + } else { + self.do_describe_statement(client, &stmt).await? + }; send_describe_response(client, &describe_response).await?; } else { return Err(PgWireError::StatementNotFound(name.to_owned())); @@ -404,7 +408,11 @@ pub trait ExtendedQueryHandler: Send + Sync { } TARGET_TYPE_BYTE_PORTAL => { if let Some(portal) = client.portal_store().get_portal(name) { - let describe_response = self.do_describe_portal(client, &portal).await?; + let describe_response = if portal.statement.statement.is_none() { + DescribePortalResponse::no_data() + } else { + self.do_describe_portal(client, &portal).await? + }; send_describe_response(client, &describe_response).await?; } else { return Err(PgWireError::PortalNotFound(name.to_owned())); @@ -489,7 +497,9 @@ pub trait ExtendedQueryHandler: Send + Sync { C::Error: Debug, PgWireError: From<>::Error>, { - let stmt = &target.statement; + let Some(stmt) = target.statement.as_ref() else { + return Ok(DescribeStatementResponse::no_data()); + }; let query_parser = self.query_parser(); let server_param_types = query_parser.get_parameter_types(stmt)?; @@ -523,7 +533,9 @@ pub trait ExtendedQueryHandler: Send + Sync { C::Error: Debug, PgWireError: From<>::Error>, { - let stmt = &target.statement.statement; + let Some(stmt) = target.statement.statement.as_ref() else { + return Ok(DescribePortalResponse::no_data()); + }; let query_parser = self.query_parser(); let result_schema = diff --git a/src/api/stmt.rs b/src/api/stmt.rs index 3cefa650..ffc50916 100644 --- a/src/api/stmt.rs +++ b/src/api/stmt.rs @@ -12,14 +12,19 @@ use super::portal::Format; use super::results::FieldInfo; use super::{ClientInfo, DEFAULT_NAME}; +pub(crate) fn is_empty_query(query: &str) -> bool { + let query = query.trim(); + query.is_empty() || query == ";" +} + /// A parsed SQL statement stored in the portal store. #[non_exhaustive] #[derive(Debug, Default, new)] pub struct StoredStatement { /// name of the statement pub id: String, - /// parsed query statement - pub statement: S, + /// parsed query statement, or none for an empty query + pub statement: Option, /// type ids of query parameters, can be empty if frontend asks backend for /// type inference pub parameter_types: Vec>, @@ -41,7 +46,11 @@ impl StoredStatement { .iter() .map(|oid| Type::from_oid(*oid)) .collect::>(); - let statement = parser.parse_sql(client, &parse.query, &types).await?; + let statement = if is_empty_query(&parse.query) { + None + } else { + Some(parser.parse_sql(client, &parse.query, &types).await?) + }; Ok(StoredStatement { id: parse .name diff --git a/tests-integration/test-server/src/main.rs b/tests-integration/test-server/src/main.rs index f77f4b81..1f0faaa1 100644 --- a/tests-integration/test-server/src/main.rs +++ b/tests-integration/test-server/src/main.rs @@ -181,7 +181,7 @@ impl ExtendedQueryHandler for DummyDatabase { where C: ClientInfo + Unpin + Send + Sync, { - let query = &portal.statement.statement; + let query = portal.statement.statement.as_ref().unwrap(); println!("extended query: {:?}", query); if query.starts_with("SELECT") { // try to parse all parameters @@ -306,7 +306,13 @@ impl ExtendedQueryHandler for DummyDatabase { C: ClientInfo + Unpin + Send + Sync, { println!("describe: {:?}", portal); - if portal.statement.statement.starts_with("SELECT") { + if portal + .statement + .statement + .as_ref() + .unwrap() + .starts_with("SELECT") + { let schema = self.schema(&portal.result_column_format); Ok(DescribePortalResponse::new(schema)) } else { diff --git a/tests/empty_extended_query.rs b/tests/empty_extended_query.rs new file mode 100644 index 00000000..cf38d162 --- /dev/null +++ b/tests/empty_extended_query.rs @@ -0,0 +1,136 @@ +use std::fmt::Debug; +use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::sync::Arc; + +use async_trait::async_trait; +use bytes::BytesMut; +use futures::Sink; +use pgwire::api::portal::Portal; +use pgwire::api::query::ExtendedQueryHandler; +use pgwire::api::results::{DescribePortalResponse, DescribeStatementResponse, Response}; +use pgwire::api::stmt::{NoopQueryParser, StoredStatement}; +use pgwire::api::{ClientInfo, DefaultClient, PgWireConnectionState}; +use pgwire::error::{PgWireError, PgWireResult}; +use pgwire::messages::extendedquery::{ + Bind, Describe, Execute, Parse, Sync as PgSync, TARGET_TYPE_BYTE_PORTAL, + TARGET_TYPE_BYTE_STATEMENT, +}; +use pgwire::messages::{DecodeContext, PgWireBackendMessage}; +use pgwire::tokio::server::PgWireMessageServerCodec; +use tokio::io::{AsyncReadExt, duplex}; +use tokio_util::codec::Framed; + +struct TestExtendedQueryHandler; + +#[async_trait] +impl ExtendedQueryHandler for TestExtendedQueryHandler { + type Statement = String; + type QueryParser = NoopQueryParser; + + fn query_parser(&self) -> Arc { + Arc::new(NoopQueryParser) + } + + async fn do_describe_statement( + &self, + _client: &mut C, + _target: &StoredStatement, + ) -> PgWireResult + where + C: ClientInfo + Unpin + Send + Sync, + { + panic!("empty query reached statement description") + } + + async fn do_describe_portal( + &self, + _client: &mut C, + _target: &Portal, + ) -> PgWireResult + where + C: ClientInfo + Unpin + Send + Sync, + { + panic!("empty query reached portal description") + } + + async fn do_query( + &self, + _client: &mut C, + _portal: &Portal, + _max_rows: usize, + ) -> PgWireResult + where + C: ClientInfo + Sink + Unpin + Send + Sync, + C::Error: Debug, + PgWireError: From<>::Error>, + { + panic!("empty query reached query execution") + } +} + +#[tokio::test] +async fn empty_extended_query() { + let address = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 5432); + let mut client_info = DefaultClient::::new(address, false); + client_info.set_state(PgWireConnectionState::ReadyForQuery); + let codec = PgWireMessageServerCodec::new(client_info); + let (server_stream, mut client_stream) = duplex(1024); + let mut server = Framed::new(server_stream, codec); + let handler = TestExtendedQueryHandler; + + handler + .on_parse(&mut server, Parse::new(None, String::new(), vec![])) + .await + .unwrap(); + handler + .on_describe(&mut server, Describe::new(TARGET_TYPE_BYTE_STATEMENT, None)) + .await + .unwrap(); + handler + .on_bind(&mut server, Bind::new(None, None, vec![], vec![], vec![])) + .await + .unwrap(); + handler + .on_describe(&mut server, Describe::new(TARGET_TYPE_BYTE_PORTAL, None)) + .await + .unwrap(); + handler + .on_execute(&mut server, Execute::new(None, 0)) + .await + .unwrap(); + handler.on_sync(&mut server, PgSync::new()).await.unwrap(); + + let mut buffer = BytesMut::new(); + let mut messages = Vec::new(); + while messages.len() < 7 { + assert_ne!(client_stream.read_buf(&mut buffer).await.unwrap(), 0); + while let Some(message) = + PgWireBackendMessage::decode(&mut buffer, &DecodeContext::default()).unwrap() + { + messages.push(message); + } + } + + assert_eq!(messages.len(), 7); + assert!(matches!( + messages[0], + PgWireBackendMessage::ParseComplete(_) + )); + match &messages[1] { + PgWireBackendMessage::ParameterDescription(description) => { + assert!(description.types.is_empty()); + } + message => panic!("unexpected message: {message:?}"), + } + assert!(matches!(messages[2], PgWireBackendMessage::NoData(_))); + assert!(matches!(messages[3], PgWireBackendMessage::BindComplete(_))); + assert!(matches!(messages[4], PgWireBackendMessage::NoData(_))); + assert!(matches!( + messages[5], + PgWireBackendMessage::EmptyQueryResponse(_) + )); + assert!(matches!( + messages[6], + PgWireBackendMessage::ReadyForQuery(_) + )); +}