aykroyd 0.3.1

//! An asynchronous, pipelined, PostgreSQL client.

use crate::client::{FromColumnIndexed, FromColumnNamed, ToParam};
use crate::postgres_common::params_iter;
use crate::query::StaticQueryText;
use crate::{error, postgres_client, FromRow, Query, QueryOne, Statement};

/// A convenience function which parses a connection string and connects to the database.
///
/// See the documentation for [`tokio_postgres::Config`] for details on the connection string format.
pub async fn connect<T>(
    config: &str,
    tls: T,
) -> Result<
    (
        Client,
        tokio_postgres::Connection<tokio_postgres::Socket, T::Stream>,
    ),
    Error,
>
where
    T: tokio_postgres::tls::MakeTlsConnect<tokio_postgres::Socket>,
{
    let (client, connection) = tokio_postgres::connect(config, tls)
        .await
        .map_err(Error::connect)?;
    Ok((client.into(), connection))
}

/// An asynchronous PostgreSQL client.
pub struct Client {
    client: tokio_postgres::Client,
    statements: std::collections::HashMap<String, tokio_postgres::Statement>,
}

postgres_client!(Client);

impl AsMut<tokio_postgres::Client> for Client {
    fn as_mut(&mut self) -> &mut tokio_postgres::Client {
        &mut self.client
    }
}

impl AsRef<tokio_postgres::Client> for Client {
    fn as_ref(&self) -> &tokio_postgres::Client {
        &self.client
    }
}

impl From<tokio_postgres::Client> for Client {
    fn from(client: tokio_postgres::Client) -> Self {
        Self::new(client)
    }
}

impl std::fmt::Debug for Client {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        f.debug_struct("Client")
            .field("client", &self.client)
            .finish()
    }
}

impl Client {
    /// Create a new `Client` from a `tokio_postgres::Client`.
    pub fn new(client: tokio_postgres::Client) -> Self {
        let statements = std::collections::HashMap::new();
        Client { client, statements }
    }

    async fn prepare_internal<S: Into<String>>(
        &mut self,
        query_text: S,
    ) -> Result<tokio_postgres::Statement, Error> {
        match self.statements.entry(query_text.into()) {
            std::collections::hash_map::Entry::Occupied(entry) => Ok(entry.get().clone()),
            std::collections::hash_map::Entry::Vacant(entry) => {
                let statement = self
                    .client
                    .prepare(entry.key())
                    .await
                    .map_err(Error::prepare)?;
                Ok(entry.insert(statement).clone())
            }
        }
    }

