chore(proxy): vendor a subset of rust-postgres (#9930)

Our rust-postgres fork is getting messy. Mostly because proxy wants more
control over the raw protocol than tokio-postgres provides. As such,
it's diverging more and more. Storage and compute also make use of
rust-postgres, but in more normal usage, thus they don't need our crazy
changes.

Idea: 
* proxy maintains their subset
* other teams use a minimal patch set against upstream rust-postgres

Reviewing this code will be difficult. To implement it, I
1. Copied tokio-postgres, postgres-protocol and postgres-types from
00940fcdb5
2. Updated their package names with the `2` suffix to make them compile
in the workspace.
3. Updated proxy to use those packages
4. Copied in the code from tokio-postgres-rustls 0.13 (with some patches
applied https://github.com/jbg/tokio-postgres-rustls/pull/32
https://github.com/jbg/tokio-postgres-rustls/pull/33)
5. Removed as much dead code as I could find in the vendored libraries
6. Updated the tokio-postgres-rustls code to use our existing channel
binding implementation

libs/proxy/tokio-postgres2/src/cancel_query.rs

@@ -0,0 +1,40 @@
+use tokio::net::TcpStream;
+
+use crate::client::SocketConfig;
+use crate::config::{Host, SslMode};
+use crate::tls::MakeTlsConnect;
+use crate::{cancel_query_raw, connect_socket, Error};
+use std::io;
+
+pub(crate) async fn cancel_query<T>(
+    config: Option<SocketConfig>,
+    ssl_mode: SslMode,
+    mut tls: T,
+    process_id: i32,
+    secret_key: i32,
+) -> Result<(), Error>
+where
+    T: MakeTlsConnect<TcpStream>,
+{
+    let config = match config {
+        Some(config) => config,
+        None => {
+            return Err(Error::connect(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "unknown host",
+            )))
+        }
+    };
+
+    let hostname = match &config.host {
+        Host::Tcp(host) => &**host,
+    };
+    let tls = tls
+        .make_tls_connect(hostname)
+        .map_err(|e| Error::tls(e.into()))?;
+
+    let socket =
+        connect_socket::connect_socket(&config.host, config.port, config.connect_timeout).await?;
+
+    cancel_query_raw::cancel_query_raw(socket, ssl_mode, tls, process_id, secret_key).await
+}

libs/proxy/tokio-postgres2/src/cancel_query_raw.rs

@@ -0,0 +1,29 @@
+use crate::config::SslMode;
+use crate::tls::TlsConnect;
+use crate::{connect_tls, Error};
+use bytes::BytesMut;
+use postgres_protocol2::message::frontend;
+use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
+
+pub async fn cancel_query_raw<S, T>(
+    stream: S,
+    mode: SslMode,
+    tls: T,
+    process_id: i32,
+    secret_key: i32,
+) -> Result<(), Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: TlsConnect<S>,
+{
+    let mut stream = connect_tls::connect_tls(stream, mode, tls).await?;
+
+    let mut buf = BytesMut::new();
+    frontend::cancel_request(process_id, secret_key, &mut buf);
+
+    stream.write_all(&buf).await.map_err(Error::io)?;
+    stream.flush().await.map_err(Error::io)?;
+    stream.shutdown().await.map_err(Error::io)?;
+
+    Ok(())
+}

libs/proxy/tokio-postgres2/src/cancel_token.rs

@@ -0,0 +1,62 @@
+use crate::config::SslMode;
+use crate::tls::TlsConnect;
+
+use crate::{cancel_query, client::SocketConfig, tls::MakeTlsConnect};
+use crate::{cancel_query_raw, Error};
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::net::TcpStream;
+
+/// The capability to request cancellation of in-progress queries on a
+/// connection.
+#[derive(Clone)]
+pub struct CancelToken {
+    pub(crate) socket_config: Option<SocketConfig>,
+    pub(crate) ssl_mode: SslMode,
+    pub(crate) process_id: i32,
+    pub(crate) secret_key: i32,
+}
+
+impl CancelToken {
+    /// Attempts to cancel the in-progress query on the connection associated
+    /// with this `CancelToken`.
+    ///
+    /// The server provides no information about whether a cancellation attempt was successful or not. An error will
+    /// only be returned if the client was unable to connect to the database.
+    ///
+    /// Cancellation is inherently racy. There is no guarantee that the
+    /// cancellation request will reach the server before the query terminates
+    /// normally, or that the connection associated with this token is still
+    /// active.
+    ///
+    /// Requires the `runtime` Cargo feature (enabled by default).
+    pub async fn cancel_query<T>(&self, tls: T) -> Result<(), Error>
+    where
+        T: MakeTlsConnect<TcpStream>,
+    {
+        cancel_query::cancel_query(
+            self.socket_config.clone(),
+            self.ssl_mode,
+            tls,
+            self.process_id,
+            self.secret_key,
+        )
+        .await
+    }
+
+    /// Like `cancel_query`, but uses a stream which is already connected to the server rather than opening a new
+    /// connection itself.
+    pub async fn cancel_query_raw<S, T>(&self, stream: S, tls: T) -> Result<(), Error>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+        T: TlsConnect<S>,
+    {
+        cancel_query_raw::cancel_query_raw(
+            stream,
+            self.ssl_mode,
+            tls,
+            self.process_id,
+            self.secret_key,
+        )
+        .await
+    }
+}

libs/proxy/tokio-postgres2/src/client.rs

@@ -0,0 +1,439 @@
+use crate::codec::{BackendMessages, FrontendMessage};
+
+use crate::config::Host;
+use crate::config::SslMode;
+use crate::connection::{Request, RequestMessages};
+
+use crate::query::RowStream;
+use crate::simple_query::SimpleQueryStream;
+
+use crate::types::{Oid, ToSql, Type};
+
+use crate::{
+    prepare, query, simple_query, slice_iter, CancelToken, Error, ReadyForQueryStatus, Row,
+    SimpleQueryMessage, Statement, ToStatement, Transaction, TransactionBuilder,
+};
+use bytes::BytesMut;
+use fallible_iterator::FallibleIterator;
+use futures_util::{future, ready, TryStreamExt};
+use parking_lot::Mutex;
+use postgres_protocol2::message::{backend::Message, frontend};
+use std::collections::HashMap;
+use std::fmt;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+use tokio::sync::mpsc;
+
+use std::time::Duration;
+
+pub struct Responses {
+    receiver: mpsc::Receiver<BackendMessages>,
+    cur: BackendMessages,
+}
+
+impl Responses {
+    pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Result<Message, Error>> {
+        loop {
+            match self.cur.next().map_err(Error::parse)? {
+                Some(Message::ErrorResponse(body)) => return Poll::Ready(Err(Error::db(body))),
+                Some(message) => return Poll::Ready(Ok(message)),
+                None => {}
+            }
+
+            match ready!(self.receiver.poll_recv(cx)) {
+                Some(messages) => self.cur = messages,
+                None => return Poll::Ready(Err(Error::closed())),
+            }
+        }
+    }
+
+    pub async fn next(&mut self) -> Result<Message, Error> {
+        future::poll_fn(|cx| self.poll_next(cx)).await
+    }
+}
+
+/// A cache of type info and prepared statements for fetching type info
+/// (corresponding to the queries in the [prepare] module).
+#[derive(Default)]
+struct CachedTypeInfo {
+    /// A statement for basic information for a type from its
+    /// OID. Corresponds to [TYPEINFO_QUERY](prepare::TYPEINFO_QUERY) (or its
+    /// fallback).
+    typeinfo: Option<Statement>,
+    /// A statement for getting information for a composite type from its OID.
+    /// Corresponds to [TYPEINFO_QUERY](prepare::TYPEINFO_COMPOSITE_QUERY).
+    typeinfo_composite: Option<Statement>,
+    /// A statement for getting information for a composite type from its OID.
+    /// Corresponds to [TYPEINFO_QUERY](prepare::TYPEINFO_COMPOSITE_QUERY) (or
+    /// its fallback).
+    typeinfo_enum: Option<Statement>,
+
+    /// Cache of types already looked up.
+    types: HashMap<Oid, Type>,
+}
+
+pub struct InnerClient {
+    sender: mpsc::UnboundedSender<Request>,
+    cached_typeinfo: Mutex<CachedTypeInfo>,
+
+    /// A buffer to use when writing out postgres commands.
+    buffer: Mutex<BytesMut>,
+}
+
+impl InnerClient {
+    pub fn send(&self, messages: RequestMessages) -> Result<Responses, Error> {
+        let (sender, receiver) = mpsc::channel(1);
+        let request = Request { messages, sender };
+        self.sender.send(request).map_err(|_| Error::closed())?;
+
+        Ok(Responses {
+            receiver,
+            cur: BackendMessages::empty(),
+        })
+    }
+
+    pub fn typeinfo(&self) -> Option<Statement> {
+        self.cached_typeinfo.lock().typeinfo.clone()
+    }
+
+    pub fn set_typeinfo(&self, statement: &Statement) {
+        self.cached_typeinfo.lock().typeinfo = Some(statement.clone());
+    }
+
+    pub fn typeinfo_composite(&self) -> Option<Statement> {
+        self.cached_typeinfo.lock().typeinfo_composite.clone()
+    }
+
+    pub fn set_typeinfo_composite(&self, statement: &Statement) {
+        self.cached_typeinfo.lock().typeinfo_composite = Some(statement.clone());
+    }
+
+    pub fn typeinfo_enum(&self) -> Option<Statement> {
+        self.cached_typeinfo.lock().typeinfo_enum.clone()
+    }
+
+    pub fn set_typeinfo_enum(&self, statement: &Statement) {
+        self.cached_typeinfo.lock().typeinfo_enum = Some(statement.clone());
+    }
+
+    pub fn type_(&self, oid: Oid) -> Option<Type> {
+        self.cached_typeinfo.lock().types.get(&oid).cloned()
+    }
+
+    pub fn set_type(&self, oid: Oid, type_: &Type) {
+        self.cached_typeinfo.lock().types.insert(oid, type_.clone());
+    }
+
+    /// Call the given function with a buffer to be used when writing out
+    /// postgres commands.
+    pub fn with_buf<F, R>(&self, f: F) -> R
+    where
+        F: FnOnce(&mut BytesMut) -> R,
+    {
+        let mut buffer = self.buffer.lock();
+        let r = f(&mut buffer);
+        buffer.clear();
+        r
+    }
+}
+
+#[derive(Clone)]
+pub(crate) struct SocketConfig {
+    pub host: Host,
+    pub port: u16,
+    pub connect_timeout: Option<Duration>,
+    // pub keepalive: Option<KeepaliveConfig>,
+}
+
+/// An asynchronous PostgreSQL client.
+///
+/// The client is one half of what is returned when a connection is established. Users interact with the database
+/// through this client object.
+pub struct Client {
+    inner: Arc<InnerClient>,
+
+    socket_config: Option<SocketConfig>,
+    ssl_mode: SslMode,
+    process_id: i32,
+    secret_key: i32,
+}
+
+impl Client {
+    pub(crate) fn new(
+        sender: mpsc::UnboundedSender<Request>,
+        ssl_mode: SslMode,
+        process_id: i32,
+        secret_key: i32,
+    ) -> Client {
+        Client {
+            inner: Arc::new(InnerClient {
+                sender,
+                cached_typeinfo: Default::default(),
+                buffer: Default::default(),
+            }),
+
+            socket_config: None,
+            ssl_mode,
+            process_id,
+            secret_key,
+        }
+    }
+
+    /// Returns process_id.
+    pub fn get_process_id(&self) -> i32 {
+        self.process_id
+    }
+
+    pub(crate) fn inner(&self) -> &Arc<InnerClient> {
+        &self.inner
+    }
+
+    pub(crate) fn set_socket_config(&mut self, socket_config: SocketConfig) {
+        self.socket_config = Some(socket_config);
+    }
+
+    /// Creates a new prepared statement.
+    ///
+    /// Prepared statements can be executed repeatedly, and may contain query parameters (indicated by `$1`, `$2`, etc),
+    /// which are set when executed. Prepared statements can only be used with the connection that created them.
+    pub async fn prepare(&self, query: &str) -> Result<Statement, Error> {
+        self.prepare_typed(query, &[]).await
+    }
+
+    /// Like `prepare`, but allows the types of query parameters to be explicitly specified.
+    ///
+    /// The list of types may be smaller than the number of parameters - the types of the remaining parameters will be
+    /// inferred. For example, `client.prepare_typed(query, &[])` is equivalent to `client.prepare(query)`.
+    pub async fn prepare_typed(
+        &self,
+        query: &str,
+        parameter_types: &[Type],
+    ) -> Result<Statement, Error> {
+        prepare::prepare(&self.inner, query, parameter_types).await
+    }
+
+    /// Executes a statement, returning a vector of the resulting rows.
+    ///
+    /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
+    /// provided, 1-indexed.
+    ///
+    /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be
+    /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
+    /// with the `prepare` method.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the number of parameters provided does not match the number expected.
+    pub async fn query<T>(
+        &self,
+        statement: &T,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<Vec<Row>, Error>
+    where
+        T: ?Sized + ToStatement,
+    {
+        self.query_raw(statement, slice_iter(params))
+            .await?
+            .try_collect()
+            .await
+    }
+
+    /// The maximally flexible version of [`query`].
+    ///
+    /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
+    /// provided, 1-indexed.
+    ///
+    /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be
+    /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
+    /// with the `prepare` method.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the number of parameters provided does not match the number expected.
+    ///
+    /// [`query`]: #method.query
+    pub async fn query_raw<'a, T, I>(&self, statement: &T, params: I) -> Result<RowStream, Error>
+    where
+        T: ?Sized + ToStatement,
+        I: IntoIterator<Item = &'a (dyn ToSql + Sync)>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        let statement = statement.__convert().into_statement(self).await?;
+        query::query(&self.inner, statement, params).await
+    }
+
+    /// Pass text directly to the Postgres backend to allow it to sort out typing itself and
+    /// to save a roundtrip
+    pub async fn query_raw_txt<S, I>(&self, statement: &str, params: I) -> Result<RowStream, Error>
+    where
+        S: AsRef<str>,
+        I: IntoIterator<Item = Option<S>>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        query::query_txt(&self.inner, statement, params).await
+    }
+
+    /// Executes a statement, returning the number of rows modified.
+    ///
+    /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
+    /// provided, 1-indexed.
+    ///
+    /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be
+    /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
+    /// with the `prepare` method.
+    ///
+    /// If the statement does not modify any rows (e.g. `SELECT`), 0 is returned.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the number of parameters provided does not match the number expected.
+    pub async fn execute<T>(
+        &self,
+        statement: &T,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<u64, Error>
+    where
+        T: ?Sized + ToStatement,
+    {
+        self.execute_raw(statement, slice_iter(params)).await
+    }
+
+    /// The maximally flexible version of [`execute`].
+    ///
+    /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
+    /// provided, 1-indexed.
+    ///
+    /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be
+    /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
+    /// with the `prepare` method.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the number of parameters provided does not match the number expected.
+    ///
+    /// [`execute`]: #method.execute
+    pub async fn execute_raw<'a, T, I>(&self, statement: &T, params: I) -> Result<u64, Error>
+    where
+        T: ?Sized + ToStatement,
+        I: IntoIterator<Item = &'a (dyn ToSql + Sync)>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        let statement = statement.__convert().into_statement(self).await?;
+        query::execute(self.inner(), statement, params).await
+    }
+
+    /// Executes a sequence of SQL statements using the simple query protocol, returning the resulting rows.
+    ///
+    /// Statements should be separated by semicolons. If an error occurs, execution of the sequence will stop at that
+    /// point. The simple query protocol returns the values in rows as strings rather than in their binary encodings,
+    /// so the associated row type doesn't work with the `FromSql` trait. Rather than simply returning a list of the
+    /// rows, this method returns a list of an enum which indicates either the completion of one of the commands,
+    /// or a row of data. This preserves the framing between the separate statements in the request.
+    ///
+    /// # Warning
+    ///
+    /// Prepared statements should be use for any query which contains user-specified data, as they provided the
+    /// functionality to safely embed that data in the request. Do not form statements via string concatenation and pass
+    /// them to this method!
+    pub async fn simple_query(&self, query: &str) -> Result<Vec<SimpleQueryMessage>, Error> {
+        self.simple_query_raw(query).await?.try_collect().await
+    }
+
+    pub(crate) async fn simple_query_raw(&self, query: &str) -> Result<SimpleQueryStream, Error> {
+        simple_query::simple_query(self.inner(), query).await
+    }
+
+    /// Executes a sequence of SQL statements using the simple query protocol.
+    ///
+    /// Statements should be separated by semicolons. If an error occurs, execution of the sequence will stop at that
+    /// point. This is intended for use when, for example, initializing a database schema.
+    ///
+    /// # Warning
+    ///
+    /// Prepared statements should be use for any query which contains user-specified data, as they provided the
+    /// functionality to safely embed that data in the request. Do not form statements via string concatenation and pass
+    /// them to this method!
+    pub async fn batch_execute(&self, query: &str) -> Result<ReadyForQueryStatus, Error> {
+        simple_query::batch_execute(self.inner(), query).await
+    }
+
+    /// Begins a new database transaction.
+    ///
+    /// The transaction will roll back by default - use the `commit` method to commit it.
+    pub async fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
+        struct RollbackIfNotDone<'me> {
+            client: &'me Client,
+            done: bool,
+        }
+
+        impl Drop for RollbackIfNotDone<'_> {
+            fn drop(&mut self) {
+                if self.done {
+                    return;
+                }
+
+                let buf = self.client.inner().with_buf(|buf| {
+                    frontend::query("ROLLBACK", buf).unwrap();
+                    buf.split().freeze()
+                });
+                let _ = self
+                    .client
+                    .inner()
+                    .send(RequestMessages::Single(FrontendMessage::Raw(buf)));
+            }
+        }
+
+        // This is done, as `Future` created by this method can be dropped after
+        // `RequestMessages` is synchronously send to the `Connection` by
+        // `batch_execute()`, but before `Responses` is asynchronously polled to
+        // completion. In that case `Transaction` won't be created and thus
+        // won't be rolled back.
+        {
+            let mut cleaner = RollbackIfNotDone {
+                client: self,
+                done: false,
+            };
+            self.batch_execute("BEGIN").await?;
+            cleaner.done = true;
+        }
+
+        Ok(Transaction::new(self))
+    }
+
+    /// Returns a builder for a transaction with custom settings.
+    ///
+    /// Unlike the `transaction` method, the builder can be used to control the transaction's isolation level and other
+    /// attributes.
+    pub fn build_transaction(&mut self) -> TransactionBuilder<'_> {
+        TransactionBuilder::new(self)
+    }
+
+    /// Constructs a cancellation token that can later be used to request cancellation of a query running on the
+    /// connection associated with this client.
+    pub fn cancel_token(&self) -> CancelToken {
+        CancelToken {
+            socket_config: self.socket_config.clone(),
+            ssl_mode: self.ssl_mode,
+            process_id: self.process_id,
+            secret_key: self.secret_key,
+        }
+    }
+
+    /// Query for type information
+    pub async fn get_type(&self, oid: Oid) -> Result<Type, Error> {
+        crate::prepare::get_type(&self.inner, oid).await
+    }
+
+    /// Determines if the connection to the server has already closed.
+    ///
+    /// In that case, all future queries will fail.
+    pub fn is_closed(&self) -> bool {
+        self.inner.sender.is_closed()
+    }
+}
+
+impl fmt::Debug for Client {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Client").finish()
+    }
+}

