diesel_uplete/

lib.rs

// SPDX-FileCopyrightText: 2024, 2025 Joseph Silva
//
// SPDX-License-Identifier: MIT OR Apache-2.0

use diesel::associations::HasTable;
use diesel::dsl;
use diesel::expression::is_aggregate;
use diesel::expression::ValidGrouping;
use diesel::pg::Pg;
use diesel::query_builder::AsQuery;
use diesel::query_builder::AstPass;
use diesel::query_builder::Query;
use diesel::query_builder::QueryFragment;
use diesel::query_builder::QueryId;
use diesel::query_dsl::methods::FilterDsl;
use diesel::query_dsl::methods::SelectDsl;
use diesel::result::Error;
use diesel::sql_types;
use diesel::Column;
use diesel::Expression;
use diesel::Queryable;
use diesel::RunQueryDsl;
use diesel::Table;
use std::any::TypeId;
use tuplex::IntoArray;

/// Set columns (each specified with [`UpleteBuilder::set_null`]) to null in the rows found by
/// `query`, and delete rows that have no remaining non-null values in non-ignored columns
///
/// # Panics
///
/// Running the built query panics if [`UpleteBuilder::set_null`] is not called at least once.
pub fn uplete<Q>(query: Q) -> UpleteBuilder<dsl::Select<Q::Query, <Q::Table as Table>::PrimaryKey>>
where
    Q: AsQuery + HasTable,
    Q::Table: Default,
    Q::Query: SelectDsl<<Q::Table as Table>::PrimaryKey>,

    // For better error messages
    UpleteBuilder<Q>: AsQuery,
{
    UpleteBuilder {
        query: query.as_query().select(Q::Table::default().primary_key()),
        columns_being_set_to_null: Vec::new(),
    }
}

pub struct UpleteBuilder<Q> {
    query: Q,
    columns_being_set_to_null: Vec<DynColumn>,
}

impl<Q: HasTable> UpleteBuilder<Q> {
    pub fn set_null<C: Column<Table = Q::Table> + Into<DynColumn>>(mut self, column: C) -> Self {
        self.columns_being_set_to_null.push(column.into());
        self
    }
}

impl<Q> AsQuery for UpleteBuilder<Q>
where
    Q: HasTable,
    Q::Table: SupportedTable,
    <Q::Table as Table>::AllColumns: IntoArray<DynColumn>,
    <<Q::Table as SupportedTable>::Key as IntoArray<DynColumn>>::Output:
        IntoIterator<Item = DynColumn>,
    <<Q::Table as SupportedTable>::AdditionalIgnoredColumns as IntoArray<DynColumn>>::Output:
        IntoIterator<Item = DynColumn>,
    <<Q::Table as Table>::AllColumns as IntoArray<DynColumn>>::Output:
        IntoIterator<Item = DynColumn>,
    Q: Clone + FilterDsl<AllNull> + FilterDsl<dsl::not<AllNull>>,
    dsl::Filter<Q, AllNull>: QueryFragment<Pg> + Send + 'static,
    dsl::Filter<Q, dsl::not<AllNull>>: QueryFragment<Pg> + Send + 'static,
{
    type Query = UpleteQuery;

    type SqlType = (sql_types::BigInt, sql_types::BigInt);

    fn as_query(self) -> Self::Query {
        let table = Q::Table::default;
        let deletion_condition = AllNull(
            Q::Table::all_columns()
                .into_array()
                .into_iter()
                .filter(|c: &DynColumn| {
                    self.columns_being_set_to_null
                        .iter()
                        .cloned()
                        .chain(<Q::Table as SupportedTable>::Key::default().into_array())
                        .chain(
                            <Q::Table as SupportedTable>::AdditionalIgnoredColumns::default()
                                .into_array(),
                        )
                        .all(|excluded_column| excluded_column.type_id != c.type_id)
                })
                .collect::<Vec<_>>(),
        );
        UpleteQuery {
            // Updated rows and deleted rows must not overlap, so updating all rows and using the returned
            // new rows to determine which ones to delete is not an option.
            //
            // https://www.postgresql.org/docs/16/queries-with.html#QUERIES-WITH-MODIFYING
            //
            // "Trying to update the same row twice in a single statement is not supported. Only one of
            // the modifications takes place, but it is not easy (and sometimes not possible) to reliably
            // predict which one. This also applies to deleting a row that was already updated in the same
            // statement: only the update is performed."
            update_subquery: Box::new(
                self.query
                    .clone()
                    .filter(dsl::not(deletion_condition.clone())),
            ),
            delete_subquery: Box::new(self.query.filter(deletion_condition)),
            table: Box::new(table()),
            key: Box::new(<Q::Table as SupportedTable>::Key::default()),
            columns_being_set_to_null: self.columns_being_set_to_null,
        }
    }
}

