Skip to main content

cognee_database/ops/
task_runs.rs

1use cognee_utils::tracing_keys::COGNEE_DB_SYSTEM;
2use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter};
3use tracing::{Span, instrument};
4use uuid::Uuid;
5
6use crate::conversions::map_sea_err;
7use crate::database_system_label;
8use crate::entities::task_run;
9use crate::types::{DatabaseError, TaskRun};
10use crate::uuid_hex;
11
12#[instrument(
13    name = "cognee.db.relational.task_runs.create_task_run",
14    level = "info",
15    skip_all,
16    fields(cognee.db.system = tracing::field::Empty),
17    err,
18)]
19pub async fn create_task_run(
20    db: &DatabaseConnection,
21    run: TaskRun,
22) -> Result<TaskRun, DatabaseError> {
23    Span::current().record(COGNEE_DB_SYSTEM, database_system_label(db));
24    task_run::ActiveModel::from(&run)
25        .insert(db)
26        .await
27        .map_err(map_sea_err)?;
28    Ok(run)
29}
30
31#[instrument(
32    name = "cognee.db.relational.task_runs.update_task_run_status",
33    level = "info",
34    skip_all,
35    fields(cognee.db.system = tracing::field::Empty),
36    err,
37)]
38pub async fn update_task_run_status(
39    db: &DatabaseConnection,
40    id: Uuid,
41    status: &str,
42) -> Result<(), DatabaseError> {
43    Span::current().record(COGNEE_DB_SYSTEM, database_system_label(db));
44    task_run::Entity::update_many()
45        .col_expr(
46            task_run::Column::Status,
47            sea_orm::sea_query::Expr::value(status),
48        )
49        .filter(task_run::Column::Id.eq(uuid_hex::to_hex(id)))
50        .exec(db)
51        .await
52        .map_err(map_sea_err)?;
53    Ok(())
54}