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> { // 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, }