pub struct UpleteQuery {
    update_subquery: Box<dyn QueryFragment<Pg> + Send + 'static>,
    delete_subquery: Box<dyn QueryFragment<Pg> + Send + 'static>,
    table: Box<dyn QueryFragment<Pg> + Send + 'static>,
    key: Box<dyn QueryFragment<Pg> + Send + 'static>,
    columns_being_set_to_null: Vec<DynColumn>,
}

impl QueryId for UpleteQuery {
    type QueryId = ();

    const HAS_STATIC_QUERY_ID: bool = false;
}

impl Query for UpleteQuery {
    type SqlType = (sql_types::BigInt, sql_types::BigInt);
}

impl QueryFragment<Pg> for UpleteQuery {
    fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> Result<(), Error> {
        assert_ne!(
            self.columns_being_set_to_null.len(),
            0,
            "`set_null` was not called"
        );

        // This is checked by require_uplete triggers
        out.push_sql("/**/");

        // Declare `update_keys` and `delete_keys` CTEs, which select primary keys
        for (prefix, subquery) in [
            ("WITH update_keys", &self.update_subquery),
            (", delete_keys", &self.delete_subquery),
        ] {
            out.push_sql(prefix);
            out.push_sql(" AS (");
            subquery.walk_ast(out.reborrow())?;
            out.push_sql(" FOR UPDATE)");
        }

        // Update rows that are referenced in `update_keys`
        out.push_sql(", update_result AS (UPDATE ");
        self.table.walk_ast(out.reborrow())?;
        let mut item_prefix = " SET ";
        for column in &self.columns_being_set_to_null {
            out.push_sql(item_prefix);
            out.push_identifier(column.name)?;
            out.push_sql(" = NULL");
            item_prefix = ",";
        }
        out.push_sql(" WHERE (");
        self.key.walk_ast(out.reborrow())?;
        out.push_sql(") = ANY (SELECT * FROM update_keys) RETURNING 1)");

        // Delete rows that are referenced in `delete_keys`
        out.push_sql(", delete_result AS (DELETE FROM ");
        self.table.walk_ast(out.reborrow())?;
        out.push_sql(" WHERE (");
        self.key.walk_ast(out.reborrow())?;
        out.push_sql(") = ANY (SELECT * FROM delete_keys) RETURNING 1)");

        // Count updated rows and deleted rows (`RETURNING 1` makes this possible)
        out.push_sql(" SELECT (SELECT count(*) FROM update_result)");
        out.push_sql(", (SELECT count(*) FROM delete_result)");

        Ok(())
    }
}

/// Appears in some trait bounds.
#[derive(Clone)]
pub struct AllNull<T = DynColumn>(Vec<T>);

impl<T> Expression for AllNull<T> {
    type SqlType = sql_types::Bool;
}

impl<T> ValidGrouping<()> for AllNull<T> {
    type IsAggregate = is_aggregate::No;
}

impl<T: QueryFragment<Pg>> QueryFragment<Pg> for AllNull<T> {
    fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> Result<(), Error> {
        // Must produce a valid expression even if `self.0` is empty
        out.push_sql("(TRUE");
        for item in &self.0 {
            out.push_sql(" AND (");
            item.walk_ast(out.reborrow())?;
            out.push_sql(" IS NULL)");
        }
        out.push_sql(")");

        Ok(())
    }
}

/// Appears in some trait bounds.
#[derive(Clone)]
pub struct DynColumn {
    type_id: TypeId,
    name: &'static str,
}

impl<T: Column + 'static> From<T> for DynColumn {
    fn from(_value: T) -> Self {
        DynColumn {
            type_id: TypeId::of::<T>(),
            name: T::NAME,
        }
    }
}

impl QueryFragment<Pg> for DynColumn {
    fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> Result<(), Error> {
        out.push_identifier(self.name)
    }
}

/// The output of an uplete.
///
/// Each row in the query given to [`uplete`] is either updated or deleted. No row is both updated and deleted.
#[derive(Queryable, PartialEq, Eq, Debug)]
pub struct UpleteCount {
    /// Number of rows that were updated.
    pub updated: i64,
    /// Number of rows that were deleted.
    pub deleted: i64,
}

impl UpleteCount {
    /// Returns `UpleteCount { updated: n, deleted: 0 }`.
    pub fn only_updated(n: i64) -> Self {
        UpleteCount {
            updated: n,
            deleted: 0,
        }
    }

