rustango 0.62.0

//! `delete` audit rows for the audited rows an FK `ON DELETE CASCADE`
//! removes, `update` rows for the ones `SET NULL` / `SET DEFAULT` changes.
//! Read before the delete, written after it for the rows it took or changed.

use std::collections::{HashMap, HashSet};
use std::sync::OnceLock;

use crate::core::{
    CountQuery, DeleteQuery, Filter, ModelEntry, ModelSchema, OnDeleteAction, Op, Relation,
    SelectQuery, SqlValue, WhereExpr,
};
use crate::sql::{ExecError, PoolTx};

use super::{emit_many_tx, pk_in, AuditOp, PendingEntry, BULK_AUDIT_CHUNK};

/// `child.fk_column` references `parent.on` with `ON DELETE CASCADE`, or
/// with `SET NULL` / `SET DEFAULT` to an audited child (`sets`).
struct Edge {
    child: &'static ModelEntry,
    fk_column: &'static str,
    on: &'static str,
    sets: bool,
}

/// The cascade edges out of `table` that lead to an audited model.
fn edges(table: &str) -> &'static [Edge] {
    static GRAPH: OnceLock<HashMap<&'static str, Vec<Edge>>> = OnceLock::new();
    GRAPH
        .get_or_init(build_graph)
        .get(table)
        .map_or(&[], Vec::as_slice)
}

/// One registry entry per table: a project model over a framework one,
/// then the first by model name, so the pick does not depend on link order.
fn canonical_entries() -> Vec<&'static ModelEntry> {
    let mut by_table: HashMap<&'static str, &'static ModelEntry> = HashMap::new();
    for entry in crate::core::inventory::iter::<ModelEntry> {
        let key = |e: &ModelEntry| (e.is_framework(), e.schema.name);
        by_table
            .entry(entry.schema.table)
            .and_modify(|kept| {
                if key(entry) < key(kept) {
                    *kept = entry;
                }
            })
            .or_insert(entry);
    }
    by_table.into_values().collect()
}

fn build_graph() -> HashMap<&'static str, Vec<Edge>> {
    let entries = canonical_entries();
    let mut graph: HashMap<&'static str, Vec<Edge>> = HashMap::new();
    for &entry in &entries {
        let audited = entry.schema.audit_track.is_some();
        for field in entry.schema.fields {
            let (Some(action), Some(Relation::Fk { to, on } | Relation::O2O { to, on })) =
                (field.fk_on_delete, field.relation)
            else {
                continue;
            };
            let sets = match action {
                OnDeleteAction::Cascade => false,
                OnDeleteAction::SetNull | OnDeleteAction::SetDefault if audited => true,
                _ => continue,
            };
            graph.entry(to).or_default().push(Edge {
                child: entry,
                fk_column: field.column,
                on,
                sets,
            });
        }
    }
    // Tables whose cascade reaches an audited model, to a fixed point.
    let mut reaches: HashSet<&str> = entries
        .iter()
        .filter(|e| e.schema.audit_track.is_some())
        .map(|e| e.schema.table)
        .collect();
    reaches.extend(
        graph
            .iter()
            .filter(|(_, es)| es.iter().any(|e| e.sets))
            .map(|(t, _)| *t),
    );
    loop {
        let grown: Vec<&'static str> = graph
            .iter()
            .filter(|(t, es)| {
                !reaches.contains(*t)
                    && es
                        .iter()
                        .any(|e| !e.sets && reaches.contains(e.child.schema.table))
            })
            .map(|(t, _)| *t)
            .collect();
        if grown.is_empty() {
            break;
        }
        reaches.extend(grown);
    }
    for es in graph.values_mut() {
        es.retain(|e| e.sets || reaches.contains(e.child.schema.table));
    }
    graph.retain(|_, es| !es.is_empty());
    graph
}

/// An audited table a truncate of `table` changes: rows it takes, or, for a
/// plain truncate, the FK it sets to NULL / default (`set`).
pub(crate) struct Reached {
    schema: &'static ModelSchema,
    set: Option<&'static str>,
    before: i64,
}