libs/proxy/tokio-postgres2/src/codec.rs

@@ -0,0 +1,109 @@
+use bytes::{Buf, Bytes, BytesMut};
+use fallible_iterator::FallibleIterator;
+use postgres_protocol2::message::backend;
+use postgres_protocol2::message::frontend::CopyData;
+use std::io;
+use tokio_util::codec::{Decoder, Encoder};
+
+pub enum FrontendMessage {
+    Raw(Bytes),
+    CopyData(CopyData<Box<dyn Buf + Send>>),
+}
+
+pub enum BackendMessage {
+    Normal {
+        messages: BackendMessages,
+        request_complete: bool,
+    },
+    Async(backend::Message),
+}
+
+pub struct BackendMessages(BytesMut);
+
+impl BackendMessages {
+    pub fn empty() -> BackendMessages {
+        BackendMessages(BytesMut::new())
+    }
+}
+
+impl FallibleIterator for BackendMessages {
+    type Item = backend::Message;
+    type Error = io::Error;
+
+    fn next(&mut self) -> io::Result<Option<backend::Message>> {
+        backend::Message::parse(&mut self.0)
+    }
+}
+
+pub struct PostgresCodec {
+    pub max_message_size: Option<usize>,
+}
+
+impl Encoder<FrontendMessage> for PostgresCodec {
+    type Error = io::Error;
+
+    fn encode(&mut self, item: FrontendMessage, dst: &mut BytesMut) -> io::Result<()> {
+        match item {
+            FrontendMessage::Raw(buf) => dst.extend_from_slice(&buf),
+            FrontendMessage::CopyData(data) => data.write(dst),
+        }
+
+        Ok(())
+    }
+}
+
+impl Decoder for PostgresCodec {
+    type Item = BackendMessage;
+    type Error = io::Error;
+
+    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<BackendMessage>, io::Error> {
+        let mut idx = 0;
+        let mut request_complete = false;
+
+        while let Some(header) = backend::Header::parse(&src[idx..])? {
+            let len = header.len() as usize + 1;
+            if src[idx..].len() < len {
+                break;
+            }
+
+            if let Some(max) = self.max_message_size {
+                if len > max {
+                    return Err(io::Error::new(
+                        io::ErrorKind::InvalidInput,
+                        "message too large",
+                    ));
+                }
+            }
+
+            match header.tag() {
+                backend::NOTICE_RESPONSE_TAG
+                | backend::NOTIFICATION_RESPONSE_TAG
+                | backend::PARAMETER_STATUS_TAG => {
+                    if idx == 0 {
+                        let message = backend::Message::parse(src)?.unwrap();
+                        return Ok(Some(BackendMessage::Async(message)));
+                    } else {
+                        break;
+                    }
+                }
+                _ => {}
+            }
+
+            idx += len;
+
+            if header.tag() == backend::READY_FOR_QUERY_TAG {
+                request_complete = true;
+                break;
+            }
+        }
+
+        if idx == 0 {
+            Ok(None)
+        } else {
+            Ok(Some(BackendMessage::Normal {
+                messages: BackendMessages(src.split_to(idx)),
+                request_complete,
+            }))
+        }
+    }
+}

libs/proxy/tokio-postgres2/src/config.rs