    /// Returns `UpleteCount { updated: 0, deleted: n }`.
    pub fn only_deleted(n: i64) -> Self {
        UpleteCount {
            updated: 0,
            deleted: n,
        }
    }
}

/// Trait that you must implement for a table in order to do upletes on it.
pub trait SupportedTable: Default + QueryFragment<Pg> + Send + 'static {
    /// A tuple of columns that, as a whole, must always has a unique value for each row. Subqueries in the generated SQL rely on this requirement being met.
    ///
    /// Each column in `Key` is also ignored. See the documentation of [`uplete`] for what that means.
    ///
    /// `Key` usually should the same as the primary key.
    ///
    /// To use only one column, use a trailing comma: `type Key = (my_table::my_column,);`
    type Key: Default + IntoArray<DynColumn> + QueryFragment<Pg> + Send + 'static;
    /// Each column in `AdditionalIgnoredColumns` is ignored. See the documentation of [`uplete`] for what that means.
    ///
    /// May be `()` if `Key` has all columns that should be ignored.
    ///
    /// To use only one column, use a trailing comma: `type Key = (my_table::my_column,);`
    type AdditionalIgnoredColumns: Default + IntoArray<DynColumn>;
}

impl<Conn, Q> RunQueryDsl<Conn> for UpleteBuilder<Q> where Self: AsQuery {}

impl<Conn> RunQueryDsl<Conn> for UpleteQuery {}

#[cfg(test)]
mod tests {
    use crate::AllNull;
    use diesel::connection::SimpleConnection;
    use diesel::debug_query;
    use diesel::insert_into;
    use diesel::pg::Pg;
    use diesel::query_builder::AsQuery;
    use diesel::query_builder::QueryId;
    use diesel::select;
    use diesel::sql_types;
    use diesel::AppearsOnTable;
    use diesel::Connection;
    use diesel::ExpressionMethods;
    use diesel::IntoSql;
    use diesel::PgConnection;
    use diesel::QueryDsl;
    use diesel::QueryResult;
    use diesel::RunQueryDsl;
    use diesel::SelectableExpression;
    use std::env;
    use std::process::Command;
    use tempfile::tempdir_in;

    impl<T, QS> AppearsOnTable<QS> for AllNull<T> {}

    impl<T, QS> SelectableExpression<QS> for AllNull<T> {}

    impl<T> QueryId for AllNull<T> {
        type QueryId = ();
        const HAS_STATIC_QUERY_ID: bool = false;
    }

    diesel::table! {
      t (id1, id2) {
        // uplete doesn't work for non-tuple primary key
        id1 -> Int4,
        id2 -> Int4,
        a -> Nullable<Int4>,
        b -> Nullable<Int4>,
      }
    }

    impl crate::SupportedTable for t::table {
        type Key = (t::id1, t::id2);
        type AdditionalIgnoredColumns = ();
    }

    fn expect_rows(
        conn: &mut PgConnection,
        expected: &[(Option<i32>, Option<i32>)],
    ) -> QueryResult<()> {
        let rows: Vec<(Option<i32>, Option<i32>)> =
            t::table.select((t::a, t::b)).order_by(t::id1).load(conn)?;
        assert_eq!(expected, &rows);

        Ok(())
    }

    // Main purpose of this test is to check accuracy of the returned `UpleteCount`, which other modules'
    // tests rely on
    fn test_count(conn: &mut PgConnection) -> QueryResult<()> {
        conn
            .batch_execute("CREATE TABLE t (id1 serial, id2 int NOT NULL DEFAULT 1, a int, b int, PRIMARY KEY (id1, id2));")?;
        expect_rows(conn, &[])?;

        insert_into(t::table)
            .values(&[
                (t::a.eq(Some(1)), t::b.eq(Some(2))),
                (t::a.eq(Some(3)), t::b.eq(None)),
                (t::a.eq(Some(4)), t::b.eq(Some(5))),
            ])
            .execute(conn)?;
        expect_rows(
            conn,
            &[(Some(1), Some(2)), (Some(3), None), (Some(4), Some(5))],
        )?;

        let count1 = crate::uplete(t::table).set_null(t::a).get_result(conn)?;
        assert_eq!(
            crate::UpleteCount {
                updated: 2,
                deleted: 1
            },
            count1
        );
        expect_rows(conn, &[(None, Some(2)), (None, Some(5))])?;

        let count2 = crate::uplete(t::table).set_null(t::b).get_result(conn)?;
        assert_eq!(crate::UpleteCount::only_deleted(2), count2);
        expect_rows(conn, &[])?;

        conn.batch_execute("DROP TABLE t;")?;

        Ok(())
    }

