cognee_database/ops/
task_runs.rs1use 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}