@@ -0,0 +1,897 @@
+//! Connection configuration.
+
+use crate::connect::connect;
+use crate::connect_raw::connect_raw;
+use crate::tls::MakeTlsConnect;
+use crate::tls::TlsConnect;
+use crate::{Client, Connection, Error};
+use std::borrow::Cow;
+use std::str;
+use std::str::FromStr;
+use std::time::Duration;
+use std::{error, fmt, iter, mem};
+use tokio::io::{AsyncRead, AsyncWrite};
+
+pub use postgres_protocol2::authentication::sasl::ScramKeys;
+use tokio::net::TcpStream;
+
+/// Properties required of a session.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+#[non_exhaustive]
+pub enum TargetSessionAttrs {
+    /// No special properties are required.
+    Any,
+    /// The session must allow writes.
+    ReadWrite,
+}
+
+/// TLS configuration.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+#[non_exhaustive]
+pub enum SslMode {
+    /// Do not use TLS.
+    Disable,
+    /// Attempt to connect with TLS but allow sessions without.
+    Prefer,
+    /// Require the use of TLS.
+    Require,
+}
+
+/// Channel binding configuration.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+#[non_exhaustive]
+pub enum ChannelBinding {
+    /// Do not use channel binding.
+    Disable,
+    /// Attempt to use channel binding but allow sessions without.
+    Prefer,
+    /// Require the use of channel binding.
+    Require,
+}
+
+/// Replication mode configuration.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+#[non_exhaustive]
+pub enum ReplicationMode {
+    /// Physical replication.
+    Physical,
+    /// Logical replication.
+    Logical,
+}
+
+/// A host specification.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum Host {
+    /// A TCP hostname.
+    Tcp(String),
+}
+
+/// Precomputed keys which may override password during auth.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum AuthKeys {
+    /// A `ClientKey` & `ServerKey` pair for `SCRAM-SHA-256`.
+    ScramSha256(ScramKeys<32>),
+}
+
+/// Connection configuration.
+///
+/// Configuration can be parsed from libpq-style connection strings. These strings come in two formats:
+///
+/// # Key-Value
+///
+/// This format consists of space-separated key-value pairs. Values which are either the empty string or contain
+/// whitespace should be wrapped in `'`. `'` and `\` characters should be backslash-escaped.
+///
+/// ## Keys
+///
+/// * `user` - The username to authenticate with. Required.
+/// * `password` - The password to authenticate with.
+/// * `dbname` - The name of the database to connect to. Defaults to the username.
+/// * `options` - Command line options used to configure the server.
+/// * `application_name` - Sets the `application_name` parameter on the server.
+/// * `sslmode` - Controls usage of TLS. If set to `disable`, TLS will not be used. If set to `prefer`, TLS will be used
+///     if available, but not used otherwise. If set to `require`, TLS will be forced to be used. Defaults to `prefer`.
+/// * `host` - The host to connect to. On Unix platforms, if the host starts with a `/` character it is treated as the
+///     path to the directory containing Unix domain sockets. Otherwise, it is treated as a hostname. Multiple hosts
+///     can be specified, separated by commas. Each host will be tried in turn when connecting. Required if connecting
+///     with the `connect` method.
+/// * `port` - The port to connect to. Multiple ports can be specified, separated by commas. The number of ports must be
+///     either 1, in which case it will be used for all hosts, or the same as the number of hosts. Defaults to 5432 if
+///     omitted or the empty string.
+/// * `connect_timeout` - The time limit in seconds applied to each socket-level connection attempt. Note that hostnames
+///     can resolve to multiple IP addresses, and this limit is applied to each address. Defaults to no timeout.
+/// * `target_session_attrs` - Specifies requirements of the session. If set to `read-write`, the client will check that
+///     the `transaction_read_write` session parameter is set to `on`. This can be used to connect to the primary server
+///     in a database cluster as opposed to the secondary read-only mirrors. Defaults to `all`.
+/// * `channel_binding` - Controls usage of channel binding in the authentication process. If set to `disable`, channel
+///     binding will not be used. If set to `prefer`, channel binding will be used if available, but not used otherwise.
+///     If set to `require`, the authentication process will fail if channel binding is not used. Defaults to `prefer`.
+///
+/// ## Examples
+///
+/// ```not_rust
+/// host=localhost user=postgres connect_timeout=10 keepalives=0
+/// ```
+///
+/// ```not_rust
+/// host=/var/lib/postgresql,localhost port=1234 user=postgres password='password with spaces'
+/// ```
+///
+/// ```not_rust
+/// host=host1,host2,host3 port=1234,,5678 user=postgres target_session_attrs=read-write
+/// ```
+///
+/// # Url
+///
+/// This format resembles a URL with a scheme of either `postgres://` or `postgresql://`. All components are optional,
+/// and the format accepts query parameters for all of the key-value pairs described in the section above. Multiple
+/// host/port pairs can be comma-separated. Unix socket paths in the host section of the URL should be percent-encoded,
+/// as the path component of the URL specifies the database name.
+///
+/// ## Examples
+///
+/// ```not_rust
+/// postgresql://user@localhost
+/// ```
+///
+/// ```not_rust
+/// postgresql://user:password@%2Fvar%2Flib%2Fpostgresql/mydb?connect_timeout=10
+/// ```
+///
+/// ```not_rust
+/// postgresql://user@host1:1234,host2,host3:5678?target_session_attrs=read-write
+/// ```
+///
+/// ```not_rust
+/// postgresql:///mydb?user=user&host=/var/lib/postgresql
+/// ```
+#[derive(Clone, PartialEq, Eq)]
+pub struct Config {
+    pub(crate) user: Option<String>,
+    pub(crate) password: Option<Vec<u8>>,
+    pub(crate) auth_keys: Option<Box<AuthKeys>>,
+    pub(crate) dbname: Option<String>,
+    pub(crate) options: Option<String>,
+    pub(crate) application_name: Option<String>,
+    pub(crate) ssl_mode: SslMode,
+    pub(crate) host: Vec<Host>,
+    pub(crate) port: Vec<u16>,
+    pub(crate) connect_timeout: Option<Duration>,
+    pub(crate) target_session_attrs: TargetSessionAttrs,
+    pub(crate) channel_binding: ChannelBinding,
+    pub(crate) replication_mode: Option<ReplicationMode>,
+    pub(crate) max_backend_message_size: Option<usize>,
+}
+
+impl Default for Config {
+    fn default() -> Config {
+        Config::new()
+    }
+}
+
+impl Config {
+    /// Creates a new configuration.
+    pub fn new() -> Config {
+        Config {
+            user: None,
+            password: None,
+            auth_keys: None,
+            dbname: None,
+            options: None,
+            application_name: None,
+            ssl_mode: SslMode::Prefer,
+            host: vec![],
+            port: vec![],
+            connect_timeout: None,
+            target_session_attrs: TargetSessionAttrs::Any,
+            channel_binding: ChannelBinding::Prefer,
+            replication_mode: None,
+            max_backend_message_size: None,
+        }
+    }
+
+    /// Sets the user to authenticate with.
+    ///
+    /// Required.
+    pub fn user(&mut self, user: &str) -> &mut Config {
+        self.user = Some(user.to_string());
+        self
+    }
+
+    /// Gets the user to authenticate with, if one has been configured with
+    /// the `user` method.
+    pub fn get_user(&self) -> Option<&str> {
+        self.user.as_deref()
+    }
+
+    /// Sets the password to authenticate with.
+    pub fn password<T>(&mut self, password: T) -> &mut Config
+    where
+        T: AsRef<[u8]>,
+    {
+        self.password = Some(password.as_ref().to_vec());
+        self
+    }
+
+    /// Gets the password to authenticate with, if one has been configured with
+    /// the `password` method.
+    pub fn get_password(&self) -> Option<&[u8]> {
+        self.password.as_deref()
+    }
+
+    /// Sets precomputed protocol-specific keys to authenticate with.
+    /// When set, this option will override `password`.
+    /// See [`AuthKeys`] for more information.
+    pub fn auth_keys(&mut self, keys: AuthKeys) -> &mut Config {
+        self.auth_keys = Some(Box::new(keys));
+        self
+    }
+
+    /// Gets precomputed protocol-specific keys to authenticate with.
+    /// if one has been configured with the `auth_keys` method.
+    pub fn get_auth_keys(&self) -> Option<AuthKeys> {
+        self.auth_keys.as_deref().copied()
+    }
+
+    /// Sets the name of the database to connect to.
+    ///
+    /// Defaults to the user.
+    pub fn dbname(&mut self, dbname: &str) -> &mut Config {
+        self.dbname = Some(dbname.to_string());
+        self
+    }
+
+    /// Gets the name of the database to connect to, if one has been configured
+    /// with the `dbname` method.
+    pub fn get_dbname(&self) -> Option<&str> {
+        self.dbname.as_deref()
+    }
+
+    /// Sets command line options used to configure the server.
+    pub fn options(&mut self, options: &str) -> &mut Config {
+        self.options = Some(options.to_string());
+        self
+    }
+
+    /// Gets the command line options used to configure the server, if the
+    /// options have been set with the `options` method.
+    pub fn get_options(&self) -> Option<&str> {
+        self.options.as_deref()
+    }
+
+    /// Sets the value of the `application_name` runtime parameter.
+    pub fn application_name(&mut self, application_name: &str) -> &mut Config {
+        self.application_name = Some(application_name.to_string());
+        self
+    }
+
+    /// Gets the value of the `application_name` runtime parameter, if it has
+    /// been set with the `application_name` method.
+    pub fn get_application_name(&self) -> Option<&str> {
+        self.application_name.as_deref()
+    }
+
+    /// Sets the SSL configuration.
+    ///
+    /// Defaults to `prefer`.
+    pub fn ssl_mode(&mut self, ssl_mode: SslMode) -> &mut Config {
+        self.ssl_mode = ssl_mode;
+        self
+    }
+
+    /// Gets the SSL configuration.
+    pub fn get_ssl_mode(&self) -> SslMode {
+        self.ssl_mode
+    }
+
+    /// Adds a host to the configuration.
+    ///
+    /// Multiple hosts can be specified by calling this method multiple times, and each will be tried in order.
+    pub fn host(&mut self, host: &str) -> &mut Config {
+        self.host.push(Host::Tcp(host.to_string()));
+        self
+    }
+
+    /// Gets the hosts that have been added to the configuration with `host`.
+    pub fn get_hosts(&self) -> &[Host] {
+        &self.host
+    }
+
+    /// Adds a port to the configuration.
+    ///
+    /// Multiple ports can be specified by calling this method multiple times. There must either be no ports, in which
+    /// case the default of 5432 is used, a single port, in which it is used for all hosts, or the same number of ports
+    /// as hosts.
+    pub fn port(&mut self, port: u16) -> &mut Config {
+        self.port.push(port);
+        self
+    }
+
+    /// Gets the ports that have been added to the configuration with `port`.
+    pub fn get_ports(&self) -> &[u16] {
+        &self.port
+    }
+
+    /// Sets the timeout applied to socket-level connection attempts.
+    ///
+    /// Note that hostnames can resolve to multiple IP addresses, and this timeout will apply to each address of each
+    /// host separately. Defaults to no limit.
+    pub fn connect_timeout(&mut self, connect_timeout: Duration) -> &mut Config {
+        self.connect_timeout = Some(connect_timeout);
+        self
+    }
+
+    /// Gets the connection timeout, if one has been set with the
+    /// `connect_timeout` method.
+    pub fn get_connect_timeout(&self) -> Option<&Duration> {
+        self.connect_timeout.as_ref()
+    }
+
+    /// Sets the requirements of the session.
+    ///
+    /// This can be used to connect to the primary server in a clustered database rather than one of the read-only
+    /// secondary servers. Defaults to `Any`.
+    pub fn target_session_attrs(
+        &mut self,
+        target_session_attrs: TargetSessionAttrs,
+    ) -> &mut Config {
+        self.target_session_attrs = target_session_attrs;
+        self
+    }
+
+    /// Gets the requirements of the session.
+    pub fn get_target_session_attrs(&self) -> TargetSessionAttrs {
+        self.target_session_attrs
+    }
+
+    /// Sets the channel binding behavior.
+    ///
+    /// Defaults to `prefer`.
+    pub fn channel_binding(&mut self, channel_binding: ChannelBinding) -> &mut Config {
+        self.channel_binding = channel_binding;
+        self
+    }
+
+    /// Gets the channel binding behavior.
+    pub fn get_channel_binding(&self) -> ChannelBinding {
+        self.channel_binding
+    }
+
+    /// Set replication mode.
+    pub fn replication_mode(&mut self, replication_mode: ReplicationMode) -> &mut Config {
+        self.replication_mode = Some(replication_mode);
+        self
+    }
+
+    /// Get replication mode.
+    pub fn get_replication_mode(&self) -> Option<ReplicationMode> {
+        self.replication_mode
+    }
+
+    /// Set limit for backend messages size.
+    pub fn max_backend_message_size(&mut self, max_backend_message_size: usize) -> &mut Config {
+        self.max_backend_message_size = Some(max_backend_message_size);
+        self
+    }
+
+    /// Get limit for backend messages size.
+    pub fn get_max_backend_message_size(&self) -> Option<usize> {
+        self.max_backend_message_size
+    }
+
+    fn param(&mut self, key: &str, value: &str) -> Result<(), Error> {
+        match key {
+            "user" => {
+                self.user(value);
+            }
+            "password" => {
+                self.password(value);
+            }
+            "dbname" => {
+                self.dbname(value);
+            }
+            "options" => {
+                self.options(value);
+            }
+            "application_name" => {
+                self.application_name(value);
+            }
+            "sslmode" => {
+                let mode = match value {
+                    "disable" => SslMode::Disable,
+                    "prefer" => SslMode::Prefer,
+                    "require" => SslMode::Require,
+                    _ => return Err(Error::config_parse(Box::new(InvalidValue("sslmode")))),
+                };
+                self.ssl_mode(mode);
+            }
+            "host" => {
+                for host in value.split(',') {
+                    self.host(host);
+                }
+            }
+            "port" => {
+                for port in value.split(',') {
+                    let port = if port.is_empty() {
+                        5432
+                    } else {
+                        port.parse()
+                            .map_err(|_| Error::config_parse(Box::new(InvalidValue("port"))))?
+                    };
+                    self.port(port);
+                }
+            }
+            "connect_timeout" => {
+                let timeout = value
+                    .parse::<i64>()
+                    .map_err(|_| Error::config_parse(Box::new(InvalidValue("connect_timeout"))))?;
+                if timeout > 0 {
+                    self.connect_timeout(Duration::from_secs(timeout as u64));
+                }
+            }
+            "target_session_attrs" => {
+                let target_session_attrs = match value {
+                    "any" => TargetSessionAttrs::Any,
+                    "read-write" => TargetSessionAttrs::ReadWrite,
+                    _ => {
+                        return Err(Error::config_parse(Box::new(InvalidValue(
+                            "target_session_attrs",
+                        ))));
+                    }
+                };
+                self.target_session_attrs(target_session_attrs);
+            }
+            "channel_binding" => {
+                let channel_binding = match value {
+                    "disable" => ChannelBinding::Disable,
+                    "prefer" => ChannelBinding::Prefer,
+                    "require" => ChannelBinding::Require,
+                    _ => {
+                        return Err(Error::config_parse(Box::new(InvalidValue(
+                            "channel_binding",
+                        ))))
+                    }
+                };
+                self.channel_binding(channel_binding);
+            }
+            "max_backend_message_size" => {
+                let limit = value.parse::<usize>().map_err(|_| {
+                    Error::config_parse(Box::new(InvalidValue("max_backend_message_size")))
+                })?;
+                if limit > 0 {
+                    self.max_backend_message_size(limit);
+                }
+            }
+            key => {
+                return Err(Error::config_parse(Box::new(UnknownOption(
+                    key.to_string(),
+                ))));
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Opens a connection to a PostgreSQL database.
+    ///
+    /// Requires the `runtime` Cargo feature (enabled by default).
+    pub async fn connect<T>(
+        &self,
+        tls: T,
+    ) -> Result<(Client, Connection<TcpStream, T::Stream>), Error>
+    where
+        T: MakeTlsConnect<TcpStream>,
+    {
+        connect(tls, self).await
+    }
+
+    /// Connects to a PostgreSQL database over an arbitrary stream.
+    ///
+    /// All of the settings other than `user`, `password`, `dbname`, `options`, and `application_name` name are ignored.
+    pub async fn connect_raw<S, T>(
+        &self,
+        stream: S,
+        tls: T,
+    ) -> Result<(Client, Connection<S, T::Stream>), Error>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+        T: TlsConnect<S>,
+    {
+        connect_raw(stream, tls, self).await
+    }
+}
+
+impl FromStr for Config {
+    type Err = Error;
+
+    fn from_str(s: &str) -> Result<Config, Error> {
+        match UrlParser::parse(s)? {
+            Some(config) => Ok(config),
+            None => Parser::parse(s),
+        }
+    }
+}
+
+// Omit password from debug output
+impl fmt::Debug for Config {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        struct Redaction {}
+        impl fmt::Debug for Redaction {
+            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+                write!(f, "_")
+            }
+        }
+
+        f.debug_struct("Config")
+            .field("user", &self.user)
+            .field("password", &self.password.as_ref().map(|_| Redaction {}))
+            .field("dbname", &self.dbname)
+            .field("options", &self.options)
+            .field("application_name", &self.application_name)
+            .field("ssl_mode", &self.ssl_mode)
+            .field("host", &self.host)
+            .field("port", &self.port)
+            .field("connect_timeout", &self.connect_timeout)
+            .field("target_session_attrs", &self.target_session_attrs)
+            .field("channel_binding", &self.channel_binding)
+            .field("replication", &self.replication_mode)
+            .finish()
+    }
+}
+
+#[derive(Debug)]
+struct UnknownOption(String);
+
+impl fmt::Display for UnknownOption {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(fmt, "unknown option `{}`", self.0)
+    }
+}
+
+impl error::Error for UnknownOption {}
+
+#[derive(Debug)]
+struct InvalidValue(&'static str);
+
+impl fmt::Display for InvalidValue {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(fmt, "invalid value for option `{}`", self.0)
+    }
+}
+
+impl error::Error for InvalidValue {}
+
+struct Parser<'a> {
+    s: &'a str,
+    it: iter::Peekable<str::CharIndices<'a>>,
+}
+
+impl<'a> Parser<'a> {
+    fn parse(s: &'a str) -> Result<Config, Error> {
+        let mut parser = Parser {
+            s,
+            it: s.char_indices().peekable(),
+        };
+
+        let mut config = Config::new();
+
+        while let Some((key, value)) = parser.parameter()? {
+            config.param(key, &value)?;
+        }
+
+        Ok(config)
+    }
+
+    fn skip_ws(&mut self) {
+        self.take_while(char::is_whitespace);
+    }
+
+    fn take_while<F>(&mut self, f: F) -> &'a str
+    where
+        F: Fn(char) -> bool,
+    {
+        let start = match self.it.peek() {
+            Some(&(i, _)) => i,
+            None => return "",
+        };
+
+        loop {
+            match self.it.peek() {
+                Some(&(_, c)) if f(c) => {
+                    self.it.next();
+                }
+                Some(&(i, _)) => return &self.s[start..i],
+                None => return &self.s[start..],
+            }
+        }
+    }
+
+    fn eat(&mut self, target: char) -> Result<(), Error> {
+        match self.it.next() {
+            Some((_, c)) if c == target => Ok(()),
+            Some((i, c)) => {
+                let m = format!(
+                    "unexpected character at byte {}: expected `{}` but got `{}`",
+                    i, target, c
+                );
+                Err(Error::config_parse(m.into()))
+            }
+            None => Err(Error::config_parse("unexpected EOF".into())),
+        }
+    }
+
+    fn eat_if(&mut self, target: char) -> bool {
+        match self.it.peek() {
+            Some(&(_, c)) if c == target => {
+                self.it.next();
+                true
+            }
+            _ => false,
+        }
+    }
+
+    fn keyword(&mut self) -> Option<&'a str> {
+        let s = self.take_while(|c| match c {
+            c if c.is_whitespace() => false,
+            '=' => false,
+            _ => true,
+        });
+
+        if s.is_empty() {
+            None
+        } else {
+            Some(s)
+        }
+    }
+
+    fn value(&mut self) -> Result<String, Error> {
+        let value = if self.eat_if('\'') {
+            let value = self.quoted_value()?;
+            self.eat('\'')?;
+            value
+        } else {
+            self.simple_value()?
+        };
+
+        Ok(value)
+    }
+
+    fn simple_value(&mut self) -> Result<String, Error> {
+        let mut value = String::new();
+
+        while let Some(&(_, c)) = self.it.peek() {
+            if c.is_whitespace() {
+                break;
+            }
+
+            self.it.next();
+            if c == '\\' {
+                if let Some((_, c2)) = self.it.next() {
+                    value.push(c2);
+                }
+            } else {
+                value.push(c);
+            }
+        }
+
+        if value.is_empty() {
+            return Err(Error::config_parse("unexpected EOF".into()));
+        }
+
+        Ok(value)
+    }
+
+    fn quoted_value(&mut self) -> Result<String, Error> {
+        let mut value = String::new();
+
+        while let Some(&(_, c)) = self.it.peek() {
+            if c == '\'' {
+                return Ok(value);
+            }
+
+            self.it.next();
+            if c == '\\' {
+                if let Some((_, c2)) = self.it.next() {
+                    value.push(c2);
+                }
+            } else {
+                value.push(c);
+            }
+        }
+
+        Err(Error::config_parse(
+            "unterminated quoted connection parameter value".into(),
+        ))
+    }
+
+    fn parameter(&mut self) -> Result<Option<(&'a str, String)>, Error> {
+        self.skip_ws();
+        let keyword = match self.keyword() {
+            Some(keyword) => keyword,
+            None => return Ok(None),
+        };
+        self.skip_ws();
+        self.eat('=')?;
+        self.skip_ws();
+        let value = self.value()?;
+
+        Ok(Some((keyword, value)))
+    }
+}
+
+// This is a pretty sloppy "URL" parser, but it matches the behavior of libpq, where things really aren't very strict
+struct UrlParser<'a> {
+    s: &'a str,
+    config: Config,
+}
+
+impl<'a> UrlParser<'a> {
+    fn parse(s: &'a str) -> Result<Option<Config>, Error> {
+        let s = match Self::remove_url_prefix(s) {
+            Some(s) => s,
+            None => return Ok(None),
+        };
+
+        let mut parser = UrlParser {
+            s,
+            config: Config::new(),
+        };
+
+        parser.parse_credentials()?;
+        parser.parse_host()?;
+        parser.parse_path()?;
+        parser.parse_params()?;
+
+        Ok(Some(parser.config))
+    }
+
+    fn remove_url_prefix(s: &str) -> Option<&str> {
+        for prefix in &["postgres://", "postgresql://"] {
+            if let Some(stripped) = s.strip_prefix(prefix) {
+                return Some(stripped);
+            }
+        }
+
+        None
+    }
+
+    fn take_until(&mut self, end: &[char]) -> Option<&'a str> {
+        match self.s.find(end) {
+            Some(pos) => {
+                let (head, tail) = self.s.split_at(pos);
+                self.s = tail;
+                Some(head)
+            }
+            None => None,
+        }
+    }
+
+    fn take_all(&mut self) -> &'a str {
+        mem::take(&mut self.s)
+    }
+
+    fn eat_byte(&mut self) {
+        self.s = &self.s[1..];
+    }
+
+    fn parse_credentials(&mut self) -> Result<(), Error> {
+        let creds = match self.take_until(&['@']) {
+            Some(creds) => creds,
+            None => return Ok(()),
+        };
+        self.eat_byte();
+
+        let mut it = creds.splitn(2, ':');
+        let user = self.decode(it.next().unwrap())?;
+        self.config.user(&user);
+
+        if let Some(password) = it.next() {
+            let password = Cow::from(percent_encoding::percent_decode(password.as_bytes()));
+            self.config.password(password);
+        }
+
+        Ok(())
+    }
+
+    fn parse_host(&mut self) -> Result<(), Error> {
+        let host = match self.take_until(&['/', '?']) {
+            Some(host) => host,
+            None => self.take_all(),
+        };
+
+        if host.is_empty() {
+            return Ok(());
+        }
+
+        for chunk in host.split(',') {
+            let (host, port) = if chunk.starts_with('[') {
+                let idx = match chunk.find(']') {
+                    Some(idx) => idx,
+                    None => return Err(Error::config_parse(InvalidValue("host").into())),
+                };
+
+                let host = &chunk[1..idx];
+                let remaining = &chunk[idx + 1..];
+                let port = if let Some(port) = remaining.strip_prefix(':') {
+                    Some(port)
+                } else if remaining.is_empty() {
+                    None
+                } else {
+                    return Err(Error::config_parse(InvalidValue("host").into()));
+                };
+
+                (host, port)
+            } else {
+                let mut it = chunk.splitn(2, ':');
+                (it.next().unwrap(), it.next())
+            };
+
+            self.host_param(host)?;
+            let port = self.decode(port.unwrap_or("5432"))?;
+            self.config.param("port", &port)?;
+        }
+
+        Ok(())
+    }
+
+    fn parse_path(&mut self) -> Result<(), Error> {
+        if !self.s.starts_with('/') {
+            return Ok(());
+        }
+        self.eat_byte();
+
+        let dbname = match self.take_until(&['?']) {
+            Some(dbname) => dbname,
+            None => self.take_all(),
+        };
+
+        if !dbname.is_empty() {
+            self.config.dbname(&self.decode(dbname)?);
+        }
+
+        Ok(())
+    }
+
+    fn parse_params(&mut self) -> Result<(), Error> {
+        if !self.s.starts_with('?') {
+            return Ok(());
+        }
+        self.eat_byte();
+
+        while !self.s.is_empty() {
+            let key = match self.take_until(&['=']) {
+                Some(key) => self.decode(key)?,
+                None => return Err(Error::config_parse("unterminated parameter".into())),
+            };
+            self.eat_byte();
+
+            let value = match self.take_until(&['&']) {
+                Some(value) => {
+                    self.eat_byte();
+                    value
+                }
+                None => self.take_all(),
+            };
+
+            if key == "host" {
+                self.host_param(value)?;
+            } else {
+                let value = self.decode(value)?;
+                self.config.param(&key, &value)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    fn host_param(&mut self, s: &str) -> Result<(), Error> {
+        let s = self.decode(s)?;
+        self.config.param("host", &s)
+    }
+
+    fn decode(&self, s: &'a str) -> Result<Cow<'a, str>, Error> {
+        percent_encoding::percent_decode(s.as_bytes())
+            .decode_utf8()
+            .map_err(|e| Error::config_parse(e.into()))
+    }
+}

libs/proxy/tokio-postgres2/src/connect.rs

@@ -0,0 +1,112 @@
+use crate::client::SocketConfig;
+use crate::config::{Host, TargetSessionAttrs};
+use crate::connect_raw::connect_raw;
+use crate::connect_socket::connect_socket;
+use crate::tls::{MakeTlsConnect, TlsConnect};
+use crate::{Client, Config, Connection, Error, SimpleQueryMessage};
+use futures_util::{future, pin_mut, Future, FutureExt, Stream};
+use std::io;
+use std::task::Poll;
+use tokio::net::TcpStream;
+
+pub async fn connect<T>(
+    mut tls: T,
+    config: &Config,
+) -> Result<(Client, Connection<TcpStream, T::Stream>), Error>
+where
+    T: MakeTlsConnect<TcpStream>,
+{
+    if config.host.is_empty() {
+        return Err(Error::config("host missing".into()));
+    }
+
+    if config.port.len() > 1 && config.port.len() != config.host.len() {
+        return Err(Error::config("invalid number of ports".into()));
+    }
+
+    let mut error = None;
+    for (i, host) in config.host.iter().enumerate() {
+        let port = config
+            .port
+            .get(i)
+            .or_else(|| config.port.first())
+            .copied()
+            .unwrap_or(5432);
+
+        let hostname = match host {
+            Host::Tcp(host) => host.as_str(),
+        };
+
+        let tls = tls
+            .make_tls_connect(hostname)
+            .map_err(|e| Error::tls(e.into()))?;
+
+        match connect_once(host, port, tls, config).await {
+            Ok((client, connection)) => return Ok((client, connection)),
+            Err(e) => error = Some(e),
+        }
+    }
+
+    Err(error.unwrap())
+}
+
+async fn connect_once<T>(
+    host: &Host,
+    port: u16,
+    tls: T,
+    config: &Config,
+) -> Result<(Client, Connection<TcpStream, T::Stream>), Error>
+where
+    T: TlsConnect<TcpStream>,
+{
+    let socket = connect_socket(host, port, config.connect_timeout).await?;
+    let (mut client, mut connection) = connect_raw(socket, tls, config).await?;
+
+    if let TargetSessionAttrs::ReadWrite = config.target_session_attrs {
+        let rows = client.simple_query_raw("SHOW transaction_read_only");
+        pin_mut!(rows);
+
+        let rows = future::poll_fn(|cx| {
+            if connection.poll_unpin(cx)?.is_ready() {
+                return Poll::Ready(Err(Error::closed()));
+            }
+
+            rows.as_mut().poll(cx)
+        })
+        .await?;
+        pin_mut!(rows);
+
+        loop {
+            let next = future::poll_fn(|cx| {
+                if connection.poll_unpin(cx)?.is_ready() {
+                    return Poll::Ready(Some(Err(Error::closed())));
+                }
+
+                rows.as_mut().poll_next(cx)
+            });
+
+            match next.await.transpose()? {
+                Some(SimpleQueryMessage::Row(row)) => {
+                    if row.try_get(0)? == Some("on") {
+                        return Err(Error::connect(io::Error::new(
+                            io::ErrorKind::PermissionDenied,
+                            "database does not allow writes",
+                        )));
+                    } else {
+                        break;
+                    }
+                }
+                Some(_) => {}
+                None => return Err(Error::unexpected_message()),
+            }
+        }
+    }
+
+    client.set_socket_config(SocketConfig {
+        host: host.clone(),
+        port,
+        connect_timeout: config.connect_timeout,
+    });
+
+    Ok((client, connection))
+}

libs/proxy/tokio-postgres2/src/connect_raw.rs

@@ -0,0 +1,359 @@
+use crate::codec::{BackendMessage, BackendMessages, FrontendMessage, PostgresCodec};
+use crate::config::{self, AuthKeys, Config, ReplicationMode};
+use crate::connect_tls::connect_tls;
+use crate::maybe_tls_stream::MaybeTlsStream;
+use crate::tls::{TlsConnect, TlsStream};
+use crate::{Client, Connection, Error};
+use bytes::BytesMut;
+use fallible_iterator::FallibleIterator;
+use futures_util::{ready, Sink, SinkExt, Stream, TryStreamExt};
+use postgres_protocol2::authentication;
+use postgres_protocol2::authentication::sasl;
+use postgres_protocol2::authentication::sasl::ScramSha256;
+use postgres_protocol2::message::backend::{AuthenticationSaslBody, Message};
+use postgres_protocol2::message::frontend;
+use std::collections::{HashMap, VecDeque};
+use std::io;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::sync::mpsc;
+use tokio_util::codec::Framed;
+
+pub struct StartupStream<S, T> {
+    inner: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
+    buf: BackendMessages,
+    delayed: VecDeque<BackendMessage>,
+}
+
+impl<S, T> Sink<FrontendMessage> for StartupStream<S, T>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    type Error = io::Error;
+
+    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.inner).poll_ready(cx)
+    }
+
+    fn start_send(mut self: Pin<&mut Self>, item: FrontendMessage) -> io::Result<()> {
+        Pin::new(&mut self.inner).start_send(item)
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.inner).poll_flush(cx)
+    }
+
+    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.inner).poll_close(cx)
+    }
+}
+
+impl<S, T> Stream for StartupStream<S, T>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    type Item = io::Result<Message>;
+
+    fn poll_next(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<io::Result<Message>>> {
+        loop {
+            match self.buf.next() {
+                Ok(Some(message)) => return Poll::Ready(Some(Ok(message))),
+                Ok(None) => {}
+                Err(e) => return Poll::Ready(Some(Err(e))),
+            }
+
+            match ready!(Pin::new(&mut self.inner).poll_next(cx)) {
+                Some(Ok(BackendMessage::Normal { messages, .. })) => self.buf = messages,
+                Some(Ok(BackendMessage::Async(message))) => return Poll::Ready(Some(Ok(message))),
+                Some(Err(e)) => return Poll::Ready(Some(Err(e))),
+                None => return Poll::Ready(None),
+            }
+        }
+    }
+}
+
+pub async fn connect_raw<S, T>(
+    stream: S,
+    tls: T,
+    config: &Config,
+) -> Result<(Client, Connection<S, T::Stream>), Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: TlsConnect<S>,
+{
+    let stream = connect_tls(stream, config.ssl_mode, tls).await?;
+
+    let mut stream = StartupStream {
+        inner: Framed::new(
+            stream,
+            PostgresCodec {
+                max_message_size: config.max_backend_message_size,
+            },
+        ),
+        buf: BackendMessages::empty(),
+        delayed: VecDeque::new(),
+    };
+
+    startup(&mut stream, config).await?;
+    authenticate(&mut stream, config).await?;
+    let (process_id, secret_key, parameters) = read_info(&mut stream).await?;
+
+    let (sender, receiver) = mpsc::unbounded_channel();
+    let client = Client::new(sender, config.ssl_mode, process_id, secret_key);
+    let connection = Connection::new(stream.inner, stream.delayed, parameters, receiver);
+
+    Ok((client, connection))
+}
+
+async fn startup<S, T>(stream: &mut StartupStream<S, T>, config: &Config) -> Result<(), Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    let mut params = vec![("client_encoding", "UTF8")];
+    if let Some(user) = &config.user {
+        params.push(("user", &**user));
+    }
+    if let Some(dbname) = &config.dbname {
+        params.push(("database", &**dbname));
+    }
+    if let Some(options) = &config.options {
+        params.push(("options", &**options));
+    }
+    if let Some(application_name) = &config.application_name {
+        params.push(("application_name", &**application_name));
+    }
+    if let Some(replication_mode) = &config.replication_mode {
+        match replication_mode {
+            ReplicationMode::Physical => params.push(("replication", "true")),
+            ReplicationMode::Logical => params.push(("replication", "database")),
+        }
+    }
+
+    let mut buf = BytesMut::new();
+    frontend::startup_message(params, &mut buf).map_err(Error::encode)?;
+
+    stream
+        .send(FrontendMessage::Raw(buf.freeze()))
+        .await
+        .map_err(Error::io)
+}
+
+async fn authenticate<S, T>(stream: &mut StartupStream<S, T>, config: &Config) -> Result<(), Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: TlsStream + Unpin,
+{
+    match stream.try_next().await.map_err(Error::io)? {
+        Some(Message::AuthenticationOk) => {
+            can_skip_channel_binding(config)?;
+            return Ok(());
+        }
+        Some(Message::AuthenticationCleartextPassword) => {
+            can_skip_channel_binding(config)?;
+
+            let pass = config
+                .password
+                .as_ref()
+                .ok_or_else(|| Error::config("password missing".into()))?;
+
+            authenticate_password(stream, pass).await?;
+        }
+        Some(Message::AuthenticationMd5Password(body)) => {
+            can_skip_channel_binding(config)?;
+
+            let user = config
+                .user
+                .as_ref()
+                .ok_or_else(|| Error::config("user missing".into()))?;
+            let pass = config
+                .password
+                .as_ref()
+                .ok_or_else(|| Error::config("password missing".into()))?;
+
+            let output = authentication::md5_hash(user.as_bytes(), pass, body.salt());
+            authenticate_password(stream, output.as_bytes()).await?;
+        }
+        Some(Message::AuthenticationSasl(body)) => {
+            authenticate_sasl(stream, body, config).await?;
+        }
+        Some(Message::AuthenticationKerberosV5)
+        | Some(Message::AuthenticationScmCredential)
+        | Some(Message::AuthenticationGss)
+        | Some(Message::AuthenticationSspi) => {
+            return Err(Error::authentication(
+                "unsupported authentication method".into(),
+            ))
+        }
+        Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
+        Some(_) => return Err(Error::unexpected_message()),
+        None => return Err(Error::closed()),
+    }
+
+    match stream.try_next().await.map_err(Error::io)? {
+        Some(Message::AuthenticationOk) => Ok(()),
+        Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
+        Some(_) => Err(Error::unexpected_message()),
+        None => Err(Error::closed()),
+    }
+}
+
+fn can_skip_channel_binding(config: &Config) -> Result<(), Error> {
+    match config.channel_binding {
+        config::ChannelBinding::Disable | config::ChannelBinding::Prefer => Ok(()),
+        config::ChannelBinding::Require => Err(Error::authentication(
+            "server did not use channel binding".into(),
+        )),
+    }
+}
+
+async fn authenticate_password<S, T>(
+    stream: &mut StartupStream<S, T>,
+    password: &[u8],
+) -> Result<(), Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    let mut buf = BytesMut::new();
+    frontend::password_message(password, &mut buf).map_err(Error::encode)?;
+
+    stream
+        .send(FrontendMessage::Raw(buf.freeze()))
+        .await
+        .map_err(Error::io)
+}
+
+async fn authenticate_sasl<S, T>(
+    stream: &mut StartupStream<S, T>,
+    body: AuthenticationSaslBody,
+    config: &Config,
+) -> Result<(), Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: TlsStream + Unpin,
+{
+    let mut has_scram = false;
+    let mut has_scram_plus = false;
+    let mut mechanisms = body.mechanisms();
+    while let Some(mechanism) = mechanisms.next().map_err(Error::parse)? {
+        match mechanism {
+            sasl::SCRAM_SHA_256 => has_scram = true,
+            sasl::SCRAM_SHA_256_PLUS => has_scram_plus = true,
+            _ => {}
+        }
+    }
+
+    let channel_binding = stream
+        .inner
+        .get_ref()
+        .channel_binding()
+        .tls_server_end_point
+        .filter(|_| config.channel_binding != config::ChannelBinding::Disable)
+        .map(sasl::ChannelBinding::tls_server_end_point);
+
+    let (channel_binding, mechanism) = if has_scram_plus {
+        match channel_binding {
+            Some(channel_binding) => (channel_binding, sasl::SCRAM_SHA_256_PLUS),
+            None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
+        }
+    } else if has_scram {
+        match channel_binding {
+            Some(_) => (sasl::ChannelBinding::unrequested(), sasl::SCRAM_SHA_256),
+            None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
+        }
+    } else {
+        return Err(Error::authentication("unsupported SASL mechanism".into()));
+    };
+
+    if mechanism != sasl::SCRAM_SHA_256_PLUS {
+        can_skip_channel_binding(config)?;
+    }
+
+    let mut scram = if let Some(AuthKeys::ScramSha256(keys)) = config.get_auth_keys() {
+        ScramSha256::new_with_keys(keys, channel_binding)
+    } else if let Some(password) = config.get_password() {
+        ScramSha256::new(password, channel_binding)
+    } else {
+        return Err(Error::config("password or auth keys missing".into()));
+    };
+
+    let mut buf = BytesMut::new();
+    frontend::sasl_initial_response(mechanism, scram.message(), &mut buf).map_err(Error::encode)?;
+    stream
+        .send(FrontendMessage::Raw(buf.freeze()))
+        .await
+        .map_err(Error::io)?;
+
+    let body = match stream.try_next().await.map_err(Error::io)? {
+        Some(Message::AuthenticationSaslContinue(body)) => body,
+        Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
+        Some(_) => return Err(Error::unexpected_message()),
+        None => return Err(Error::closed()),
+    };
+
+    scram
+        .update(body.data())
+        .await
+        .map_err(|e| Error::authentication(e.into()))?;
+
+    let mut buf = BytesMut::new();
+    frontend::sasl_response(scram.message(), &mut buf).map_err(Error::encode)?;
+    stream
+        .send(FrontendMessage::Raw(buf.freeze()))
+        .await
+        .map_err(Error::io)?;
+
+    let body = match stream.try_next().await.map_err(Error::io)? {
+        Some(Message::AuthenticationSaslFinal(body)) => body,
+        Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
+        Some(_) => return Err(Error::unexpected_message()),
+        None => return Err(Error::closed()),
+    };
+
+    scram
+        .finish(body.data())
+        .map_err(|e| Error::authentication(e.into()))?;
+
+    Ok(())
+}
+
+async fn read_info<S, T>(
+    stream: &mut StartupStream<S, T>,
+) -> Result<(i32, i32, HashMap<String, String>), Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    let mut process_id = 0;
+    let mut secret_key = 0;
+    let mut parameters = HashMap::new();
+
+    loop {
+        match stream.try_next().await.map_err(Error::io)? {
+            Some(Message::BackendKeyData(body)) => {
+                process_id = body.process_id();
+                secret_key = body.secret_key();
+            }
+            Some(Message::ParameterStatus(body)) => {
+                parameters.insert(
+                    body.name().map_err(Error::parse)?.to_string(),
+                    body.value().map_err(Error::parse)?.to_string(),
+                );
+            }
+            Some(msg @ Message::NoticeResponse(_)) => {
+                stream.delayed.push_back(BackendMessage::Async(msg))
+            }
+            Some(Message::ReadyForQuery(_)) => return Ok((process_id, secret_key, parameters)),
+            Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
+            Some(_) => return Err(Error::unexpected_message()),
+            None => return Err(Error::closed()),
+        }
+    }
+}

