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}