Skip to main content

briefcase_core/storage/
mod.rs

1use crate::models::{DecisionSnapshot, Snapshot};
2#[cfg(feature = "async")]
3use async_trait::async_trait;
4use chrono::{DateTime, Utc};
5use std::collections::HashMap;
6use thiserror::Error;
7
8pub mod buffered;
9#[cfg(feature = "lakefs-storage")]
10pub mod lakefs;
11#[cfg(feature = "sqlite-storage")]
12pub mod sqlite;
13pub mod sync;
14#[cfg(feature = "vcs-storage")]
15pub mod vcs;
16
17#[cfg(feature = "lakefs-storage")]
18pub use lakefs::{
19    ActionHook, BranchInfo, CommitInfo, DiffEntry, LakeFSAction, LakeFSBackend, LakeFSConfig,
20    MergeResult, ObjectStats, RepositoryInfo,
21};
22#[cfg(feature = "sqlite-storage")]
23pub use sqlite::SqliteBackend;
24#[cfg(feature = "sqlite-storage")]
25pub use sync::SyncSqliteBackend;
26pub use sync::{MemoryStorageBackend, SyncStorageBackend};
27
28#[cfg(feature = "async")]
29#[async_trait]
30pub trait StorageBackend: Send + Sync {
31    /// Save a snapshot, return its ID
32    async fn save(&self, snapshot: &Snapshot) -> Result<String, StorageError>;
33
34    /// Save a single decision snapshot
35    async fn save_decision(&self, decision: &DecisionSnapshot) -> Result<String, StorageError>;
36
37    /// Load a snapshot by ID
38    async fn load(&self, snapshot_id: &str) -> Result<Snapshot, StorageError>;
39
40    /// Load a decision by ID
41    async fn load_decision(&self, decision_id: &str) -> Result<DecisionSnapshot, StorageError>;
42
43    /// Query snapshots with filters
44    async fn query(&self, query: SnapshotQuery) -> Result<Vec<Snapshot>, StorageError>;
45
46    /// Delete a snapshot
47    async fn delete(&self, snapshot_id: &str) -> Result<bool, StorageError>;
48
49    /// Flush pending writes (for batching backends)
50    async fn flush(&self) -> Result<FlushResult, StorageError>;
51
52    /// Check health/connectivity
53    async fn health_check(&self) -> Result<bool, StorageError>;
54}
55
56#[derive(Debug, Clone, Default)]
57pub struct SnapshotQuery {
58    pub function_name: Option<String>,
59    pub module_name: Option<String>,
60    pub model_name: Option<String>,
61    pub start_time: Option<DateTime<Utc>>,
62    pub end_time: Option<DateTime<Utc>>,
63    pub tags: Option<HashMap<String, String>>,
64    pub limit: Option<usize>,
65    pub offset: Option<usize>,
66}
67
68impl SnapshotQuery {
69    pub fn new() -> Self {
70        Self::default()
71    }
72
73    pub fn with_function_name(mut self, function_name: impl Into<String>) -> Self {
74        self.function_name = Some(function_name.into());
75        self
76    }
77
78    pub fn with_module_name(mut self, module_name: impl Into<String>) -> Self {
79        self.module_name = Some(module_name.into());
80        self
81    }
82
83    pub fn with_model_name(mut self, model_name: impl Into<String>) -> Self {
84        self.model_name = Some(model_name.into());
85        self
86    }
87
88    pub fn with_time_range(mut self, start: DateTime<Utc>, end: DateTime<Utc>) -> Self {
89        self.start_time = Some(start);
90        self.end_time = Some(end);
91        self
92    }
93
94    pub fn with_tag(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
95        if self.tags.is_none() {
96            self.tags = Some(HashMap::new());
97        }
98        self.tags.as_mut().unwrap().insert(key.into(), value.into());
99        self
100    }
101
102    pub fn with_limit(mut self, limit: usize) -> Self {
103        self.limit = Some(limit);
104        self
105    }
106
107    pub fn with_offset(mut self, offset: usize) -> Self {
108        self.offset = Some(offset);
109        self
110    }
111}
112
113#[derive(Debug, Clone)]
114pub struct FlushResult {
115    pub snapshots_written: usize,
116    pub bytes_written: usize,
117    pub checkpoint_id: Option<String>, // LakeFS commit ID
118}
119
120#[derive(Error, Debug, Clone, PartialEq)]
121pub enum StorageError {
122    #[error("Not found: {0}")]
123    NotFound(String),
124    #[error("Connection error: {0}")]
125    ConnectionError(String),
126    #[error("Serialization error: {0}")]
127    SerializationError(String),
128    #[error("Permission denied: {0}")]
129    PermissionDenied(String),
130    #[error("Quota exceeded")]
131    QuotaExceeded,
132    #[error("Invalid query: {0}")]
133    InvalidQuery(String),
134    #[error("IO error: {0}")]
135    IoError(String),
136}
137
138impl From<serde_json::Error> for StorageError {
139    fn from(err: serde_json::Error) -> Self {
140        StorageError::SerializationError(err.to_string())
141    }
142}
143
144impl From<std::io::Error> for StorageError {
145    fn from(err: std::io::Error) -> Self {
146        StorageError::IoError(err.to_string())
147    }
148}