libs/proxy/tokio-postgres2/src/connect_socket.rs

@@ -0,0 +1,65 @@
+use crate::config::Host;
+use crate::Error;
+use std::future::Future;
+use std::io;
+use std::time::Duration;
+use tokio::net::{self, TcpStream};
+use tokio::time;
+
+pub(crate) async fn connect_socket(
+    host: &Host,
+    port: u16,
+    connect_timeout: Option<Duration>,
+) -> Result<TcpStream, Error> {
+    match host {
+        Host::Tcp(host) => {
+            let addrs = net::lookup_host((&**host, port))
+                .await
+                .map_err(Error::connect)?;
+
+            let mut last_err = None;
+
+            for addr in addrs {
+                let stream =
+                    match connect_with_timeout(TcpStream::connect(addr), connect_timeout).await {
+                        Ok(stream) => stream,
+                        Err(e) => {
+                            last_err = Some(e);
+                            continue;
+                        }
+                    };
+
+                stream.set_nodelay(true).map_err(Error::connect)?;
+
+                return Ok(stream);
+            }
+
+            Err(last_err.unwrap_or_else(|| {
+                Error::connect(io::Error::new(
+                    io::ErrorKind::InvalidInput,
+                    "could not resolve any addresses",
+                ))
+            }))
+        }
+    }
+}
+
+async fn connect_with_timeout<F, T>(connect: F, timeout: Option<Duration>) -> Result<T, Error>
+where
+    F: Future<Output = io::Result<T>>,
+{
+    match timeout {
+        Some(timeout) => match time::timeout(timeout, connect).await {
+            Ok(Ok(socket)) => Ok(socket),
+            Ok(Err(e)) => Err(Error::connect(e)),
+            Err(_) => Err(Error::connect(io::Error::new(
+                io::ErrorKind::TimedOut,
+                "connection timed out",
+            ))),
+        },
+        None => match connect.await {
+            Ok(socket) => Ok(socket),
+            Err(e) => Err(Error::connect(e)),
+        },
+    }
+}

libs/proxy/tokio-postgres2/src/connect_tls.rs

@@ -0,0 +1,48 @@
+use crate::config::SslMode;
+use crate::maybe_tls_stream::MaybeTlsStream;
+use crate::tls::private::ForcePrivateApi;
+use crate::tls::TlsConnect;
+use crate::Error;
+use bytes::BytesMut;
+use postgres_protocol2::message::frontend;
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
+
+pub async fn connect_tls<S, T>(
+    mut stream: S,
+    mode: SslMode,
+    tls: T,
+) -> Result<MaybeTlsStream<S, T::Stream>, Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: TlsConnect<S>,
+{
+    match mode {
+        SslMode::Disable => return Ok(MaybeTlsStream::Raw(stream)),
+        SslMode::Prefer if !tls.can_connect(ForcePrivateApi) => {
+            return Ok(MaybeTlsStream::Raw(stream))
+        }
+        SslMode::Prefer | SslMode::Require => {}
+    }
+
+    let mut buf = BytesMut::new();
+    frontend::ssl_request(&mut buf);
+    stream.write_all(&buf).await.map_err(Error::io)?;
+
+    let mut buf = [0];
+    stream.read_exact(&mut buf).await.map_err(Error::io)?;
+
+    if buf[0] != b'S' {
+        if SslMode::Require == mode {
+            return Err(Error::tls("server does not support TLS".into()));
+        } else {
+            return Ok(MaybeTlsStream::Raw(stream));
+        }
+    }
+
+    let stream = tls
+        .connect(stream)
+        .await
+        .map_err(|e| Error::tls(e.into()))?;
+
+    Ok(MaybeTlsStream::Tls(stream))
+}

libs/proxy/tokio-postgres2/src/connection.rs

