Handle empty extended queries

Cargo.toml

@@ -156,6 +156,10 @@ members = [
   "tests-integration/test-server",
 ]
 
+[[test]]
+name = "empty_extended_query"
+required-features = ["server-api"]
+
 [[example]]
 name = "server"
 required-features = ["server-api-aws-lc-rs"]

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,9 @@ async fn handle_fetch(
         portal.state().lock().await.deref(),
         pgwire::api::portal::PortalExecutionState::Initial
     ) {
-        let inner_query = &portal.statement.statement;
+        let Some(inner_query) = portal.statement.statement.as_ref() else {
+            return Ok(vec![Response::EmptyQuery]);
+        };
         println!("  -> Lazy execution of: {}", inner_query);
         let response = execute_inner_query(inner_query)?;
         portal.start(response).await;

examples/sqlite.rs

@@ -210,7 +210,9 @@ impl ExtendedQueryHandler for SqliteBackend {
         C: ClientInfo + Unpin + Send + Sync,
     {
         let conn = self.conn.lock().unwrap();
-        let query = &portal.statement.statement;
+        let Some(query) = portal.statement.statement.as_ref() else {
+            return Ok(Response::EmptyQuery);
+        };
         let mut stmt = conn
             .prepare_cached(query)
             .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
@@ -249,8 +251,11 @@ impl ExtendedQueryHandler for SqliteBackend {
             .iter()
             .map(|t| t.clone().unwrap_or(Type::UNKNOWN))
             .collect();
+        let Some(statement) = stmt.statement.as_ref() else {
+            return Ok(DescribeStatementResponse::no_data());
+        };
         let stmt = conn
-            .prepare_cached(&stmt.statement)
+            .prepare_cached(statement)
             .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
         row_desc_from_stmt(&stmt, &Format::UnifiedBinary)
             .map(|fields| DescribeStatementResponse::new(param_types, fields))
@@ -265,8 +270,11 @@ impl ExtendedQueryHandler for SqliteBackend {
         C: ClientInfo + Unpin + Send + Sync,
     {
         let conn = self.conn.lock().unwrap();
+        let Some(statement) = portal.statement.statement.as_ref() else {
+            return Ok(DescribePortalResponse::no_data());
+        };
         let stmt = conn
-            .prepare_cached(&portal.statement.statement)
+            .prepare_cached(statement)
             .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
         row_desc_from_stmt(&stmt, &portal.result_column_format).map(DescribePortalResponse::new)
     }

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<C>(client: &mut C) -> Option<oneshot::Receiver<()>>
 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,15 +396,23 @@ 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()));
                 }
             }
             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<<C as Sink<PgWireBackendMessage>>::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<<C as Sink<PgWireBackendMessage>>::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 =

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<S> {
     /// 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<S>,
     /// type ids of query parameters, can be empty if frontend asks backend for
     /// type inference
     pub parameter_types: Vec<Option<Type>>,
@@ -41,7 +46,11 @@ impl<S> StoredStatement<S> {
             .iter()
             .map(|oid| Type::from_oid(*oid))
             .collect::<Vec<_>>();
-        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

tests-integration/test-server/src/main.rs

@@ -181,7 +181,9 @@ impl ExtendedQueryHandler for DummyDatabase {
     where
         C: ClientInfo + Unpin + Send + Sync,
     {
-        let query = &portal.statement.statement;
+        let Some(query) = portal.statement.statement.as_ref() else {
+            return Ok(Response::EmptyQuery);
+        };
         println!("extended query: {:?}", query);
         if query.starts_with("SELECT") {
             // try to parse all parameters
@@ -306,7 +308,10 @@ impl ExtendedQueryHandler for DummyDatabase {
         C: ClientInfo + Unpin + Send + Sync,
     {
         println!("describe: {:?}", portal);
-        if portal.statement.statement.starts_with("SELECT") {
+        let Some(statement) = portal.statement.statement.as_ref() else {
+            return Ok(DescribePortalResponse::no_data());
+        };
+        if statement.starts_with("SELECT") {
             let schema = self.schema(&portal.result_column_format);
             Ok(DescribePortalResponse::new(schema))
         } else {

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<Self::QueryParser> {
+        Arc::new(NoopQueryParser)
+    }
+
+    async fn do_describe_statement<C>(
+        &self,
+        _client: &mut C,
+        _target: &StoredStatement<Self::Statement>,
+    ) -> PgWireResult<DescribeStatementResponse>
+    where
+        C: ClientInfo + Unpin + Send + Sync,
+    {
+        panic!("empty query reached statement description")
+    }
+
+    async fn do_describe_portal<C>(
+        &self,
+        _client: &mut C,
+        _target: &Portal<Self::Statement>,
+    ) -> PgWireResult<DescribePortalResponse>
+    where
+        C: ClientInfo + Unpin + Send + Sync,
+    {
+        panic!("empty query reached portal description")
+    }
+
+    async fn do_query<C>(
+        &self,
+        _client: &mut C,
+        _portal: &Portal<Self::Statement>,
+        _max_rows: usize,
+    ) -> PgWireResult<Response>
+    where
+        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
+        C::Error: Debug,
+        PgWireError: From<<C as Sink<PgWireBackendMessage>>::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::<String>::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(_)
+    ));
+}