/// The audited tables a truncate of `table` reaches besides it: a cascading
/// one empties every FK child, a plain `DELETE` only follows
/// `ON DELETE CASCADE` and sets the `SET NULL` / `SET DEFAULT` children.
fn truncate_reach(
    table: &str,
    every_fk: bool,
) -> Vec<(&'static ModelSchema, Option<&'static str>)> {
    let entries = canonical_entries();
    let mut out = Vec::new();
    let mut done: HashSet<&str> = HashSet::from([table]);
    let mut frontier = vec![table.to_owned()];
    let mut sets = Vec::new();
    while let Some(parent) = frontier.pop() {
        for entry in &entries {
            let schema = entry.schema;
            let refs = |f: &&crate::core::FieldSchema| matches!(f.relation, Some(Relation::Fk { to, .. } | Relation::O2O { to, .. }) if to == parent);
            let by_field = schema
                .fields
                .iter()
                .filter(refs)
                .any(|f| every_fk || f.fk_on_delete == Some(OnDeleteAction::Cascade));
            let by_composite =
                every_fk && schema.composite_relations.iter().any(|r| r.to == parent);
            if !every_fk && schema.audit_track.is_some() {
                sets.extend(
                    schema
                        .fields
                        .iter()
                        .filter(refs)
                        .filter(|f| {
                            matches!(
                                f.fk_on_delete,
                                Some(OnDeleteAction::SetNull | OnDeleteAction::SetDefault)
                            )
                        })
                        .map(|f| (schema, f.column)),
                );
            }
            if (by_field || by_composite) && done.insert(schema.table) {
                if schema.audit_track.is_some() {
                    out.push((schema, None));
                }
                frontier.push(schema.table.to_owned());
            }
        }
    }
    // A child the truncate also deletes from is counted as removed.
    for (schema, column) in sets {
        if !done.contains(schema.table) {
            out.push((schema, Some(column)));
        }
    }
    out
}

/// The audited tables a truncate of `table` reaches through its FKs, with
/// their counts; taken before it runs, for [`emit_truncated`].
pub(crate) async fn truncate_counts(
    tx: &mut PoolTx<'_>,
    table: &str,
    mode: crate::sql::TruncateMode,
) -> Result<Vec<Reached>, ExecError> {
    let every_fk = mode == crate::sql::TruncateMode::Cascade;
    let mut out = Vec::new();
    for (schema, set) in truncate_reach(table, every_fk) {
        let before = row_count(tx, schema, set).await?;
        out.push(Reached {
            schema,
            set,
            before,
        });
    }
    Ok(out)
}

/// `schema`'s rows, or those whose `set` FK is not NULL.
async fn row_count(
    tx: &mut PoolTx<'_>,
    schema: &'static ModelSchema,
    set: Option<&'static str>,
) -> Result<i64, ExecError> {
    let mut select = SelectQuery::new(schema);
    if let Some(column) = set {
        select = select.where_clause(WhereExpr::Predicate(Filter::new(
            column,
            Op::IsNull,
            SqlValue::Bool(false),
        )));
    }
    crate::sql::count_rows_on(&mut *tx, &CountQuery::from_select(select)).await
}

/// One bulk entry per audited table the truncate of `table` changed: a
/// `delete` for rows taken, an `update` for an FK it set.
pub(crate) async fn emit_truncated(
    tx: &mut PoolTx<'_>,
    table: &str,
    before: Vec<Reached>,
) -> Result<(), ExecError> {
    let mut entries = Vec::new();
    for r in before {
        if r.before > 0 && row_count(tx, r.schema, r.set).await? < r.before {
            let (operation, changes) = match r.set {
                None => (
                    AuditOp::Delete,
                    serde_json::json!({ "bulk": "truncate", "cascade_from": table }),
                ),
                Some(column) => (
                    AuditOp::Update,
                    serde_json::json!({ "bulk": "truncate", "cascade_from": table, "set": column }),
                ),
            };
            entries.push(PendingEntry {
                entity_table: r.schema.table,
                entity_pk: String::new(),
                operation,
                source: super::current_source(),
                changes,
            });
        }
    }
    emit_many_tx(tx, &entries).await
}