@@ -0,0 +1,323 @@
+use crate::codec::{BackendMessage, BackendMessages, FrontendMessage, PostgresCodec};
+use crate::error::DbError;
+use crate::maybe_tls_stream::MaybeTlsStream;
+use crate::{AsyncMessage, Error, Notification};
+use bytes::BytesMut;
+use fallible_iterator::FallibleIterator;
+use futures_util::{ready, Sink, Stream};
+use log::{info, trace};
+use postgres_protocol2::message::backend::Message;
+use postgres_protocol2::message::frontend;
+use std::collections::{HashMap, VecDeque};
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::sync::mpsc;
+use tokio_util::codec::Framed;
+use tokio_util::sync::PollSender;
+
+pub enum RequestMessages {
+    Single(FrontendMessage),
+}
+
+pub struct Request {
+    pub messages: RequestMessages,
+    pub sender: mpsc::Sender<BackendMessages>,
+}
+
+pub struct Response {
+    sender: PollSender<BackendMessages>,
+}
+
+#[derive(PartialEq, Debug)]
+enum State {
+    Active,
+    Terminating,
+    Closing,
+}
+
+/// A connection to a PostgreSQL database.
+///
+/// This is one half of what is returned when a new connection is established. It performs the actual IO with the
+/// server, and should generally be spawned off onto an executor to run in the background.
+///
+/// `Connection` implements `Future`, and only resolves when the connection is closed, either because a fatal error has
+/// occurred, or because its associated `Client` has dropped and all outstanding work has completed.
+#[must_use = "futures do nothing unless polled"]
+pub struct Connection<S, T> {
+    /// HACK: we need this in the Neon Proxy.
+    pub stream: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
+    /// HACK: we need this in the Neon Proxy to forward params.
+    pub parameters: HashMap<String, String>,
+    receiver: mpsc::UnboundedReceiver<Request>,
+    pending_request: Option<RequestMessages>,
+    pending_responses: VecDeque<BackendMessage>,
+    responses: VecDeque<Response>,
+    state: State,
+}
+
+impl<S, T> Connection<S, T>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    pub(crate) fn new(
+        stream: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
+        pending_responses: VecDeque<BackendMessage>,
+        parameters: HashMap<String, String>,
+        receiver: mpsc::UnboundedReceiver<Request>,
+    ) -> Connection<S, T> {
+        Connection {
+            stream,
+            parameters,
+            receiver,
+            pending_request: None,
+            pending_responses,
+            responses: VecDeque::new(),
+            state: State::Active,
+        }
+    }
+
+    fn poll_response(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<BackendMessage, Error>>> {
+        if let Some(message) = self.pending_responses.pop_front() {
+            trace!("retrying pending response");
+            return Poll::Ready(Some(Ok(message)));
+        }
+
+        Pin::new(&mut self.stream)
+            .poll_next(cx)
+            .map(|o| o.map(|r| r.map_err(Error::io)))
+    }
+
+    fn poll_read(&mut self, cx: &mut Context<'_>) -> Result<Option<AsyncMessage>, Error> {
+        if self.state != State::Active {
+            trace!("poll_read: done");
+            return Ok(None);
+        }
+
+        loop {
+            let message = match self.poll_response(cx)? {
+                Poll::Ready(Some(message)) => message,
+                Poll::Ready(None) => return Err(Error::closed()),
+                Poll::Pending => {
+                    trace!("poll_read: waiting on response");
+                    return Ok(None);
+                }
+            };
+
+            let (mut messages, request_complete) = match message {
+                BackendMessage::Async(Message::NoticeResponse(body)) => {
+                    let error = DbError::parse(&mut body.fields()).map_err(Error::parse)?;
+                    return Ok(Some(AsyncMessage::Notice(error)));
+                }
+                BackendMessage::Async(Message::NotificationResponse(body)) => {
+                    let notification = Notification {
+                        process_id: body.process_id(),
+                        channel: body.channel().map_err(Error::parse)?.to_string(),
+                        payload: body.message().map_err(Error::parse)?.to_string(),
+                    };
+                    return Ok(Some(AsyncMessage::Notification(notification)));
+                }
+                BackendMessage::Async(Message::ParameterStatus(body)) => {
+                    self.parameters.insert(
+                        body.name().map_err(Error::parse)?.to_string(),
+                        body.value().map_err(Error::parse)?.to_string(),
+                    );
+                    continue;
+                }
+                BackendMessage::Async(_) => unreachable!(),
+                BackendMessage::Normal {
+                    messages,
+                    request_complete,
+                } => (messages, request_complete),
+            };
+
+            let mut response = match self.responses.pop_front() {
+                Some(response) => response,
+                None => match messages.next().map_err(Error::parse)? {
+                    Some(Message::ErrorResponse(error)) => return Err(Error::db(error)),
+                    _ => return Err(Error::unexpected_message()),
+                },
+            };
+
+            match response.sender.poll_reserve(cx) {
+                Poll::Ready(Ok(())) => {
+                    let _ = response.sender.send_item(messages);
+                    if !request_complete {
+                        self.responses.push_front(response);
+                    }
+                }
+                Poll::Ready(Err(_)) => {
+                    // we need to keep paging through the rest of the messages even if the receiver's hung up
+                    if !request_complete {
+                        self.responses.push_front(response);
+                    }
+                }
+                Poll::Pending => {
+                    self.responses.push_front(response);
+                    self.pending_responses.push_back(BackendMessage::Normal {
+                        messages,
+                        request_complete,
+                    });
+                    trace!("poll_read: waiting on sender");
+                    return Ok(None);
+                }
+            }
+        }
+    }
+
+    fn poll_request(&mut self, cx: &mut Context<'_>) -> Poll<Option<RequestMessages>> {
+        if let Some(messages) = self.pending_request.take() {
+            trace!("retrying pending request");
+            return Poll::Ready(Some(messages));
+        }
+
+        if self.receiver.is_closed() {
+            return Poll::Ready(None);
+        }
+
+        match self.receiver.poll_recv(cx) {
+            Poll::Ready(Some(request)) => {
+                trace!("polled new request");
+                self.responses.push_back(Response {
+                    sender: PollSender::new(request.sender),
+                });
+                Poll::Ready(Some(request.messages))
+            }
+            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+
+    fn poll_write(&mut self, cx: &mut Context<'_>) -> Result<bool, Error> {
+        loop {
+            if self.state == State::Closing {
+                trace!("poll_write: done");
+                return Ok(false);
+            }
+
+            if Pin::new(&mut self.stream)
+                .poll_ready(cx)
+                .map_err(Error::io)?
+                .is_pending()
+            {
+                trace!("poll_write: waiting on socket");
+                return Ok(false);
+            }
+
+            let request = match self.poll_request(cx) {
+                Poll::Ready(Some(request)) => request,
+                Poll::Ready(None) if self.responses.is_empty() && self.state == State::Active => {
+                    trace!("poll_write: at eof, terminating");
+                    self.state = State::Terminating;
+                    let mut request = BytesMut::new();
+                    frontend::terminate(&mut request);
+                    RequestMessages::Single(FrontendMessage::Raw(request.freeze()))
+                }
+                Poll::Ready(None) => {
+                    trace!(
+                        "poll_write: at eof, pending responses {}",
+                        self.responses.len()
+                    );
+                    return Ok(true);
+                }
+                Poll::Pending => {
+                    trace!("poll_write: waiting on request");
+                    return Ok(true);
+                }
+            };
+
+            match request {
+                RequestMessages::Single(request) => {
+                    Pin::new(&mut self.stream)
+                        .start_send(request)
+                        .map_err(Error::io)?;
+                    if self.state == State::Terminating {
+                        trace!("poll_write: sent eof, closing");
+                        self.state = State::Closing;
+                    }
+                }
+            }
+        }
+    }
+
+    fn poll_flush(&mut self, cx: &mut Context<'_>) -> Result<(), Error> {
+        match Pin::new(&mut self.stream)
+            .poll_flush(cx)
+            .map_err(Error::io)?
+        {
+            Poll::Ready(()) => trace!("poll_flush: flushed"),
+            Poll::Pending => trace!("poll_flush: waiting on socket"),
+        }
+        Ok(())
+    }
+
+    fn poll_shutdown(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
+        if self.state != State::Closing {
+            return Poll::Pending;
+        }
+
+        match Pin::new(&mut self.stream)
+            .poll_close(cx)
+            .map_err(Error::io)?
+        {
+            Poll::Ready(()) => {
+                trace!("poll_shutdown: complete");
+                Poll::Ready(Ok(()))
+            }
+            Poll::Pending => {
+                trace!("poll_shutdown: waiting on socket");
+                Poll::Pending
+            }
+        }
+    }
+
+    /// Returns the value of a runtime parameter for this connection.
+    pub fn parameter(&self, name: &str) -> Option<&str> {
+        self.parameters.get(name).map(|s| &**s)
+    }
+
+    /// Polls for asynchronous messages from the server.
+    ///
+    /// The server can send notices as well as notifications asynchronously to the client. Applications that wish to
+    /// examine those messages should use this method to drive the connection rather than its `Future` implementation.
+    pub fn poll_message(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<AsyncMessage, Error>>> {
+        let message = self.poll_read(cx)?;
+        let want_flush = self.poll_write(cx)?;
+        if want_flush {
+            self.poll_flush(cx)?;
+        }
+        match message {
+            Some(message) => Poll::Ready(Some(Ok(message))),
+            None => match self.poll_shutdown(cx) {
+                Poll::Ready(Ok(())) => Poll::Ready(None),
+                Poll::Ready(Err(e)) => Poll::Ready(Some(Err(e))),
+                Poll::Pending => Poll::Pending,
+            },
+        }
+    }
+}
+
+impl<S, T> Future for Connection<S, T>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    type Output = Result<(), Error>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
+        while let Some(message) = ready!(self.poll_message(cx)?) {
+            if let AsyncMessage::Notice(notice) = message {
+                info!("{}: {}", notice.severity(), notice.message());
+            }
+        }
+        Poll::Ready(Ok(()))
+    }
+}

libs/proxy/tokio-postgres2/src/error/mod.rs

@@ -0,0 +1,501 @@
+//! Errors.
+
+use fallible_iterator::FallibleIterator;
+use postgres_protocol2::message::backend::{ErrorFields, ErrorResponseBody};
+use std::error::{self, Error as _Error};
+use std::fmt;
+use std::io;
+
+pub use self::sqlstate::*;
+
+#[allow(clippy::unreadable_literal)]
+mod sqlstate;
+
+/// The severity of a Postgres error or notice.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum Severity {
+    /// PANIC
+    Panic,
+    /// FATAL
+    Fatal,
+    /// ERROR
+    Error,
+    /// WARNING
+    Warning,
+    /// NOTICE
+    Notice,
+    /// DEBUG
+    Debug,
+    /// INFO
+    Info,
+    /// LOG
+    Log,
+}
+
+impl fmt::Display for Severity {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let s = match *self {
+            Severity::Panic => "PANIC",
+            Severity::Fatal => "FATAL",
+            Severity::Error => "ERROR",
+            Severity::Warning => "WARNING",
+            Severity::Notice => "NOTICE",
+            Severity::Debug => "DEBUG",
+            Severity::Info => "INFO",
+            Severity::Log => "LOG",
+        };
+        fmt.write_str(s)
+    }
+}
+
+impl Severity {
+    fn from_str(s: &str) -> Option<Severity> {
+        match s {
+            "PANIC" => Some(Severity::Panic),
+            "FATAL" => Some(Severity::Fatal),
+            "ERROR" => Some(Severity::Error),
+            "WARNING" => Some(Severity::Warning),
+            "NOTICE" => Some(Severity::Notice),
+            "DEBUG" => Some(Severity::Debug),
+            "INFO" => Some(Severity::Info),
+            "LOG" => Some(Severity::Log),
+            _ => None,
+        }
+    }
+}
+
+/// A Postgres error or notice.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct DbError {
+    severity: String,
+    parsed_severity: Option<Severity>,
+    code: SqlState,
+    message: String,
+    detail: Option<String>,
+    hint: Option<String>,
+    position: Option<ErrorPosition>,
+    where_: Option<String>,
+    schema: Option<String>,
+    table: Option<String>,
+    column: Option<String>,
+    datatype: Option<String>,
+    constraint: Option<String>,
+    file: Option<String>,
+    line: Option<u32>,
+    routine: Option<String>,
+}
+
+impl DbError {
+    pub(crate) fn parse(fields: &mut ErrorFields<'_>) -> io::Result<DbError> {
+        let mut severity = None;
+        let mut parsed_severity = None;
+        let mut code = None;
+        let mut message = None;
+        let mut detail = None;
+        let mut hint = None;
+        let mut normal_position = None;
+        let mut internal_position = None;
+        let mut internal_query = None;
+        let mut where_ = None;
+        let mut schema = None;
+        let mut table = None;
+        let mut column = None;
+        let mut datatype = None;
+        let mut constraint = None;
+        let mut file = None;
+        let mut line = None;
+        let mut routine = None;
+
+        while let Some(field) = fields.next()? {
+            match field.type_() {
+                b'S' => severity = Some(field.value().to_owned()),
+                b'C' => code = Some(SqlState::from_code(field.value())),
+                b'M' => message = Some(field.value().to_owned()),
+                b'D' => detail = Some(field.value().to_owned()),
+                b'H' => hint = Some(field.value().to_owned()),
+                b'P' => {
+                    normal_position = Some(field.value().parse::<u32>().map_err(|_| {
+                        io::Error::new(
+                            io::ErrorKind::InvalidInput,
+                            "`P` field did not contain an integer",
+                        )
+                    })?);
+                }
+                b'p' => {
+                    internal_position = Some(field.value().parse::<u32>().map_err(|_| {
+                        io::Error::new(
+                            io::ErrorKind::InvalidInput,
+                            "`p` field did not contain an integer",
+                        )
+                    })?);
+                }
+                b'q' => internal_query = Some(field.value().to_owned()),
+                b'W' => where_ = Some(field.value().to_owned()),
+                b's' => schema = Some(field.value().to_owned()),
+                b't' => table = Some(field.value().to_owned()),
+                b'c' => column = Some(field.value().to_owned()),
+                b'd' => datatype = Some(field.value().to_owned()),
+                b'n' => constraint = Some(field.value().to_owned()),
+                b'F' => file = Some(field.value().to_owned()),
+                b'L' => {
+                    line = Some(field.value().parse::<u32>().map_err(|_| {
+                        io::Error::new(
+                            io::ErrorKind::InvalidInput,
+                            "`L` field did not contain an integer",
+                        )
+                    })?);
+                }
+                b'R' => routine = Some(field.value().to_owned()),
+                b'V' => {
+                    parsed_severity = Some(Severity::from_str(field.value()).ok_or_else(|| {
+                        io::Error::new(
+                            io::ErrorKind::InvalidInput,
+                            "`V` field contained an invalid value",
+                        )
+                    })?);
+                }
+                _ => {}
+            }
+        }
+
+        Ok(DbError {
+            severity: severity
+                .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "`S` field missing"))?,
+            parsed_severity,
+            code: code
+                .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "`C` field missing"))?,
+            message: message
+                .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "`M` field missing"))?,
+            detail,
+            hint,
+            position: match normal_position {
+                Some(position) => Some(ErrorPosition::Original(position)),
+                None => match internal_position {
+                    Some(position) => Some(ErrorPosition::Internal {
+                        position,
+                        query: internal_query.ok_or_else(|| {
+                            io::Error::new(
+                                io::ErrorKind::InvalidInput,
+                                "`q` field missing but `p` field present",
+                            )
+                        })?,
+                    }),
+                    None => None,
+                },
+            },
+            where_,
+            schema,
+            table,
+            column,
+            datatype,
+            constraint,
+            file,
+            line,
+            routine,
+        })
+    }
+
+    /// The field contents are ERROR, FATAL, or PANIC (in an error message),
+    /// or WARNING, NOTICE, DEBUG, INFO, or LOG (in a notice message), or a
+    /// localized translation of one of these.
+    pub fn severity(&self) -> &str {
+        &self.severity
+    }
+
+    /// A parsed, nonlocalized version of `severity`. (PostgreSQL 9.6+)
+    pub fn parsed_severity(&self) -> Option<Severity> {
+        self.parsed_severity
+    }
+
+    /// The SQLSTATE code for the error.
+    pub fn code(&self) -> &SqlState {
+        &self.code
+    }
+
+    /// The primary human-readable error message.
+    ///
+    /// This should be accurate but terse (typically one line).
+    pub fn message(&self) -> &str {
+        &self.message
+    }
+
+    /// An optional secondary error message carrying more detail about the
+    /// problem.
+    ///
+    /// Might run to multiple lines.
+    pub fn detail(&self) -> Option<&str> {
+        self.detail.as_deref()
+    }
+
+    /// An optional suggestion what to do about the problem.
+    ///
+    /// This is intended to differ from `detail` in that it offers advice
+    /// (potentially inappropriate) rather than hard facts. Might run to
+    /// multiple lines.
+    pub fn hint(&self) -> Option<&str> {
+        self.hint.as_deref()
+    }
+
+    /// An optional error cursor position into either the original query string
+    /// or an internally generated query.
+    pub fn position(&self) -> Option<&ErrorPosition> {
+        self.position.as_ref()
+    }
+
+    /// An indication of the context in which the error occurred.
+    ///
+    /// Presently this includes a call stack traceback of active procedural
+    /// language functions and internally-generated queries. The trace is one
+    /// entry per line, most recent first.
+    pub fn where_(&self) -> Option<&str> {
+        self.where_.as_deref()
+    }
+
+    /// If the error was associated with a specific database object, the name
+    /// of the schema containing that object, if any. (PostgreSQL 9.3+)
+    pub fn schema(&self) -> Option<&str> {
+        self.schema.as_deref()
+    }
+
+    /// If the error was associated with a specific table, the name of the
+    /// table. (Refer to the schema name field for the name of the table's
+    /// schema.) (PostgreSQL 9.3+)
+    pub fn table(&self) -> Option<&str> {
+        self.table.as_deref()
+    }
+
+    /// If the error was associated with a specific table column, the name of
+    /// the column.
+    ///
+    /// (Refer to the schema and table name fields to identify the table.)
+    /// (PostgreSQL 9.3+)
+    pub fn column(&self) -> Option<&str> {
+        self.column.as_deref()
+    }
+
+    /// If the error was associated with a specific data type, the name of the
+    /// data type. (Refer to the schema name field for the name of the data
+    /// type's schema.) (PostgreSQL 9.3+)
+    pub fn datatype(&self) -> Option<&str> {
+        self.datatype.as_deref()
+    }
+
+    /// If the error was associated with a specific constraint, the name of the
+    /// constraint.
+    ///
+    /// Refer to fields listed above for the associated table or domain.
+    /// (For this purpose, indexes are treated as constraints, even if they
+    /// weren't created with constraint syntax.) (PostgreSQL 9.3+)
+    pub fn constraint(&self) -> Option<&str> {
+        self.constraint.as_deref()
+    }
+
+    /// The file name of the source-code location where the error was reported.
+    pub fn file(&self) -> Option<&str> {
+        self.file.as_deref()
+    }
+
+    /// The line number of the source-code location where the error was
+    /// reported.
+    pub fn line(&self) -> Option<u32> {
+        self.line
+    }
+
+    /// The name of the source-code routine reporting the error.
+    pub fn routine(&self) -> Option<&str> {
+        self.routine.as_deref()
+    }
+}
+
+impl fmt::Display for DbError {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(fmt, "{}: {}", self.severity, self.message)?;
+        if let Some(detail) = &self.detail {
+            write!(fmt, "\nDETAIL: {}", detail)?;
+        }
+        if let Some(hint) = &self.hint {
+            write!(fmt, "\nHINT: {}", hint)?;
+        }
+        Ok(())
+    }
+}
+
+impl error::Error for DbError {}
+
+/// Represents the position of an error in a query.
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub enum ErrorPosition {
+    /// A position in the original query.
+    Original(u32),
+    /// A position in an internally generated query.
+    Internal {
+        /// The byte position.
+        position: u32,
+        /// A query generated by the Postgres server.
+        query: String,
+    },
+}
+
+#[derive(Debug, PartialEq)]
+enum Kind {
+    Io,
+    UnexpectedMessage,
+    Tls,
+    ToSql(usize),
+    FromSql(usize),
+    Column(String),
+    Closed,
+    Db,
+    Parse,
+    Encode,
+    Authentication,
+    ConfigParse,
+    Config,
+    Connect,
+    Timeout,
+}
+
+struct ErrorInner {
+    kind: Kind,
+    cause: Option<Box<dyn error::Error + Sync + Send>>,
+}
+
+/// An error communicating with the Postgres server.
+pub struct Error(Box<ErrorInner>);
+
+impl fmt::Debug for Error {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt.debug_struct("Error")
+            .field("kind", &self.0.kind)
+            .field("cause", &self.0.cause)
+            .finish()
+    }
+}
+
+impl fmt::Display for Error {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match &self.0.kind {
+            Kind::Io => fmt.write_str("error communicating with the server")?,
+            Kind::UnexpectedMessage => fmt.write_str("unexpected message from server")?,
+            Kind::Tls => fmt.write_str("error performing TLS handshake")?,
+            Kind::ToSql(idx) => write!(fmt, "error serializing parameter {}", idx)?,
+            Kind::FromSql(idx) => write!(fmt, "error deserializing column {}", idx)?,
+            Kind::Column(column) => write!(fmt, "invalid column `{}`", column)?,
+            Kind::Closed => fmt.write_str("connection closed")?,
+            Kind::Db => fmt.write_str("db error")?,
+            Kind::Parse => fmt.write_str("error parsing response from server")?,
+            Kind::Encode => fmt.write_str("error encoding message to server")?,
+            Kind::Authentication => fmt.write_str("authentication error")?,
+            Kind::ConfigParse => fmt.write_str("invalid connection string")?,
+            Kind::Config => fmt.write_str("invalid configuration")?,
+            Kind::Connect => fmt.write_str("error connecting to server")?,
+            Kind::Timeout => fmt.write_str("timeout waiting for server")?,
+        };
+        if let Some(ref cause) = self.0.cause {
+            write!(fmt, ": {}", cause)?;
+        }
+        Ok(())
+    }
+}
+
+impl error::Error for Error {
+    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
+        self.0.cause.as_ref().map(|e| &**e as _)
+    }
+}
+
+impl Error {
+    /// Consumes the error, returning its cause.
+    pub fn into_source(self) -> Option<Box<dyn error::Error + Sync + Send>> {
+        self.0.cause
+    }
+
+    /// Returns the source of this error if it was a `DbError`.
+    ///
+    /// This is a simple convenience method.
+    pub fn as_db_error(&self) -> Option<&DbError> {
+        self.source().and_then(|e| e.downcast_ref::<DbError>())
+    }
+
+    /// Determines if the error was associated with closed connection.
+    pub fn is_closed(&self) -> bool {
+        self.0.kind == Kind::Closed
+    }
+
+    /// Returns the SQLSTATE error code associated with the error.
+    ///
+    /// This is a convenience method that downcasts the cause to a `DbError` and returns its code.
+    pub fn code(&self) -> Option<&SqlState> {
+        self.as_db_error().map(DbError::code)
+    }
+
+    fn new(kind: Kind, cause: Option<Box<dyn error::Error + Sync + Send>>) -> Error {
+        Error(Box::new(ErrorInner { kind, cause }))
+    }
+
+    pub(crate) fn closed() -> Error {
+        Error::new(Kind::Closed, None)
+    }
+
+    pub(crate) fn unexpected_message() -> Error {
+        Error::new(Kind::UnexpectedMessage, None)
+    }
+
+    #[allow(clippy::needless_pass_by_value)]
+    pub(crate) fn db(error: ErrorResponseBody) -> Error {
+        match DbError::parse(&mut error.fields()) {
+            Ok(e) => Error::new(Kind::Db, Some(Box::new(e))),
+            Err(e) => Error::new(Kind::Parse, Some(Box::new(e))),
+        }
+    }
+
+    pub(crate) fn parse(e: io::Error) -> Error {
+        Error::new(Kind::Parse, Some(Box::new(e)))
+    }
+
+    pub(crate) fn encode(e: io::Error) -> Error {
+        Error::new(Kind::Encode, Some(Box::new(e)))
+    }
+
+    #[allow(clippy::wrong_self_convention)]
+    pub(crate) fn to_sql(e: Box<dyn error::Error + Sync + Send>, idx: usize) -> Error {
+        Error::new(Kind::ToSql(idx), Some(e))
+    }
+
+    pub(crate) fn from_sql(e: Box<dyn error::Error + Sync + Send>, idx: usize) -> Error {
+        Error::new(Kind::FromSql(idx), Some(e))
+    }
+
+    pub(crate) fn column(column: String) -> Error {
+        Error::new(Kind::Column(column), None)
+    }
+
+    pub(crate) fn tls(e: Box<dyn error::Error + Sync + Send>) -> Error {
+        Error::new(Kind::Tls, Some(e))
+    }
+
+    pub(crate) fn io(e: io::Error) -> Error {
+        Error::new(Kind::Io, Some(Box::new(e)))
+    }
+
+    pub(crate) fn authentication(e: Box<dyn error::Error + Sync + Send>) -> Error {
+        Error::new(Kind::Authentication, Some(e))
+    }
+
+    pub(crate) fn config_parse(e: Box<dyn error::Error + Sync + Send>) -> Error {
+        Error::new(Kind::ConfigParse, Some(e))
+    }
+
+    pub(crate) fn config(e: Box<dyn error::Error + Sync + Send>) -> Error {
+        Error::new(Kind::Config, Some(e))
+    }
+
+    pub(crate) fn connect(e: io::Error) -> Error {
+        Error::new(Kind::Connect, Some(Box::new(e)))
+    }
+
+    #[doc(hidden)]
+    pub fn __private_api_timeout() -> Error {
+        Error::new(Kind::Timeout, None)
+    }
+}

