mirror of
https://github.com/neondatabase/neon.git
synced 2026-01-15 17:32:56 +00:00
6768a71c86
## Problem For #11992 I realised we need to get the type info before executing the query. This is important to know how to decode rows with custom types, eg the following query: ```sql CREATE TYPE foo AS ENUM ('foo','bar','baz'); SELECT ARRAY['foo'::foo, 'bar'::foo, 'baz'::foo] AS data; ``` Getting that to work was harder that it seems. The original tokio-postgres setup has a split between `Client` and `Connection`, where messages are passed between. Because multiple clients were supported, each client message included a dedicated response channel. Each request would be terminated by the `ReadyForQuery` message. The flow I opted to use for parsing types early would not trigger a `ReadyForQuery`. The flow is as follows: ``` PARSE "" // parse the user provided query DESCRIBE "" // describe the query, returning param/result type oids FLUSH // force postgres to flush the responses early // wait for descriptions // check if we know the types, if we don't then // setup the typeinfo query and execute it against each OID: PARSE typeinfo // prepare our typeinfo query DESCRIBE typeinfo FLUSH // force postgres to flush the responses early // wait for typeinfo statement // for each OID we don't know: BIND typeinfo EXECUTE FLUSH // wait for type info, might reveal more OIDs to inspect // close the typeinfo query, we cache the OID->type map and this is kinder to pgbouncer. CLOSE typeinfo // finally once we know all the OIDs: BIND "" // bind the user provided query - already parsed - to the user provided params EXECUTE // run the user provided query SYNC // commit the transaction ``` ## Summary of changes Please review commit by commit. The main challenge was allowing one query to issue multiple sub-queries. To do this I first made sure that the client could fully own the connection, which required removing any shared client state. I then had to replace the way responses are sent to the client, by using only a single permanent channel. This required some additional effort to track which query is being processed. Lastly I had to modify the query/typeinfo functions to not issue `sync` commands, so it would fit into the desired flow above. To note: the flow above does force an extra roundtrip into each query. I don't know yet if this has a measurable latency overhead.
122 lines
3.1 KiB
Rust
122 lines
3.1 KiB
Rust
//! An asynchronous, pipelined, PostgreSQL client.
|
|
#![warn(clippy::all)]
|
|
|
|
use postgres_protocol2::message::backend::ReadyForQueryBody;
|
|
|
|
pub use crate::cancel_token::CancelToken;
|
|
pub use crate::client::{Client, SocketConfig};
|
|
pub use crate::config::Config;
|
|
pub use crate::connect_raw::RawConnection;
|
|
pub use crate::connection::Connection;
|
|
use crate::error::DbError;
|
|
pub use crate::error::Error;
|
|
pub use crate::generic_client::GenericClient;
|
|
pub use crate::query::RowStream;
|
|
pub use crate::row::{Row, SimpleQueryRow};
|
|
pub use crate::simple_query::SimpleQueryStream;
|
|
pub use crate::statement::{Column, Statement};
|
|
pub use crate::tls::NoTls;
|
|
pub use crate::transaction::Transaction;
|
|
pub use crate::transaction_builder::{IsolationLevel, TransactionBuilder};
|
|
|
|
/// After executing a query, the connection will be in one of these states
|
|
#[derive(Clone, Copy, Debug, PartialEq)]
|
|
#[repr(u8)]
|
|
pub enum ReadyForQueryStatus {
|
|
/// Connection state is unknown
|
|
Unknown,
|
|
/// Connection is idle (no transactions)
|
|
Idle = b'I',
|
|
/// Connection is in a transaction block
|
|
Transaction = b'T',
|
|
/// Connection is in a failed transaction block
|
|
FailedTransaction = b'E',
|
|
}
|
|
|
|
impl From<ReadyForQueryBody> for ReadyForQueryStatus {
|
|
fn from(value: ReadyForQueryBody) -> Self {
|
|
match value.status() {
|
|
b'I' => Self::Idle,
|
|
b'T' => Self::Transaction,
|
|
b'E' => Self::FailedTransaction,
|
|
_ => Self::Unknown,
|
|
}
|
|
}
|
|
}
|
|
|
|
mod cancel_query;
|
|
mod cancel_query_raw;
|
|
mod cancel_token;
|
|
mod client;
|
|
mod codec;
|
|
pub mod config;
|
|
mod connect;
|
|
mod connect_raw;
|
|
mod connect_socket;
|
|
mod connect_tls;
|
|
mod connection;
|
|
pub mod error;
|
|
mod generic_client;
|
|
pub mod maybe_tls_stream;
|
|
mod prepare;
|
|
mod query;
|
|
pub mod row;
|
|
mod simple_query;
|
|
mod statement;
|
|
pub mod tls;
|
|
mod transaction;
|
|
mod transaction_builder;
|
|
pub mod types;
|
|
|
|
/// An asynchronous notification.
|
|
#[derive(Clone, Debug)]
|
|
pub struct Notification {
|
|
process_id: i32,
|
|
channel: String,
|
|
payload: String,
|
|
}
|
|
|
|
impl Notification {
|
|
/// The process ID of the notifying backend process.
|
|
pub fn process_id(&self) -> i32 {
|
|
self.process_id
|
|
}
|
|
|
|
/// The name of the channel that the notify has been raised on.
|
|
pub fn channel(&self) -> &str {
|
|
&self.channel
|
|
}
|
|
|
|
/// The "payload" string passed from the notifying process.
|
|
pub fn payload(&self) -> &str {
|
|
&self.payload
|
|
}
|
|
}
|
|
|
|
/// An asynchronous message from the server.
|
|
#[allow(clippy::large_enum_variant)]
|
|
#[derive(Debug, Clone)]
|
|
#[non_exhaustive]
|
|
pub enum AsyncMessage {
|
|
/// A notice.
|
|
///
|
|
/// Notices use the same format as errors, but aren't "errors" per-se.
|
|
Notice(DbError),
|
|
/// A notification.
|
|
///
|
|
/// Connections can subscribe to notifications with the `LISTEN` command.
|
|
Notification(Notification),
|
|
}
|
|
|
|
/// Message returned by the `SimpleQuery` stream.
|
|
#[derive(Debug)]
|
|
#[non_exhaustive]
|
|
pub enum SimpleQueryMessage {
|
|
/// A row of data.
|
|
Row(SimpleQueryRow),
|
|
/// A statement in the query has completed.
|
|
///
|
|
/// The number of rows modified or selected is returned.
|
|
CommandComplete(u64),
|
|
}
|