solstice-ci/crates/orchestrator/src/persist.rs

use miette::{IntoDiagnostic as _, Result};
use sea_orm::{Database, DatabaseConnection, sea_query::{OnConflict, Expr}, Query, EntityTrait, DbBackend, Statement};
use sea_orm::sea_query::{InsertStatement, PostgresQueryBuilder, ColumnRef};
use time::OffsetDateTime;
use uuid::Uuid;

pub async fn maybe_init_db(url: &str, run_migrations: bool) -> Result<Option<DatabaseConnection>> {
    // Only enable DB when explicitly requested to avoid forcing Postgres locally.
    let enabled = std::env::var("ORCH_ENABLE_DB").ok().map(|v| v == "1" || v.eq_ignore_ascii_case("true")).unwrap_or(false);
    if !enabled { return Ok(None); }
    let db = Database::connect(url).await.into_diagnostic()?;
    if run_migrations {
        migration::Migrator::up(&db, None).await.into_diagnostic()?;
    }
    Ok(Some(db))
}

pub async fn record_job_state(
    db: &DatabaseConnection,
    request_id: Uuid,
    repo_url: &str,
    commit_sha: &str,
    runs_on: Option<&str>,
    state: &str,
) -> Result<()> {
    let now = OffsetDateTime::now_utc();
    // INSERT ... ON CONFLICT (request_id) DO UPDATE SET state=EXCLUDED.state, updated_at=EXCLUDED.updated_at
    let mut insert: InsertStatement = sea_orm::sea_query::Query::insert()
        .into_table(Jobs::Table)
        .columns([
            Jobs::RequestId,
            Jobs::RepoUrl,
            Jobs::CommitSha,
            Jobs::RunsOn,
            Jobs::State,
            Jobs::CreatedAt,
            Jobs::UpdatedAt,
        ])
        .values_panic([
            Expr::val(request_id),
            Expr::val(repo_url),
            Expr::val(commit_sha),
            Expr::val(runs_on.map(|s| s.to_string())),
            Expr::val(state),
            Expr::val(now),
            Expr::val(now),
        ])
        .on_conflict(
            OnConflict::column(Jobs::RequestId)
                .update_columns([Jobs::State, Jobs::UpdatedAt])
                .to_owned(),
        )
        .to_owned();

    let sql = insert.build(PostgresQueryBuilder);
    db.execute(Statement::from_string(DbBackend::Postgres, sql)).await.into_diagnostic()?;
    Ok(())
}

#[derive(sea_orm::sea_query::Iden)]
enum Jobs {
    #[iden = "jobs"]
    Table,
    #[iden = "request_id"]
    RequestId,
    #[iden = "repo_url"]
    RepoUrl,
    #[iden = "commit_sha"]
    CommitSha,
    #[iden = "runs_on"]
    RunsOn,
    #[iden = "state"]
    State,
    #[iden = "created_at"]
    CreatedAt,
    #[iden = "updated_at"]
    UpdatedAt,
}