    /// Creates a new prepared statement.
    ///
    /// Everything required to prepare the statement is available on the
    /// type argument, so no runtime input is needed:
    ///
    /// ```no_run
    /// # async fn xmain() -> Result<(), aykroyd::tokio_postgres::Error> {
    /// # use aykroyd::{Query, FromRow};
    /// # use aykroyd::tokio_postgres::connect;
    /// # use tokio_postgres::NoTls;
    /// # #[derive(FromRow)]
    /// # pub struct Customer;
    /// #[derive(Query)]
    /// #[aykroyd(row(Customer), text = "
    ///     SELECT id, first, last FROM customers WHERE first = $1
    /// ")]
    /// pub struct GetCustomersByFirstName<'a>(&'a str);
    ///
    /// let (mut client, conn) = connect("host=localhost user=postgres", NoTls).await?;
    ///
    /// // Prepare the query in the database.
    /// client.prepare::<GetCustomersByFirstName>().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn prepare<S: StaticQueryText>(&mut self) -> Result<(), Error> {
        self.prepare_internal(S::QUERY_TEXT).await?;
        Ok(())
    }

    /// Executes a statement, returning the resulting rows.
    ///
    /// We'll prepare the statement first if we haven't yet.
    ///
    /// ```no_run
    /// # async fn xmain() -> Result<(), aykroyd::tokio_postgres::Error> {
    /// # use aykroyd::{Query, FromRow};
    /// # use aykroyd::tokio_postgres::connect;
    /// # use tokio_postgres::NoTls;
    /// # #[derive(FromRow)]
    /// # pub struct Customer {
    /// #   id: i32,
    /// #   first: String,
    /// #   last: String,
    /// # }
    /// #[derive(Query)]
    /// #[aykroyd(row(Customer), text = "
    ///     SELECT id, first, last FROM customers WHERE first = $1
    /// ")]
    /// pub struct GetCustomersByFirstName<'a>(&'a str);
    ///
    /// let (mut client, conn) = connect("host=localhost user=postgres", NoTls).await?;
    ///
    /// // Run the query and iterate over the results.
    /// for customer in client.query(&GetCustomersByFirstName("Sammy")).await? {
    ///     println!("Got customer {} {} with id {}", customer.first, customer.last, customer.id);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn query<Q: Query<Self>>(&mut self, query: &Q) -> Result<Vec<Q::Row>, Error> {
        use futures_core::stream::Stream;
        use futures_util::{pin_mut, TryStreamExt};

        let params = params_iter::ParamsIter::from_params(query.to_params());
        let statement = self.prepare_internal(query.query_text()).await?;

        let rows = self
            .client
            .query_raw(&statement, params)
            .await
            .map_err(Error::query)?;

        let mut res = Vec::with_capacity(rows.size_hint().0);
        pin_mut!(rows);
        while let Some(row) = rows.try_next().await.map_err(Error::query)? {
            res.push(FromRow::from_row(&row)?);
        }

        Ok(res)
    }

    /// Executes a statement which returns a single row, returning it.
    ///
    /// Returns an error if the query does not return exactly one row.  We'll prepare the statement first if we haven't yet.
    ///
    /// ```no_run
    /// # async fn xmain() -> Result<(), aykroyd::tokio_postgres::Error> {
    /// # use aykroyd::{QueryOne, FromRow};
    /// # use aykroyd::tokio_postgres::connect;
    /// # use tokio_postgres::NoTls;
    /// # #[derive(FromRow)]
    /// # pub struct Customer {
    /// #   id: i32,
    /// #   first: String,
    /// #   last: String,
    /// # }
    /// #[derive(QueryOne)]
    /// #[aykroyd(row(Customer), text = "
    ///     SELECT id, first, last FROM customers WHERE id = $1
    /// ")]
    /// pub struct GetCustomerById(i32);
    ///
    /// let (mut client, conn) = connect("host=localhost user=postgres", NoTls).await?;
    ///
    /// // Run the query returning a single row.
    /// let customer = client.query_one(&GetCustomerById(42)).await?;
    /// println!("Got customer {} {} with id {}", customer.first, customer.last, customer.id);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn query_one<Q: QueryOne<Self>>(&mut self, query: &Q) -> Result<Q::Row, Error> {
        let params = query.to_params();
        let params = params.as_ref().map(AsRef::as_ref).unwrap_or(&[][..]);
        let statement = self.prepare_internal(query.query_text()).await?;

        let row = self
            .client
            .query_one(&statement, params)
            .await
            .map_err(Error::query)?;

        FromRow::from_row(&row)
    }

    /// Executes a statement which returns zero or one rows, returning it.
    ///
    /// Returns an error if the query returns more than one row.  We'll prepare the statement first if we haven't yet.
    ///
    /// ```no_run
    /// # async fn xmain() -> Result<(), aykroyd::tokio_postgres::Error> {
    /// # use aykroyd::{QueryOne, FromRow};
    /// # use aykroyd::tokio_postgres::connect;
    /// # use tokio_postgres::NoTls;
    /// # #[derive(FromRow)]
    /// # pub struct Customer {
    /// #   id: i32,
    /// #   first: String,
    /// #   last: String,
    /// # }
    /// #[derive(QueryOne)]
    /// #[aykroyd(row(Customer), text = "
    ///     SELECT id, first, last FROM customers WHERE id = $1
    /// ")]
    /// pub struct GetCustomerById(i32);
    ///
    /// let (mut client, conn) = connect("host=localhost user=postgres", NoTls).await?;
    ///
    /// // Run the query, possibly returning a single row.
    /// if let Some(customer) = client.query_opt(&GetCustomerById(42)).await? {
    ///     println!("Got customer {} {} with id {}", customer.first, customer.last, customer.id);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn query_opt<Q: QueryOne<Self>>(
        &mut self,
        query: &Q,
    ) -> Result<Option<Q::Row>, Error> {
        let params = query.to_params();
        let params = params.as_ref().map(AsRef::as_ref).unwrap_or(&[][..]);
        let statement = self.prepare_internal(query.query_text()).await?;

        let row = self
            .client
            .query_opt(&statement, params)
            .await
            .map_err(Error::query)?;

        row.map(|row| FromRow::from_row(&row)).transpose()
    }

    /// Executes a statement, returning the number of rows modified.
    ///
    /// If the statement does not modify any rows (e.g. SELECT), 0 is returned.  We'll prepare the statement first if we haven't yet.
    ///
    /// ```no_run
    /// # async fn xmain() -> Result<(), aykroyd::tokio_postgres::Error> {
    /// # use aykroyd::{Statement};
    /// # use aykroyd::tokio_postgres::connect;
    /// # use tokio_postgres::NoTls;
    /// #[derive(Statement)]
    /// #[aykroyd(text = "
    ///     UPDATE customers SET first = $2, last = $3 WHERE id = $1
    /// ")]
    /// pub struct UpdateCustomerName<'a>(i32, &'a str, &'a str);
    ///
    /// let (mut client, conn) = connect("host=localhost user=postgres", NoTls).await?;
    ///
    /// // Execute the statement, returning the number of rows modified.
    /// let rows_affected = client.execute(&UpdateCustomerName(42, "Anakin", "Skywalker")).await?;
    /// assert_eq!(rows_affected, 1);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn execute<S: Statement<Self>>(&mut self, statement: &S) -> Result<u64, Error> {
        let params = statement.to_params();
        let params = params.as_ref().map(AsRef::as_ref).unwrap_or(&[][..]);
        let statement = self.prepare_internal(statement.query_text()).await?;

        let rows_affected = self
            .client
            .execute(&statement, params)
            .await
            .map_err(Error::query)?;

        Ok(rows_affected)
    }

    /// 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> {
        Ok(Transaction {
            txn: self
                .client
                .transaction()
                .await
                .map_err(Error::transaction)?,
            statements: &mut self.statements,
        })
    }
}