libs/proxy/tokio-postgres2/src/generic_client.rs

@@ -0,0 +1,64 @@
+use crate::query::RowStream;
+use crate::types::Type;
+use crate::{Client, Error, Transaction};
+use async_trait::async_trait;
+use postgres_protocol2::Oid;
+
+mod private {
+    pub trait Sealed {}
+}
+
+/// A trait allowing abstraction over connections and transactions.
+///
+/// This trait is "sealed", and cannot be implemented outside of this crate.
+#[async_trait]
+pub trait GenericClient: private::Sealed {
+    /// Like `Client::query_raw_txt`.
+    async fn query_raw_txt<S, I>(&self, statement: &str, params: I) -> Result<RowStream, Error>
+    where
+        S: AsRef<str> + Sync + Send,
+        I: IntoIterator<Item = Option<S>> + Sync + Send,
+        I::IntoIter: ExactSizeIterator + Sync + Send;
+
+    /// Query for type information
+    async fn get_type(&self, oid: Oid) -> Result<Type, Error>;
+}
+
+impl private::Sealed for Client {}
+
+#[async_trait]
+impl GenericClient for Client {
+    async fn query_raw_txt<S, I>(&self, statement: &str, params: I) -> Result<RowStream, Error>
+    where
+        S: AsRef<str> + Sync + Send,
+        I: IntoIterator<Item = Option<S>> + Sync + Send,
+        I::IntoIter: ExactSizeIterator + Sync + Send,
+    {
+        self.query_raw_txt(statement, params).await
+    }
+
+    /// Query for type information
+    async fn get_type(&self, oid: Oid) -> Result<Type, Error> {
+        self.get_type(oid).await
+    }
+}
+
+impl private::Sealed for Transaction<'_> {}
+
+#[async_trait]
+#[allow(clippy::needless_lifetimes)]
+impl GenericClient for Transaction<'_> {
+    async fn query_raw_txt<S, I>(&self, statement: &str, params: I) -> Result<RowStream, Error>
+    where
+        S: AsRef<str> + Sync + Send,
+        I: IntoIterator<Item = Option<S>> + Sync + Send,
+        I::IntoIter: ExactSizeIterator + Sync + Send,
+    {
+        self.query_raw_txt(statement, params).await
+    }
+
+    /// Query for type information
+    async fn get_type(&self, oid: Oid) -> Result<Type, Error> {
+        self.client().get_type(oid).await
+    }
+}

libs/proxy/tokio-postgres2/src/lib.rs

@@ -0,0 +1,148 @@
+//! An asynchronous, pipelined, PostgreSQL client.
+#![warn(rust_2018_idioms, clippy::all, missing_docs)]
+
+pub use crate::cancel_token::CancelToken;
+pub use crate::client::Client;
+pub use crate::config::Config;
+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};
+use crate::tls::MakeTlsConnect;
+pub use crate::tls::NoTls;
+pub use crate::to_statement::ToStatement;
+pub use crate::transaction::Transaction;
+pub use crate::transaction_builder::{IsolationLevel, TransactionBuilder};
+use crate::types::ToSql;
+use postgres_protocol2::message::backend::ReadyForQueryBody;
+use tokio::net::TcpStream;
+
+/// 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 to_statement;
+mod transaction;
+mod transaction_builder;
+pub mod types;
+
+/// A convenience function which parses a connection string and connects to the database.
+///
+/// See the documentation for [`Config`] for details on the connection string format.
+///
+/// Requires the `runtime` Cargo feature (enabled by default).
+///
+/// [`Config`]: config/struct.Config.html
+pub async fn connect<T>(
+    config: &str,
+    tls: T,
+) -> Result<(Client, Connection<TcpStream, T::Stream>), Error>
+where
+    T: MakeTlsConnect<TcpStream>,
+{
+    let config = config.parse::<Config>()?;
+    config.connect(tls).await
+}
+
+/// 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),
+}
+
+fn slice_iter<'a>(
+    s: &'a [&'a (dyn ToSql + Sync)],
+) -> impl ExactSizeIterator<Item = &'a (dyn ToSql + Sync)> + 'a {
+    s.iter().map(|s| *s as _)
+}

libs/proxy/tokio-postgres2/src/maybe_tls_stream.rs

@@ -0,0 +1,77 @@
+//! MaybeTlsStream.
+//!
+//! Represents a stream that may or may not be encrypted with TLS.
+use crate::tls::{ChannelBinding, TlsStream};
+use std::io;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+
+/// A stream that may or may not be encrypted with TLS.
+pub enum MaybeTlsStream<S, T> {
+    /// An unencrypted stream.
+    Raw(S),
+    /// An encrypted stream.
+    Tls(T),
+}
+
+impl<S, T> AsyncRead for MaybeTlsStream<S, T>
+where
+    S: AsyncRead + Unpin,
+    T: AsyncRead + Unpin,
+{
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        match &mut *self {
+            MaybeTlsStream::Raw(s) => Pin::new(s).poll_read(cx, buf),
+            MaybeTlsStream::Tls(s) => Pin::new(s).poll_read(cx, buf),
+        }
+    }
+}
+
+impl<S, T> AsyncWrite for MaybeTlsStream<S, T>
+where
+    S: AsyncWrite + Unpin,
+    T: AsyncWrite + Unpin,
+{
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        match &mut *self {
+            MaybeTlsStream::Raw(s) => Pin::new(s).poll_write(cx, buf),
+            MaybeTlsStream::Tls(s) => Pin::new(s).poll_write(cx, buf),
+        }
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        match &mut *self {
+            MaybeTlsStream::Raw(s) => Pin::new(s).poll_flush(cx),
+            MaybeTlsStream::Tls(s) => Pin::new(s).poll_flush(cx),
+        }
+    }
+
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        match &mut *self {
+            MaybeTlsStream::Raw(s) => Pin::new(s).poll_shutdown(cx),
+            MaybeTlsStream::Tls(s) => Pin::new(s).poll_shutdown(cx),
+        }
+    }
+}
+
+impl<S, T> TlsStream for MaybeTlsStream<S, T>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: TlsStream + Unpin,
+{
+    fn channel_binding(&self) -> ChannelBinding {
+        match self {
+            MaybeTlsStream::Raw(_) => ChannelBinding::none(),
+            MaybeTlsStream::Tls(s) => s.channel_binding(),
+        }
+    }
+}

libs/proxy/tokio-postgres2/src/prepare.rs

@@ -0,0 +1,262 @@
+use crate::client::InnerClient;
+use crate::codec::FrontendMessage;
+use crate::connection::RequestMessages;
+use crate::error::SqlState;
+use crate::types::{Field, Kind, Oid, Type};
+use crate::{query, slice_iter};
+use crate::{Column, Error, Statement};
+use bytes::Bytes;
+use fallible_iterator::FallibleIterator;
+use futures_util::{pin_mut, TryStreamExt};
+use log::debug;
+use postgres_protocol2::message::backend::Message;
+use postgres_protocol2::message::frontend;
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+
+pub(crate) const TYPEINFO_QUERY: &str = "\
+SELECT t.typname, t.typtype, t.typelem, r.rngsubtype, t.typbasetype, n.nspname, t.typrelid
+FROM pg_catalog.pg_type t
+LEFT OUTER JOIN pg_catalog.pg_range r ON r.rngtypid = t.oid
+INNER JOIN pg_catalog.pg_namespace n ON t.typnamespace = n.oid
+WHERE t.oid = $1
+";
+
+// Range types weren't added until Postgres 9.2, so pg_range may not exist
+const TYPEINFO_FALLBACK_QUERY: &str = "\
+SELECT t.typname, t.typtype, t.typelem, NULL::OID, t.typbasetype, n.nspname, t.typrelid
+FROM pg_catalog.pg_type t
+INNER JOIN pg_catalog.pg_namespace n ON t.typnamespace = n.oid
+WHERE t.oid = $1
+";
+
+const TYPEINFO_ENUM_QUERY: &str = "\
+SELECT enumlabel
+FROM pg_catalog.pg_enum
+WHERE enumtypid = $1
+ORDER BY enumsortorder
+";
+
+// Postgres 9.0 didn't have enumsortorder
+const TYPEINFO_ENUM_FALLBACK_QUERY: &str = "\
+SELECT enumlabel
+FROM pg_catalog.pg_enum
+WHERE enumtypid = $1
+ORDER BY oid
+";
+
+pub(crate) const TYPEINFO_COMPOSITE_QUERY: &str = "\
+SELECT attname, atttypid
+FROM pg_catalog.pg_attribute
+WHERE attrelid = $1
+AND NOT attisdropped
+AND attnum > 0
+ORDER BY attnum
+";
+
+static NEXT_ID: AtomicUsize = AtomicUsize::new(0);
+
+pub async fn prepare(
+    client: &Arc<InnerClient>,
+    query: &str,
+    types: &[Type],
+) -> Result<Statement, Error> {
+    let name = format!("s{}", NEXT_ID.fetch_add(1, Ordering::SeqCst));
+    let buf = encode(client, &name, query, types)?;
+    let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+
+    match responses.next().await? {
+        Message::ParseComplete => {}
+        _ => return Err(Error::unexpected_message()),
+    }
+
+    let parameter_description = match responses.next().await? {
+        Message::ParameterDescription(body) => body,
+        _ => return Err(Error::unexpected_message()),
+    };
+
+    let row_description = match responses.next().await? {
+        Message::RowDescription(body) => Some(body),
+        Message::NoData => None,
+        _ => return Err(Error::unexpected_message()),
+    };
+
+    let mut parameters = vec![];
+    let mut it = parameter_description.parameters();
+    while let Some(oid) = it.next().map_err(Error::parse)? {
+        let type_ = get_type(client, oid).await?;
+        parameters.push(type_);
+    }
+
+    let mut columns = vec![];
+    if let Some(row_description) = row_description {
+        let mut it = row_description.fields();
+        while let Some(field) = it.next().map_err(Error::parse)? {
+            let type_ = get_type(client, field.type_oid()).await?;
+            let column = Column::new(field.name().to_string(), type_, field);
+            columns.push(column);
+        }
+    }
+
+    Ok(Statement::new(client, name, parameters, columns))
+}
+
+fn prepare_rec<'a>(
+    client: &'a Arc<InnerClient>,
+    query: &'a str,
+    types: &'a [Type],
+) -> Pin<Box<dyn Future<Output = Result<Statement, Error>> + 'a + Send>> {
+    Box::pin(prepare(client, query, types))
+}
+
+fn encode(client: &InnerClient, name: &str, query: &str, types: &[Type]) -> Result<Bytes, Error> {
+    if types.is_empty() {
+        debug!("preparing query {}: {}", name, query);
+    } else {
+        debug!("preparing query {} with types {:?}: {}", name, types, query);
+    }
+
+    client.with_buf(|buf| {
+        frontend::parse(name, query, types.iter().map(Type::oid), buf).map_err(Error::encode)?;
+        frontend::describe(b'S', name, buf).map_err(Error::encode)?;
+        frontend::sync(buf);
+        Ok(buf.split().freeze())
+    })
+}
+
+pub async fn get_type(client: &Arc<InnerClient>, oid: Oid) -> Result<Type, Error> {
+    if let Some(type_) = Type::from_oid(oid) {
+        return Ok(type_);
+    }
+
+    if let Some(type_) = client.type_(oid) {
+        return Ok(type_);
+    }
+
+    let stmt = typeinfo_statement(client).await?;
+
+    let rows = query::query(client, stmt, slice_iter(&[&oid])).await?;
+    pin_mut!(rows);
+
+    let row = match rows.try_next().await? {
+        Some(row) => row,
+        None => return Err(Error::unexpected_message()),
+    };
+
+    let name: String = row.try_get(0)?;
+    let type_: i8 = row.try_get(1)?;
+    let elem_oid: Oid = row.try_get(2)?;
+    let rngsubtype: Option<Oid> = row.try_get(3)?;
+    let basetype: Oid = row.try_get(4)?;
+    let schema: String = row.try_get(5)?;
+    let relid: Oid = row.try_get(6)?;
+
+    let kind = if type_ == b'e' as i8 {
+        let variants = get_enum_variants(client, oid).await?;
+        Kind::Enum(variants)
+    } else if type_ == b'p' as i8 {
+        Kind::Pseudo
+    } else if basetype != 0 {
+        let type_ = get_type_rec(client, basetype).await?;
+        Kind::Domain(type_)
+    } else if elem_oid != 0 {
+        let type_ = get_type_rec(client, elem_oid).await?;
+        Kind::Array(type_)
+    } else if relid != 0 {
+        let fields = get_composite_fields(client, relid).await?;
+        Kind::Composite(fields)
+    } else if let Some(rngsubtype) = rngsubtype {
+        let type_ = get_type_rec(client, rngsubtype).await?;
+        Kind::Range(type_)
+    } else {
+        Kind::Simple
+    };
+
+    let type_ = Type::new(name, oid, kind, schema);
+    client.set_type(oid, &type_);
+
+    Ok(type_)
+}
+
+fn get_type_rec<'a>(
+    client: &'a Arc<InnerClient>,
+    oid: Oid,
+) -> Pin<Box<dyn Future<Output = Result<Type, Error>> + Send + 'a>> {
+    Box::pin(get_type(client, oid))
+}
+
+async fn typeinfo_statement(client: &Arc<InnerClient>) -> Result<Statement, Error> {
+    if let Some(stmt) = client.typeinfo() {
+        return Ok(stmt);
+    }
+
+    let stmt = match prepare_rec(client, TYPEINFO_QUERY, &[]).await {
+        Ok(stmt) => stmt,
+        Err(ref e) if e.code() == Some(&SqlState::UNDEFINED_TABLE) => {
+            prepare_rec(client, TYPEINFO_FALLBACK_QUERY, &[]).await?
+        }
+        Err(e) => return Err(e),
+    };
+
+    client.set_typeinfo(&stmt);
+    Ok(stmt)
+}
+
+async fn get_enum_variants(client: &Arc<InnerClient>, oid: Oid) -> Result<Vec<String>, Error> {
+    let stmt = typeinfo_enum_statement(client).await?;
+
+    query::query(client, stmt, slice_iter(&[&oid]))
+        .await?
+        .and_then(|row| async move { row.try_get(0) })
+        .try_collect()
+        .await
+}
+
+async fn typeinfo_enum_statement(client: &Arc<InnerClient>) -> Result<Statement, Error> {
+    if let Some(stmt) = client.typeinfo_enum() {
+        return Ok(stmt);
+    }
+
+    let stmt = match prepare_rec(client, TYPEINFO_ENUM_QUERY, &[]).await {
+        Ok(stmt) => stmt,
+        Err(ref e) if e.code() == Some(&SqlState::UNDEFINED_COLUMN) => {
+            prepare_rec(client, TYPEINFO_ENUM_FALLBACK_QUERY, &[]).await?
+        }
+        Err(e) => return Err(e),
+    };
+
+    client.set_typeinfo_enum(&stmt);
+    Ok(stmt)
+}
+
+async fn get_composite_fields(client: &Arc<InnerClient>, oid: Oid) -> Result<Vec<Field>, Error> {
+    let stmt = typeinfo_composite_statement(client).await?;
+
+    let rows = query::query(client, stmt, slice_iter(&[&oid]))
+        .await?
+        .try_collect::<Vec<_>>()
+        .await?;
+
+    let mut fields = vec![];
+    for row in rows {
+        let name = row.try_get(0)?;
+        let oid = row.try_get(1)?;
+        let type_ = get_type_rec(client, oid).await?;
+        fields.push(Field::new(name, type_));
+    }
+
+    Ok(fields)
+}
+
+async fn typeinfo_composite_statement(client: &Arc<InnerClient>) -> Result<Statement, Error> {
+    if let Some(stmt) = client.typeinfo_composite() {
+        return Ok(stmt);
+    }
+
+    let stmt = prepare_rec(client, TYPEINFO_COMPOSITE_QUERY, &[]).await?;
+
+    client.set_typeinfo_composite(&stmt);
+    Ok(stmt)
+}

libs/proxy/tokio-postgres2/src/query.rs

