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

use std::collections::HashMap;
use std::fmt;
use std::net::IpAddr;
use std::task::{Context, Poll};
use std::time::Duration;

use bytes::BytesMut;
use fallible_iterator::FallibleIterator;
use futures_util::{TryStreamExt, future, ready};
use postgres_protocol2::message::backend::Message;
use postgres_protocol2::message::frontend;
use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;

use crate::cancel_token::RawCancelToken;
use crate::codec::{BackendMessages, FrontendMessage, RecordNotices};
use crate::config::{Host, SslMode};
use crate::query::RowStream;
use crate::simple_query::SimpleQueryStream;
use crate::types::{Oid, Type};
use crate::{CancelToken, Error, ReadyForQueryStatus, SimpleQueryMessage, query, simple_query};

pub struct Responses {
    /// new messages from conn
    receiver: mpsc::Receiver<BackendMessages>,
    /// current batch of messages
    cur: BackendMessages,
    /// number of total queries sent.
    waiting: usize,
    /// number of ReadyForQuery messages received.
    received: usize,

    /// The last query status we received.
    last_status: ReadyForQueryStatus,
}

impl Responses {
    pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Result<Message, Error>> {
        loop {
            // get the next saved message
            if let Some(message) = self.cur.next().map_err(Error::parse)? {
                let received = self.received;

                // increase the query head if this is the last message.
                if let Message::ReadyForQuery(ref status) = message {
                    self.last_status = (*status).into();
                    self.received += 1;
                }

                // check if the client has skipped this query.
                if received + 1 < self.waiting {
                    // grab the next message.
                    continue;
                }

                // convenience: turn the error messaage into a proper error.
                let res = match message {
                    Message::ErrorResponse(body) => Err(Error::db(body)),
                    message => Ok(message),
                };
                return Poll::Ready(res);
            }

            // get the next batch of messages.
            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
    }

    pub async fn wait_until_ready(&mut self) -> Result<ReadyForQueryStatus, Error> {
        while self.received < self.waiting {
            if let Message::ReadyForQuery(status) = self.next().await? {
                return Ok(status.into());
            }
        }
        Ok(self.last_status)
    }
}

/// A cache of type info and prepared statements for fetching type info
/// (corresponding to the queries in the [crate::prepare] module).
#[derive(Default)]
pub(crate) struct CachedTypeInfo {
    /// Cache of types already looked up.
    pub(crate) types: HashMap<Oid, Type>,
}

pub struct InnerClient {
    sender: mpsc::UnboundedSender<FrontendMessage>,
    responses: Responses,

    /// A buffer to use when writing out postgres commands.
    buffer: BytesMut,
}

impl InnerClient {
    pub fn start(&mut self) -> Result<PartialQuery<'_>, Error> {
        self.responses.waiting += 1;
        Ok(PartialQuery(Some(self)))
    }

    pub fn send_simple_query(&mut self, query: &str) -> Result<&mut Responses, Error> {
        self.responses.waiting += 1;

        self.buffer.clear();
        // simple queries do not need sync.
        frontend::query(query, &mut self.buffer).map_err(Error::encode)?;
        let buf = self.buffer.split().freeze();
        self.send_message(FrontendMessage::Raw(buf))
    }

    fn send_message(&mut self, messages: FrontendMessage) -> Result<&mut Responses, Error> {
        self.sender.send(messages).map_err(|_| Error::closed())?;
        Ok(&mut self.responses)
    }
}

pub struct PartialQuery<'a>(Option<&'a mut InnerClient>);

impl Drop for PartialQuery<'_> {
    fn drop(&mut self) {
        if let Some(client) = self.0.take() {
            client.buffer.clear();
            frontend::sync(&mut client.buffer);
            let buf = client.buffer.split().freeze();
            let _ = client.send_message(FrontendMessage::Raw(buf));
        }
    }
}

impl<'a> PartialQuery<'a> {
    pub fn send_with_flush<F>(&mut self, f: F) -> Result<&mut Responses, Error>
    where
        F: FnOnce(&mut BytesMut) -> Result<(), Error>,
    {
        let client = self.0.as_deref_mut().unwrap();

        client.buffer.clear();
        f(&mut client.buffer)?;
        frontend::flush(&mut client.buffer);
        let buf = client.buffer.split().freeze();
        client.send_message(FrontendMessage::Raw(buf))
    }

