postrust_core/db/

sqlx_impl.rs

//! [`Db`] over a sqlx pool: the driver the server and the Lambda adapter use.

use super::{Db, DbError, FetchedThen, RawRows, Tx};
use crate::error::DatabaseError;
use crate::row_json::{row_to_json, row_to_values};
use postrust_sql::SqlParam;
use sqlx::postgres::{PgArguments, PgPool, PgRow};
use sqlx::query::Query;
use sqlx::{AssertSqlSafe, Postgres, Row, ValueRef};

/// A transaction on a pooled connection, which goes back to the pool when this
/// is committed or dropped.
///
/// sqlx cannot pipeline, so [`Tx::queue`]d statements run one after another
/// ahead of the next statement -- the round trips they always took. On a pool
/// next to the database each is a fraction of a millisecond.
pub struct SqlxTx {
    tx: sqlx::Transaction<'static, Postgres>,
    queued: Vec<(String, Vec<SqlParam>)>,
}

impl Db for PgPool {
    type Tx<'a> = SqlxTx;

    async fn begin(&self) -> Result<SqlxTx, DbError> {
        Ok(SqlxTx {
            tx: PgPool::begin(self).await?,
            queued: Vec::new(),
        })
    }
}

impl SqlxTx {
    /// Run what was queued, in order.
    async fn flush(&mut self) -> Result<(), DbError> {
        for (sql, params) in std::mem::take(&mut self.queued) {
            query(&sql, &params)
                .execute(&mut *self.tx)
                .await
                .map_err(|e| DbError::from(e).queued())?;
        }
        Ok(())
    }
}

impl Tx for SqlxTx {
    async fn execute(&mut self, sql: &str, params: &[SqlParam]) -> Result<u64, DbError> {
        self.flush().await?;
        Ok(query(sql, params)
            .execute(&mut *self.tx)
            .await?
            .rows_affected())
    }

    async fn fetch_json(
        &mut self,
        sql: &str,
        params: &[SqlParam],
    ) -> Result<Vec<serde_json::Value>, DbError> {
        // `into_iter` matters for large result sets: each row's buffers are
        // freed as soon as it has been converted, rather than the whole
        // `Vec<PgRow>` staying alive alongside the whole `Vec<Value>`.
        self.flush().await?;
        Ok(fetch_all(&mut self.tx, sql, params)
            .await?
            .into_iter()
            .map(|row| row_to_json(&row))
            .collect())
    }

    async fn fetch_values(
        &mut self,
        sql: &str,
        params: &[SqlParam],
    ) -> Result<Vec<Vec<serde_json::Value>>, DbError> {
        self.flush().await?;
        Ok(fetch_all(&mut self.tx, sql, params)
            .await?
            .into_iter()
            .map(|row| row_to_values(&row))
            .collect())
    }

    async fn fetch_raw(
        &mut self,
        sql: &str,
        params: &[SqlParam],
        count_column: &str,
    ) -> Result<RawRows, DbError> {
        self.flush().await?;
        let rows = fetch_all(&mut self.tx, sql, params).await?;
        let first = rows.first().and_then(|row| {
            let value = row.try_get_raw(0).ok()?;
            if value.is_null() {
                return None;
            }
            Some(value.as_bytes().ok()?.to_vec())
        });
        let count = rows
            .first()
            .and_then(|row| row.try_get::<i64, _>(count_column).ok());
        Ok(RawRows::new(rows.len(), first, count))
    }

    async fn commit(mut self) -> Result<(), DbError> {
        self.flush().await?;
        Ok(self.tx.commit().await?)
    }

    fn queue(&mut self, sql: &str, params: &[SqlParam]) {
        self.queued.push((sql.to_string(), params.to_vec()));
    }

    async fn fetch_json_then(&mut self, sql: &str, params: &[SqlParam], then: &str) -> FetchedThen {
        let rows = self.fetch_json(sql, params).await?;
        let then = self.fetch_values(then, &[]).await?;
        Ok((rows, then))
    }

    async fn fetch_json_and_commit(
        mut self,
        sql: &str,
        params: &[SqlParam],
    ) -> Result<Vec<serde_json::Value>, DbError> {
        let rows = self.fetch_json(sql, params).await?;
        self.commit().await?;
        Ok(rows)
    }
}

impl crate::session::Session<'_> {
    /// Run this session's statements on a sqlx connection, in order.
    ///
    /// For code that holds a sqlx transaction of its own rather than a
    /// [`Tx`] -- GraphQL's resolvers, a hand-written route -- and has to make
    /// it the caller's. Pass the transaction (`&mut *tx`): every setting is
    /// transaction-local, so it is gone when the transaction ends and the
    /// connection goes back to the pool as it came. On a connection outside
    /// a transaction each statement would be its own, and nothing would stay.
    ///
    /// Parameters are bound as [`SqlxTx`] binds them, so a session means the
    /// same thing to both. A failure here is reported as a queued statement's,
    /// as it is when [`Tx::queue`] runs the same statements.
    pub async fn apply_sqlx(&self, conn: &mut sqlx::PgConnection) -> Result<(), DbError> {
        for statement in self.statements() {
            query(&statement.sql, &statement.params)
                .execute(&mut *conn)
                .await
                .map_err(|e| DbError::from(e).queued())?;
        }
        Ok(())
    }
}

