chroma_types/
operation.rs

1use super::ConversionError;
2use crate::chroma_proto;
3use chroma_error::{ChromaError, ErrorCodes};
4use thiserror::Error;
5
6#[derive(Clone, Copy, Debug, PartialEq)]
7pub enum Operation {
8    Add,
9    Update,
10    Upsert,
11    Delete,
12    BackfillFn,
13}
14
15#[derive(Clone, Copy, Debug, PartialEq)]
16pub enum MaterializedLogOperation {
17    // Set when the record is initially read from the segment
18    // before it is processed based on state of the log.
19    Initial,
20    // Set for records that don't exist in the segment and
21    // have been encountered for the first time in the log.
22    AddNew,
23    // Assume there is a record in the segment and in the log
24    // there is a DEL followed by an ADD for the same id.
25    // In this case, the overwriteexisting state is set for
26    // the record. Easy to construct other cases.
27    OverwriteExisting,
28    // Set for entries that are present in the segment and
29    // have been updated/upserted in the log.
30    UpdateExisting,
31    // Set for entries that are present in the segment and
32    // have been deleted in the log.
33    DeleteExisting,
34}
35
36#[derive(Error, Debug)]
37pub enum OperationConversionError {
38    #[error("Invalid operation, valid operations are: Add, Upsert, Update, Delete")]
39    InvalidOperation,
40    #[error(transparent)]
41    DecodeError(#[from] ConversionError),
42}
43
44impl_base_convert_error!(OperationConversionError, {
45    OperationConversionError::InvalidOperation => ErrorCodes::InvalidArgument,
46});
47
48impl TryFrom<chroma_proto::Operation> for Operation {
49    type Error = OperationConversionError;
50
51    fn try_from(op: chroma_proto::Operation) -> Result<Self, Self::Error> {
52        match op {
53            chroma_proto::Operation::Add => Ok(Operation::Add),
54            chroma_proto::Operation::Upsert => Ok(Operation::Upsert),
55            chroma_proto::Operation::Update => Ok(Operation::Update),
56            chroma_proto::Operation::Delete => Ok(Operation::Delete),
57            chroma_proto::Operation::BackfillFn => Ok(Operation::BackfillFn),
58        }
59    }
60}
61
62impl TryFrom<i32> for Operation {
63    type Error = OperationConversionError;
64
65    fn try_from(op: i32) -> Result<Self, Self::Error> {
66        let maybe_op = chroma_proto::Operation::try_from(op);
67        match maybe_op {
68            Ok(op) => match op {
69                chroma_proto::Operation::Add => Ok(Operation::Add),
70                chroma_proto::Operation::Upsert => Ok(Operation::Upsert),
71                chroma_proto::Operation::Update => Ok(Operation::Update),
72                chroma_proto::Operation::Delete => Ok(Operation::Delete),
73                chroma_proto::Operation::BackfillFn => Ok(Operation::BackfillFn),
74            },
75            Err(_) => Err(OperationConversionError::DecodeError(
76                ConversionError::DecodeError,
77            )),
78        }
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85    use crate::chroma_proto;
86
87    #[test]
88    fn test_operation_try_from() {
89        let proto_op = chroma_proto::Operation::Add;
90        let converted_op: Operation = proto_op.try_into().unwrap();
91        assert_eq!(converted_op, Operation::Add);
92    }
93}