    pub fn send_with_sync<F>(mut self, f: F) -> Result<&'a mut Responses, Error>
    where
        F: FnOnce(&mut BytesMut) -> Result<(), Error>,
    {
        let client = self.0.as_deref_mut().unwrap();

        client.buffer.clear();
        f(&mut client.buffer)?;
        frontend::sync(&mut client.buffer);
        let buf = client.buffer.split().freeze();
        let _ = client.send_message(FrontendMessage::Raw(buf));

        Ok(&mut self.0.take().unwrap().responses)
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct SocketConfig {
    pub host_addr: Option<IpAddr>,
    pub host: Host,
    pub port: u16,
    pub connect_timeout: Option<Duration>,
}

/// 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: InnerClient,
    cached_typeinfo: CachedTypeInfo,

    socket_config: SocketConfig,
    ssl_mode: SslMode,
    process_id: i32,
    secret_key: i32,
}

impl Client {
    pub(crate) fn new(
        sender: mpsc::UnboundedSender<FrontendMessage>,
        receiver: mpsc::Receiver<BackendMessages>,
        socket_config: SocketConfig,
        ssl_mode: SslMode,
        process_id: i32,
        secret_key: i32,
    ) -> Client {
        Client {
            inner: InnerClient {
                sender,
                responses: Responses {
                    receiver,
                    cur: BackendMessages::empty(),
                    waiting: 0,
                    received: 0,
                    // new connections are always idle.
                    last_status: ReadyForQueryStatus::Idle,
                },
                buffer: Default::default(),
            },
            cached_typeinfo: Default::default(),

            socket_config,
            ssl_mode,
            process_id,
            secret_key,
        }
    }

    /// Returns process_id.
    pub fn get_process_id(&self) -> i32 {
        self.process_id
    }

    pub(crate) fn inner_mut(&mut self) -> &mut InnerClient {
        &mut self.inner
    }

    pub fn record_notices(&mut self, limit: usize) -> mpsc::UnboundedReceiver<Box<str>> {
        let (tx, rx) = mpsc::unbounded_channel();

        let notices = RecordNotices { sender: tx, limit };
        self.inner
            .sender
            .send(FrontendMessage::RecordNotices(notices))
            .ok();

        rx
    }

    /// Wait until this connection has no more active queries.
    pub async fn wait_until_ready(&mut self) -> Result<ReadyForQueryStatus, Error> {
        self.inner_mut().responses.wait_until_ready().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>(
        &mut self,
        statement: &str,
        params: I,
    ) -> Result<RowStream<'_>, Error>
    where
        S: AsRef<str>,
        I: IntoIterator<Item = Option<S>>,
        I::IntoIter: ExactSizeIterator,
    {
        query::query_txt(
            &mut self.inner,
            &mut self.cached_typeinfo,
            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(&mut self, query: &str) -> Result<Vec<SimpleQueryMessage>, Error> {
        self.simple_query_raw(query).await?.try_collect().await
    }

    pub(crate) async fn simple_query_raw(
        &mut self,
        query: &str,
    ) -> Result<SimpleQueryStream<'_>, Error> {
        simple_query::simple_query(self.inner_mut(), 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(&mut self, query: &str) -> Result<ReadyForQueryStatus, Error> {
        simple_query::batch_execute(self.inner_mut(), query).await
    }

    /// Similar to `discard_all`, but it does not clear any query plans
    ///
    /// This runs in the background, so it can be executed without `await`ing.
    pub fn reset_session_background(&mut self) -> Result<(), Error> {
        // "CLOSE ALL": closes any cursors
        // "SET SESSION AUTHORIZATION DEFAULT": resets the current_user back to the session_user
        // "RESET ALL": resets any GUCs back to their session defaults.
        // "DEALLOCATE ALL": deallocates any prepared statements
        // "UNLISTEN *": stops listening on all channels
        // "SELECT pg_advisory_unlock_all();": unlocks all advisory locks
        // "DISCARD TEMP;": drops all temporary tables
        // "DISCARD SEQUENCES;": deallocates all cached sequence state

        let _responses = self.inner_mut().send_simple_query(
            "ROLLBACK;
            CLOSE ALL;
            SET SESSION AUTHORIZATION DEFAULT;
            RESET ALL;
            DEALLOCATE ALL;
            UNLISTEN *;
            SELECT pg_advisory_unlock_all();
            DISCARD TEMP;
            DISCARD SEQUENCES;",
        )?;

        Ok(())
    }

    /// 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(),
            raw: RawCancelToken {
                ssl_mode: self.ssl_mode,
                process_id: self.process_id,
                secret_key: self.secret_key,
            },
        }
    }

    /// 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()
    }
}