ecad_processor/archive/
temp_manager.rs

1use crate::error::{ProcessingError, Result};
2use std::collections::HashMap;
3use std::fs::File;
4use std::io::{BufWriter, Write};
5use std::path::{Path, PathBuf};
6use tempfile::TempDir;
7use zip::ZipArchive;
8
9pub struct TempFileManager {
10    temp_dir: TempDir,
11    extracted_files: HashMap<String, PathBuf>,
12}
13
14impl TempFileManager {
15    pub fn new() -> Result<Self> {
16        let temp_dir = TempDir::new().map_err(|e| {
17            ProcessingError::Io(std::io::Error::new(
18                std::io::ErrorKind::Other,
19                format!("Failed to create temporary directory: {}", e),
20            ))
21        })?;
22
23        Ok(Self {
24            temp_dir,
25            extracted_files: HashMap::new(),
26        })
27    }
28
29    pub fn temp_dir_path(&self) -> &Path {
30        self.temp_dir.path()
31    }
32
33    pub fn extract_file(&mut self, zip_path: &Path, file_name: &str) -> Result<PathBuf> {
34        // Check if already extracted
35        if let Some(path) = self.extracted_files.get(file_name) {
36            return Ok(path.clone());
37        }
38
39        let file = File::open(zip_path)?;
40        let mut archive = ZipArchive::new(file)?;
41
42        // Find the file in the archive
43        let mut zip_file = archive.by_name(file_name).map_err(|_| {
44            ProcessingError::InvalidFormat(format!(
45                "File '{}' not found in archive '{}'",
46                file_name,
47                zip_path.display()
48            ))
49        })?;
50
51        // Create destination path
52        let dest_path = self.temp_dir.path().join(file_name);
53
54        // Extract the file
55        let mut dest_file = File::create(&dest_path)?;
56        let mut writer = BufWriter::new(&mut dest_file);
57        std::io::copy(&mut zip_file, &mut writer)?;
58        writer.flush()?;
59
60        // Track the extracted file
61        self.extracted_files
62            .insert(file_name.to_string(), dest_path.clone());
63
64        Ok(dest_path)
65    }
66
67    pub fn extract_files_matching_pattern(
68        &mut self,
69        zip_path: &Path,
70        pattern: &str,
71    ) -> Result<Vec<PathBuf>> {
72        let file = File::open(zip_path)?;
73        let mut archive = ZipArchive::new(file)?;
74        let mut extracted_paths = Vec::new();
75
76        // Iterate through all files in the archive
77        for i in 0..archive.len() {
78            let mut zip_file = archive.by_index(i)?;
79            let file_name = zip_file.name().to_string();
80
81            // Check if file matches pattern
82            if file_name.contains(pattern) {
83                // Skip if already extracted
84                if self.extracted_files.contains_key(&file_name) {
85                    extracted_paths.push(self.extracted_files[&file_name].clone());
86                    continue;
87                }
88
89                // Create destination path
90                let dest_path = self.temp_dir.path().join(&file_name);
91
92                // Create parent directories if needed
93                if let Some(parent) = dest_path.parent() {
94                    std::fs::create_dir_all(parent)?;
95                }
96
97                // Extract the file
98                let mut dest_file = File::create(&dest_path)?;
99                let mut writer = BufWriter::new(&mut dest_file);
100                std::io::copy(&mut zip_file, &mut writer)?;
101                writer.flush()?;
102
103                // Track the extracted file
104                self.extracted_files.insert(file_name, dest_path.clone());
105                extracted_paths.push(dest_path);
106            }
107        }
108
109        Ok(extracted_paths)
110    }
111
112    pub fn extract_metadata_files(&mut self, zip_path: &Path) -> Result<HashMap<String, PathBuf>> {
113        let metadata_files = [
114            "stations.txt",
115            "elements.txt",
116            "metadata.txt",
117            "sources.txt",
118        ];
119        let mut extracted = HashMap::new();
120
121        for file_name in &metadata_files {
122            if let Ok(path) = self.extract_file(zip_path, file_name) {
123                extracted.insert(file_name.to_string(), path);
124            }
125        }
126
127        if extracted.is_empty() {
128            return Err(ProcessingError::InvalidFormat(
129                "No metadata files found in archive".to_string(),
130            ));
131        }
132
133        Ok(extracted)
134    }
135
136    pub fn get_extracted_file(&self, file_name: &str) -> Option<&PathBuf> {
137        self.extracted_files.get(file_name)
138    }
139
140    pub fn list_extracted_files(&self) -> Vec<&String> {
141        self.extracted_files.keys().collect()
142    }
143
144    pub fn cleanup(&mut self) -> Result<()> {
145        self.extracted_files.clear();
146
147        // TempDir automatically cleans up on drop, so we don't need to call close() explicitly
148        // The temporary directory will be cleaned up when TempDir is dropped
149        Ok(())
150    }
151
152    pub fn estimate_extraction_size(&self, zip_path: &Path) -> Result<u64> {
153        let file = File::open(zip_path)?;
154        let mut archive = ZipArchive::new(file)?;
155        let mut total_size = 0u64;
156
157        for i in 0..archive.len() {
158            let zip_file = archive.by_index(i)?;
159            total_size += zip_file.size();
160        }
161
162        Ok(total_size)
163    }
164}
165
166impl Drop for TempFileManager {
167    fn drop(&mut self) {
168        if let Err(e) = self.cleanup() {
169            eprintln!("Warning: Failed to cleanup temporary files: {}", e);
170        }
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177    use std::io::Write;
178    use tempfile::NamedTempFile;
179    use zip::{CompressionMethod, ZipWriter};
180
181    fn create_test_zip() -> Result<NamedTempFile> {
182        let file = NamedTempFile::new()?;
183        {
184            let mut zip = ZipWriter::new(&file);
185
186            // Add stations.txt
187            zip.start_file(
188                "stations.txt",
189                zip::write::FileOptions::default().compression_method(CompressionMethod::Stored),
190            )?;
191            zip.write_all(
192                b"STAID,STANAME,CN,LAT,LON,HGHT\n257,TEST STATION,GB,+51:30:00,-000:07:00,100\n",
193            )?;
194
195            // Add elements.txt
196            zip.start_file(
197                "elements.txt",
198                zip::write::FileOptions::default().compression_method(CompressionMethod::Stored),
199            )?;
200            zip.write_all(b"ELEID,DESC,UNIT\nTX1,Maximum temperature,0.1 C\n")?;
201
202            // Add a data file
203            zip.start_file(
204                "TX_STAID000257.txt",
205                zip::write::FileOptions::default().compression_method(CompressionMethod::Stored),
206            )?;
207            zip.write_all(b"Header\n101,20230101,125,0\n")?;
208
209            zip.finish()?;
210        } // zip goes out of scope here
211        Ok(file)
212    }
213
214    #[test]
215    fn test_temp_file_manager_creation() -> Result<()> {
216        let manager = TempFileManager::new()?;
217        assert!(manager.temp_dir_path().exists());
218        Ok(())
219    }
220
221    #[test]
222    fn test_extract_file() -> Result<()> {
223        let test_zip = create_test_zip()?;
224        let mut manager = TempFileManager::new()?;
225
226        let extracted_path = manager.extract_file(test_zip.path(), "stations.txt")?;
227        assert!(extracted_path.exists());
228
229        let content = std::fs::read_to_string(&extracted_path)?;
230        assert!(content.contains("TEST STATION"));
231
232        Ok(())
233    }
234
235    #[test]
236    fn test_extract_metadata_files() -> Result<()> {
237        let test_zip = create_test_zip()?;
238        let mut manager = TempFileManager::new()?;
239
240        let metadata_files = manager.extract_metadata_files(test_zip.path())?;
241        assert!(metadata_files.contains_key("stations.txt"));
242        assert!(metadata_files.contains_key("elements.txt"));
243
244        Ok(())
245    }
246
247    #[test]
248    fn test_extract_files_matching_pattern() -> Result<()> {
249        let test_zip = create_test_zip()?;
250        let mut manager = TempFileManager::new()?;
251
252        let data_files = manager.extract_files_matching_pattern(test_zip.path(), "STAID")?;
253        assert_eq!(data_files.len(), 1);
254        assert!(data_files[0]
255            .file_name()
256            .unwrap()
257            .to_str()
258            .unwrap()
259            .contains("TX_STAID"));
260
261        Ok(())
262    }
263
264    #[test]
265    fn test_already_extracted_file() -> Result<()> {
266        let test_zip = create_test_zip()?;
267        let mut manager = TempFileManager::new()?;
268
269        // Extract once
270        let path1 = manager.extract_file(test_zip.path(), "stations.txt")?;
271        // Extract again - should return same path
272        let path2 = manager.extract_file(test_zip.path(), "stations.txt")?;
273
274        assert_eq!(path1, path2);
275        assert_eq!(manager.list_extracted_files().len(), 1);
276
277        Ok(())
278    }
279}
ecad_processor/archive/temp_manager.rs

ecad_processor/archive/
temp_manager.rs
