llm_optimizer_processor/storage/
wrapper.rs

1//! Storage wrapper for dynamic dispatch
2//!
3//! This module provides an enum-based wrapper that allows storing different
4//! storage backend types while maintaining the ability to use generic methods.
5
6use async_trait::async_trait;
7use chrono::{DateTime, Utc};
8use serde::{de::DeserializeOwned, Serialize};
9
10use super::error::StorageResult;
11use super::postgresql::PostgreSQLStorage;
12use super::redis::RedisStorage;
13use super::sled_backend::SledStorage;
14use super::traits::Storage;
15use super::types::{
16    Batch, Query, StorageBackend, StorageEntry, StorageHealth, StorageMetadata, StorageStats,
17    Transaction,
18};
19
20/// Enum wrapper for different storage backend implementations
21///
22/// This allows dynamic dispatch while preserving the ability to call
23/// generic methods that require `Self: Sized`.
24#[derive(Debug)]
25pub enum AnyStorage {
26    PostgreSQL(PostgreSQLStorage),
27    Redis(RedisStorage),
28    Sled(SledStorage),
29}
30
31impl AnyStorage {
32    /// Create from PostgreSQL storage
33    pub fn from_postgresql(storage: PostgreSQLStorage) -> Self {
34        Self::PostgreSQL(storage)
35    }
36
37    /// Create from Redis storage
38    pub fn from_redis(storage: RedisStorage) -> Self {
39        Self::Redis(storage)
40    }
41
42    /// Create from Sled storage
43    pub fn from_sled(storage: SledStorage) -> Self {
44        Self::Sled(storage)
45    }
46}
47
48#[async_trait]
49impl Storage for AnyStorage {
50    fn backend(&self) -> StorageBackend {
51        match self {
52            Self::PostgreSQL(s) => s.backend(),
53            Self::Redis(s) => s.backend(),
54            Self::Sled(s) => s.backend(),
55        }
56    }
57
58    async fn metadata(&self) -> StorageResult<StorageMetadata> {
59        match self {
60            Self::PostgreSQL(s) => s.metadata().await,
61            Self::Redis(s) => s.metadata().await,
62            Self::Sled(s) => s.metadata().await,
63        }
64    }
65
66    async fn health_check(&self) -> StorageResult<StorageHealth> {
67        match self {
68            Self::PostgreSQL(s) => s.health_check().await,
69            Self::Redis(s) => s.health_check().await,
70            Self::Sled(s) => s.health_check().await,
71        }
72    }
73
74    async fn stats(&self) -> StorageResult<StorageStats> {
75        match self {
76            Self::PostgreSQL(s) => s.stats().await,
77            Self::Redis(s) => s.stats().await,
78            Self::Sled(s) => s.stats().await,
79        }
80    }
81
82    async fn insert<T: Serialize + DeserializeOwned + Send + Sync>(
83        &self,
84        table: &str,
85        key: &str,
86        value: &T,
87    ) -> StorageResult<()> {
88        match self {
89            Self::PostgreSQL(s) => s.insert(table, key, value).await,
90            Self::Redis(s) => s.insert(table, key, value).await,
91            Self::Sled(s) => s.insert(table, key, value).await,
92        }
93    }
94
95    async fn insert_with_ttl<T: Serialize + Send + Sync>(
96        &self,
97        table: &str,
98        key: &str,
99        value: &T,
100        ttl_seconds: i64,
101    ) -> StorageResult<()> {
102        match self {
103            Self::PostgreSQL(s) => s.insert_with_ttl(table, key, value, ttl_seconds).await,
104            Self::Redis(s) => s.insert_with_ttl(table, key, value, ttl_seconds).await,
105            Self::Sled(s) => s.insert_with_ttl(table, key, value, ttl_seconds).await,
106        }
107    }
108
109    async fn get<T: DeserializeOwned + Serialize + Send + Sync>(
110        &self,
111        table: &str,
112        key: &str,
113    ) -> StorageResult<Option<StorageEntry<T>>> {
114        match self {
115            Self::PostgreSQL(s) => s.get(table, key).await,
116            Self::Redis(s) => s.get(table, key).await,
117            Self::Sled(s) => s.get(table, key).await,
118        }
119    }
120
121    async fn update<T: Serialize + Send + Sync>(
122        &self,
123        table: &str,
124        key: &str,
125        value: &T,
126    ) -> StorageResult<()> {
127        match self {
128            Self::PostgreSQL(s) => s.update(table, key, value).await,
129            Self::Redis(s) => s.update(table, key, value).await,
130            Self::Sled(s) => s.update(table, key, value).await,
131        }
132    }
133
134    async fn update_versioned<T: Serialize + Send + Sync>(
135        &self,
136        table: &str,
137        key: &str,
138        value: &T,
139        expected_version: u64,
140    ) -> StorageResult<()> {
141        match self {
142            Self::PostgreSQL(s) => s.update_versioned(table, key, value, expected_version).await,
143            Self::Redis(s) => s.update_versioned(table, key, value, expected_version).await,
144            Self::Sled(s) => s.update_versioned(table, key, value, expected_version).await,
145        }
146    }
147
148    async fn delete(&self, table: &str, key: &str) -> StorageResult<()> {
149        match self {
150            Self::PostgreSQL(s) => s.delete(table, key).await,
151            Self::Redis(s) => s.delete(table, key).await,
152            Self::Sled(s) => s.delete(table, key).await,
153        }
154    }
155
156    async fn exists(&self, table: &str, key: &str) -> StorageResult<bool> {
157        match self {
158            Self::PostgreSQL(s) => s.exists(table, key).await,
159            Self::Redis(s) => s.exists(table, key).await,
160            Self::Sled(s) => s.exists(table, key).await,
161        }
162    }
163
164    async fn execute_batch(&self, batch: Batch) -> StorageResult<()> {
165        match self {
166            Self::PostgreSQL(s) => s.execute_batch(batch).await,
167            Self::Redis(s) => s.execute_batch(batch).await,
168            Self::Sled(s) => s.execute_batch(batch).await,
169        }
170    }
171
172    async fn query<T: DeserializeOwned + Serialize + Send + Sync + 'static>(
173        &self,
174        query: Query,
175    ) -> StorageResult<Vec<StorageEntry<T>>> {
176        match self {
177            Self::PostgreSQL(s) => s.query(query).await,
178            Self::Redis(s) => s.query(query).await,
179            Self::Sled(s) => s.query(query).await,
180        }
181    }
182
183    async fn count(&self, query: Query) -> StorageResult<usize> {
184        match self {
185            Self::PostgreSQL(s) => s.count(query).await,
186            Self::Redis(s) => s.count(query).await,
187            Self::Sled(s) => s.count(query).await,
188        }
189    }
190
191    async fn begin_transaction(&self) -> StorageResult<Transaction> {
192        match self {
193            Self::PostgreSQL(s) => s.begin_transaction().await,
194            Self::Redis(s) => s.begin_transaction().await,
195            Self::Sled(s) => s.begin_transaction().await,
196        }
197    }
198
199    async fn commit_transaction(&self, transaction: Transaction) -> StorageResult<()> {
200        match self {
201            Self::PostgreSQL(s) => s.commit_transaction(transaction).await,
202            Self::Redis(s) => s.commit_transaction(transaction).await,
203            Self::Sled(s) => s.commit_transaction(transaction).await,
204        }
205    }
206
207    async fn rollback_transaction(&self, transaction: Transaction) -> StorageResult<()> {
208        match self {
209            Self::PostgreSQL(s) => s.rollback_transaction(transaction).await,
210            Self::Redis(s) => s.rollback_transaction(transaction).await,
211            Self::Sled(s) => s.rollback_transaction(transaction).await,
212        }
213    }
214
215    async fn scan_prefix(&self, table: &str, prefix: &str) -> StorageResult<Vec<String>> {
216        match self {
217            Self::PostgreSQL(s) => s.scan_prefix(table, prefix).await,
218            Self::Redis(s) => s.scan_prefix(table, prefix).await,
219            Self::Sled(s) => s.scan_prefix(table, prefix).await,
220        }
221    }
222
223    async fn get_multi<T: DeserializeOwned + Serialize + Send + Sync>(
224        &self,
225        table: &str,
226        keys: Vec<String>,
227    ) -> StorageResult<Vec<Option<StorageEntry<T>>>> {
228        match self {
229            Self::PostgreSQL(s) => s.get_multi(table, keys).await,
230            Self::Redis(s) => s.get_multi(table, keys).await,
231            Self::Sled(s) => s.get_multi(table, keys).await,
232        }
233    }
234
235    async fn delete_multi(&self, table: &str, keys: Vec<String>) -> StorageResult<()> {
236        match self {
237            Self::PostgreSQL(s) => s.delete_multi(table, keys).await,
238            Self::Redis(s) => s.delete_multi(table, keys).await,
239            Self::Sled(s) => s.delete_multi(table, keys).await,
240        }
241    }
242
243    async fn compare_and_swap<T: Serialize + Send + Sync>(
244        &self,
245        table: &str,
246        key: &str,
247        expected_version: u64,
248        new_value: &T,
249    ) -> StorageResult<bool> {
250        match self {
251            Self::PostgreSQL(s) => s.compare_and_swap(table, key, expected_version, new_value).await,
252            Self::Redis(s) => s.compare_and_swap(table, key, expected_version, new_value).await,
253            Self::Sled(s) => s.compare_and_swap(table, key, expected_version, new_value).await,
254        }
255    }
256
257    async fn get_modified_since<T: DeserializeOwned + Serialize + Send + Sync + 'static>(
258        &self,
259        table: &str,
260        since: DateTime<Utc>,
261    ) -> StorageResult<Vec<StorageEntry<T>>> {
262        match self {
263            Self::PostgreSQL(s) => s.get_modified_since(table, since).await,
264            Self::Redis(s) => s.get_modified_since(table, since).await,
265            Self::Sled(s) => s.get_modified_since(table, since).await,
266        }
267    }
268
269    async fn get_expiring_before<T: DeserializeOwned + Serialize + Send + Sync + 'static>(
270        &self,
271        table: &str,
272        before: DateTime<Utc>,
273    ) -> StorageResult<Vec<StorageEntry<T>>> {
274        match self {
275            Self::PostgreSQL(s) => s.get_expiring_before(table, before).await,
276            Self::Redis(s) => s.get_expiring_before(table, before).await,
277            Self::Sled(s) => s.get_expiring_before(table, before).await,
278        }
279    }
280
281    async fn cleanup_expired(&self, table: &str) -> StorageResult<usize> {
282        match self {
283            Self::PostgreSQL(s) => s.cleanup_expired(table).await,
284            Self::Redis(s) => s.cleanup_expired(table).await,
285            Self::Sled(s) => s.cleanup_expired(table).await,
286        }
287    }
288
289    async fn compact(&self) -> StorageResult<()> {
290        match self {
291            Self::PostgreSQL(s) => s.compact().await,
292            Self::Redis(s) => s.compact().await,
293            Self::Sled(s) => s.compact().await,
294        }
295    }
296
297    async fn flush(&self) -> StorageResult<()> {
298        match self {
299            Self::PostgreSQL(s) => s.flush().await,
300            Self::Redis(s) => s.flush().await,
301            Self::Sled(s) => s.flush().await,
302        }
303    }
304
305    async fn initialize(&mut self) -> StorageResult<()> {
306        match self {
307            Self::PostgreSQL(s) => s.initialize().await,
308            Self::Redis(s) => s.initialize().await,
309            Self::Sled(s) => s.initialize().await,
310        }
311    }
312
313    async fn shutdown(&mut self) -> StorageResult<()> {
314        match self {
315            Self::PostgreSQL(s) => s.shutdown().await,
316            Self::Redis(s) => s.shutdown().await,
317            Self::Sled(s) => s.shutdown().await,
318        }
319    }
320}
llm_optimizer_processor/storage/wrapper.rs

llm_optimizer_processor/storage/
wrapper.rs
