Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions examples/cursor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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));

Expand Down Expand Up @@ -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;
Expand Down
6 changes: 3 additions & 3 deletions examples/sqlite.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)))?;
Expand Down Expand Up @@ -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))
Expand All @@ -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)
}
Expand Down
34 changes: 23 additions & 11 deletions src/api/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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,
Expand Down Expand Up @@ -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
) {
Expand Down Expand Up @@ -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()));
Expand Down Expand Up @@ -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)?;
Expand Down Expand Up @@ -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 =
Expand Down
15 changes: 12 additions & 3 deletions src/api/stmt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>>,
Expand All @@ -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
Expand Down
10 changes: 8 additions & 2 deletions tests-integration/test-server/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down
136 changes: 136 additions & 0 deletions tests/empty_extended_query.rs
Original file line number Diff line number Diff line change
@@ -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(_)
));
}
Loading