/// Whether deleting `model` rows may cascade to an audited model.
pub(crate) fn reaches_audited(model: &ModelSchema) -> bool {
    !edges(model.table).is_empty()
}

/// Audited rows a delete will cascade to, with their entries, and the
/// audited rows whose FK it will set.
struct Pending {
    removed: Vec<(&'static ModelSchema, Vec<(SqlValue, PendingEntry)>)>,
    set: Vec<SetRows>,
}

/// Audited `child` rows `pks` whose `fk_column` holds one of `keys`.
struct SetRows {
    child: &'static ModelEntry,
    fk_column: &'static str,
    keys: Vec<SqlValue>,
    pks: Vec<SqlValue>,
}

/// Delete `query`'s rows in chunks of roots, each with the audited rows it
/// cascades to: read them locked, delete the chunk, emit for the rows gone.
/// Memory holds one chunk's cascade, plus the keys of the rows seen.
pub(crate) async fn delete_with_audit(
    tx: &mut PoolTx<'_>,
    query: &DeleteQuery,
) -> Result<u64, ExecError> {
    let mut seen: HashSet<(&'static str, String)> = HashSet::new();
    let Some(pk) = query.model.primary_key() else {
        return delete_chunk(tx, query, &mut seen).await;
    };
    // Locked up front, so the chunks cover exactly the rows `query` matches.
    let roots = column_values(tx, query.model, query.where_clause.clone(), pk.column).await?;
    // An audited model's own rows are audited by its delete path.
    if query.model.audit_track.is_some() {
        seen.extend(roots.iter().map(|k| (query.model.table, format!("{k:?}"))));
    }
    let mut affected = 0;
    for chunk in roots.chunks(BULK_AUDIT_CHUNK) {
        let part = DeleteQuery {
            model: query.model,
            where_clause: pk_in(query.model, chunk.to_vec())?,
        };
        affected += delete_chunk(tx, &part, &mut seen).await?;
    }
    Ok(affected)
}

async fn delete_chunk(
    tx: &mut PoolTx<'_>,
    query: &DeleteQuery,
    seen: &mut HashSet<(&'static str, String)>,
) -> Result<u64, ExecError> {
    let pending = snapshot(tx, query, seen).await?;
    let plan = crate::sql::compile_delete_plan(tx.dialect(), query)?;
    let affected = crate::sql::run_delete_statements_tx(tx, plan).await?;
    emit_set(tx, pending.set, seen).await?;
    emit_removed(tx, pending.removed).await?;
    Ok(affected)
}

/// Read, locked, every audited row deleting `query`'s rows would cascade to.
async fn snapshot(
    tx: &mut PoolTx<'_>,
    query: &DeleteQuery,
    seen: &mut HashSet<(&'static str, String)>,
) -> Result<Pending, ExecError> {
    let mut out: Vec<(&'static ModelSchema, Vec<(SqlValue, PendingEntry)>)> = Vec::new();
    let mut set = Vec::new();
    let mut frontier = vec![(query.model, query.where_clause.clone())];
    while let Some((model, rows)) = frontier.pop() {
        for edge in edges(model.table) {
            let keys = column_values(tx, model, rows.clone(), edge.on).await?;
            let child = edge.child.schema;
            let pk = child
                .primary_key()
                .ok_or(ExecError::MissingPrimaryKey { table: child.table })?;
            for chunk in keys.chunks(BULK_AUDIT_CHUNK) {
                let child_rows = WhereExpr::Predicate(Filter::new(
                    edge.fk_column,
                    Op::In,
                    SqlValue::List(chunk.to_vec()),
                ));
                if edge.sets {
                    let pks = column_values(tx, child, child_rows, pk.column).await?;
                    if !pks.is_empty() {
                        set.push(SetRows {
                            child: edge.child,
                            fk_column: edge.fk_column,
                            keys: chunk.to_vec(),
                            pks,
                        });
                    }
                    continue;
                }
                let found: Vec<(SqlValue, Option<PendingEntry>)> = if child.audit_track.is_some() {
                    let read = edge
                        .child
                        .audited_cascade()
                        .ok_or(ExecError::AuditUnsupported {
                            table: child.table,
                            reason: "its registry entry has no cascade recorder",
                        })?;
                    let rows = read(tx, child_rows).await?;
                    rows.into_iter().map(|(k, e)| (k, Some(e))).collect()
                } else {
                    let pks = column_values(tx, child, child_rows, pk.column).await?;
                    pks.into_iter().map(|k| (k, None)).collect()
                };
                let mut fresh = Vec::new();
                let mut entries = Vec::new();
                for (key, entry) in found {
                    if seen.insert((child.table, format!("{key:?}"))) {
                        fresh.push(key.clone());
                        entries.extend(entry.map(|e| (key, e)));
                    }
                }
                if !entries.is_empty() {
                    out.push((child, entries));
                }
                if reaches_audited(child) {
                    for part in fresh.chunks(BULK_AUDIT_CHUNK) {
                        frontier.push((child, pk_in(child, part.to_vec())?));
                    }
                }
            }
        }
    }
    Ok(Pending { removed: out, set })
}

/// Emit an `update` entry for each audited row the delete set the FK of:
/// still there (no other edge removed it), no longer pointing at a key.
async fn emit_set(
    tx: &mut PoolTx<'_>,
    set: Vec<SetRows>,
    seen: &mut HashSet<(&'static str, String)>,
) -> Result<(), ExecError> {
    for rows in set {
        let child = rows.child.schema;
        let read = rows
            .child
            .audited_cascade()
            .ok_or(ExecError::AuditUnsupported {
                table: child.table,
                reason: "its registry entry has no cascade recorder",
            })?;
        let pk = child
            .primary_key()
            .ok_or(ExecError::MissingPrimaryKey { table: child.table })?;
        let mut entries = Vec::new();
        for part in rows.pks.chunks(BULK_AUDIT_CHUNK) {
            // A missing FK (or SQLite with `foreign_keys` off) sets nothing.
            let unchanged = WhereExpr::And(vec![
                pk_in(child, part.to_vec())?,
                WhereExpr::Predicate(Filter::new(
                    rows.fk_column,
                    Op::In,
                    SqlValue::List(rows.keys.clone()),
                )),
            ]);
            let left: HashSet<String> = column_values(tx, child, unchanged, pk.column)
                .await?
                .iter()
                .map(|k| format!("{k:?}"))
                .collect();
            for (key, mut entry) in read(tx, pk_in(child, part.to_vec())?).await? {
                let key = format!("{key:?}");
                if !left.contains(&key) && seen.insert((child.table, format!("set {key}"))) {
                    entry.operation = AuditOp::Update;
                    entries.push(entry);
                }
            }
        }
        emit_many_tx(tx, &entries).await?;
    }
    Ok(())
}

/// Emit the entries of the snapshot rows the delete removed.
async fn emit_removed(
    tx: &mut PoolTx<'_>,
    removed: Vec<(&'static ModelSchema, Vec<(SqlValue, PendingEntry)>)>,
) -> Result<(), ExecError> {
    for (model, rows) in removed {
        let pk = model
            .primary_key()
            .ok_or(ExecError::MissingPrimaryKey { table: model.table })?;
        let keys: Vec<SqlValue> = rows.iter().map(|(k, _)| k.clone()).collect();
        // A missing FK (or SQLite with `foreign_keys` off) cascades nothing.
        let mut left = HashSet::new();
        for part in keys.chunks(BULK_AUDIT_CHUNK) {
            let found = column_values(tx, model, pk_in(model, part.to_vec())?, pk.column).await?;
            left.extend(found.iter().map(|k| format!("{k:?}")));
        }
        let entries: Vec<PendingEntry> = rows
            .into_iter()
            .filter(|(k, _)| !left.contains(&format!("{k:?}")))
            .map(|(_, e)| e)
            .collect();
        emit_many_tx(tx, &entries).await?;
    }
    Ok(())
}

/// `column` of `model`'s rows matching `rows`, locked on PG / MySQL.
async fn column_values(
    tx: &mut PoolTx<'_>,
    model: &'static ModelSchema,
    rows: WhereExpr,
    column: &'static str,
) -> Result<Vec<SqlValue>, ExecError> {
    let mut select = SelectQuery::new(model).where_clause(rows);
    select.projection = Some(vec![column]);
    select.lock_mode = Some(crate::core::LockMode {
        silent_on_sqlite: true,
        ..crate::core::LockMode::default()
    });
    let rows = crate::sql::fetch_values_list_on(tx, &select).await?;
    Ok(rows
        .into_iter()
        .filter_map(|r| r.into_iter().next())
        .collect())
}

#[cfg(all(test, feature = "tenancy"))]
mod tx_source_tests {
    use crate::audit::{fetch_for_entity_pool, with_tenant_source, writing_to_tenant, AuditSource};
    use crate::core::Model as _;
    #[cfg(feature = "postgres")]
    use crate::sql::CounterPool as _;
    use crate::sql::Pool;

    #[derive(crate::Model, Debug, Clone)]
    #[rustango(table = "audit_txsrc_shelf", app = "audit_txsrc")]
    #[allow(dead_code)]
    pub struct Shelf {
        #[rustango(primary_key)]
        pub id: i64,
    }

    #[derive(crate::Model, Debug, Clone)]
    #[rustango(
        table = "audit_txsrc_book",
        app = "audit_txsrc",
        audit(track = "title")
    )]
    #[allow(dead_code)]
    pub struct Book {
        #[rustango(primary_key)]
        pub id: i64,
        #[rustango(on_delete = "cascade")]
        pub shelf: crate::sql::ForeignKey<Shelf>,
        #[rustango(max_length = 32)]
        pub title: String,
    }

    async fn setup(pool: &Pool) {
        crate::testkit::matrix::drop_table(pool, Book::SCHEMA.table).await;
        crate::testkit::matrix::fresh_table::<Shelf>(pool).await;
        crate::testkit::matrix::fresh_table::<Book>(pool).await;
        crate::audit::ensure_table_pool(pool).await.unwrap();
        crate::audit::AuditLog::delete_where("entity_table", Book::SCHEMA.table, pool)
            .await
            .unwrap();
        Shelf { id: 1 }.insert(pool).await.unwrap();
        let mut book = Book {
            id: 1,
            shelf: crate::sql::ForeignKey::unloaded(1_i64),
            title: "b".into(),
        };
        book.insert(pool).await.unwrap();
    }

    /// How the test deletes the shelf.
    #[derive(Clone, Copy)]
    enum Via {
        DeleteTx,
        DeleteOnTx,
        ConnBegin,
        DeleteOnConn,
        DeleteOnTenantConn,
    }

    /// Deletes the shelf as tenant `a`'s user, scoped to `pool`.
    async fn cascade_source(pool: &Pool, via: Via) -> Option<String> {
        let user = AuditSource::User { id: "7".into() };
        let query = || {
            Shelf::objects()
                .filter("id", 1_i64)
                .compile_delete()
                .unwrap()
        };
        with_tenant_source(user, "a".into(), async {
            writing_to_tenant("a".into(), Some(pool), async {
                match via {
                    Via::DeleteTx | Via::DeleteOnTx => {
                        let mut tx = crate::sql::transaction_pool(pool).await.unwrap();
                        if matches!(via, Via::DeleteTx) {
                            crate::sql::delete_tx(&mut tx, &query()).await.unwrap();
                        } else {
                            let qs = Shelf::objects().filter("id", 1_i64);
                            qs.delete_on(&mut tx).await.unwrap();
                        }
                        tx.commit().await.unwrap();
                    }
                    Via::ConnBegin => {
                        let mut conn = pool.acquire().await.unwrap();
                        let mut tx = conn.begin().await.unwrap();
                        crate::sql::delete_tx(&mut tx, &query()).await.unwrap();
                        tx.commit().await.unwrap();
                    }
                    Via::DeleteOnConn => {
                        let mut conn = pool.acquire().await.unwrap();
                        let qs = Shelf::objects().filter("id", 1_i64);
                        qs.delete_on(&mut conn).await.unwrap();
                    }
                    Via::DeleteOnTenantConn => delete_on_tenant_conn(pool).await,
                }
            })
            .await;
        })
        .await;
        let rows = fetch_for_entity_pool(pool, Book::SCHEMA.table, "1")
            .await
            .unwrap();
        rows.first().map(|e| e.source.clone())
    }

    /// `delete_on` through a database-mode tenant connection out of `pool`.
    async fn delete_on_tenant_conn(pool: &Pool) {
        #[cfg(any(feature = "mysql", feature = "sqlite"))]
        use crate::tenancy::DatabaseConn;
        #[cfg(any(feature = "mysql", feature = "sqlite"))]
        let qs = Shelf::objects().filter("id", 1_i64);
        match pool {
            #[cfg(feature = "mysql")]
            Pool::Mysql(p) => {
                let mut conn = DatabaseConn::new(p.acquire().await.unwrap(), pool);
                qs.delete_on(&mut conn).await.unwrap();
            }
            #[cfg(feature = "sqlite")]
            Pool::Sqlite(p) => {
                let mut conn = DatabaseConn::new(p.acquire().await.unwrap(), pool);
                qs.delete_on(&mut conn).await.unwrap();
            }
            #[allow(unreachable_patterns)]
            _ => unreachable!(),
        }
    }

    async fn delete_on_tenant_conn_keeps_the_tenant_source(pool: &Pool) {
        // Postgres tenants run `delete_on` on the raw sqlx connection.
        if pool.dialect().name() == "postgres" {
            return;
        }
        let source = cascade_source(pool, Via::DeleteOnTenantConn).await;
        assert_eq!(source.as_deref(), Some("user:7"));
    }

    /// A raw sqlx executor cannot audit the cascade, so the delete refuses.
    async fn raw_executor_cascade_is_refused(pool: &Pool) {
        #[cfg(not(feature = "postgres"))]
        let _ = pool;
        #[cfg(feature = "postgres")]
        #[allow(irrefutable_let_patterns)]
        if let Pool::Postgres(pg) = pool {
            let res = Shelf::objects().filter("id", 1_i64).delete_on(pg).await;
            assert!(
                matches!(res, Err(crate::sql::ExecError::AuditUnsupported { .. })),
                "{res:?}"
            );
            assert_eq!(Book::objects().count(pool).await.unwrap(), 1);
        }
    }

    async fn delete_tx_keeps_the_tenant_source(pool: &Pool) {
        let source = cascade_source(pool, Via::DeleteTx).await;
        assert_eq!(source.as_deref(), Some("user:7"));
    }

    async fn delete_on_tx_keeps_the_tenant_source(pool: &Pool) {
        let source = cascade_source(pool, Via::DeleteOnTx).await;
        assert_eq!(source.as_deref(), Some("user:7"));
    }

    async fn conn_begin_keeps_the_tenant_source(pool: &Pool) {
        let source = cascade_source(pool, Via::ConnBegin).await;
        assert_eq!(source.as_deref(), Some("user:7"));
    }

    async fn delete_on_conn_keeps_the_tenant_source(pool: &Pool) {
        let source = cascade_source(pool, Via::DeleteOnConn).await;
        assert_eq!(source.as_deref(), Some("user:7"));
    }

    crate::tri_dialect_test!(
        setup: setup,
        scenarios: [
            raw_executor_cascade_is_refused,
            delete_tx_keeps_the_tenant_source,
            delete_on_tx_keeps_the_tenant_source,
            conn_begin_keeps_the_tenant_source,
            delete_on_conn_keeps_the_tenant_source,
            delete_on_tenant_conn_keeps_the_tenant_source,
        ],
    );
}