@@ -0,0 +1,340 @@
+use crate::client::{InnerClient, Responses};
+use crate::codec::FrontendMessage;
+use crate::connection::RequestMessages;
+use crate::types::IsNull;
+use crate::{Column, Error, ReadyForQueryStatus, Row, Statement};
+use bytes::{BufMut, Bytes, BytesMut};
+use fallible_iterator::FallibleIterator;
+use futures_util::{ready, Stream};
+use log::{debug, log_enabled, Level};
+use pin_project_lite::pin_project;
+use postgres_protocol2::message::backend::Message;
+use postgres_protocol2::message::frontend;
+use postgres_types2::{Format, ToSql, Type};
+use std::fmt;
+use std::marker::PhantomPinned;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+struct BorrowToSqlParamsDebug<'a>(&'a [&'a (dyn ToSql + Sync)]);
+
+impl fmt::Debug for BorrowToSqlParamsDebug<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_list().entries(self.0.iter()).finish()
+    }
+}
+
+pub async fn query<'a, I>(
+    client: &InnerClient,
+    statement: Statement,
+    params: I,
+) -> Result<RowStream, Error>
+where
+    I: IntoIterator<Item = &'a (dyn ToSql + Sync)>,
+    I::IntoIter: ExactSizeIterator,
+{
+    let buf = if log_enabled!(Level::Debug) {
+        let params = params.into_iter().collect::<Vec<_>>();
+        debug!(
+            "executing statement {} with parameters: {:?}",
+            statement.name(),
+            BorrowToSqlParamsDebug(params.as_slice()),
+        );
+        encode(client, &statement, params)?
+    } else {
+        encode(client, &statement, params)?
+    };
+    let responses = start(client, buf).await?;
+    Ok(RowStream {
+        statement,
+        responses,
+        command_tag: None,
+        status: ReadyForQueryStatus::Unknown,
+        output_format: Format::Binary,
+        _p: PhantomPinned,
+    })
+}
+
+pub async fn query_txt<S, I>(
+    client: &Arc<InnerClient>,
+    query: &str,
+    params: I,
+) -> Result<RowStream, Error>
+where
+    S: AsRef<str>,
+    I: IntoIterator<Item = Option<S>>,
+    I::IntoIter: ExactSizeIterator,
+{
+    let params = params.into_iter();
+
+    let buf = client.with_buf(|buf| {
+        frontend::parse(
+            "",                 // unnamed prepared statement
+            query,              // query to parse
+            std::iter::empty(), // give no type info
+            buf,
+        )
+        .map_err(Error::encode)?;
+        frontend::describe(b'S', "", buf).map_err(Error::encode)?;
+        // Bind, pass params as text, retrieve as binary
+        match frontend::bind(
+            "",                 // empty string selects the unnamed portal
+            "",                 // unnamed prepared statement
+            std::iter::empty(), // all parameters use the default format (text)
+            params,
+            |param, buf| match param {
+                Some(param) => {
+                    buf.put_slice(param.as_ref().as_bytes());
+                    Ok(postgres_protocol2::IsNull::No)
+                }
+                None => Ok(postgres_protocol2::IsNull::Yes),
+            },
+            Some(0), // all text
+            buf,
+        ) {
+            Ok(()) => Ok(()),
+            Err(frontend::BindError::Conversion(e)) => Err(Error::to_sql(e, 0)),
+            Err(frontend::BindError::Serialization(e)) => Err(Error::encode(e)),
+        }?;
+
+        // Execute
+        frontend::execute("", 0, buf).map_err(Error::encode)?;
+        // Sync
+        frontend::sync(buf);
+
+        Ok(buf.split().freeze())
+    })?;
+
+    // now read the responses
+    let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+
+    match responses.next().await? {
+        Message::ParseComplete => {}
+        _ => return Err(Error::unexpected_message()),
+    }
+
+    let parameter_description = match responses.next().await? {
+        Message::ParameterDescription(body) => body,
+        _ => return Err(Error::unexpected_message()),
+    };
+
+    let row_description = match responses.next().await? {
+        Message::RowDescription(body) => Some(body),
+        Message::NoData => None,
+        _ => return Err(Error::unexpected_message()),
+    };
+
+    match responses.next().await? {
+        Message::BindComplete => {}
+        _ => return Err(Error::unexpected_message()),
+    }
+
+    let mut parameters = vec![];
+    let mut it = parameter_description.parameters();
+    while let Some(oid) = it.next().map_err(Error::parse)? {
+        let type_ = Type::from_oid(oid).unwrap_or(Type::UNKNOWN);
+        parameters.push(type_);
+    }
+
+    let mut columns = vec![];
+    if let Some(row_description) = row_description {
+        let mut it = row_description.fields();
+        while let Some(field) = it.next().map_err(Error::parse)? {
+            let type_ = Type::from_oid(field.type_oid()).unwrap_or(Type::UNKNOWN);
+            let column = Column::new(field.name().to_string(), type_, field);
+            columns.push(column);
+        }
+    }
+
+    Ok(RowStream {
+        statement: Statement::new_anonymous(parameters, columns),
+        responses,
+        command_tag: None,
+        status: ReadyForQueryStatus::Unknown,
+        output_format: Format::Text,
+        _p: PhantomPinned,
+    })
+}
+
+pub async fn execute<'a, I>(
+    client: &InnerClient,
+    statement: Statement,
+    params: I,
+) -> Result<u64, Error>
+where
+    I: IntoIterator<Item = &'a (dyn ToSql + Sync)>,
+    I::IntoIter: ExactSizeIterator,
+{
+    let buf = if log_enabled!(Level::Debug) {
+        let params = params.into_iter().collect::<Vec<_>>();
+        debug!(
+            "executing statement {} with parameters: {:?}",
+            statement.name(),
+            BorrowToSqlParamsDebug(params.as_slice()),
+        );
+        encode(client, &statement, params)?
+    } else {
+        encode(client, &statement, params)?
+    };
+    let mut responses = start(client, buf).await?;
+
+    let mut rows = 0;
+    loop {
+        match responses.next().await? {
+            Message::DataRow(_) => {}
+            Message::CommandComplete(body) => {
+                rows = body
+                    .tag()
+                    .map_err(Error::parse)?
+                    .rsplit(' ')
+                    .next()
+                    .unwrap()
+                    .parse()
+                    .unwrap_or(0);
+            }
+            Message::EmptyQueryResponse => rows = 0,
+            Message::ReadyForQuery(_) => return Ok(rows),
+            _ => return Err(Error::unexpected_message()),
+        }
+    }
+}
+
+async fn start(client: &InnerClient, buf: Bytes) -> Result<Responses, Error> {
+    let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+
+    match responses.next().await? {
+        Message::BindComplete => {}
+        _ => return Err(Error::unexpected_message()),
+    }
+
+    Ok(responses)
+}
+
+pub fn encode<'a, I>(client: &InnerClient, statement: &Statement, params: I) -> Result<Bytes, Error>
+where
+    I: IntoIterator<Item = &'a (dyn ToSql + Sync)>,
+    I::IntoIter: ExactSizeIterator,
+{
+    client.with_buf(|buf| {
+        encode_bind(statement, params, "", buf)?;
+        frontend::execute("", 0, buf).map_err(Error::encode)?;
+        frontend::sync(buf);
+        Ok(buf.split().freeze())
+    })
+}
+
+pub fn encode_bind<'a, I>(
+    statement: &Statement,
+    params: I,
+    portal: &str,
+    buf: &mut BytesMut,
+) -> Result<(), Error>
+where
+    I: IntoIterator<Item = &'a (dyn ToSql + Sync)>,
+    I::IntoIter: ExactSizeIterator,
+{
+    let param_types = statement.params();
+    let params = params.into_iter();
+
+    assert!(
+        param_types.len() == params.len(),
+        "expected {} parameters but got {}",
+        param_types.len(),
+        params.len()
+    );
+
+    let (param_formats, params): (Vec<_>, Vec<_>) = params
+        .zip(param_types.iter())
+        .map(|(p, ty)| (p.encode_format(ty) as i16, p))
+        .unzip();
+
+    let params = params.into_iter();
+
+    let mut error_idx = 0;
+    let r = frontend::bind(
+        portal,
+        statement.name(),
+        param_formats,
+        params.zip(param_types).enumerate(),
+        |(idx, (param, ty)), buf| match param.to_sql_checked(ty, buf) {
+            Ok(IsNull::No) => Ok(postgres_protocol2::IsNull::No),
+            Ok(IsNull::Yes) => Ok(postgres_protocol2::IsNull::Yes),
+            Err(e) => {
+                error_idx = idx;
+                Err(e)
+            }
+        },
+        Some(1),
+        buf,
+    );
+    match r {
+        Ok(()) => Ok(()),
+        Err(frontend::BindError::Conversion(e)) => Err(Error::to_sql(e, error_idx)),
+        Err(frontend::BindError::Serialization(e)) => Err(Error::encode(e)),
+    }
+}
+
+pin_project! {
+    /// A stream of table rows.
+    pub struct RowStream {
+        statement: Statement,
+        responses: Responses,
+        command_tag: Option<String>,
+        output_format: Format,
+        status: ReadyForQueryStatus,
+        #[pin]
+        _p: PhantomPinned,
+    }
+}
+
+impl Stream for RowStream {
+    type Item = Result<Row, Error>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+        loop {
+            match ready!(this.responses.poll_next(cx)?) {
+                Message::DataRow(body) => {
+                    return Poll::Ready(Some(Ok(Row::new(
+                        this.statement.clone(),
+                        body,
+                        *this.output_format,
+                    )?)))
+                }
+                Message::EmptyQueryResponse | Message::PortalSuspended => {}
+                Message::CommandComplete(body) => {
+                    if let Ok(tag) = body.tag() {
+                        *this.command_tag = Some(tag.to_string());
+                    }
+                }
+                Message::ReadyForQuery(status) => {
+                    *this.status = status.into();
+                    return Poll::Ready(None);
+                }
+                _ => return Poll::Ready(Some(Err(Error::unexpected_message()))),
+            }
+        }
+    }
+}
+
+impl RowStream {
+    /// Returns information about the columns of data in the row.
+    pub fn columns(&self) -> &[Column] {
+        self.statement.columns()
+    }
+
+    /// Returns the command tag of this query.
+    ///
+    /// This is only available after the stream has been exhausted.
+    pub fn command_tag(&self) -> Option<String> {
+        self.command_tag.clone()
+    }
+
+    /// Returns if the connection is ready for querying, with the status of the connection.
+    ///
+    /// This might be available only after the stream has been exhausted.
+    pub fn ready_status(&self) -> ReadyForQueryStatus {
+        self.status
+    }
+}

libs/proxy/tokio-postgres2/src/row.rs

@@ -0,0 +1,300 @@
+//! Rows.
+
+use crate::row::sealed::{AsName, Sealed};
+use crate::simple_query::SimpleColumn;
+use crate::statement::Column;
+use crate::types::{FromSql, Type, WrongType};
+use crate::{Error, Statement};
+use fallible_iterator::FallibleIterator;
+use postgres_protocol2::message::backend::DataRowBody;
+use postgres_types2::{Format, WrongFormat};
+use std::fmt;
+use std::ops::Range;
+use std::str;
+use std::sync::Arc;
+
+mod sealed {
+    pub trait Sealed {}
+
+    pub trait AsName {
+        fn as_name(&self) -> &str;
+    }
+}
+
+impl AsName for Column {
+    fn as_name(&self) -> &str {
+        self.name()
+    }
+}
+
+impl AsName for String {
+    fn as_name(&self) -> &str {
+        self
+    }
+}
+
+/// A trait implemented by types that can index into columns of a row.
+///
+/// This cannot be implemented outside of this crate.
+pub trait RowIndex: Sealed {
+    #[doc(hidden)]
+    fn __idx<T>(&self, columns: &[T]) -> Option<usize>
+    where
+        T: AsName;
+}
+
+impl Sealed for usize {}
+
+impl RowIndex for usize {
+    #[inline]
+    fn __idx<T>(&self, columns: &[T]) -> Option<usize>
+    where
+        T: AsName,
+    {
+        if *self >= columns.len() {
+            None
+        } else {
+            Some(*self)
+        }
+    }
+}
+
+impl Sealed for str {}
+
+impl RowIndex for str {
+    #[inline]
+    fn __idx<T>(&self, columns: &[T]) -> Option<usize>
+    where
+        T: AsName,
+    {
+        if let Some(idx) = columns.iter().position(|d| d.as_name() == self) {
+            return Some(idx);
+        };
+
+        // FIXME ASCII-only case insensitivity isn't really the right thing to
+        // do. Postgres itself uses a dubious wrapper around tolower and JDBC
+        // uses the US locale.
+        columns
+            .iter()
+            .position(|d| d.as_name().eq_ignore_ascii_case(self))
+    }
+}
+
+impl<T> Sealed for &T where T: ?Sized + Sealed {}
+
+impl<T> RowIndex for &T
+where
+    T: ?Sized + RowIndex,
+{
+    #[inline]
+    fn __idx<U>(&self, columns: &[U]) -> Option<usize>
+    where
+        U: AsName,
+    {
+        T::__idx(*self, columns)
+    }
+}
+
+/// A row of data returned from the database by a query.
+pub struct Row {
+    statement: Statement,
+    output_format: Format,
+    body: DataRowBody,
+    ranges: Vec<Option<Range<usize>>>,
+}
+
+impl fmt::Debug for Row {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Row")
+            .field("columns", &self.columns())
+            .finish()
+    }
+}
+
+impl Row {
+    pub(crate) fn new(
+        statement: Statement,
+        body: DataRowBody,
+        output_format: Format,
+    ) -> Result<Row, Error> {
+        let ranges = body.ranges().collect().map_err(Error::parse)?;
+        Ok(Row {
+            statement,
+            body,
+            ranges,
+            output_format,
+        })
+    }
+
+    /// Returns information about the columns of data in the row.
+    pub fn columns(&self) -> &[Column] {
+        self.statement.columns()
+    }
+
+    /// Determines if the row contains no values.
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    /// Returns the number of values in the row.
+    pub fn len(&self) -> usize {
+        self.columns().len()
+    }
+
+    /// Deserializes a value from the row.
+    ///
+    /// The value can be specified either by its numeric index in the row, or by its column name.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the index is out of bounds or if the value cannot be converted to the specified type.
+    pub fn get<'a, I, T>(&'a self, idx: I) -> T
+    where
+        I: RowIndex + fmt::Display,
+        T: FromSql<'a>,
+    {
+        match self.get_inner(&idx) {
+            Ok(ok) => ok,
+            Err(err) => panic!("error retrieving column {}: {}", idx, err),
+        }
+    }
+
+    /// Like `Row::get`, but returns a `Result` rather than panicking.
+    pub fn try_get<'a, I, T>(&'a self, idx: I) -> Result<T, Error>
+    where
+        I: RowIndex + fmt::Display,
+        T: FromSql<'a>,
+    {
+        self.get_inner(&idx)
+    }
+
+    fn get_inner<'a, I, T>(&'a self, idx: &I) -> Result<T, Error>
+    where
+        I: RowIndex + fmt::Display,
+        T: FromSql<'a>,
+    {
+        let idx = match idx.__idx(self.columns()) {
+            Some(idx) => idx,
+            None => return Err(Error::column(idx.to_string())),
+        };
+
+        let ty = self.columns()[idx].type_();
+        if !T::accepts(ty) {
+            return Err(Error::from_sql(
+                Box::new(WrongType::new::<T>(ty.clone())),
+                idx,
+            ));
+        }
+
+        FromSql::from_sql_nullable(ty, self.col_buffer(idx)).map_err(|e| Error::from_sql(e, idx))
+    }
+
+    /// Get the raw bytes for the column at the given index.
+    fn col_buffer(&self, idx: usize) -> Option<&[u8]> {
+        let range = self.ranges.get(idx)?.to_owned()?;
+        Some(&self.body.buffer()[range])
+    }
+
+    /// Interpret the column at the given index as text
+    ///
+    /// Useful when using query_raw_txt() which sets text transfer mode
+    pub fn as_text(&self, idx: usize) -> Result<Option<&str>, Error> {
+        if self.output_format == Format::Text {
+            match self.col_buffer(idx) {
+                Some(raw) => {
+                    FromSql::from_sql(&Type::TEXT, raw).map_err(|e| Error::from_sql(e, idx))
+                }
+                None => Ok(None),
+            }
+        } else {
+            Err(Error::from_sql(Box::new(WrongFormat {}), idx))
+        }
+    }
+
+    /// Row byte size
+    pub fn body_len(&self) -> usize {
+        self.body.buffer().len()
+    }
+}
+
+impl AsName for SimpleColumn {
+    fn as_name(&self) -> &str {
+        self.name()
+    }
+}
+
+/// A row of data returned from the database by a simple query.
+#[derive(Debug)]
+pub struct SimpleQueryRow {
+    columns: Arc<[SimpleColumn]>,
+    body: DataRowBody,
+    ranges: Vec<Option<Range<usize>>>,
+}
+
+impl SimpleQueryRow {
+    #[allow(clippy::new_ret_no_self)]
+    pub(crate) fn new(
+        columns: Arc<[SimpleColumn]>,
+        body: DataRowBody,
+    ) -> Result<SimpleQueryRow, Error> {
+        let ranges = body.ranges().collect().map_err(Error::parse)?;
+        Ok(SimpleQueryRow {
+            columns,
+            body,
+            ranges,
+        })
+    }
+
+    /// Returns information about the columns of data in the row.
+    pub fn columns(&self) -> &[SimpleColumn] {
+        &self.columns
+    }
+
+    /// Determines if the row contains no values.
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    /// Returns the number of values in the row.
+    pub fn len(&self) -> usize {
+        self.columns.len()
+    }
+
+    /// Returns a value from the row.
+    ///
+    /// The value can be specified either by its numeric index in the row, or by its column name.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the index is out of bounds or if the value cannot be converted to the specified type.
+    pub fn get<I>(&self, idx: I) -> Option<&str>
+    where
+        I: RowIndex + fmt::Display,
+    {
+        match self.get_inner(&idx) {
+            Ok(ok) => ok,
+            Err(err) => panic!("error retrieving column {}: {}", idx, err),
+        }
+    }
+
+    /// Like `SimpleQueryRow::get`, but returns a `Result` rather than panicking.
+    pub fn try_get<I>(&self, idx: I) -> Result<Option<&str>, Error>
+    where
+        I: RowIndex + fmt::Display,
+    {
+        self.get_inner(&idx)
+    }
+
+    fn get_inner<I>(&self, idx: &I) -> Result<Option<&str>, Error>
+    where
+        I: RowIndex + fmt::Display,
+    {
+        let idx = match idx.__idx(&self.columns) {
+            Some(idx) => idx,
+            None => return Err(Error::column(idx.to_string())),
+        };
+
+        let buf = self.ranges[idx].clone().map(|r| &self.body.buffer()[r]);
+        FromSql::from_sql_nullable(&Type::TEXT, buf).map_err(|e| Error::from_sql(e, idx))
+    }
+}

libs/proxy/tokio-postgres2/src/simple_query.rs

@@ -0,0 +1,142 @@
+use crate::client::{InnerClient, Responses};
+use crate::codec::FrontendMessage;
+use crate::connection::RequestMessages;
+use crate::{Error, ReadyForQueryStatus, SimpleQueryMessage, SimpleQueryRow};
+use bytes::Bytes;
+use fallible_iterator::FallibleIterator;
+use futures_util::{ready, Stream};
+use log::debug;
+use pin_project_lite::pin_project;
+use postgres_protocol2::message::backend::Message;
+use postgres_protocol2::message::frontend;
+use std::marker::PhantomPinned;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+/// Information about a column of a single query row.
+#[derive(Debug)]
+pub struct SimpleColumn {
+    name: String,
+}
+
+impl SimpleColumn {
+    pub(crate) fn new(name: String) -> SimpleColumn {
+        SimpleColumn { name }
+    }
+
+    /// Returns the name of the column.
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+}
+
+pub async fn simple_query(client: &InnerClient, query: &str) -> Result<SimpleQueryStream, Error> {
+    debug!("executing simple query: {}", query);
+
+    let buf = encode(client, query)?;
+    let responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+
+    Ok(SimpleQueryStream {
+        responses,
+        columns: None,
+        status: ReadyForQueryStatus::Unknown,
+        _p: PhantomPinned,
+    })
+}
+
+pub async fn batch_execute(
+    client: &InnerClient,
+    query: &str,
+) -> Result<ReadyForQueryStatus, Error> {
+    debug!("executing statement batch: {}", query);
+
+    let buf = encode(client, query)?;
+    let mut responses = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)))?;
+
+    loop {
+        match responses.next().await? {
+            Message::ReadyForQuery(status) => return Ok(status.into()),
+            Message::CommandComplete(_)
+            | Message::EmptyQueryResponse
+            | Message::RowDescription(_)
+            | Message::DataRow(_) => {}
+            _ => return Err(Error::unexpected_message()),
+        }
+    }
+}
+
+pub(crate) fn encode(client: &InnerClient, query: &str) -> Result<Bytes, Error> {
+    client.with_buf(|buf| {
+        frontend::query(query, buf).map_err(Error::encode)?;
+        Ok(buf.split().freeze())
+    })
+}
+
+pin_project! {
+    /// A stream of simple query results.
+    pub struct SimpleQueryStream {
+        responses: Responses,
+        columns: Option<Arc<[SimpleColumn]>>,
+        status: ReadyForQueryStatus,
+        #[pin]
+        _p: PhantomPinned,
+    }
+}
+
+impl SimpleQueryStream {
+    /// Returns if the connection is ready for querying, with the status of the connection.
+    ///
+    /// This might be available only after the stream has been exhausted.
+    pub fn ready_status(&self) -> ReadyForQueryStatus {
+        self.status
+    }
+}
+
+impl Stream for SimpleQueryStream {
+    type Item = Result<SimpleQueryMessage, Error>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+        loop {
+            match ready!(this.responses.poll_next(cx)?) {
+                Message::CommandComplete(body) => {
+                    let rows = body
+                        .tag()
+                        .map_err(Error::parse)?
+                        .rsplit(' ')
+                        .next()
+                        .unwrap()
+                        .parse()
+                        .unwrap_or(0);
+                    return Poll::Ready(Some(Ok(SimpleQueryMessage::CommandComplete(rows))));
+                }
+                Message::EmptyQueryResponse => {
+                    return Poll::Ready(Some(Ok(SimpleQueryMessage::CommandComplete(0))));
+                }
+                Message::RowDescription(body) => {
+                    let columns = body
+                        .fields()
+                        .map(|f| Ok(SimpleColumn::new(f.name().to_string())))
+                        .collect::<Vec<_>>()
+                        .map_err(Error::parse)?
+                        .into();
+
+                    *this.columns = Some(columns);
+                }
+                Message::DataRow(body) => {
+                    let row = match &this.columns {
+                        Some(columns) => SimpleQueryRow::new(columns.clone(), body)?,
+                        None => return Poll::Ready(Some(Err(Error::unexpected_message()))),
+                    };
+                    return Poll::Ready(Some(Ok(SimpleQueryMessage::Row(row))));
+                }
+                Message::ReadyForQuery(s) => {
+                    *this.status = s.into();
+                    return Poll::Ready(None);
+                }
+                _ => return Poll::Ready(Some(Err(Error::unexpected_message()))),
+            }
+        }
+    }
+}

