chroma_types/
flush.rs

1use super::{CollectionUuid, ConversionError};
2use crate::{
3    chroma_proto::{FilePaths, FlushCollectionCompactionResponse, FlushSegmentCompactionInfo},
4    SegmentUuid,
5};
6use chroma_error::{ChromaError, ErrorCodes};
7use std::collections::HashMap;
8use thiserror::Error;
9use uuid::Uuid;
10
11#[derive(Debug, Clone)]
12pub struct SegmentFlushInfo {
13    pub segment_id: SegmentUuid,
14    pub file_paths: HashMap<String, Vec<String>>,
15}
16
17impl TryInto<FlushSegmentCompactionInfo> for &SegmentFlushInfo {
18    type Error = SegmentFlushInfoConversionError;
19
20    fn try_into(self) -> Result<FlushSegmentCompactionInfo, Self::Error> {
21        let mut file_paths = HashMap::new();
22        for (key, value) in self.file_paths.clone() {
23            file_paths.insert(key, FilePaths { paths: value });
24        }
25
26        Ok(FlushSegmentCompactionInfo {
27            segment_id: self.segment_id.to_string(),
28            file_paths,
29        })
30    }
31}
32
33#[derive(Error, Debug)]
34pub enum SegmentFlushInfoConversionError {
35    #[error("Invalid segment id, valid UUID required")]
36    InvalidSegmentId,
37    #[error(transparent)]
38    DecodeError(#[from] ConversionError),
39}
40
41#[derive(Debug)]
42pub struct FlushCompactionResponse {
43    pub collection_id: CollectionUuid,
44    pub collection_version: i32,
45    pub last_compaction_time: i64,
46}
47
48impl FlushCompactionResponse {
49    pub fn new(
50        collection_id: CollectionUuid,
51        collection_version: i32,
52        last_compaction_time: i64,
53    ) -> Self {
54        FlushCompactionResponse {
55            collection_id,
56            collection_version,
57            last_compaction_time,
58        }
59    }
60}
61
62impl TryFrom<FlushCollectionCompactionResponse> for FlushCompactionResponse {
63    type Error = FlushCompactionResponseConversionError;
64
65    fn try_from(value: FlushCollectionCompactionResponse) -> Result<Self, Self::Error> {
66        let id = Uuid::parse_str(&value.collection_id)
67            .map_err(|_| FlushCompactionResponseConversionError::InvalidUuid)?;
68        Ok(FlushCompactionResponse {
69            collection_id: CollectionUuid(id),
70            collection_version: value.collection_version,
71            last_compaction_time: value.last_compaction_time,
72        })
73    }
74}
75
76#[derive(Error, Debug)]
77pub enum FlushCompactionResponseConversionError {
78    #[error(transparent)]
79    DecodeError(#[from] ConversionError),
80    #[error("Invalid collection id, valid UUID required")]
81    InvalidUuid,
82}
83
84impl ChromaError for FlushCompactionResponseConversionError {
85    fn code(&self) -> ErrorCodes {
86        match self {
87            FlushCompactionResponseConversionError::InvalidUuid => ErrorCodes::InvalidArgument,
88            FlushCompactionResponseConversionError::DecodeError(e) => e.code(),
89        }
90    }
91}