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: if total_objects > 0 {
361                total_size / total_objects
362            } else {
363                0
364            },
365        })
366    }
367}
368
369/// Sync mode for dataset synchronization
370#[derive(Debug, Clone)]
371pub enum SyncMode {
372    Upload,
373    Download,
374    Bidirectional,
375}
376
377/// Sync result
378#[derive(Debug, Clone, Default)]
379pub struct SyncResult {
380    pub uploaded: Vec<String>,
381    pub downloaded: Vec<String>,
382    pub errors: Vec<String>,
383}
384
385/// Storage metrics
386#[derive(Debug, Clone)]
387pub struct StorageMetrics {
388    pub total_size_bytes: u64,
389    pub total_objects: u64,
390    pub file_types: HashMap<String, usize>,
391    pub average_file_size: u64,
392}
393
394impl fmt::Display for StorageMetrics {
395    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
396        writeln!(f, "Storage Metrics:")?;
397        writeln!(
398            f,
399            "  Total Size: {:.2} MB",
400            self.total_size_bytes as f64 / 1024.0 / 1024.0
401        )?;
402        writeln!(f, "  Total Objects: {}", self.total_objects)?;
403        writeln!(
404            f,
405            "  Average File Size: {:.2} KB",
406            self.average_file_size as f64 / 1024.0
407        )?;
408        writeln!(f, "  File Types:")?;
409        for (ext, count) in &self.file_types {
410            writeln!(f, "    .{ext}: {count}")?;
411        }
412        Ok(())
413    }
414}
415
416#[allow(non_snake_case)]
417#[cfg(test)]
418mod tests {
419    use super::*;
420    use std::fs;
421
422    #[test]
423    fn test_cloud_storage_config() {
424        let config = CloudStorageConfig {
425            provider: CloudProvider::AWS,
426            bucket: "test-bucket".to_string(),
427            ..Default::default()
428        };
429
430        assert_eq!(config.provider, CloudProvider::AWS);
431        assert_eq!(config.bucket, "test-bucket");
432        assert_eq!(config.region, Some("us-east-1".to_string()));
433        assert!(config.use_ssl);
434    }
435
436    #[test]
437    fn test_cloud_provider_display() {
438        assert_eq!(CloudProvider::AWS.to_string(), "aws");
439        assert_eq!(CloudProvider::GoogleCloud.to_string(), "gcp");
440        assert_eq!(CloudProvider::Azure.to_string(), "azure");
441        assert_eq!(CloudProvider::MinIO.to_string(), "minio");
442        assert_eq!(
443            CloudProvider::Custom("test".to_string()).to_string(),
444            "test"
445        );
446    }
447
448    #[test]
449    fn test_mock_client_upload_download() {
450        let client = MockCloudStorageClient::new();
451        let test_data = b"hello world";
452
453        // Test upload
454        let url = client
455            .upload("test-key", test_data)
456            .expect("operation should succeed");
457        assert_eq!(url, "mock://bucket/test-key");
458
459        // Test download
460        let downloaded = client
461            .download("test-key")
462            .expect("operation should succeed");
463        assert_eq!(downloaded, test_data);
464
465        // Test exists
466        assert!(client.exists("test-key").expect("operation should succeed"));
467        assert!(!client
468            .exists("nonexistent-key")
469            .expect("operation should succeed"));
470    }
471
472    #[test]
473    fn test_mock_client_metadata() {
474        let client = MockCloudStorageClient::new();
475        let test_data = b"hello world";
476
477        client
478            .upload("test-key", test_data)
479            .expect("operation should succeed");
480
481        let metadata = client
482            .get_metadata("test-key")
483            .expect("operation should succeed");
484        assert_eq!(metadata.size, test_data.len() as u64);
485        assert_eq!(metadata.etag, Some("mock-etag-test-key".to_string()));
486        assert_eq!(
487            metadata.content_type,
488            Some("application/octet-stream".to_string())
489        );
490    }
491
492    #[test]
493    fn test_mock_client_list_objects() {
494        let client = MockCloudStorageClient::new();
495
496        client
497            .upload("data/file1.txt", b"content1")
498            .expect("operation should succeed");
499        client
500            .upload("data/file2.txt", b"content2")
501            .expect("operation should succeed");
502        client
503            .upload("other/file3.txt", b"content3")
504            .expect("operation should succeed");
505
506        let objects = client
507            .list_objects("data/")
508            .expect("operation should succeed");
509        assert_eq!(objects.len(), 2);
510        assert!(objects.contains(&"data/file1.txt".to_string()));
511        assert!(objects.contains(&"data/file2.txt".to_string()));
512    }
513
514    #[test]
515    fn test_mock_client_delete() {
516        let client = MockCloudStorageClient::new();
517
518        client
519            .upload("test-key", b"hello")
520            .expect("operation should succeed");
521        assert!(client.exists("test-key").expect("operation should succeed"));
522
523        client.delete("test-key").expect("operation should succeed");
524        assert!(!client.exists("test-key").expect("operation should succeed"));
525    }
526
527    #[test]
528    fn test_cloud_storage_factory() {
529        let config = CloudStorageConfig {
530            provider: CloudProvider::AWS,
531            bucket: "test-bucket".to_string(),
532            ..Default::default()
533        };
534
535        let client = CloudStorageFactory::create_client(&config).expect("operation should succeed");
536
537        // Test that we can use the client
538        client
539            .upload("test", b"data")
540            .expect("operation should succeed");
541        let downloaded = client.download("test").expect("operation should succeed");
542        assert_eq!(downloaded, b"data");
543    }
544
545    #[test]
546    fn test_storage_metrics_display() {
547        let mut file_types = HashMap::new();
548        file_types.insert("txt".to_string(), 5);
549        file_types.insert("csv".to_string(), 3);
550
551        let metrics = StorageMetrics {
552            total_size_bytes: 1_048_576, // 1 MB
553            total_objects: 8,
554            file_types,
555            average_file_size: 131_072, // 128 KB
556        };
557
558        let display = metrics.to_string();
559        assert!(display.contains("Total Size: 1.00 MB"));
560        assert!(display.contains("Total Objects: 8"));
561        assert!(display.contains("Average File Size: 128.00 KB"));
562        assert!(display.contains(".txt: 5"));
563        assert!(display.contains(".csv: 3"));
564    }
565
566    #[test]
567    fn test_sync_result_default() {
568        let result = SyncResult::default();
569        assert!(result.uploaded.is_empty());
570        assert!(result.downloaded.is_empty());
571        assert!(result.errors.is_empty());
572    }
573
574    #[test]
575    fn test_file_upload_download() {
576        let client = MockCloudStorageClient::new();
577        let temp_dir = tempfile::tempdir().expect("operation should succeed");
578        let file_path = temp_dir.path().join("test.txt");
579
580        // Create test file
581        fs::write(&file_path, b"test content").expect("operation should succeed");
582
583        // Upload file
584        let url = client
585            .upload_file(
586                "test.txt",
587                file_path.to_str().expect("operation should succeed"),
588            )
589            .expect("operation should succeed");
590        assert_eq!(url, "mock://bucket/test.txt");
591
592        // Download file
593        let download_path = temp_dir.path().join("downloaded.txt");
594        client
595            .download_file(
596                "test.txt",
597                download_path.to_str().expect("operation should succeed"),
598            )
599            .expect("operation should succeed");
600
601        // Verify content
602        let downloaded_content = fs::read(&download_path).expect("operation should succeed");
603        assert_eq!(downloaded_content, b"test content");
604    }
605
606    #[test]
607    fn test_calculate_storage_metrics() {
608        let client = MockCloudStorageClient::new();
609
610        // Upload test files
611        client
612            .upload("data/file1.txt", b"hello")
613            .expect("operation should succeed");
614        client
615            .upload("data/file2.csv", b"world")
616            .expect("operation should succeed");
617        client
618            .upload("data/file3.txt", b"test")
619            .expect("operation should succeed");
620
621        let metrics = CloudStorageUtils::calculate_storage_metrics(&client, "data/")
622            .expect("operation should succeed");
623
624        assert_eq!(metrics.total_objects, 3);
625        assert_eq!(metrics.total_size_bytes, 14); // 5 + 5 + 4
626        assert_eq!(metrics.file_types.get("txt"), Some(&2));
627        assert_eq!(metrics.file_types.get("csv"), Some(&1));
628    }
629
630    #[test]
631    fn test_batch_upload() {
632        let client = MockCloudStorageClient::new();
633        let temp_dir = tempfile::tempdir().expect("operation should succeed");
634
635        // Create test files
636        let file1_path = temp_dir.path().join("file1.txt");
637        let file2_path = temp_dir.path().join("file2.txt");
638        fs::write(&file1_path, b"content1").expect("operation should succeed");
639        fs::write(&file2_path, b"content2").expect("operation should succeed");
640
641        let files = vec![
642            (
643                file1_path
644                    .to_str()
645                    .expect("operation should succeed")
646                    .to_string(),
647                "batch/file1.txt".to_string(),
648            ),
649            (
650                file2_path
651                    .to_str()
652                    .expect("operation should succeed")
653                    .to_string(),
654                "batch/file2.txt".to_string(),
655            ),
656        ];
657
658        let results =
659            CloudStorageUtils::batch_upload(&client, &files).expect("operation should succeed");
660
661        assert_eq!(results.len(), 2);
662        assert_eq!(results[0], "mock://bucket/batch/file1.txt");
663        assert_eq!(results[1], "mock://bucket/batch/file2.txt");
664
665        // Verify uploads
666        let content1 = client
667            .download("batch/file1.txt")
668            .expect("operation should succeed");
669        let content2 = client
670            .download("batch/file2.txt")
671            .expect("operation should succeed");
672        assert_eq!(content1, b"content1");
673        assert_eq!(content2, b"content2");
674    }
675}
sklears_utils/cloud_storage.rs

sklears_utils/
cloud_storage.rs