    fn expected_sql(check_null: &str, set_null: &str) -> String {
        let with_queries = {
            let key = r#""t"."id1", "t"."id2""#;
            let t = r#""t""#;

            let update_keys =
                format!("SELECT {key} FROM {t} WHERE  NOT (({check_null})) FOR UPDATE");
            let delete_keys = format!("SELECT {key} FROM {t} WHERE ({check_null}) FOR UPDATE");
            let update_result = format!(
          "UPDATE {t} SET {set_null} WHERE ({key}) = ANY (SELECT * FROM update_keys) RETURNING 1"
        );
            let delete_result = format!(
                "DELETE FROM {t} WHERE ({key}) = ANY (SELECT * FROM delete_keys) RETURNING 1"
            );

            format!("update_keys AS ({update_keys}), delete_keys AS ({delete_keys}), update_result AS ({update_result}), delete_result AS ({delete_result})")
        };
        let update_count = "SELECT count(*) FROM update_result";
        let delete_count = "SELECT count(*) FROM delete_result";

        format!(r#"/**/WITH {with_queries} SELECT ({update_count}), ({delete_count}) -- binds: []"#)
    }

    #[test]
    fn test_generated_sql() {
        // Unlike the `get_result` method, `debug_query` does not automatically call `as_query`
        assert_eq!(
            debug_query::<Pg, _>(&crate::uplete(t::table).set_null(t::b).as_query()).to_string(),
            expected_sql(r#"TRUE AND ("a" IS NULL)"#, r#""b" = NULL"#)
        );
        assert_eq!(
            debug_query::<Pg, _>(
                &crate::uplete(t::table)
                    .set_null(t::a)
                    .set_null(t::b)
                    .as_query()
            )
            .to_string(),
            expected_sql(r#"TRUE"#, r#""a" = NULL,"b" = NULL"#)
        );
    }

    #[test]
    fn test_count_methods() {
        assert_eq!(
            crate::UpleteCount::only_updated(1),
            crate::UpleteCount {
                updated: 1,
                deleted: 0
            }
        );
        assert_eq!(
            crate::UpleteCount::only_deleted(1),
            crate::UpleteCount {
                updated: 0,
                deleted: 1
            }
        );
    }

    fn test_all_null(conn: &mut PgConnection) -> QueryResult<()> {
        let some = Some(1).into_sql::<sql_types::Nullable<sql_types::Integer>>();
        let none = None::<i32>.into_sql::<sql_types::Nullable<sql_types::Integer>>();

        // Allows type inference for `vec![]`
        let mut all_null = |items| select(AllNull(items)).get_result::<bool>(conn);

        assert!(all_null(vec![])?);
        assert!(all_null(vec![none])?);
        assert!(all_null(vec![none, none])?);
        assert!(all_null(vec![none, none, none])?);
        assert!(!all_null(vec![some])?);
        assert!(!all_null(vec![some, none])?);
        assert!(!all_null(vec![none, some, none])?);

        Ok(())
    }

    #[test]
    fn test_db_stuff() -> QueryResult<()> {
        let user_specific_tmp = env::var("XDG_RUNTIME_DIR").unwrap();
        let tempdir = tempdir_in(&user_specific_tmp).unwrap();
        let tempdir_path = tempdir.path().to_str().unwrap();
        let pgdata = tempdir.path().join("pgdata");
        let envs = [("PGDATA", &pgdata)];
        assert!(Command::new("initdb")
            .envs(envs)
            .args(["--username=postgres", "--auth=trust", "--no-sync"])
            .status()
            .unwrap()
            .success());
        let mut postgres_process = Command::new("postgres")
            .envs(envs)
            .args([
                "-c",
                "listen_addresses=",
                "-c",
                &format!("unix_socket_directories={}", tempdir_path),
                "-c",
                "fsync=off",
                "-c",
                "logging_collector=off",
                "-c",
                "port=1",
            ])
            .spawn()
            .unwrap();
        // TODO: kill postgres on failure
        std::thread::sleep(std::time::Duration::from_secs(3));
        let mut conn = PgConnection::establish(&format!(
            "postgresql://postgres@{}:1",
            tempdir_path.replace('/', "%2F")
        ))
        .unwrap();
        test_count(&mut conn)?;
        test_all_null(&mut conn)?;
        postgres_process.kill().unwrap();
        tempdir.close().unwrap();
        Ok(())
    }
}