use anyhow::Result;
use chrono::Utc;
use sea_orm::entity::prelude::*;
use sea_orm::{ActiveValue::Set, QueryOrder};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Serialize, Deserialize)]
#[sea_orm(rs_type = "String", db_type = "String(StringLen::None)")]
pub enum TaskRepoStatus {
#[sea_orm(string_value = "pending")]
Pending,
#[sea_orm(string_value = "in_progress")]
InProgress,
#[sea_orm(string_value = "done")]
Done,
#[sea_orm(string_value = "failed")]
Failed,
}
impl std::fmt::Display for TaskRepoStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
TaskRepoStatus::Pending => write!(f, "pending"),
TaskRepoStatus::InProgress => write!(f, "in_progress"),
TaskRepoStatus::Done => write!(f, "done"),
TaskRepoStatus::Failed => write!(f, "failed"),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "task_repos")]
pub struct Model {
#[sea_orm(primary_key)]
pub id: i32,
pub task_id: i32,
pub repo_id: i32,
pub status: TaskRepoStatus,
pub branch_name: Option<String>,
pub pr_url: Option<String>,
pub error_message: Option<String>,
pub started_at: Option<String>,
pub completed_at: Option<String>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {}
impl ActiveModelBehavior for ActiveModel {}
impl Model {
pub async fn create(db: &DatabaseConnection, task_id: i32, repo_id: i32) -> Result<()> {
let existing = Entity::find()
.filter(Column::TaskId.eq(task_id))
.filter(Column::RepoId.eq(repo_id))
.one(db)
.await?;
if existing.is_none() {
let active = ActiveModel {
task_id: Set(task_id),
repo_id: Set(repo_id),
status: Set(TaskRepoStatus::Pending),
..Default::default()
};
active.insert(db).await?;
}
Ok(())
}
pub async fn find_by_task(db: &DatabaseConnection, task_id: i32) -> Result<Vec<Model>> {
let rows = Entity::find()
.filter(Column::TaskId.eq(task_id))
.order_by_asc(Column::Id)
.all(db)
.await?;
Ok(rows)
}
pub async fn set_in_progress(db: &DatabaseConnection, id: i32) -> Result<()> {
let now = Utc::now().to_rfc3339();
Entity::update_many()
.col_expr(Column::Status, Expr::value(TaskRepoStatus::InProgress.to_string()))
.col_expr(Column::StartedAt, Expr::value(now))
.filter(Column::Id.eq(id))
.exec(db)
.await?;
Ok(())
}
pub async fn set_done(
db: &DatabaseConnection,
id: i32,
branch_name: &str,
pr_url: &str,
) -> Result<()> {
let now = Utc::now().to_rfc3339();
Entity::update_many()
.col_expr(Column::Status, Expr::value(TaskRepoStatus::Done.to_string()))
.col_expr(Column::BranchName, Expr::value(branch_name))
.col_expr(Column::PrUrl, Expr::value(pr_url))
.col_expr(Column::CompletedAt, Expr::value(now))
.filter(Column::Id.eq(id))
.exec(db)
.await?;
Ok(())
}
pub async fn set_failed(db: &DatabaseConnection, id: i32, error: &str) -> Result<()> {
let now = Utc::now().to_rfc3339();
Entity::update_many()
.col_expr(Column::Status, Expr::value(TaskRepoStatus::Failed.to_string()))
.col_expr(Column::ErrorMessage, Expr::value(error))
.col_expr(Column::CompletedAt, Expr::value(now))
.filter(Column::Id.eq(id))
.exec(db)
.await?;
Ok(())
}
}