libs/proxy/tokio-postgres2/src/statement.rs

@@ -0,0 +1,157 @@
+use crate::client::InnerClient;
+use crate::codec::FrontendMessage;
+use crate::connection::RequestMessages;
+use crate::types::Type;
+use postgres_protocol2::{
+    message::{backend::Field, frontend},
+    Oid,
+};
+use std::{
+    fmt,
+    sync::{Arc, Weak},
+};
+
+struct StatementInner {
+    client: Weak<InnerClient>,
+    name: String,
+    params: Vec<Type>,
+    columns: Vec<Column>,
+}
+
+impl Drop for StatementInner {
+    fn drop(&mut self) {
+        if let Some(client) = self.client.upgrade() {
+            let buf = client.with_buf(|buf| {
+                frontend::close(b'S', &self.name, buf).unwrap();
+                frontend::sync(buf);
+                buf.split().freeze()
+            });
+            let _ = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)));
+        }
+    }
+}
+
+/// A prepared statement.
+///
+/// Prepared statements can only be used with the connection that created them.
+#[derive(Clone)]
+pub struct Statement(Arc<StatementInner>);
+
+impl Statement {
+    pub(crate) fn new(
+        inner: &Arc<InnerClient>,
+        name: String,
+        params: Vec<Type>,
+        columns: Vec<Column>,
+    ) -> Statement {
+        Statement(Arc::new(StatementInner {
+            client: Arc::downgrade(inner),
+            name,
+            params,
+            columns,
+        }))
+    }
+
+    pub(crate) fn new_anonymous(params: Vec<Type>, columns: Vec<Column>) -> Statement {
+        Statement(Arc::new(StatementInner {
+            client: Weak::new(),
+            name: String::new(),
+            params,
+            columns,
+        }))
+    }
+
+    pub(crate) fn name(&self) -> &str {
+        &self.0.name
+    }
+
+    /// Returns the expected types of the statement's parameters.
+    pub fn params(&self) -> &[Type] {
+        &self.0.params
+    }
+
+    /// Returns information about the columns returned when the statement is queried.
+    pub fn columns(&self) -> &[Column] {
+        &self.0.columns
+    }
+}
+
+/// Information about a column of a query.
+pub struct Column {
+    name: String,
+    type_: Type,
+
+    // raw fields from RowDescription
+    table_oid: Oid,
+    column_id: i16,
+    format: i16,
+
+    // that better be stored in self.type_, but that is more radical refactoring
+    type_oid: Oid,
+    type_size: i16,
+    type_modifier: i32,
+}
+
+impl Column {
+    pub(crate) fn new(name: String, type_: Type, raw_field: Field<'_>) -> Column {
+        Column {
+            name,
+            type_,
+            table_oid: raw_field.table_oid(),
+            column_id: raw_field.column_id(),
+            format: raw_field.format(),
+            type_oid: raw_field.type_oid(),
+            type_size: raw_field.type_size(),
+            type_modifier: raw_field.type_modifier(),
+        }
+    }
+
+    /// Returns the name of the column.
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+
+    /// Returns the type of the column.
+    pub fn type_(&self) -> &Type {
+        &self.type_
+    }
+
+    /// Returns the table OID of the column.
+    pub fn table_oid(&self) -> Oid {
+        self.table_oid
+    }
+
+    /// Returns the column ID of the column.
+    pub fn column_id(&self) -> i16 {
+        self.column_id
+    }
+
+    /// Returns the format of the column.
+    pub fn format(&self) -> i16 {
+        self.format
+    }
+
+    /// Returns the type OID of the column.
+    pub fn type_oid(&self) -> Oid {
+        self.type_oid
+    }
+
+    /// Returns the type size of the column.
+    pub fn type_size(&self) -> i16 {
+        self.type_size
+    }
+
+    /// Returns the type modifier of the column.
+    pub fn type_modifier(&self) -> i32 {
+        self.type_modifier
+    }
+}
+
+impl fmt::Debug for Column {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt.debug_struct("Column")
+            .field("name", &self.name)
+            .field("type", &self.type_)
+            .finish()
+    }
+}

libs/proxy/tokio-postgres2/src/tls.rs

@@ -0,0 +1,162 @@
+//! TLS support.
+
+use std::error::Error;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use std::{fmt, io};
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+
+pub(crate) mod private {
+    pub struct ForcePrivateApi;
+}
+
+/// Channel binding information returned from a TLS handshake.
+pub struct ChannelBinding {
+    pub(crate) tls_server_end_point: Option<Vec<u8>>,
+}
+
+impl ChannelBinding {
+    /// Creates a `ChannelBinding` containing no information.
+    pub fn none() -> ChannelBinding {
+        ChannelBinding {
+            tls_server_end_point: None,
+        }
+    }
+
+    /// Creates a `ChannelBinding` containing `tls-server-end-point` channel binding information.
+    pub fn tls_server_end_point(tls_server_end_point: Vec<u8>) -> ChannelBinding {
+        ChannelBinding {
+            tls_server_end_point: Some(tls_server_end_point),
+        }
+    }
+}
+
+/// A constructor of `TlsConnect`ors.
+///
+/// Requires the `runtime` Cargo feature (enabled by default).
+pub trait MakeTlsConnect<S> {
+    /// The stream type created by the `TlsConnect` implementation.
+    type Stream: TlsStream + Unpin;
+    /// The `TlsConnect` implementation created by this type.
+    type TlsConnect: TlsConnect<S, Stream = Self::Stream>;
+    /// The error type returned by the `TlsConnect` implementation.
+    type Error: Into<Box<dyn Error + Sync + Send>>;
+
+    /// Creates a new `TlsConnect`or.
+    ///
+    /// The domain name is provided for certificate verification and SNI.
+    fn make_tls_connect(&mut self, domain: &str) -> Result<Self::TlsConnect, Self::Error>;
+}
+
+/// An asynchronous function wrapping a stream in a TLS session.
+pub trait TlsConnect<S> {
+    /// The stream returned by the future.
+    type Stream: TlsStream + Unpin;
+    /// The error returned by the future.
+    type Error: Into<Box<dyn Error + Sync + Send>>;
+    /// The future returned by the connector.
+    type Future: Future<Output = Result<Self::Stream, Self::Error>>;
+
+    /// Returns a future performing a TLS handshake over the stream.
+    fn connect(self, stream: S) -> Self::Future;
+
+    #[doc(hidden)]
+    fn can_connect(&self, _: private::ForcePrivateApi) -> bool {
+        true
+    }
+}
+
+/// A TLS-wrapped connection to a PostgreSQL database.
+pub trait TlsStream: AsyncRead + AsyncWrite {
+    /// Returns channel binding information for the session.
+    fn channel_binding(&self) -> ChannelBinding;
+}
+
+/// A `MakeTlsConnect` and `TlsConnect` implementation which simply returns an error.
+///
+/// This can be used when `sslmode` is `none` or `prefer`.
+#[derive(Debug, Copy, Clone)]
+pub struct NoTls;
+
+impl<S> MakeTlsConnect<S> for NoTls {
+    type Stream = NoTlsStream;
+    type TlsConnect = NoTls;
+    type Error = NoTlsError;
+
+    fn make_tls_connect(&mut self, _: &str) -> Result<NoTls, NoTlsError> {
+        Ok(NoTls)
+    }
+}
+
+impl<S> TlsConnect<S> for NoTls {
+    type Stream = NoTlsStream;
+    type Error = NoTlsError;
+    type Future = NoTlsFuture;
+
+    fn connect(self, _: S) -> NoTlsFuture {
+        NoTlsFuture(())
+    }
+
+    fn can_connect(&self, _: private::ForcePrivateApi) -> bool {
+        false
+    }
+}
+
+/// The future returned by `NoTls`.
+pub struct NoTlsFuture(());
+
+impl Future for NoTlsFuture {
+    type Output = Result<NoTlsStream, NoTlsError>;
+
+    fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
+        Poll::Ready(Err(NoTlsError(())))
+    }
+}
+
+/// The TLS "stream" type produced by the `NoTls` connector.
+///
+/// Since `NoTls` doesn't support TLS, this type is uninhabited.
+pub enum NoTlsStream {}
+
+impl AsyncRead for NoTlsStream {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        _: &mut Context<'_>,
+        _: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        match *self {}
+    }
+}
+
+impl AsyncWrite for NoTlsStream {
+    fn poll_write(self: Pin<&mut Self>, _: &mut Context<'_>, _: &[u8]) -> Poll<io::Result<usize>> {
+        match *self {}
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
+        match *self {}
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
+        match *self {}
+    }
+}
+
+impl TlsStream for NoTlsStream {
+    fn channel_binding(&self) -> ChannelBinding {
+        match *self {}
+    }
+}
+
+/// The error returned by `NoTls`.
+#[derive(Debug)]
+pub struct NoTlsError(());
+
+impl fmt::Display for NoTlsError {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt.write_str("no TLS implementation configured")
+    }
+}
+
+impl Error for NoTlsError {}

libs/proxy/tokio-postgres2/src/to_statement.rs

@@ -0,0 +1,57 @@
+use crate::to_statement::private::{Sealed, ToStatementType};
+use crate::Statement;
+
+mod private {
+    use crate::{Client, Error, Statement};
+
+    pub trait Sealed {}
+
+    pub enum ToStatementType<'a> {
+        Statement(&'a Statement),
+        Query(&'a str),
+    }
+
+    impl<'a> ToStatementType<'a> {
+        pub async fn into_statement(self, client: &Client) -> Result<Statement, Error> {
+            match self {
+                ToStatementType::Statement(s) => Ok(s.clone()),
+                ToStatementType::Query(s) => client.prepare(s).await,
+            }
+        }
+    }
+}
+
+/// A trait abstracting over prepared and unprepared statements.
+///
+/// Many methods are generic over this bound, so that they support both a raw query string as well as a statement which
+/// was prepared previously.
+///
+/// This trait is "sealed" and cannot be implemented by anything outside this crate.
+pub trait ToStatement: Sealed {
+    #[doc(hidden)]
+    fn __convert(&self) -> ToStatementType<'_>;
+}
+
+impl ToStatement for Statement {
+    fn __convert(&self) -> ToStatementType<'_> {
+        ToStatementType::Statement(self)
+    }
+}
+
+impl Sealed for Statement {}
+
+impl ToStatement for str {
+    fn __convert(&self) -> ToStatementType<'_> {
+        ToStatementType::Query(self)
+    }
+}
+
+impl Sealed for str {}
+
+impl ToStatement for String {
+    fn __convert(&self) -> ToStatementType<'_> {
+        ToStatementType::Query(self)
+    }
+}
+
+impl Sealed for String {}

libs/proxy/tokio-postgres2/src/transaction.rs

@@ -0,0 +1,74 @@
+use crate::codec::FrontendMessage;
+use crate::connection::RequestMessages;
+use crate::query::RowStream;
+use crate::{CancelToken, Client, Error, ReadyForQueryStatus};
+use postgres_protocol2::message::frontend;
+
+/// A representation of a PostgreSQL database transaction.
+///
+/// Transactions will implicitly roll back when dropped. Use the `commit` method to commit the changes made in the
+/// transaction. Transactions can be nested, with inner transactions implemented via safepoints.
+pub struct Transaction<'a> {
+    client: &'a mut Client,
+    done: bool,
+}
+
+impl Drop for Transaction<'_> {
+    fn drop(&mut self) {
+        if self.done {
+            return;
+        }
+
+        let buf = self.client.inner().with_buf(|buf| {
+            frontend::query("ROLLBACK", buf).unwrap();
+            buf.split().freeze()
+        });
+        let _ = self
+            .client
+            .inner()
+            .send(RequestMessages::Single(FrontendMessage::Raw(buf)));
+    }
+}
+
+impl<'a> Transaction<'a> {
+    pub(crate) fn new(client: &'a mut Client) -> Transaction<'a> {
+        Transaction {
+            client,
+            done: false,
+        }
+    }
+
+    /// Consumes the transaction, committing all changes made within it.
+    pub async fn commit(mut self) -> Result<ReadyForQueryStatus, Error> {
+        self.done = true;
+        self.client.batch_execute("COMMIT").await
+    }
+
+    /// Rolls the transaction back, discarding all changes made within it.
+    ///
+    /// This is equivalent to `Transaction`'s `Drop` implementation, but provides any error encountered to the caller.
+    pub async fn rollback(mut self) -> Result<ReadyForQueryStatus, Error> {
+        self.done = true;
+        self.client.batch_execute("ROLLBACK").await
+    }
+
+    /// Like `Client::query_raw_txt`.
+    pub async fn query_raw_txt<S, I>(&self, statement: &str, params: I) -> Result<RowStream, Error>
+    where
+        S: AsRef<str>,
+        I: IntoIterator<Item = Option<S>>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        self.client.query_raw_txt(statement, params).await
+    }
+
+    /// Like `Client::cancel_token`.
+    pub fn cancel_token(&self) -> CancelToken {
+        self.client.cancel_token()
+    }
+
+    /// Returns a reference to the underlying `Client`.
+    pub fn client(&self) -> &Client {
+        self.client
+    }
+}

libs/proxy/tokio-postgres2/src/transaction_builder.rs

@@ -0,0 +1,113 @@
+use crate::{Client, Error, Transaction};
+
+/// The isolation level of a database transaction.
+#[derive(Debug, Copy, Clone)]
+#[non_exhaustive]
+pub enum IsolationLevel {
+    /// Equivalent to `ReadCommitted`.
+    ReadUncommitted,
+
+    /// An individual statement in the transaction will see rows committed before it began.
+    ReadCommitted,
+
+    /// All statements in the transaction will see the same view of rows committed before the first query in the
+    /// transaction.
+    RepeatableRead,
+
+    /// The reads and writes in this transaction must be able to be committed as an atomic "unit" with respect to reads
+    /// and writes of all other concurrent serializable transactions without interleaving.
+    Serializable,
+}
+
+/// A builder for database transactions.
+pub struct TransactionBuilder<'a> {
+    client: &'a mut Client,
+    isolation_level: Option<IsolationLevel>,
+    read_only: Option<bool>,
+    deferrable: Option<bool>,
+}
+
+impl<'a> TransactionBuilder<'a> {
+    pub(crate) fn new(client: &'a mut Client) -> TransactionBuilder<'a> {
+        TransactionBuilder {
+            client,
+            isolation_level: None,
+            read_only: None,
+            deferrable: None,
+        }
+    }
+
+    /// Sets the isolation level of the transaction.
+    pub fn isolation_level(mut self, isolation_level: IsolationLevel) -> Self {
+        self.isolation_level = Some(isolation_level);
+        self
+    }
+
+    /// Sets the access mode of the transaction.
+    pub fn read_only(mut self, read_only: bool) -> Self {
+        self.read_only = Some(read_only);
+        self
+    }
+
+    /// Sets the deferrability of the transaction.
+    ///
+    /// If the transaction is also serializable and read only, creation of the transaction may block, but when it
+    /// completes the transaction is able to run with less overhead and a guarantee that it will not be aborted due to
+    /// serialization failure.
+    pub fn deferrable(mut self, deferrable: bool) -> Self {
+        self.deferrable = Some(deferrable);
+        self
+    }
+
+    /// Begins the transaction.
+    ///
+    /// The transaction will roll back by default - use the `commit` method to commit it.
+    pub async fn start(self) -> Result<Transaction<'a>, Error> {
+        let mut query = "START TRANSACTION".to_string();
+        let mut first = true;
+
+        if let Some(level) = self.isolation_level {
+            first = false;
+
+            query.push_str(" ISOLATION LEVEL ");
+            let level = match level {
+                IsolationLevel::ReadUncommitted => "READ UNCOMMITTED",
+                IsolationLevel::ReadCommitted => "READ COMMITTED",
+                IsolationLevel::RepeatableRead => "REPEATABLE READ",
+                IsolationLevel::Serializable => "SERIALIZABLE",
+            };
+            query.push_str(level);
+        }
+
+        if let Some(read_only) = self.read_only {
+            if !first {
+                query.push(',');
+            }
+            first = false;
+
+            let s = if read_only {
+                " READ ONLY"
+            } else {
+                " READ WRITE"
+            };
+            query.push_str(s);
+        }
+
+        if let Some(deferrable) = self.deferrable {
+            if !first {
+                query.push(',');
+            }
+
+            let s = if deferrable {
+                " DEFERRABLE"
+            } else {
+                " NOT DEFERRABLE"
+            };
+            query.push_str(s);
+        }
+
+        self.client.batch_execute(&query).await?;
+
+        Ok(Transaction::new(self.client))
+    }
+}

libs/proxy/tokio-postgres2/src/types.rs

@@ -0,0 +1,6 @@
+//! Types.
+//!
+//! This module is a reexport of the `postgres_types` crate.
+
+#[doc(inline)]
+pub use postgres_types2::*;