litellm_rs/core/cache/cloud/
gcs.rs

1//! Google Cloud Storage cache backend implementation
2//!
3//! Provides a cache backend using Google Cloud Storage for persistent, distributed caching.
4
5#[cfg(feature = "s3")]
6mod implementation {
7    use async_trait::async_trait;
8    use object_store::ObjectStore;
9    use object_store::ObjectStoreExt;
10    use object_store::gcp::GoogleCloudStorageBuilder;
11    use object_store::path::Path;
12    use serde::{Serialize, de::DeserializeOwned};
13    use std::sync::Arc;
14    use std::time::Duration;
15    use tracing::{debug, error, warn};
16
17    use crate::core::cache::cloud::{CacheMetadata, CloudCache, CloudCacheConfig};
18    use crate::core::cache::types::CacheKey;
19    use crate::utils::error::gateway_error::{GatewayError, Result};
20
21    /// GCS cache configuration
22    #[derive(Debug, Clone, Default)]
23    pub struct GcsCacheConfig {
24        /// Base cloud cache config
25        pub base: CloudCacheConfig,
26        /// Service account key path
27        pub service_account_path: Option<String>,
28        /// Service account key JSON content
29        pub service_account_key: Option<String>,
30    }
31
32    impl GcsCacheConfig {
33        /// Create a new GCS cache configuration
34        pub fn new(bucket: impl Into<String>) -> Self {
35            Self {
36                base: CloudCacheConfig::new(bucket),
37                ..Default::default()
38            }
39        }
40
41        /// Set the service account key path
42        pub fn service_account_path(mut self, path: impl Into<String>) -> Self {
43            self.service_account_path = Some(path.into());
44            self
45        }
46
47        /// Set the service account key JSON content
48        pub fn service_account_key(mut self, key: impl Into<String>) -> Self {
49            self.service_account_key = Some(key.into());
50            self
51        }
52
53        /// Set the key prefix
54        pub fn prefix(mut self, prefix: impl Into<String>) -> Self {
55            self.base.prefix = prefix.into();
56            self
57        }
58
59        /// Set the default TTL
60        pub fn default_ttl(mut self, ttl: Duration) -> Self {
61            self.base.default_ttl = ttl;
62            self
63        }
64
65        /// Create from environment variables
66        pub fn from_env() -> Option<Self> {
67            let bucket = std::env::var("GCS_CACHE_BUCKET").ok()?;
68            Some(Self {
69                base: CloudCacheConfig::new(bucket),
70                service_account_path: std::env::var("GOOGLE_APPLICATION_CREDENTIALS").ok(),
71                service_account_key: std::env::var("GCS_SERVICE_ACCOUNT_KEY").ok(),
72            })
73        }
74    }
75
76    /// Google Cloud Storage cache backend
77    pub struct GcsCache {
78        store: Arc<dyn ObjectStore>,
79        config: GcsCacheConfig,
80    }
81
82    impl GcsCache {
83        /// Create a new GCS cache
84        pub async fn new(config: GcsCacheConfig) -> Result<Self> {
85            let mut builder =
86                GoogleCloudStorageBuilder::new().with_bucket_name(&config.base.bucket);
87
88            if let Some(ref key_path) = config.service_account_path {
89                builder = builder.with_service_account_path(key_path);
90            }
91
92            if let Some(ref key) = config.service_account_key {
93                builder = builder.with_service_account_key(key);
94            }
95
96            let store = builder.build().map_err(|e| {
97                GatewayError::Internal(format!("Failed to create GCS client: {}", e))
98            })?;
99
100            Ok(Self {
101                store: Arc::new(store),
102                config,
103            })
104        }
105
106        /// Create from environment variables
107        pub async fn from_env() -> Result<Self> {
108            let config = GcsCacheConfig::from_env().ok_or_else(|| {
109                GatewayError::Config("GCS_CACHE_BUCKET not set in environment".to_string())
110            })?;
111            Self::new(config).await
112        }
113
114        /// Build the full GCS path for a cache key
115        fn build_path(&self, key: &CacheKey) -> Path {
116            Path::from(format!("{}{}", self.config.base.prefix, key.as_str()))
117        }
118
119        /// Build the metadata path for a cache key
120        fn build_metadata_path(&self, key: &CacheKey) -> Path {
121            Path::from(format!("{}{}._meta", self.config.base.prefix, key.as_str()))
122        }
123    }
124
125    #[async_trait]
126    impl CloudCache for GcsCache {
127        async fn get<T: DeserializeOwned + Send>(&self, key: &CacheKey) -> Result<Option<T>> {
128            let path = self.build_path(key);
129            let meta_path = self.build_metadata_path(key);
130            debug!(key = %path, "Reading from GCS cache");
131
132            // First check metadata for expiration
133            match self.store.get(&meta_path).await {
134                Ok(meta_result) => {
135                    let meta_bytes = meta_result.bytes().await.map_err(|e| {
136                        GatewayError::Internal(format!("Failed to read metadata: {}", e))
137                    })?;
138
139                    let metadata: CacheMetadata =
140                        serde_json::from_slice(&meta_bytes).map_err(|e| {
141                            GatewayError::Internal(format!("Failed to parse metadata: {}", e))
142                        })?;
143
144                    if metadata.is_expired() {
145                        debug!(key = %path, "Cache entry expired");
146                        let _ = self.delete(key).await;
147                        return Ok(None);
148                    }
149                }
150                Err(object_store::Error::NotFound { .. }) => {
151                    debug!(key = %path, "Cache miss - no metadata");
152                    return Ok(None);
153                }
154                Err(err) => {
155                    warn!(key = %path, error = %err, "Failed to read metadata");
156                    return Ok(None);
157                }
158            }
159
160            // Read the actual value
161            match self.store.get(&path).await {
162                Ok(result) => {
163                    let bytes = result.bytes().await.map_err(|e| {
164                        GatewayError::Internal(format!("Failed to read body: {}", e))
165                    })?;
166
167                    let value: T = serde_json::from_slice(&bytes).map_err(|e| {
168                        GatewayError::Internal(format!("Failed to deserialize: {}", e))
169                    })?;
170
171                    debug!(key = %path, "GCS cache hit");
172                    Ok(Some(value))
173                }
174                Err(object_store::Error::NotFound { .. }) => {
175                    debug!(key = %path, "Cache miss");
176                    Ok(None)
177                }
178                Err(err) => {
179                    error!(key = %path, error = %err, "Failed to read from GCS");
180                    Err(GatewayError::Internal(format!("GCS read error: {}", err)))
181                }
182            }
183        }
184
185        async fn set<T: Serialize + Send + Sync>(
186            &self,
187            key: &CacheKey,
188            value: &T,
189            ttl: Duration,
190        ) -> Result<()> {
191            let path = self.build_path(key);
192            let meta_path = self.build_metadata_path(key);
193            debug!(key = %path, ttl_secs = ttl.as_secs(), "Writing to GCS cache");
194
195            // Serialize the value
196            let bytes = serde_json::to_vec(value)
197                .map_err(|e| GatewayError::Internal(format!("Failed to serialize: {}", e)))?;
198
199            // Create metadata
200            let metadata = CacheMetadata::new(ttl, bytes.len(), false);
201            let meta_bytes = serde_json::to_vec(&metadata).map_err(|e| {
202                GatewayError::Internal(format!("Failed to serialize metadata: {}", e))
203            })?;
204
205            // Write the value
206            self.store
207                .put(&path, bytes.into())
208                .await
209                .map_err(|e| GatewayError::Internal(format!("Failed to write to GCS: {}", e)))?;
210
211            // Write metadata
212            self.store
213                .put(&meta_path, meta_bytes.into())
214                .await
215                .map_err(|e| GatewayError::Internal(format!("Failed to write metadata: {}", e)))?;
216
217            debug!(key = %path, "GCS cache write successful");
218            Ok(())
219        }
220
221        async fn delete(&self, key: &CacheKey) -> Result<bool> {
222            let path = self.build_path(key);
223            let meta_path = self.build_metadata_path(key);
224            debug!(key = %path, "Deleting from GCS cache");
225
226            let value_result = self.store.delete(&path).await;
227            let meta_result = self.store.delete(&meta_path).await;
228
229            // Check if at least one deletion succeeded
230            let deleted = value_result.is_ok() || meta_result.is_ok();
231
232            if deleted {
233                debug!(key = %path, "GCS cache delete successful");
234            }
235
236            Ok(deleted)
237        }
238
239        async fn exists(&self, key: &CacheKey) -> Result<bool> {
240            let path = self.build_path(key);
241
242            match self.store.head(&path).await {
243                Ok(_) => Ok(true),
244                Err(object_store::Error::NotFound { .. }) => Ok(false),
245                Err(err) => Err(GatewayError::Internal(format!("GCS head error: {}", err))),
246            }
247        }
248
249        async fn list_keys(&self, prefix: &str) -> Result<Vec<String>> {
250            use futures::StreamExt;
251
252            let full_prefix = Path::from(format!("{}{}", self.config.base.prefix, prefix));
253
254            let mut keys = Vec::new();
255            let mut stream = self.store.list(Some(&full_prefix));
256
257            while let Some(result) = stream.next().await {
258                match result {
259                    Ok(meta) => {
260                        let key = meta.location.to_string();
261                        // Remove prefix and filter out metadata keys
262                        if !key.ends_with("._meta") {
263                            if let Some(stripped) = key.strip_prefix(&self.config.base.prefix) {
264                                keys.push(stripped.to_string());
265                            } else {
266                                keys.push(key);
267                            }
268                        }
269                    }
270                    Err(err) => {
271                        warn!(error = %err, "Error listing GCS objects");
272                    }
273                }
274            }
275
276            Ok(keys)
277        }
278
279        async fn clear(&self) -> Result<()> {
280            warn!("Clearing all GCS cache entries");
281
282            let keys = self.list_keys("").await?;
283
284            for key in keys {
285                let cache_key = CacheKey::new(key);
286                let _ = self.delete(&cache_key).await;
287            }
288
289            Ok(())
290        }
291
292        fn name(&self) -> &'static str {
293            "gcs"
294        }
295    }
296
297    impl std::fmt::Debug for GcsCache {
298        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
299            f.debug_struct("GcsCache")
300                .field("bucket", &self.config.base.bucket)
301                .field("prefix", &self.config.base.prefix)
302                .finish()
303        }
304    }
305}
306
307#[cfg(feature = "s3")]
308pub use implementation::*;
309
310// Stub implementation when feature is disabled
311#[cfg(not(feature = "s3"))]
312mod stub {
313    use crate::utils::error::gateway_error::{GatewayError, Result};
314
315    /// GCS cache configuration (stub)
316    #[derive(Debug, Clone, Default)]
317    pub struct GcsCacheConfig {
318        pub bucket: String,
319    }
320
321    impl GcsCacheConfig {
322        pub fn new(bucket: impl Into<String>) -> Self {
323            Self {
324                bucket: bucket.into(),
325            }
326        }
327
328        pub fn from_env() -> Option<Self> {
329            None
330        }
331    }
332
333    /// GCS cache (stub)
334    #[derive(Debug)]
335    pub struct GcsCache;
336
337    impl GcsCache {
338        pub async fn new(_config: GcsCacheConfig) -> Result<Self> {
339            Err(GatewayError::Config(
340                "GCS cache support not enabled. Enable the 's3' feature.".to_string(),
341            ))
342        }
343
344        pub async fn from_env() -> Result<Self> {
345            Err(GatewayError::Config(
346                "GCS cache support not enabled. Enable the 's3' feature.".to_string(),
347            ))
348        }
349    }
350}
351
352#[cfg(not(feature = "s3"))]
353pub use stub::*;
354
355#[cfg(all(test, feature = "s3"))]
356mod tests {
357    use super::*;
358
359    #[test]
360    fn test_gcs_cache_config_builder() {
361        let config = GcsCacheConfig::new("my-bucket")
362            .service_account_path("/path/to/key.json")
363            .prefix("cache/");
364
365        assert_eq!(config.base.bucket, "my-bucket");
366        assert_eq!(
367            config.service_account_path,
368            Some("/path/to/key.json".to_string())
369        );
370        assert_eq!(config.base.prefix, "cache/");
371    }
372}
litellm_rs/core/cache/cloud/gcs.rs

litellm_rs/core/cache/cloud/
gcs.rs
