f_ck/engine/
cached_engine.rs1use 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 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 let query_json = serde_json::to_string(query)?;
33
34 let _execution_result = execute_query_cached(&self.db, query_json, source_files);
36
37 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 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 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 self.db = DatabaseImpl::default();
95 }
96
97 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 let _metadata = source_metadata(&self.db, source_file);
103 }
104 Ok(())
105 }
106
107 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, 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 fs::write(&file_path, "id,name\n1,test").unwrap();
162
163 let _source_file = engine.get_or_create_source_file(&file_path).unwrap();
165 assert_eq!(engine.source_cache.len(), 1);
166
167 engine.invalidate_source(&file_path);
169 assert_eq!(engine.source_cache.len(), 0);
170 }
171}