Skip to main content

f_ck/engine/
cached_engine.rs

1use crate::dsl::QueryPlan;
2use crate::engine::{salsa_db::*, JoinEngine};
3use anyhow::Result;
4use polars::lazy::frame::LazyFrame;
5use std::collections::HashMap;
6use std::path::{Path, PathBuf};
7use std::time::SystemTime;
8
9pub struct CachedEngine {
10    db: DatabaseImpl,
11    source_cache: HashMap<PathBuf, SourceFile>,
12}
13
14impl CachedEngine {
15    pub fn new() -> Self {
16        Self {
17            db: DatabaseImpl::default(),
18            source_cache: HashMap::new(),
19        }
20    }
21
22    pub fn execute_query_cached(&mut self, query: &QueryPlan) -> Result<LazyFrame> {
23        // Convert sources to SourceFiles and register them with Salsa
24        let mut source_files = Vec::new();
25        
26        for source in &query.sources {
27            let source_file = self.get_or_create_source_file(&source.path)?;
28            source_files.push(source_file);
29        }
30
31        // Serialize the query plan for caching
32        let query_json = serde_json::to_string(query)?;
33        
34        // Execute through Salsa (this will use incremental computation)
35        let _execution_result = execute_query_cached(&self.db, query_json, source_files);
36        
37        // For now, fall back to the regular engine for the actual execution
38        // In a full implementation, we'd cache the LazyFrame itself
39        JoinEngine::execute_query(query)
40    }
41
42    fn get_or_create_source_file(&mut self, path: &PathBuf) -> Result<SourceFile> {
43        if let Some(&existing) = self.source_cache.get(path) {
44            // Check if file has been modified
45            let metadata = std::fs::metadata(path)?;
46            let current_mtime = metadata.modified().unwrap_or(SystemTime::UNIX_EPOCH);
47            
48            if existing.modified_time(&self.db) >= current_mtime {
49                return Ok(existing);
50            }
51        }
52
53        // File doesn't exist in cache or has been modified
54        let content_hash = self.compute_content_hash(path)?;
55        let metadata = std::fs::metadata(path)?;
56        let modified_time = metadata.modified().unwrap_or(SystemTime::UNIX_EPOCH);
57
58        let source_file = SourceFile::new(
59            &mut self.db,
60            path.clone(),
61            content_hash,
62            modified_time,
63        );
64
65        self.source_cache.insert(path.clone(), source_file);
66        Ok(source_file)
67    }
68
69    fn compute_content_hash(&self, path: &Path) -> Result<u64> {
70        use std::collections::hash_map::DefaultHasher;
71        use std::hash::{Hash, Hasher};
72        use std::fs;
73
74        let content = fs::read(path)?;
75        let mut hasher = DefaultHasher::new();
76        content.hash(&mut hasher);
77        Ok(hasher.finish())
78    }
79
80    pub fn invalidate_source(&mut self, path: &PathBuf) {
81        self.source_cache.remove(path);
82        invalidate_source_in_db(&mut self.db, path);
83    }
84
85    pub fn cache_stats(&self) -> HashMap<String, usize> {
86        let mut stats = get_cache_stats(&self.db);
87        stats.insert("source_files_cached".to_string(), self.source_cache.len());
88        stats
89    }
90
91    pub fn clear_cache(&mut self) {
92        self.source_cache.clear();
93        // Reset the database - in practice, you might want more granular control
94        self.db = DatabaseImpl::default();
95    }
96
97    /// Warm up the cache by pre-loading source metadata
98    pub fn warmup_sources(&mut self, paths: &[PathBuf]) -> Result<()> {
99        for path in paths {
100            let source_file = self.get_or_create_source_file(path)?;
101            // Trigger metadata computation
102            let _metadata = source_metadata(&self.db, source_file);
103        }
104        Ok(())
105    }
106
107    /// Get detailed information about a source file from cache
108    pub fn get_source_info(&mut self, path: &PathBuf) -> Result<SourceInfo> {
109        let source_file = self.get_or_create_source_file(path)?;
110        let metadata_str = source_metadata(&self.db, source_file);
111        let parsed_id = parsed_source_id(&self.db, source_file);
112
113        Ok(SourceInfo {
114            path: path.clone(),
115            content_hash: source_file.content_hash(&self.db),
116            modified_time: source_file.modified_time(&self.db),
117            schema: metadata_str,
118            row_count: 0, // TODO: Extract from metadata string
119            file_size: std::fs::metadata(path)?.len(),
120            dataframe_id: parsed_id,
121        })
122    }
123}
124
125impl Default for CachedEngine {
126    fn default() -> Self {
127        Self::new()
128    }
129}
130
131#[derive(Debug)]
132pub struct SourceInfo {
133    pub path: PathBuf,
134    pub content_hash: u64,
135    pub modified_time: SystemTime,
136    pub schema: String,
137    pub row_count: usize,
138    pub file_size: u64,
139    pub dataframe_id: String,
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145    use std::fs;
146    use tempfile::tempdir;
147
148    #[test]
149    fn test_cached_engine_creation() {
150        let engine = CachedEngine::new();
151        assert_eq!(engine.source_cache.len(), 0);
152    }
153
154    #[test]
155    fn test_cache_invalidation() {
156        let mut engine = CachedEngine::new();
157        let temp_dir = tempdir().unwrap();
158        let file_path = temp_dir.path().join("test.csv");
159        
160        // Create a test file
161        fs::write(&file_path, "id,name\n1,test").unwrap();
162        
163        // Load it into cache
164        let _source_file = engine.get_or_create_source_file(&file_path).unwrap();
165        assert_eq!(engine.source_cache.len(), 1);
166        
167        // Invalidate cache
168        engine.invalidate_source(&file_path);
169        assert_eq!(engine.source_cache.len(), 0);
170    }
171}