async fn fetch_all(
    tx: &mut sqlx::Transaction<'static, Postgres>,
    sql: &str,
    params: &[SqlParam],
) -> Result<Vec<PgRow>, sqlx::Error> {
    query(sql, params).fetch_all(&mut **tx).await
}

/// The statement to run: `sql`, with each parameter bound with the type sqlx
/// gives its Rust value, and the SQL's own casts deciding the rest.
///
/// Every statement this driver runs comes through here, which is why the
/// driver wraps its SQL in `AssertSqlSafe` here and nowhere else. sqlx 0.9
/// accepts only a `&'static str` as SQL unless told otherwise, as a speed bump
/// against `format!`-ing user input into a query. The query builder in
/// `crate::query` assembles statements at runtime, but every value a request
/// carries arrives in `params` and is bound, never spliced in. What is spliced
/// in is SQL the builder wrote itself; identifiers, each quoted by
/// `postrust_sql::escape_ident`; literals the builder quotes itself, by
/// `postrust_sql::quote_literal`; and the type name of a `::` cast in a
/// select, which the parser admits only as a single type name -- letters,
/// digits, `_` and `$` -- or one of the SQL standard's multi-word names, in
/// the standard's spelling (see `parse_cast_type` in
/// `crate::api_request::query_params`). The assertion is that, and nothing
/// more. The schema cache's two runtime statements, in
/// `schema_cache::queries::enum_values`, quote their identifiers the same way
/// and make the same assertion for themselves.
fn query(sql: &str, params: &[SqlParam]) -> Query<'static, Postgres, PgArguments> {
    let mut query = sqlx::query(AssertSqlSafe(sql));
    for param in params {
        query = match param {
            SqlParam::Null => query.bind(None::<String>),
            SqlParam::Bool(b) => query.bind(b),
            SqlParam::Int(n) => query.bind(n),
            SqlParam::Float(f) => query.bind(f),
            SqlParam::Text(s) => query.bind(s),
            SqlParam::Bytes(b) => query.bind(b),
            SqlParam::Json(j) => query.bind(j),
            SqlParam::Uuid(u) => query.bind(u),
            SqlParam::Timestamp(t) => query.bind(t),
            SqlParam::Array(arr) => {
                // Convert array to Vec<String> for text arrays
                let strings: Vec<String> = arr
                    .iter()
                    .map(|p| match p {
                        SqlParam::Text(s) => s.clone(),
                        SqlParam::Int(n) => n.to_string(),
                        SqlParam::Bool(b) => b.to_string(),
                        other => format!("{:?}", other),
                    })
                    .collect();
                query.bind(strings)
            }
        };
    }

    query
}

impl From<sqlx::Error> for DbError {
    fn from(e: sqlx::Error) -> Self {
        let message = e.to_string();
        match e {
            sqlx::Error::Database(db_err) => {
                // Phrased here rather than taken from sqlx's `Display`, which
                // from 0.9 appends " at line N" -- the line of PostgreSQL's
                // own C source that raised the error (its `L` field), not a
                // line of the statement. It means nothing to whoever reads the
                // message, differs between PostgreSQL builds, and is not what
                // the socket driver says. This is the 0.8 wording, word for
                // word.
                let message = format!("error returned from database: {}", db_err.message());
                // Try to downcast to Postgres-specific error for additional details
                let (details, hint) = db_err
                    .try_downcast_ref::<sqlx::postgres::PgDatabaseError>()
                    .map(|pg_err| {
                        (
                            pg_err.detail().map(String::from),
                            pg_err.hint().map(String::from),
                        )
                    })
                    .unwrap_or((None, None));

                DbError::database(
                    message,
                    DatabaseError {
                        code: db_err.code().map(|c| c.to_string()).unwrap_or_default(),
                        message: db_err.message().to_string(),
                        details,
                        hint,
                        constraint: db_err.constraint().map(|s| s.to_string()),
                        table: db_err.table().map(|s| s.to_string()),
                        column: None,
                    },
                )
            }
            // The connection, not a statement on it: it broke, it could not be
            // had from the pool in time, or the pool is gone.
            sqlx::Error::Io(_)
            | sqlx::Error::Tls(_)
            | sqlx::Error::PoolTimedOut
            | sqlx::Error::PoolClosed
            | sqlx::Error::WorkerCrashed => DbError::connection(message),
            _ => DbError::other(message),
        }
    }
}