1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use std::str::FromStr;
4use std::time::SystemTime;
5use uuid::Uuid;
6
7use crate::CollectionUuid;
8
9define_uuid_newtype!(
10 JobId,
13 new_v4
14);
15
16impl From<CollectionUuid> for JobId {
18 fn from(collection_uuid: CollectionUuid) -> Self {
19 JobId(collection_uuid.0)
20 }
21}
22
23impl From<AttachedFunctionUuid> for JobId {
24 fn from(attached_function_uuid: AttachedFunctionUuid) -> Self {
25 JobId(attached_function_uuid.0)
26 }
27}
28
29define_uuid_newtype!(
30 #[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
32 AttachedFunctionUuid,
33 new_v4
34);
35
36define_uuid_newtype!(
37 NonceUuid,
39 now_v7
40);
41
42fn default_systemtime() -> SystemTime {
45 SystemTime::UNIX_EPOCH
46}
47
48#[derive(Clone, Debug, Deserialize, Serialize)]
49pub struct AttachedFunction {
50 pub id: AttachedFunctionUuid,
52 pub name: String,
54 pub function_id: uuid::Uuid,
56 pub input_collection_id: CollectionUuid,
58 pub output_collection_name: String,
60 pub output_collection_id: Option<CollectionUuid>,
62 pub params: Option<String>,
64 pub tenant_id: String,
66 pub database_id: String,
68 #[serde(skip, default)]
70 pub last_run: Option<SystemTime>,
71 #[serde(skip, default = "default_systemtime")]
73 pub next_run: SystemTime,
74 pub completion_offset: u64,
76 pub min_records_for_invocation: u64,
78 #[serde(skip, default)]
80 pub is_deleted: bool,
81 #[serde(default = "default_systemtime")]
83 pub created_at: SystemTime,
84 #[serde(default = "default_systemtime")]
86 pub updated_at: SystemTime,
87 pub next_nonce: NonceUuid,
89 pub lowest_live_nonce: Option<NonceUuid>,
94}
95
96#[derive(Clone, Debug, Deserialize, Serialize)]
98pub struct ScheduleEntry {
99 pub collection_id: CollectionUuid,
100 pub attached_function_id: Uuid,
101 pub attached_function_run_nonce: NonceUuid,
102 pub when_to_run: Option<DateTime<Utc>>,
103 pub lowest_live_nonce: Option<Uuid>,
106}
107
108impl TryFrom<crate::chroma_proto::ScheduleEntry> for ScheduleEntry {
109 type Error = ScheduleEntryConversionError;
110
111 fn try_from(proto: crate::chroma_proto::ScheduleEntry) -> Result<Self, Self::Error> {
112 let collection_id = proto
113 .collection_id
114 .ok_or(ScheduleEntryConversionError::MissingField(
115 "collection_id".to_string(),
116 ))
117 .and_then(|id| {
118 CollectionUuid::from_str(&id).map_err(|_| {
119 ScheduleEntryConversionError::InvalidUuid("collection_id".to_string())
120 })
121 })?;
122
123 let attached_function_id = proto
124 .attached_function_id
125 .ok_or(ScheduleEntryConversionError::MissingField(
126 "attached_function_id".to_string(),
127 ))
128 .and_then(|id| {
129 Uuid::parse_str(&id).map_err(|_| {
130 ScheduleEntryConversionError::InvalidUuid("attached_function_id".to_string())
131 })
132 })?;
133
134 let attached_function_run_nonce = proto
135 .run_nonce
136 .ok_or(ScheduleEntryConversionError::MissingField(
137 "run_nonce".to_string(),
138 ))
139 .and_then(|nonce| {
140 Uuid::parse_str(&nonce)
141 .map(NonceUuid)
142 .map_err(|_| ScheduleEntryConversionError::InvalidUuid("run_nonce".to_string()))
143 })?;
144
145 let when_to_run = proto
146 .when_to_run
147 .and_then(|ms| DateTime::from_timestamp_millis(ms as i64));
148
149 let lowest_live_nonce = proto
150 .lowest_live_nonce
151 .as_ref()
152 .and_then(|nonce_str| Uuid::parse_str(nonce_str).ok());
153
154 Ok(ScheduleEntry {
155 collection_id,
156 attached_function_id,
157 attached_function_run_nonce,
158 when_to_run,
159 lowest_live_nonce,
160 })
161 }
162}
163
164#[derive(Debug, thiserror::Error)]
165pub enum ScheduleEntryConversionError {
166 #[error("Missing required field: {0}")]
167 MissingField(String),
168 #[error("Invalid UUID for field: {0}")]
169 InvalidUuid(String),
170}