/// An asynchronous PostgreSQL database transaction.
///
/// Transactions will implicitly roll back by default when dropped. Use the
/// `commit` method to commit the changes made in the transaction.
pub struct Transaction<'a> {
    txn: tokio_postgres::Transaction<'a>,
    statements: &'a mut std::collections::HashMap<String, tokio_postgres::Statement>,
}

impl<'a> AsMut<tokio_postgres::Transaction<'a>> for Transaction<'a> {
    fn as_mut(&mut self) -> &mut tokio_postgres::Transaction<'a> {
        &mut self.txn
    }
}

impl<'a> AsRef<tokio_postgres::Transaction<'a>> for Transaction<'a> {
    fn as_ref(&self) -> &tokio_postgres::Transaction<'a> {
        &self.txn
    }
}

impl<'a> std::fmt::Debug for Transaction<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        f.debug_struct("Transaction").finish()
    }
}

impl<'a> Transaction<'a> {
    async fn prepare_internal<S: Into<String>>(
        &mut self,
        query_text: S,
    ) -> Result<tokio_postgres::Statement, Error> {
        match self.statements.entry(query_text.into()) {
            std::collections::hash_map::Entry::Occupied(entry) => Ok(entry.get().clone()),
            std::collections::hash_map::Entry::Vacant(entry) => {
                let statement = self
                    .txn
                    .prepare(entry.key())
                    .await
                    .map_err(Error::prepare)?;
                Ok(entry.insert(statement).clone())
            }
        }
    }

    /// Consumes the transaction, committing all changes made within it.
    pub async fn commit(self) -> Result<(), Error> {
        self.txn.commit().await.map_err(Error::transaction)
    }

