chroma_types/
task.rs

1use serde::{Deserialize, Serialize};
2use std::time::SystemTime;
3
4use crate::CollectionUuid;
5
6fn prost_value_to_json(v: &prost_types::Value) -> serde_json::Value {
7    match &v.kind {
8        Some(prost_types::value::Kind::NullValue(_)) => serde_json::Value::Null,
9        Some(prost_types::value::Kind::NumberValue(n)) => serde_json::json!(*n),
10        Some(prost_types::value::Kind::StringValue(s)) => serde_json::Value::String(s.clone()),
11        Some(prost_types::value::Kind::BoolValue(b)) => serde_json::Value::Bool(*b),
12        Some(prost_types::value::Kind::StructValue(s)) => prost_struct_to_json(s),
13        Some(prost_types::value::Kind::ListValue(l)) => {
14            serde_json::Value::Array(l.values.iter().map(prost_value_to_json).collect())
15        }
16        None => serde_json::Value::Null,
17    }
18}
19
20fn prost_struct_to_json(s: &prost_types::Struct) -> serde_json::Value {
21    let map: serde_json::Map<String, serde_json::Value> = s
22        .fields
23        .iter()
24        .map(|(k, v)| (k.clone(), prost_value_to_json(v)))
25        .collect();
26    serde_json::Value::Object(map)
27}
28
29define_uuid_newtype!(
30    /// JobId is a wrapper around Uuid to provide a unified type for job identifiers.
31    /// Jobs can be either collection compaction jobs or task execution jobs.
32    JobId,
33    new_v4
34);
35
36// Custom From implementations for JobId
37impl From<CollectionUuid> for JobId {
38    fn from(collection_uuid: CollectionUuid) -> Self {
39        JobId(collection_uuid.0)
40    }
41}
42
43impl From<AttachedFunctionUuid> for JobId {
44    fn from(attached_function_uuid: AttachedFunctionUuid) -> Self {
45        JobId(attached_function_uuid.0)
46    }
47}
48
49define_uuid_newtype!(
50    /// AttachedFunctionUuid is a wrapper around Uuid to provide a type for attached function identifiers.
51    #[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
52    AttachedFunctionUuid,
53    new_v4
54);
55
56/// AttachedFunction represents an asynchronous function that is triggered by collection writes
57/// to map records from a source collection to a target collection.
58fn default_systemtime() -> SystemTime {
59    SystemTime::UNIX_EPOCH
60}
61
62#[derive(Clone, Debug, Deserialize, Serialize)]
63pub struct AttachedFunction {
64    /// Unique identifier for the attached function
65    pub id: AttachedFunctionUuid,
66    /// Human-readable name for the attached function instance
67    pub name: String,
68    /// UUID of the function/built-in definition this attached function uses
69    pub function_id: uuid::Uuid,
70    /// Source collection that triggers the attached function
71    pub input_collection_id: CollectionUuid,
72    /// Name of target collection where attached function output is stored
73    pub output_collection_name: String,
74    /// ID of the output collection (lazily filled in after creation)
75    pub output_collection_id: Option<CollectionUuid>,
76    /// Optional JSON parameters for the function
77    pub params: Option<String>,
78    /// Tenant ID this attached function belongs to
79    pub tenant_id: String,
80    /// Database ID this attached function belongs to
81    pub database_id: String,
82    /// Timestamp of the last successful function run
83    #[serde(skip, default)]
84    pub last_run: Option<SystemTime>,
85    /// Completion offset: the WAL position up to which the attached function has processed records
86    pub completion_offset: u64,
87    /// Minimum number of new records required before the attached function runs again
88    pub min_records_for_invocation: u64,
89    /// Whether the attached function has been soft-deleted
90    #[serde(skip, default)]
91    pub is_deleted: bool,
92    /// Whether the attached function runs asynchronously
93    #[serde(skip, default)]
94    pub is_async: bool,
95    /// Timestamp when the attached function was created
96    #[serde(default = "default_systemtime")]
97    pub created_at: SystemTime,
98    /// Timestamp when the attached function was last updated
99    #[serde(default = "default_systemtime")]
100    pub updated_at: SystemTime,
101    // is_ready is a column in the database, but not in the struct because
102    // it is not meant to be used in rust code. If it is false, rust code
103    // should never even see it.
104}
105
106#[derive(Debug, thiserror::Error)]
107pub enum AttachedFunctionConversionError {
108    #[error("Invalid UUID: {0}")]
109    InvalidUuid(String),
110}
111
112impl TryFrom<crate::chroma_proto::AttachedFunction> for AttachedFunction {
113    type Error = AttachedFunctionConversionError;
114
115    fn try_from(
116        attached_function: crate::chroma_proto::AttachedFunction,
117    ) -> Result<Self, Self::Error> {
118        // Parse attached_function_id
119        let attached_function_id = attached_function
120            .id
121            .parse::<AttachedFunctionUuid>()
122            .map_err(|_| {
123                AttachedFunctionConversionError::InvalidUuid("attached_function_id".to_string())
124            })?;
125
126        // Parse function_id
127        let function_id = attached_function
128            .function_id
129            .parse::<uuid::Uuid>()
130            .map_err(|_| AttachedFunctionConversionError::InvalidUuid("function_id".to_string()))?;
131
132        // Parse input_collection_id
133        let input_collection_id = attached_function
134            .input_collection_id
135            .parse::<CollectionUuid>()
136            .map_err(|_| {
137                AttachedFunctionConversionError::InvalidUuid("input_collection_id".to_string())
138            })?;
139
140        // Parse output_collection_id if available
141        let output_collection_id = attached_function
142            .output_collection_id
143            .map(|id| id.parse::<CollectionUuid>())
144            .transpose()
145            .map_err(|_| {
146                AttachedFunctionConversionError::InvalidUuid("output_collection_id".to_string())
147            })?;
148
149        let params = attached_function
150            .params
151            .as_ref()
152            .map(|s| serde_json::to_string(&prost_struct_to_json(s)))
153            .transpose()
154            .unwrap_or(None);
155
156        // Parse timestamps
157        let created_at = std::time::SystemTime::UNIX_EPOCH
158            + std::time::Duration::from_micros(attached_function.created_at);
159        let updated_at = std::time::SystemTime::UNIX_EPOCH
160            + std::time::Duration::from_micros(attached_function.updated_at);
161
162        Ok(AttachedFunction {
163            id: attached_function_id,
164            name: attached_function.name,
165            function_id,
166            input_collection_id,
167            output_collection_name: attached_function.output_collection_name,
168            output_collection_id,
169            params,
170            tenant_id: attached_function.tenant_id,
171            database_id: attached_function.database_id,
172            last_run: None, // Not available in proto
173            completion_offset: attached_function.completion_offset,
174            min_records_for_invocation: attached_function.min_records_for_invocation,
175            is_deleted: false, // Not available in proto, would need to be fetched separately
176            is_async: attached_function.is_async,
177            created_at,
178            updated_at,
179        })
180    }
181}
chroma_types/task.rs

chroma_types/
task.rs
