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