crank/crates/crank-registry/src/migrations/schema_guard_v9.rs

use sqlx::{PgConnection, Row, query};

use super::authority::MigrationError;
use super::schema_guard::{
    column_exists, constraint_expression, normalize_definition, schema_error,
};

pub(super) async fn validate_v9_absent(
    connection: &mut PgConnection,
) -> Result<bool, MigrationError> {
    let columns_present = column_exists(connection, "agents", "catalog_revision").await?
        || column_exists(connection, "published_agents", "catalog_revision").await?;
    let triggers_present = trigger_exists(
        connection,
        "agent_versions",
        "agent_versions_immutable_guard",
    )
    .await?
        || trigger_exists(
            connection,
            "agent_operation_bindings",
            "agent_operation_bindings_immutable_guard",
        )
        .await?
        || trigger_exists(
            connection,
            "published_agents",
            "published_agents_monotonic_guard",
        )
        .await?;
    Ok(!columns_present && !triggers_present)
}

pub(super) async fn validate_v9_agent_catalog_lifecycle(
    connection: &mut PgConnection,
) -> Result<(), MigrationError> {
    for (table, column) in [
        ("agents", "catalog_revision"),
        ("published_agents", "catalog_revision"),
    ] {
        let row = query(
            "select data_type, is_nullable
             from information_schema.columns
             where table_schema = current_schema()
               and table_name = $1
               and column_name = $2",
        )
        .bind(table)
        .bind(column)
        .fetch_optional(&mut *connection)
        .await
        .map_err(|_| MigrationError::storage("preflight.schema"))?;
        let valid = row.is_some_and(|row| {
            row.try_get::<String, _>("data_type").ok().as_deref() == Some("bigint")
                && row.try_get::<String, _>("is_nullable").ok().as_deref() == Some("NO")
        });
        if !valid {
            return Err(schema_error(9));
        }
    }

    let agents_revision =
        constraint_expression(connection, "agents", "agents_catalog_revision_check").await?;
    if normalize_definition(&agents_revision) != "catalog_revision>=0" {
        return Err(schema_error(9));
    }
    let published_revision = constraint_expression(
        connection,
        "published_agents",
        "published_agents_catalog_revision_check",
    )
    .await?;
    if normalize_definition(&published_revision) != "catalog_revision>0" {
        return Err(schema_error(9));
    }

    for (table, trigger, expected_events, expected_function, function_snippets) in [
        (
            "agent_versions",
            "agent_versions_immutable_guard",
            &["before", "update", "delete"][..],
            "crank_reject_published_agent_version_mutation",
            &["old.status='published'", "returnold", "returnnew"][..],
        ),
        (
            "agent_operation_bindings",
            "agent_operation_bindings_immutable_guard",
            &["before", "insert", "update", "delete"][..],
            "crank_reject_published_agent_binding_mutation",
            &[
                "coalescenew.agent_id,old.agent_id",
                "bound_status='published'",
                "returnold",
                "returnnew",
            ][..],
        ),
        (
            "published_agents",
            "published_agents_monotonic_guard",
            &["before", "insert", "update", "delete"][..],
            "crank_reject_published_agent_pointer_rewind",
            &[
                "tg_op='delete'",
                "new.catalog_revision<=old.catalog_revision",
                "new.version<old.version",
            ][..],
        ),
    ] {
        let Some(trigger_definition) = trigger_definition(connection, table, trigger).await? else {
            return Err(schema_error(9));
        };
        let normalized_trigger = normalize_definition(&trigger_definition);
        if expected_events
            .iter()
            .any(|event| !normalized_trigger.contains(event))
            || !normalized_trigger.contains(expected_function)
        {
            return Err(schema_error(9));
        }
        let function_definition = function_definition(connection, expected_function).await?;
        let normalized_function = normalize_definition(&function_definition);
        if function_snippets
            .iter()
            .any(|snippet| !normalized_function.contains(snippet))
        {
            return Err(schema_error(9));
        }
    }
    Ok(())
}

async fn trigger_exists(
    connection: &mut PgConnection,
    table: &str,
    trigger: &str,
) -> Result<bool, MigrationError> {
    query(
        "select exists (
            select 1
            from pg_catalog.pg_trigger trg
            join pg_catalog.pg_class t on t.oid = trg.tgrelid
            join pg_catalog.pg_namespace n on n.oid = t.relnamespace
            where n.nspname = current_schema()
              and t.relname = $1
              and trg.tgname = $2
              and not trg.tgisinternal
         ) as present",
    )
    .bind(table)
    .bind(trigger)
    .fetch_one(connection)
    .await
    .map_err(|_| MigrationError::storage("preflight.schema"))?
    .try_get::<bool, _>("present")
    .map_err(|_| MigrationError::storage("preflight.schema"))
}

async fn trigger_definition(
    connection: &mut PgConnection,
    table: &str,
    trigger: &str,
) -> Result<Option<String>, MigrationError> {
    let row = query(
        "select pg_get_triggerdef(trg.oid) as definition
         from pg_catalog.pg_trigger trg
         join pg_catalog.pg_class t on t.oid = trg.tgrelid
         join pg_catalog.pg_namespace n on n.oid = t.relnamespace
         where n.nspname = current_schema()
           and t.relname = $1
           and trg.tgname = $2
           and not trg.tgisinternal",
    )
    .bind(table)
    .bind(trigger)
    .fetch_optional(&mut *connection)
    .await
    .map_err(|_| MigrationError::storage("preflight.schema"))?;
    row.map(|row| row.try_get::<String, _>("definition"))
        .transpose()
        .map_err(|_| MigrationError::storage("preflight.schema"))
}

async fn function_definition(
    connection: &mut PgConnection,
    function_name: &str,
) -> Result<String, MigrationError> {
    query(
        "select pg_get_functiondef(p.oid) as definition
         from pg_catalog.pg_proc p
         join pg_catalog.pg_namespace n on n.oid = p.pronamespace
         where n.nspname = current_schema()
           and p.proname = $1",
    )
    .bind(function_name)
    .fetch_optional(connection)
    .await
    .map_err(|_| MigrationError::storage("preflight.schema"))?
    .ok_or_else(|| schema_error(9))?
    .try_get::<String, _>("definition")
    .map_err(|_| MigrationError::storage("preflight.schema"))
}