    /// 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(self) -> Result<(), Error> {
        self.txn.rollback().await.map_err(Error::transaction)
    }

    /// Creates a new prepared statement.
    ///
    /// Everything required to prepare the statement is available on the
    /// type argument, so no runtime input is needed:
    ///
    /// ```no_run
    /// # async fn xmain() -> Result<(), aykroyd::tokio_postgres::Error> {
    /// # use aykroyd::{Query, FromRow};
    /// # use aykroyd::tokio_postgres::connect;
    /// # use tokio_postgres::NoTls;
    /// # #[derive(FromRow)]
    /// # pub struct Customer;
    /// #[derive(Query)]
    /// #[aykroyd(row(Customer), text = "
    ///     SELECT id, first, last FROM customers WHERE first = $1
    /// ")]
    /// pub struct GetCustomersByFirstName<'a>(&'a str);
    ///
    /// let (mut client, conn) = connect("host=localhost user=postgres", NoTls).await?;
    /// let mut txn = client.transaction().await?;
    ///
    /// // Prepare the query in the database.
    /// txn.prepare::<GetCustomersByFirstName>().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn prepare<S: StaticQueryText>(&mut self) -> Result<(), Error> {
        self.prepare_internal(S::QUERY_TEXT).await?;
        Ok(())
    }

    /// Executes a statement, returning the resulting rows.
    ///
    /// We'll prepare the statement first if we haven't yet.
    ///
    /// ```no_run
    /// # async fn xmain() -> Result<(), aykroyd::tokio_postgres::Error> {
    /// # use aykroyd::{Query, FromRow};
    /// # use aykroyd::tokio_postgres::connect;
    /// # use tokio_postgres::NoTls;
    /// # #[derive(FromRow)]
    /// # pub struct Customer {
    /// #   id: i32,
    /// #   first: String,
    /// #   last: String,
    /// # }
    /// #[derive(Query)]
    /// #[aykroyd(row(Customer), text = "
    ///     SELECT id, first, last FROM customers WHERE first = $1
    /// ")]
    /// pub struct GetCustomersByFirstName<'a>(&'a str);
    ///
    /// let (mut client, conn) = connect("host=localhost user=postgres", NoTls).await?;
    /// let mut txn = client.transaction().await?;
    ///
    /// // Run the query and iterate over the results.
    /// for customer in txn.query(&GetCustomersByFirstName("Sammy")).await? {
    ///     println!("Got customer {} {} with id {}", customer.first, customer.last, customer.id);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn query<Q: Query<Client>>(&mut self, query: &Q) -> Result<Vec<Q::Row>, Error> {
        use futures_core::stream::Stream;
        use futures_util::{pin_mut, TryStreamExt};

        let params = params_iter::ParamsIter::from_params(query.to_params());
        let statement = self.prepare_internal(query.query_text()).await?;

        let rows = self
            .txn
            .query_raw(&statement, params)
            .await
            .map_err(Error::query)?;

        let mut res = Vec::with_capacity(rows.size_hint().0);
        pin_mut!(rows);
        while let Some(row) = rows.try_next().await.map_err(Error::query)? {
            res.push(FromRow::from_row(&row)?);
        }

        Ok(res)
    }

    /// Executes a statement which returns a single row, returning it.
    ///
    /// Returns an error if the query does not return exactly one row.  We'll prepare the statement first if we haven't yet.
    ///
    /// ```no_run
    /// # async fn xmain() -> Result<(), aykroyd::tokio_postgres::Error> {
    /// # use aykroyd::{QueryOne, FromRow};
    /// # use aykroyd::tokio_postgres::connect;
    /// # use tokio_postgres::NoTls;
    /// # #[derive(FromRow)]
    /// # pub struct Customer {
    /// #   id: i32,
    /// #   first: String,
    /// #   last: String,
    /// # }
    /// #[derive(QueryOne)]
    /// #[aykroyd(row(Customer), text = "
    ///     SELECT id, first, last FROM customers WHERE id = $1
    /// ")]
    /// pub struct GetCustomerById(i32);
    ///
    /// let (mut client, conn) = connect("host=localhost user=postgres", NoTls).await?;
    /// let mut txn = client.transaction().await?;
    ///
    /// // Run the query returning a single row.
    /// let customer = txn.query_one(&GetCustomerById(42)).await?;
    /// println!("Got customer {} {} with id {}", customer.first, customer.last, customer.id);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn query_one<Q: QueryOne<Client>>(&mut self, query: &Q) -> Result<Q::Row, Error> {
        let params = query.to_params();
        let params = params.as_ref().map(AsRef::as_ref).unwrap_or(&[][..]);
        let statement = self.prepare_internal(query.query_text()).await?;

        let row = self
            .txn
            .query_one(&statement, params)
            .await
            .map_err(Error::query)?;

        FromRow::from_row(&row)
    }

    /// Executes a statement which returns zero or one rows, returning it.
    ///
    /// Returns an error if the query returns more than one row.  We'll prepare the statement first if we haven't yet.
    ///
    /// ```no_run
    /// # async fn xmain() -> Result<(), aykroyd::tokio_postgres::Error> {
    /// # use aykroyd::{QueryOne, FromRow};
    /// # use aykroyd::tokio_postgres::connect;
    /// # use tokio_postgres::NoTls;
    /// # #[derive(FromRow)]
    /// # pub struct Customer {
    /// #   id: i32,
    /// #   first: String,
    /// #   last: String,
    /// # }
    /// #[derive(QueryOne)]
    /// #[aykroyd(row(Customer), text = "
    ///     SELECT id, first, last FROM customers WHERE id = $1
    /// ")]
    /// pub struct GetCustomerById(i32);
    ///
    /// let (mut client, conn) = connect("host=localhost user=postgres", NoTls).await?;
    /// let mut txn = client.transaction().await?;
    ///
    /// // Run the query, possibly returning a single row.
    /// if let Some(customer) = txn.query_opt(&GetCustomerById(42)).await? {
    ///     println!("Got customer {} {} with id {}", customer.first, customer.last, customer.id);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn query_opt<Q: QueryOne<Client>>(
        &mut self,
        query: &Q,
    ) -> Result<Option<Q::Row>, Error> {
        let params = query.to_params();
        let params = params.as_ref().map(AsRef::as_ref).unwrap_or(&[][..]);
        let statement = self.prepare_internal(query.query_text()).await?;

        let row = self
            .txn
            .query_opt(&statement, params)
            .await
            .map_err(Error::query)?;

        row.map(|row| FromRow::from_row(&row)).transpose()
    }

    /// Executes a statement, returning the number of rows modified.
    ///
    /// If the statement does not modify any rows (e.g. SELECT), 0 is returned.  We'll prepare the statement first if we haven't yet.
    ///
    /// ```no_run
    /// # async fn xmain() -> Result<(), aykroyd::tokio_postgres::Error> {
    /// # use aykroyd::{Statement};
    /// # use aykroyd::tokio_postgres::connect;
    /// # use tokio_postgres::NoTls;
    /// #[derive(Statement)]
    /// #[aykroyd(text = "
    ///     UPDATE customers SET first = $2, last = $3 WHERE id = $1
    /// ")]
    /// pub struct UpdateCustomerName<'a>(i32, &'a str, &'a str);
    ///
    /// let (mut client, conn) = connect("host=localhost user=postgres", NoTls).await?;
    /// let mut txn = client.transaction().await?;
    ///
    /// // Execute the statement, returning the number of rows modified.
    /// let rows_affected = txn.execute(&UpdateCustomerName(42, "Anakin", "Skywalker")).await?;
    /// assert_eq!(rows_affected, 1);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn execute<S: Statement<Client>>(&mut self, statement: &S) -> Result<u64, Error> {
        let params = statement.to_params();
        let params = params.as_ref().map(AsRef::as_ref).unwrap_or(&[][..]);
        let statement = self.prepare_internal(statement.query_text()).await?;

        let rows_affected = self
            .txn
            .execute(&statement, params)
            .await
            .map_err(Error::query)?;

        Ok(rows_affected)
    }
}

