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