sklears_utils/
cloud_storage.rs

1//! Cloud storage utilities for machine learning data processing
2//!
3//! This module provides unified interfaces for working with cloud storage services
4//! including AWS S3, Google Cloud Storage, and Azure Blob Storage.
5
6use crate::{UtilsError, UtilsResult};
7use std::collections::HashMap;
8use std::fmt;
9
10/// Cloud storage configuration
11#[derive(Debug, Clone)]
12pub struct CloudStorageConfig {
13    pub provider: CloudProvider,
14    pub endpoint: Option<String>,
15    pub region: Option<String>,
16    pub access_key: Option<String>,
17    pub secret_key: Option<String>,
18    pub bucket: String,
19    pub timeout_seconds: Option<u64>,
20    pub use_ssl: bool,
21    pub custom_headers: HashMap<String, String>,
22}
23
24/// Supported cloud storage providers
25#[derive(Debug, Clone, PartialEq)]
26pub enum CloudProvider {
27    AWS,
28    GoogleCloud,
29    Azure,
30    MinIO,
31    Custom(String),
32}
33
34impl fmt::Display for CloudProvider {
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        match self {
37            CloudProvider::AWS => write!(f, "aws"),
38            CloudProvider::GoogleCloud => write!(f, "gcp"),
39            CloudProvider::Azure => write!(f, "azure"),
40            CloudProvider::MinIO => write!(f, "minio"),
41            CloudProvider::Custom(name) => write!(f, "{name}"),
42        }
43    }
44}
45
46impl Default for CloudStorageConfig {
47    fn default() -> Self {
48        Self {
49            provider: CloudProvider::AWS,
50            endpoint: None,
51            region: Some("us-east-1".to_string()),
52            access_key: None,
53            secret_key: None,
54            bucket: String::new(),
55            timeout_seconds: Some(30),
56            use_ssl: true,
57            custom_headers: HashMap::new(),
58        }
59    }
60}
61
62/// Cloud storage client trait
63pub trait CloudStorageClient {
64    /// Upload data to cloud storage
65    fn upload(&self, key: &str, data: &[u8]) -> UtilsResult<String>;
66
67    /// Download data from cloud storage
68    fn download(&self, key: &str) -> UtilsResult<Vec<u8>>;
69
70    /// Delete object from cloud storage
71    fn delete(&self, key: &str) -> UtilsResult<()>;
72
73    /// List objects with prefix
74    fn list_objects(&self, prefix: &str) -> UtilsResult<Vec<String>>;
75
76    /// Check if object exists
77    fn exists(&self, key: &str) -> UtilsResult<bool>;
78
79    /// Get object metadata
80    fn get_metadata(&self, key: &str) -> UtilsResult<ObjectMetadata>;
81
82    /// Upload file from local path
83    fn upload_file(&self, key: &str, local_path: &str) -> UtilsResult<String>;
84
85    /// Download file to local path
86    fn download_file(&self, key: &str, local_path: &str) -> UtilsResult<()>;
87}
88
89/// Object metadata
90#[derive(Debug, Clone)]
91pub struct ObjectMetadata {
92    pub size: u64,
93    pub etag: Option<String>,
94    pub content_type: Option<String>,
95    pub last_modified: Option<String>,
96    pub custom_metadata: HashMap<String, String>,
97}
98
99/// Mock cloud storage client for testing
100pub struct MockCloudStorageClient {
101    storage: std::sync::Arc<std::sync::Mutex<HashMap<String, Vec<u8>>>>,
102    metadata: std::sync::Arc<std::sync::Mutex<HashMap<String, ObjectMetadata>>>,
103}
104
105impl Default for MockCloudStorageClient {
106    fn default() -> Self {
107        Self::new()
108    }
109}
110
111impl MockCloudStorageClient {
112    pub fn new() -> Self {
113        Self {
114            storage: std::sync::Arc::new(std::sync::Mutex::new(HashMap::new())),
115            metadata: std::sync::Arc::new(std::sync::Mutex::new(HashMap::new())),
116        }
117    }
118}
119
120impl CloudStorageClient for MockCloudStorageClient {
121    fn upload(&self, key: &str, data: &[u8]) -> UtilsResult<String> {
122        let mut storage = self.storage.lock().expect("operation should succeed");
123        let mut metadata = self.metadata.lock().expect("operation should succeed");
124
125        storage.insert(key.to_string(), data.to_vec());
126        metadata.insert(
127            key.to_string(),
128            ObjectMetadata {
129                size: data.len() as u64,
130                etag: Some(format!("mock-etag-{key}")),
131                content_type: Some("application/octet-stream".to_string()),
132                last_modified: Some(chrono::Utc::now().to_rfc3339()),
133                custom_metadata: HashMap::new(),
134            },
135        );
136
137        Ok(format!("mock://bucket/{key}"))
138    }
139
140    fn download(&self, key: &str) -> UtilsResult<Vec<u8>> {
141        let storage = self.storage.lock().expect("operation should succeed");
142        storage
143            .get(key)
144            .cloned()
145            .ok_or_else(|| UtilsError::InvalidParameter(format!("Object not found: {key}")))
146    }
147
148    fn delete(&self, key: &str) -> UtilsResult<()> {
149        let mut storage = self.storage.lock().expect("operation should succeed");
150        let mut metadata = self.metadata.lock().expect("operation should succeed");
151
152        storage.remove(key);
153        metadata.remove(key);
154        Ok(())
155    }
156
157    fn list_objects(&self, prefix: &str) -> UtilsResult<Vec<String>> {
158        let storage = self.storage.lock().expect("operation should succeed");
159        let objects: Vec<String> = storage
160            .keys()
161            .filter(|key| key.starts_with(prefix))
162            .cloned()
163            .collect();
164        Ok(objects)
165    }
166
167    fn exists(&self, key: &str) -> UtilsResult<bool> {
168        let storage = self.storage.lock().expect("operation should succeed");
169        Ok(storage.contains_key(key))
170    }
171
172    fn get_metadata(&self, key: &str) -> UtilsResult<ObjectMetadata> {
173        let metadata = self.metadata.lock().expect("operation should succeed");
174        metadata
175            .get(key)
176            .cloned()
177            .ok_or_else(|| UtilsError::InvalidParameter(format!("Object not found: {key}")))
178    }
179
180    fn upload_file(&self, key: &str, local_path: &str) -> UtilsResult<String> {
181        let data = std::fs::read(local_path)
182            .map_err(|e| UtilsError::InvalidParameter(format!("Failed to read file: {e}")))?;
183        self.upload(key, &data)
184    }
185
186    fn download_file(&self, key: &str, local_path: &str) -> UtilsResult<()> {
187        let data = self.download(key)?;
188        std::fs::write(local_path, data)
189            .map_err(|e| UtilsError::InvalidParameter(format!("Failed to write file: {e}")))?;
190        Ok(())
191    }
192}
193
194/// Cloud storage factory
195pub struct CloudStorageFactory;
196
197impl CloudStorageFactory {
198    /// Create a cloud storage client based on configuration
199    pub fn create_client(config: &CloudStorageConfig) -> UtilsResult<Box<dyn CloudStorageClient>> {
200        match config.provider {
201            CloudProvider::AWS => {
202                // In a real implementation, this would create an AWS S3 client
203                // For now, we'll use the mock client
204                Ok(Box::new(MockCloudStorageClient::new()))
205            }
206            CloudProvider::GoogleCloud => {
207                // In a real implementation, this would create a GCS client
208                Ok(Box::new(MockCloudStorageClient::new()))
209            }
210            CloudProvider::Azure => {
211                // In a real implementation, this would create an Azure Blob client
212                Ok(Box::new(MockCloudStorageClient::new()))
213            }
214            CloudProvider::MinIO => {
215                // In a real implementation, this would create a MinIO client
216                Ok(Box::new(MockCloudStorageClient::new()))
217            }
218            CloudProvider::Custom(_) => {
219                // For custom providers, use mock client
220                Ok(Box::new(MockCloudStorageClient::new()))
221            }
222        }
223    }
224}
225
226/// Cloud storage utilities for ML data processing
227pub struct CloudStorageUtils;
228
229impl CloudStorageUtils {
230    /// Upload ML dataset to cloud storage
231    pub fn upload_dataset(
232        client: &dyn CloudStorageClient,
233        dataset_path: &str,
234        key_prefix: &str,
235    ) -> UtilsResult<Vec<String>> {
236        let mut uploaded_keys = Vec::new();
237
238        // Read dataset directory
239        let entries = std::fs::read_dir(dataset_path)
240            .map_err(|e| UtilsError::InvalidParameter(format!("Failed to read directory: {e}")))?;
241
242        for entry in entries {
243            let entry = entry
244                .map_err(|e| UtilsError::InvalidParameter(format!("Failed to read entry: {e}")))?;
245            let path = entry.path();
246
247            if path.is_file() {
248                let filename = path
249                    .file_name()
250                    .expect("operation should succeed")
251                    .to_str()
252                    .expect("operation should succeed");
253                let key = format!("{key_prefix}/{filename}");
254                let local_path = path.to_str().expect("operation should succeed");
255
256                client.upload_file(&key, local_path)?;
257                uploaded_keys.push(key);
258            }
259        }
260
261        Ok(uploaded_keys)
262    }
263
264    /// Download ML dataset from cloud storage
265    pub fn download_dataset(
266        client: &dyn CloudStorageClient,
267        key_prefix: &str,
268        local_path: &str,
269    ) -> UtilsResult<Vec<String>> {
270        let objects = client.list_objects(key_prefix)?;
271        let mut downloaded_files = Vec::new();
272
273        // Create local directory if it doesn't exist
274        std::fs::create_dir_all(local_path).map_err(|e| {
275            UtilsError::InvalidParameter(format!("Failed to create directory: {e}"))
276        })?;
277
278        for object_key in objects {
279            let filename = object_key.split('/').next_back().unwrap_or(&object_key);
280            let local_file_path = format!("{local_path}/{filename}");
281
282            client.download_file(&object_key, &local_file_path)?;
283            downloaded_files.push(local_file_path);
284        }
285
286        Ok(downloaded_files)
287    }
288
289    /// Sync local dataset with cloud storage
290    pub fn sync_dataset(
291        client: &dyn CloudStorageClient,
292        local_path: &str,
293        key_prefix: &str,
294        sync_mode: SyncMode,
295    ) -> UtilsResult<SyncResult> {
296        let mut result = SyncResult::default();
297
298        match sync_mode {
299            SyncMode::Upload => {
300                let uploaded = Self::upload_dataset(client, local_path, key_prefix)?;
301                result.uploaded = uploaded;
302            }
303            SyncMode::Download => {
304                let downloaded = Self::download_dataset(client, key_prefix, local_path)?;
305                result.downloaded = downloaded;
306            }
307            SyncMode::Bidirectional => {
308                // Simple bidirectional sync: upload first, then download
309                let uploaded = Self::upload_dataset(client, local_path, key_prefix)?;
310                let downloaded = Self::download_dataset(client, key_prefix, local_path)?;
311                result.uploaded = uploaded;
312                result.downloaded = downloaded;
313            }
314        }
315
316        Ok(result)
317    }
318
319    /// Batch upload multiple files with metadata
320    pub fn batch_upload(
321        client: &dyn CloudStorageClient,
322        files: &[(String, String)], // (local_path, key)
323    ) -> UtilsResult<Vec<String>> {
324        let mut uploaded_keys = Vec::new();
325
326        for (local_path, key) in files {
327            let result = client.upload_file(key, local_path)?;
328            uploaded_keys.push(result);
329        }
330
331        Ok(uploaded_keys)
332    }
333
334    /// Calculate storage metrics for ML datasets
335    pub fn calculate_storage_metrics(
336        client: &dyn CloudStorageClient,
337        key_prefix: &str,
338    ) -> UtilsResult<StorageMetrics> {
339        let objects = client.list_objects(key_prefix)?;
340        let mut total_size = 0;
341        let mut total_objects = 0;
342        let mut file_types = HashMap::new();
343
344        for object_key in objects {
345            if let Ok(metadata) = client.get_metadata(&object_key) {
346                total_size += metadata.size;
347                total_objects += 1;
348
349                // Extract file extension
350                if let Some(ext) = object_key.split('.').next_back() {
351                    *file_types.entry(ext.to_string()).or_insert(0) += 1;
352                }
353            }
354        }
355
356        Ok(StorageMetrics {
357            total_size_bytes: total_size,
358            total_objects,
359            file_types,
360            average_file_size: total_size.checked_div(total_objects).unwrap_or(0),
361        })
362    }
363}
364
365/// Sync mode for dataset synchronization
366#[derive(Debug, Clone)]
367pub enum SyncMode {
368    Upload,
369    Download,
370    Bidirectional,
371}
372
373/// Sync result
374#[derive(Debug, Clone, Default)]
375pub struct SyncResult {
376    pub uploaded: Vec<String>,
377    pub downloaded: Vec<String>,
378    pub errors: Vec<String>,
379}
380
381/// Storage metrics
382#[derive(Debug, Clone)]
383pub struct StorageMetrics {
384    pub total_size_bytes: u64,
385    pub total_objects: u64,
386    pub file_types: HashMap<String, usize>,
387    pub average_file_size: u64,
388}
389
390impl fmt::Display for StorageMetrics {
391    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
392        writeln!(f, "Storage Metrics:")?;
393        writeln!(
394            f,
395            "  Total Size: {:.2} MB",
396            self.total_size_bytes as f64 / 1024.0 / 1024.0
397        )?;
398        writeln!(f, "  Total Objects: {}", self.total_objects)?;
399        writeln!(
400            f,
401            "  Average File Size: {:.2} KB",
402            self.average_file_size as f64 / 1024.0
403        )?;
404        writeln!(f, "  File Types:")?;
405        for (ext, count) in &self.file_types {
406            writeln!(f, "    .{ext}: {count}")?;
407        }
408        Ok(())
409    }
410}
411
412#[allow(non_snake_case)]
413#[cfg(test)]
414mod tests {
415    use super::*;
416    use std::fs;
417
418    #[test]
419    fn test_cloud_storage_config() {
420        let config = CloudStorageConfig {
421            provider: CloudProvider::AWS,
422            bucket: "test-bucket".to_string(),
423            ..Default::default()
424        };
425
426        assert_eq!(config.provider, CloudProvider::AWS);
427        assert_eq!(config.bucket, "test-bucket");
428        assert_eq!(config.region, Some("us-east-1".to_string()));
429        assert!(config.use_ssl);
430    }
431
432    #[test]
433    fn test_cloud_provider_display() {
434        assert_eq!(CloudProvider::AWS.to_string(), "aws");
435        assert_eq!(CloudProvider::GoogleCloud.to_string(), "gcp");
436        assert_eq!(CloudProvider::Azure.to_string(), "azure");
437        assert_eq!(CloudProvider::MinIO.to_string(), "minio");
438        assert_eq!(
439            CloudProvider::Custom("test".to_string()).to_string(),
440            "test"
441        );
442    }
443
444    #[test]
445    fn test_mock_client_upload_download() {
446        let client = MockCloudStorageClient::new();
447        let test_data = b"hello world";
448
449        // Test upload
450        let url = client
451            .upload("test-key", test_data)
452            .expect("operation should succeed");
453        assert_eq!(url, "mock://bucket/test-key");
454
455        // Test download
456        let downloaded = client
457            .download("test-key")
458            .expect("operation should succeed");
459        assert_eq!(downloaded, test_data);
460
461        // Test exists
462        assert!(client.exists("test-key").expect("operation should succeed"));
463        assert!(!client
464            .exists("nonexistent-key")
465            .expect("operation should succeed"));
466    }
467
468    #[test]
469    fn test_mock_client_metadata() {
470        let client = MockCloudStorageClient::new();
471        let test_data = b"hello world";
472
473        client
474            .upload("test-key", test_data)
475            .expect("operation should succeed");
476
477        let metadata = client
478            .get_metadata("test-key")
479            .expect("operation should succeed");
480        assert_eq!(metadata.size, test_data.len() as u64);
481        assert_eq!(metadata.etag, Some("mock-etag-test-key".to_string()));
482        assert_eq!(
483            metadata.content_type,
484            Some("application/octet-stream".to_string())
485        );
486    }
487
488    #[test]
489    fn test_mock_client_list_objects() {
490        let client = MockCloudStorageClient::new();
491
492        client
493            .upload("data/file1.txt", b"content1")
494            .expect("operation should succeed");
495        client
496            .upload("data/file2.txt", b"content2")
497            .expect("operation should succeed");
498        client
499            .upload("other/file3.txt", b"content3")
500            .expect("operation should succeed");
501
502        let objects = client
503            .list_objects("data/")
504            .expect("operation should succeed");
505        assert_eq!(objects.len(), 2);
506        assert!(objects.contains(&"data/file1.txt".to_string()));
507        assert!(objects.contains(&"data/file2.txt".to_string()));
508    }
509
510    #[test]
511    fn test_mock_client_delete() {
512        let client = MockCloudStorageClient::new();
513
514        client
515            .upload("test-key", b"hello")
516            .expect("operation should succeed");
517        assert!(client.exists("test-key").expect("operation should succeed"));
518
519        client.delete("test-key").expect("operation should succeed");
520        assert!(!client.exists("test-key").expect("operation should succeed"));
521    }
522
523    #[test]
524    fn test_cloud_storage_factory() {
525        let config = CloudStorageConfig {
526            provider: CloudProvider::AWS,
527            bucket: "test-bucket".to_string(),
528            ..Default::default()
529        };
530
531        let client = CloudStorageFactory::create_client(&config).expect("operation should succeed");
532
533        // Test that we can use the client
534        client
535            .upload("test", b"data")
536            .expect("operation should succeed");
537        let downloaded = client.download("test").expect("operation should succeed");
538        assert_eq!(downloaded, b"data");
539    }
540
541    #[test]
542    fn test_storage_metrics_display() {
543        let mut file_types = HashMap::new();
544        file_types.insert("txt".to_string(), 5);
545        file_types.insert("csv".to_string(), 3);
546
547        let metrics = StorageMetrics {
548            total_size_bytes: 1_048_576, // 1 MB
549            total_objects: 8,
550            file_types,
551            average_file_size: 131_072, // 128 KB
552        };
553
554        let display = metrics.to_string();
555        assert!(display.contains("Total Size: 1.00 MB"));
556        assert!(display.contains("Total Objects: 8"));
557        assert!(display.contains("Average File Size: 128.00 KB"));
558        assert!(display.contains(".txt: 5"));
559        assert!(display.contains(".csv: 3"));
560    }
561
562    #[test]
563    fn test_sync_result_default() {
564        let result = SyncResult::default();
565        assert!(result.uploaded.is_empty());
566        assert!(result.downloaded.is_empty());
567        assert!(result.errors.is_empty());
568    }
569
570    #[test]
571    fn test_file_upload_download() {
572        let client = MockCloudStorageClient::new();
573        let temp_dir = tempfile::tempdir().expect("operation should succeed");
574        let file_path = temp_dir.path().join("test.txt");
575
576        // Create test file
577        fs::write(&file_path, b"test content").expect("operation should succeed");
578
579        // Upload file
580        let url = client
581            .upload_file(
582                "test.txt",
583                file_path.to_str().expect("operation should succeed"),
584            )
585            .expect("operation should succeed");
586        assert_eq!(url, "mock://bucket/test.txt");
587
588        // Download file
589        let download_path = temp_dir.path().join("downloaded.txt");
590        client
591            .download_file(
592                "test.txt",
593                download_path.to_str().expect("operation should succeed"),
594            )
595            .expect("operation should succeed");
596
597        // Verify content
598        let downloaded_content = fs::read(&download_path).expect("operation should succeed");
599        assert_eq!(downloaded_content, b"test content");
600    }
601
602    #[test]
603    fn test_calculate_storage_metrics() {
604        let client = MockCloudStorageClient::new();
605
606        // Upload test files
607        client
608            .upload("data/file1.txt", b"hello")
609            .expect("operation should succeed");
610        client
611            .upload("data/file2.csv", b"world")
612            .expect("operation should succeed");
613        client
614            .upload("data/file3.txt", b"test")
615            .expect("operation should succeed");
616
617        let metrics = CloudStorageUtils::calculate_storage_metrics(&client, "data/")
618            .expect("operation should succeed");
619
620        assert_eq!(metrics.total_objects, 3);
621        assert_eq!(metrics.total_size_bytes, 14); // 5 + 5 + 4
622        assert_eq!(metrics.file_types.get("txt"), Some(&2));
623        assert_eq!(metrics.file_types.get("csv"), Some(&1));
624    }
625
626    #[test]
627    fn test_batch_upload() {
628        let client = MockCloudStorageClient::new();
629        let temp_dir = tempfile::tempdir().expect("operation should succeed");
630
631        // Create test files
632        let file1_path = temp_dir.path().join("file1.txt");
633        let file2_path = temp_dir.path().join("file2.txt");
634        fs::write(&file1_path, b"content1").expect("operation should succeed");
635        fs::write(&file2_path, b"content2").expect("operation should succeed");
636
637        let files = vec![
638            (
639                file1_path
640                    .to_str()
641                    .expect("operation should succeed")
642                    .to_string(),
643                "batch/file1.txt".to_string(),
644            ),
645            (
646                file2_path
647                    .to_str()
648                    .expect("operation should succeed")
649                    .to_string(),
650                "batch/file2.txt".to_string(),
651            ),
652        ];
653
654        let results =
655            CloudStorageUtils::batch_upload(&client, &files).expect("operation should succeed");
656
657        assert_eq!(results.len(), 2);
658        assert_eq!(results[0], "mock://bucket/batch/file1.txt");
659        assert_eq!(results[1], "mock://bucket/batch/file2.txt");
660
661        // Verify uploads
662        let content1 = client
663            .download("batch/file1.txt")
664            .expect("operation should succeed");
665        let content2 = client
666            .download("batch/file2.txt")
667            .expect("operation should succeed");
668        assert_eq!(content1, b"content1");
669        assert_eq!(content2, b"content2");
670    }
671}
sklears_utils/cloud_storage.rs

sklears_utils/
cloud_storage.rs