// TODO: not derive support
#[cfg(all(test, feature = "derive"))]
mod test {
    use super::*;

    use tokio_postgres::NoTls;

    #[derive(Statement)]
    #[aykroyd(
        text = "CREATE TABLE test_tokio_postgres (id SERIAL PRIMARY KEY, label TEXT NOT NULL)"
    )]
    struct CreateTodos;

    #[derive(Statement)]
    #[aykroyd(text = "DROP TABLE IF EXISTS test_tokio_postgres")]
    struct DropTodos;

    #[derive(Statement)]
    #[aykroyd(text = "INSERT INTO test_tokio_postgres (label) VALUES ($1)")]
    struct InsertTodo<'a>(&'a str);

    #[derive(Query)]
    #[aykroyd(row((i32, String)), text = "SELECT id, label FROM test_tokio_postgres")]
    struct GetAllTodos;

    #[tokio::test]
    async fn end_to_end() {
        const TODO_TEXT: &str = "get things done, please!";

        let (mut client, connection) = connect(
            "host=localhost user=aykroyd_test password=aykroyd_test",
            NoTls,
        )
        .await
        .unwrap();

        tokio::spawn(async move {
            if let Err(e) = connection.await {
                eprintln!("connection error: {e}");
            }
        });

        client.execute(&DropTodos).await.unwrap();

        client.execute(&CreateTodos).await.unwrap();

        client.execute(&InsertTodo(TODO_TEXT)).await.unwrap();

        let todos = client.query(&GetAllTodos).await.unwrap();
        assert_eq!(1, todos.len());
        assert_eq!(TODO_TEXT, todos[0].1);

        client.execute(&DropTodos).await.unwrap();
    }
}