1use crate::SourceRow;
2use crate::filter::{FilterCondition, matches_all as matches_filter_all};
3use crate::index::SourceIndex;
4use crate::index_builder::index_path_for_source;
5use crate::memory::InMemorySource;
6use crate::{SourceDefinition, SourceReader, open_source_reader};
7use anyhow::{Context, Result};
8use apif_twoq_cache::TwoQCache;
9use apif_utils::FileUtils;
10use serde_json::Value;
11use std::collections::HashMap;
12use std::path::{Path, PathBuf};
13use std::sync::{Arc, Mutex};
14use tracing::warn;
15
16const DIMENSION_CACHE_HOT: usize = 2048;
20const DIMENSION_CACHE_COLD: usize = 8192;
21
22const ENV_DIMENSION_MEMORY_BUDGET: &str = "GRPCTESTIFY_DIMENSION_MEMORY_BUDGET";
23const MAX_DIMENSION_MEMORY_BUDGET: u64 = 512 * 1024 * 1024;
24const MIN_DIMENSION_MEMORY_BUDGET: u64 = 32 * 1024 * 1024;
25
26fn resolve_dimension_budget() -> u64 {
27 if let Ok(val) = std::env::var(ENV_DIMENSION_MEMORY_BUDGET)
28 && !val.is_empty()
29 && let Ok(bytes) = parse_bytes(&val)
30 {
31 return bytes;
32 }
33
34 let mut sys = sysinfo::System::new_with_specifics(sysinfo::RefreshKind::nothing());
35 sys.refresh_memory();
36 let available = sys.available_memory();
37
38 if available == 0 {
39 return MIN_DIMENSION_MEMORY_BUDGET;
40 }
41
42 (available / 2).clamp(MIN_DIMENSION_MEMORY_BUDGET, MAX_DIMENSION_MEMORY_BUDGET)
43}
44
45fn parse_bytes(s: &str) -> Result<u64> {
46 let s = s.trim_ascii().to_ascii_lowercase();
47 let split_pos = s
48 .char_indices()
49 .find(|(_, c)| !c.is_ascii_digit() && *c != '.')
50 .map(|(i, _)| i)
51 .unwrap_or(s.len());
52 let num_str = &s[..split_pos];
53 let unit = s[split_pos..].trim_ascii();
54 let num: f64 = num_str
55 .parse()
56 .map_err(|_| anyhow::anyhow!("invalid number: {num_str}"))?;
57 let bytes = match unit {
58 "" | "b" => num,
59 "kb" | "k" => num * 1024.0,
60 "mb" | "m" => num * 1024.0 * 1024.0,
61 "gb" | "g" => num * 1024.0 * 1024.0 * 1024.0,
62 other => anyhow::bail!("unknown unit: {other} (use kb, mb, gb)"),
63 };
64 Ok(bytes as u64)
65}
66
67pub enum DimensionSource {
68 Memory(Arc<InMemorySource>),
69 Indexed(Box<IndexedDimension>),
70}
71
72pub struct IndexedDimension {
73 pub index: Arc<SourceIndex>,
74 pub mmap: memmap2::Mmap,
75 pub cache: Mutex<TwoQCache<String, SourceRow>>,
76}
77
78impl DimensionSource {
79 fn lookup_row(&self, key: &str) -> Result<Option<SourceRow>> {
80 match self {
81 DimensionSource::Memory(mem) => Ok(mem.lookup(key).cloned()),
82 DimensionSource::Indexed(idx) => {
83 let mut cache = idx.cache.lock().expect("cache mutex poisoned");
84 if let Some(row) = cache.get(&key.to_string()) {
85 return Ok(Some(row.clone()));
86 }
87 let Some(line) = idx.index.lookup_row_from_mmap(idx.mmap.as_ref(), key)? else {
88 return Ok(None);
89 };
90 let row = SourceRow::from_csv_line(&line);
91 cache.insert(key.to_string(), row.clone());
92 Ok(Some(row))
93 }
94 }
95 }
96
97 fn lookup_all(&self, key: &str) -> Vec<SourceRow> {
100 match self {
101 DimensionSource::Memory(mem) => mem
102 .iter()
103 .filter(|(k, _)| k.as_str() == key)
104 .map(|(_, row)| row.clone())
105 .collect(),
106 DimensionSource::Indexed(_) => Vec::new(),
107 }
108 }
109}
110
111#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
112pub enum RuntimeFallbackPolicy {
113 #[default]
114 Skip,
115 ScanSource,
116 Error,
117}
118
119#[derive(Debug, Clone, Default)]
120pub struct SourceFallbackEvent {
121 pub source_name: String,
122 pub key: String,
123 pub reason: FallbackReason,
124 pub fallback_type: FallbackType,
125}
126
127#[derive(Debug, Clone, Default)]
128pub enum FallbackReason {
129 #[default]
130 IndexLookupMiss,
131 IndexCorrupted,
132 IndexOutOfSync,
133 TypeMismatch,
134}
135
136#[derive(Debug, Clone, Default)]
137pub enum FallbackType {
138 #[default]
139 None,
140 ScanSource,
141 Error,
142}
143
144struct DimensionJoin {
145 source_name: String,
146 foreign_key: String,
147 join_type: super::definition::JoinType,
148}
149
150struct CrossProductState {
153 row: SourceRow,
154 cross_matches: Vec<Vec<SourceRow>>,
155 cross_indices: Vec<usize>,
156}
157
158#[derive(Clone)]
159struct DimTask {
160 name: String,
161 resolved_path: PathBuf,
162 key_col: String,
163 file_size: u64,
164 def: SourceDefinition,
165}
166
167fn load_dimension_source(
168 def: &SourceDefinition,
169 document_path: &Path,
170 resolved_path: &Path,
171 key_col: &str,
172) -> Result<DimensionSource> {
173 let effective_key = if key_col.is_empty() {
174 let reader = open_source_reader(def, document_path)
175 .with_context(|| format!("failed to open dimension source '{}'", def.file))?;
176 reader.headers().first().cloned().unwrap_or_default()
177 } else {
178 key_col.to_string()
179 };
180
181 let index =
182 load_or_build_index_with_key(def, document_path, &effective_key).with_context(|| {
183 format!(
184 "failed to build/load index for dimension '{}'",
185 resolved_path.display()
186 )
187 })?;
188 let file = std::fs::File::open(resolved_path)
189 .with_context(|| format!("failed to open dimension file: {}", resolved_path.display()))?;
190 let mmap = unsafe { memmap2::Mmap::map(&file) }
191 .with_context(|| format!("failed to mmap dimension file: {}", resolved_path.display()))?;
192 Ok(DimensionSource::Indexed(Box::new(IndexedDimension {
193 index: Arc::new(index),
194 mmap,
195 cache: Mutex::new(TwoQCache::new(DIMENSION_CACHE_HOT, DIMENSION_CACHE_COLD)),
196 })))
197}
198
199fn load_dimension_in_memory(
200 def: &SourceDefinition,
201 document_path: &Path,
202 resolved_path: &Path,
203 key_col: &str,
204) -> Result<DimensionSource> {
205 let mut reader = open_source_reader(def, document_path)
206 .with_context(|| format!("failed to open dimension source '{}'", def.file))?;
207 let effective_key = if key_col.is_empty() {
208 reader.headers().first().cloned().unwrap_or_default()
209 } else {
210 key_col.to_string()
211 };
212 let mem = InMemorySource::load(&mut *reader, &effective_key)
213 .with_context(|| format!("failed to load dimension '{}'", resolved_path.display()))?;
214 let mem = if let Some(ref filter) = def.filter {
216 let filtered = mem.filter(filter);
217 Arc::new(filtered)
218 } else {
219 Arc::new(mem)
220 };
221 Ok(DimensionSource::Memory(mem))
222}
223
224pub struct SourceDrivenConfig {
225 pub primary: Arc<Mutex<Box<dyn SourceReader>>>,
226 pub primary_name: String,
227 pub dimensions: HashMap<String, DimensionSource>,
228 pub resolved_paths: HashMap<String, PathBuf>,
229 dim_joins: Vec<DimensionJoin>,
230 primary_filter: Vec<FilterCondition>,
231 pub load_stats: DimLoadStats,
232 pub runtime_stats: SourceRuntimeStats,
233 pub fallback_policy: RuntimeFallbackPolicy,
234 cross_product_state: std::sync::Mutex<Option<CrossProductState>>,
235 pub loaded_at: std::time::Instant,
236 pub current_row: std::sync::atomic::AtomicU64,
237}
238
239#[derive(Debug, Clone, Default)]
240pub struct DimLoadStats {
241 pub in_memory_count: usize,
242 pub indexed_count: usize,
243 pub total_file_bytes: u64,
244 pub index_build_ms: u64,
245}
246
247#[derive(Debug)]
251pub struct SourceRuntimeStats {
252 pub dimension_lookups: std::sync::atomic::AtomicU64,
253 pub dimension_hits: std::sync::atomic::AtomicU64,
254 pub dimension_misses: std::sync::atomic::AtomicU64,
255 pub in_memory_lookups: std::sync::atomic::AtomicU64,
256 pub indexed_lookups: std::sync::atomic::AtomicU64,
257 pub index_fallbacks: std::sync::atomic::AtomicU64,
258}
259
260#[derive(Debug, Clone, Default)]
262pub struct RuntimeStatsSnapshot {
263 pub dimension_lookups: u64,
264 pub dimension_hits: u64,
265 pub dimension_misses: u64,
266 pub in_memory_lookups: u64,
267 pub indexed_lookups: u64,
268 pub index_fallbacks: u64,
269}
270
271impl SourceRuntimeStats {
272 pub fn snapshot(&self) -> RuntimeStatsSnapshot {
274 use std::sync::atomic::Ordering::Relaxed;
275 RuntimeStatsSnapshot {
276 dimension_lookups: self.dimension_lookups.load(Relaxed),
277 dimension_hits: self.dimension_hits.load(Relaxed),
278 dimension_misses: self.dimension_misses.load(Relaxed),
279 in_memory_lookups: self.in_memory_lookups.load(Relaxed),
280 indexed_lookups: self.indexed_lookups.load(Relaxed),
281 index_fallbacks: self.index_fallbacks.load(Relaxed),
282 }
283 }
284}
285
286impl Default for SourceRuntimeStats {
287 fn default() -> Self {
288 Self {
289 dimension_lookups: std::sync::atomic::AtomicU64::new(0),
290 dimension_hits: std::sync::atomic::AtomicU64::new(0),
291 dimension_misses: std::sync::atomic::AtomicU64::new(0),
292 in_memory_lookups: std::sync::atomic::AtomicU64::new(0),
293 indexed_lookups: std::sync::atomic::AtomicU64::new(0),
294 index_fallbacks: std::sync::atomic::AtomicU64::new(0),
295 }
296 }
297}
298
299impl Clone for SourceRuntimeStats {
300 fn clone(&self) -> Self {
301 Self {
302 dimension_lookups: std::sync::atomic::AtomicU64::new(
303 self.dimension_lookups
304 .load(std::sync::atomic::Ordering::Relaxed),
305 ),
306 dimension_hits: std::sync::atomic::AtomicU64::new(
307 self.dimension_hits
308 .load(std::sync::atomic::Ordering::Relaxed),
309 ),
310 dimension_misses: std::sync::atomic::AtomicU64::new(
311 self.dimension_misses
312 .load(std::sync::atomic::Ordering::Relaxed),
313 ),
314 in_memory_lookups: std::sync::atomic::AtomicU64::new(
315 self.in_memory_lookups
316 .load(std::sync::atomic::Ordering::Relaxed),
317 ),
318 indexed_lookups: std::sync::atomic::AtomicU64::new(
319 self.indexed_lookups
320 .load(std::sync::atomic::Ordering::Relaxed),
321 ),
322 index_fallbacks: std::sync::atomic::AtomicU64::new(
323 self.index_fallbacks
324 .load(std::sync::atomic::Ordering::Relaxed),
325 ),
326 }
327 }
328}
329
330impl SourceRuntimeStats {
331 pub fn record_lookup(&self, _source_name: &str, found: bool, is_indexed: bool) {
332 use std::sync::atomic::Ordering;
333 self.dimension_lookups.fetch_add(1, Ordering::Relaxed);
334 if found {
335 self.dimension_hits.fetch_add(1, Ordering::Relaxed);
336 } else {
337 self.dimension_misses.fetch_add(1, Ordering::Relaxed);
338 }
339 if is_indexed {
340 self.indexed_lookups.fetch_add(1, Ordering::Relaxed);
341 } else {
342 self.in_memory_lookups.fetch_add(1, Ordering::Relaxed);
343 }
344 }
345
346 pub fn record_fallback(&self) {
347 use std::sync::atomic::Ordering;
348 self.index_fallbacks.fetch_add(1, Ordering::Relaxed);
349 }
350}
351
352impl SourceDrivenConfig {
353 pub fn prepare(definitions: &[SourceDefinition], document_path: &Path) -> Result<Option<Self>> {
354 if definitions.is_empty() {
355 return Ok(None);
356 }
357
358 let primary_def = &definitions[0];
359 let primary_name = primary_def
360 .name
361 .clone()
362 .unwrap_or_else(|| "primary".to_string());
363
364 let primary_reader = open_source_reader(primary_def, document_path)
365 .with_context(|| format!("failed to open primary source '{}'", primary_def.file))?;
366 let primary_filter = primary_def.filter.clone().unwrap_or_default();
367
368 let mut dimensions = HashMap::new();
369 let mut resolved_paths = HashMap::new();
370 let mut dim_joins = Vec::new();
371 let mut dim_tasks: Vec<DimTask> = Vec::new();
372
373 for def in &definitions[1..] {
374 let dim_name = def.name.clone().unwrap_or_else(|| "dim".to_string());
375
376 let resolved = FileUtils::resolve_relative_path(document_path, &def.file);
377 let file_size = std::fs::metadata(&resolved).map(|m| m.len()).unwrap_or(0);
378
379 let key_col = def
380 .indexed_by
381 .as_ref()
382 .map(|v| v.join(super::index::COMPOSITE_KEY_SEPARATOR))
383 .unwrap_or_default();
384
385 dim_joins.push(DimensionJoin {
386 source_name: dim_name.clone(),
387 foreign_key: key_col.clone(),
388 join_type: def.join_type_or_default(),
389 });
390
391 resolved_paths.insert(dim_name.clone(), resolved.clone());
392 dim_tasks.push(DimTask {
393 name: dim_name,
394 resolved_path: resolved,
395 key_col,
396 file_size,
397 def: def.clone(),
398 });
399 }
400
401 let memory_bb = resolve_dimension_budget();
402 let mut in_memory: Vec<DimTask> = Vec::new();
403 let mut too_large: Vec<DimTask> = Vec::new();
404 let mut total_file_bytes = 0u64;
405 for task in dim_tasks {
406 total_file_bytes += task.file_size;
407 if task.file_size <= memory_bb {
408 in_memory.push(task);
409 } else {
410 too_large.push(task);
411 }
412 }
413 in_memory.sort_by_key(|t| t.file_size);
414
415 let task_count = in_memory.len() + too_large.len();
416 let all_tasks: Vec<DimTask> = in_memory.iter().chain(too_large.iter()).cloned().collect();
417 let stats = Arc::new(std::sync::Mutex::new((0usize, 0usize, 0u64)));
418
419 let results: Vec<(String, Result<DimensionSource>)> = if task_count <= 1 {
420 all_tasks
421 .into_iter()
422 .map(|t| {
423 let start = std::time::Instant::now();
424 let src = if t.file_size <= memory_bb {
425 load_dimension_in_memory(
426 &t.def,
427 document_path,
428 &t.resolved_path,
429 &t.key_col,
430 )
431 } else {
432 load_dimension_source(&t.def, document_path, &t.resolved_path, &t.key_col)
433 };
434 let elapsed = start.elapsed().as_millis() as u64;
435 let mut s = stats.lock().expect("stats mutex should not be poisoned");
436 if t.file_size <= memory_bb {
437 s.0 += 1;
438 } else {
439 s.1 += 1;
440 }
441 s.2 += elapsed;
442 (t.name, src)
443 })
444 .collect()
445 } else {
446 std::thread::scope(|s| {
447 all_tasks
448 .into_iter()
449 .map(|t| {
450 let doc_path = document_path.to_path_buf();
451 let mem_budget = memory_bb;
452 let stats = Arc::clone(&stats);
453 s.spawn(move || {
454 let start = std::time::Instant::now();
455 let src = if t.file_size <= mem_budget {
456 load_dimension_in_memory(
457 &t.def,
458 &doc_path,
459 &t.resolved_path,
460 &t.key_col,
461 )
462 } else {
463 load_dimension_source(
464 &t.def,
465 &doc_path,
466 &t.resolved_path,
467 &t.key_col,
468 )
469 };
470 let elapsed = start.elapsed().as_millis() as u64;
471 let mut ss = stats.lock().expect("stats mutex should not be poisoned");
472 if t.file_size <= mem_budget {
473 ss.0 += 1;
474 } else {
475 ss.1 += 1;
476 }
477 ss.2 += elapsed;
478 (t.name, src)
479 })
480 })
481 .collect::<Vec<_>>()
482 .into_iter()
483 .map(|h| h.join().expect("dimension load thread panicked"))
484 .collect()
485 })
486 };
487
488 let (in_memory_count, indexed_count, index_build_ms) =
489 *stats.lock().expect("stats mutex should not be poisoned");
490
491 for (name, result) in results {
492 match result {
493 Ok(ds) => {
494 dimensions.insert(name, ds);
495 }
496 Err(e) => {
497 return Err(e).with_context(|| format!("failed to load dimension '{}'", name));
498 }
499 }
500 }
501
502 Ok(Some(Self {
503 primary: Arc::new(Mutex::new(primary_reader)),
504 primary_name,
505 dimensions,
506 resolved_paths,
507 dim_joins,
508 primary_filter,
509 cross_product_state: std::sync::Mutex::new(None),
510 loaded_at: std::time::Instant::now(),
511 current_row: std::sync::atomic::AtomicU64::new(0),
512 load_stats: DimLoadStats {
513 in_memory_count,
514 indexed_count,
515 total_file_bytes,
516 index_build_ms,
517 },
518 runtime_stats: SourceRuntimeStats::default(),
519 fallback_policy: RuntimeFallbackPolicy::default(),
520 }))
521 }
522
523 pub fn next_row_variables(&self) -> Result<Option<HashMap<String, Value>>> {
524 {
526 let state_guard = self
527 .cross_product_state
528 .lock()
529 .expect("cross_product_state mutex should not be poisoned");
530 if state_guard.is_some() {
531 drop(state_guard);
532 return self.next_cross_product_combination();
533 }
534 }
535
536 let mut reader = self.primary.lock().map_err(|e| anyhow::anyhow!("{e}"))?;
537 let row = loop {
538 match reader.next_row()? {
539 Some(r) => {
540 if self.primary_filter.is_empty()
541 || matches_filter_all(&r, &self.primary_filter)
542 {
543 break r;
544 }
545 }
546 None => return Ok(None),
547 }
548 };
549 drop(reader);
550
551 let mut vars = self.build_primary_vars(&row);
552
553 let inner_missing = self.dim_joins.iter().any(|j| {
555 j.join_type == super::definition::JoinType::Inner
556 && row
557 .get(&j.foreign_key)
558 .is_some_and(|fk| self.dimension_lookup(&j.source_name, fk).is_none())
559 });
560 if inner_missing {
561 return self.next_row_variables();
562 }
563
564 for join in &self.dim_joins {
566 if join.join_type != super::definition::JoinType::Left {
567 continue;
568 }
569 if let Some(fk_val) = row.get(&join.foreign_key)
570 && let Some(dim_row) = self.dimension_lookup(&join.source_name, fk_val)
571 {
572 for col in dim_row.columns() {
573 if let Some(val) = dim_row.get(col) {
574 vars.insert(
575 format!("{}.{}", join.source_name, col),
576 Value::String(val.to_string()),
577 );
578 }
579 }
580 }
581 }
582
583 let has_cross = self
585 .dim_joins
586 .iter()
587 .any(|j| j.join_type == super::definition::JoinType::Cross);
588 if has_cross {
589 let mut cross_matches: Vec<Vec<SourceRow>> = Vec::new();
590 for join in &self.dim_joins {
591 if join.join_type != super::definition::JoinType::Cross {
592 continue;
593 }
594 if let Some(fk_val) = row.get(&join.foreign_key) {
595 if let Some(all_rows) = self.dimension_lookup_all(&join.source_name, fk_val) {
596 cross_matches.push(all_rows);
597 } else {
598 cross_matches.push(Vec::new());
599 }
600 } else {
601 cross_matches.push(Vec::new());
602 }
603 }
604
605 *self
606 .cross_product_state
607 .lock()
608 .expect("cross_product_state mutex should not be poisoned") =
609 Some(CrossProductState {
610 row,
611 cross_matches,
612 cross_indices: vec![
613 0;
614 self.dim_joins
615 .iter()
616 .filter(|j| j.join_type == super::definition::JoinType::Cross)
617 .count()
618 ],
619 });
620 return self.next_cross_product_combination();
621 }
622
623 self.current_row
624 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
625 Ok(Some(vars))
626 }
627
628 fn next_cross_product_combination(&self) -> Result<Option<HashMap<String, Value>>> {
630 let mut state_guard = self
631 .cross_product_state
632 .lock()
633 .expect("cross_product_state mutex should not be poisoned");
634 let state = state_guard.as_ref().unwrap();
635 let row = &state.row;
636 let mut vars = self.build_primary_vars(row);
637
638 let mut cross_idx = 0;
640 for join in &self.dim_joins {
641 if join.join_type != super::definition::JoinType::Cross {
642 continue;
643 }
644 if let Some(matches) = state.cross_matches.get(cross_idx) {
645 let idx = state.cross_indices[cross_idx];
646 if idx < matches.len() {
647 let dim_row = &matches[idx];
648 for col in dim_row.columns() {
649 if let Some(val) = dim_row.get(col) {
650 vars.insert(
651 format!("{}.{}", join.source_name, col),
652 Value::String(val.to_string()),
653 );
654 }
655 }
656 }
657 }
658 cross_idx += 1;
659 }
660
661 let cross_count = self
663 .dim_joins
664 .iter()
665 .filter(|j| j.join_type == super::definition::JoinType::Cross)
666 .count();
667 let mut new_indices = state.cross_indices.clone();
668 let mut advanced = false;
669
670 for i in (0..cross_count).rev() {
671 let max = state.cross_matches[i].len();
672 if max == 0 {
673 continue;
674 }
675 if new_indices[i] + 1 < max {
676 new_indices[i] += 1;
677 advanced = true;
678 break;
679 } else {
680 new_indices[i] = 0;
681 }
682 }
683
684 if advanced {
685 *state_guard = Some(CrossProductState {
686 row: row.clone(),
687 cross_matches: state.cross_matches.clone(),
688 cross_indices: new_indices,
689 });
690 } else {
691 *state_guard = None;
692 }
693
694 self.current_row
695 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
696 Ok(Some(vars))
697 }
698
699 fn build_primary_vars(&self, row: &SourceRow) -> HashMap<String, Value> {
700 let mut vars = HashMap::new();
701 for col in row.columns() {
702 if let Some(val) = row.get(col) {
703 vars.insert(
704 format!("{}.{}", self.primary_name, col),
705 Value::String(val.to_string()),
706 );
707 }
708 }
709 vars
710 }
711
712 pub fn dimension_lookup(&self, source_name: &str, key: &str) -> Option<SourceRow> {
713 let dim = self.dimensions.get(source_name)?;
714 let is_indexed = matches!(dim, DimensionSource::Indexed(_));
715 let result = dim.lookup_row(key).ok().flatten();
716 self.runtime_stats
717 .record_lookup(source_name, result.is_some(), is_indexed);
718 result
719 }
720
721 pub fn dimension_lookup_all(&self, source_name: &str, key: &str) -> Option<Vec<SourceRow>> {
722 let dim = self.dimensions.get(source_name)?;
723 let rows = dim.lookup_all(key);
724 if rows.is_empty() { None } else { Some(rows) }
725 }
726
727 pub fn build_dimension_variables(
728 &self,
729 row: &SourceRow,
730 joins: &[(String, String, String)],
731 ) -> HashMap<String, Value> {
732 let mut vars = HashMap::new();
733
734 for (dim_name, local_key_col, _remote_key_col) in joins {
735 if let Some(key_val) = row.get(local_key_col)
736 && let Some(dim_row) = self.dimension_lookup(dim_name, key_val)
737 {
738 for col in dim_row.columns() {
739 if let Some(val) = dim_row.get(col) {
740 vars.insert(format!("{dim_name}.{col}"), Value::String(val.to_string()));
741 }
742 }
743 }
744 }
745
746 vars
747 }
748
749 pub fn primary_headers(&self) -> Vec<String> {
750 let reader = self.primary.lock().ok();
751 match reader {
752 Some(r) => r.headers().to_vec(),
753 None => Vec::new(),
754 }
755 }
756}
757
758fn load_or_build_index_with_key(
759 def: &SourceDefinition,
760 document_path: &Path,
761 key_col: &str,
762) -> Result<SourceIndex> {
763 let source_path = FileUtils::resolve_relative_path(document_path, &def.file);
764 let idx_path = index_path_for_source(&source_path, key_col);
765
766 if idx_path.exists() {
767 match SourceIndex::read_from_file(&idx_path) {
768 Ok(index) => {
769 let idx_meta = std::fs::metadata(&idx_path);
770 let src_meta = std::fs::metadata(&source_path);
771 if let (Ok(im), Ok(sm)) = (idx_meta, src_meta)
772 && let (Ok(it), Ok(st)) = (im.modified(), sm.modified())
773 && it >= st
774 {
775 return Ok(index);
776 }
777 }
778 Err(e) => {
779 if is_corruption_error(&e) {
780 warn!(
781 "Index corrupted (checksum mismatch), rebuilding: {}. Error: {}",
782 idx_path.display(),
783 e
784 );
785 let _ = std::fs::remove_file(&idx_path);
786 }
787 }
788 }
789 }
790
791 let mut reader = open_source_reader(def, document_path)
792 .with_context(|| format!("failed to open source for indexing: {}", def.file))?;
793
794 let mut index = SourceIndex::new(key_col);
795 let header_line = read_first_line(&source_path)?;
796 let mut byte_offset = header_line.len() as u64 + 1;
797 let mut row_count = 0u64;
798
799 while let Some(row) = reader.next_row()? {
800 let key_val = row.get(key_col).ok_or_else(|| {
801 anyhow::anyhow!("column '{}' not found in row {}", key_col, row_count)
802 })?;
803 let row_bytes = estimate_row_size(&row);
804 index
805 .insert(key_val.to_string(), byte_offset, row_bytes)
806 .with_context(|| format!("failed to insert key '{}' at row {}", key_val, row_count))?;
807 byte_offset += row_bytes as u64 + 1;
808 row_count += 1;
809 }
810
811 let parent = idx_path.parent().unwrap_or(Path::new("."));
812 std::fs::create_dir_all(parent).ok();
813 let mut index_mut = index;
814 index_mut
815 .write_to_file(&idx_path)
816 .with_context(|| format!("failed to write index to {}", idx_path.display()))?;
817
818 SourceIndex::read_from_file(&idx_path)
819}
820
821fn is_corruption_error(e: &anyhow::Error) -> bool {
822 let msg = e.to_string();
823 msg.contains("corrupted")
824 || msg.contains("checksum mismatch")
825 || msg.contains("invalid index file")
826}
827
828fn read_first_line(path: &Path) -> Result<String> {
829 use std::io::{BufRead, BufReader};
830 let file = std::fs::File::open(path)?;
831 let mut reader = BufReader::new(file);
832 let mut line = String::new();
833 reader.read_line(&mut line)?;
834 Ok(line)
835}
836
837fn estimate_row_size(row: &SourceRow) -> u32 {
838 let mut size = 0u32;
839 for col in row.columns() {
840 size += col.len() as u32 + 1;
841 }
842 for val in row.values() {
843 size += val.len() as u32;
844 }
845 size + row.columns().len().saturating_sub(1) as u32
846}
847
848#[cfg(test)]
849mod tests {
850 use super::*;
851 #[cfg(not(miri))]
852 use std::io::Write;
853
854 #[cfg(not(miri))]
855 fn create_temp_csv(dir: &Path, name: &str, content: &str) -> PathBuf {
856 let path = dir.join(name);
857 let mut f = std::fs::File::create(&path).unwrap();
858 f.write_all(content.as_bytes()).unwrap();
859 path
860 }
861
862 #[test]
863 fn no_definitions_returns_none() {
864 let result = SourceDrivenConfig::prepare(&[], Path::new("test.gctf")).unwrap();
865 assert!(result.is_none());
866 }
867
868 #[test]
869 #[cfg(not(miri))]
870 fn primary_only_no_dimensions() {
871 let dir = std::env::temp_dir().join("gctf_driven_test");
872 std::fs::create_dir_all(&dir).unwrap();
873 create_temp_csv(&dir, "users.csv", "id,name,age\n1,Alice,30\n2,Bob,25\n");
874
875 let defs: Vec<SourceDefinition> =
876 serde_yaml_ng::from_str("- file: users.csv\n name: users\n").unwrap();
877
878 let doc_path = dir.join("test.gctf");
879 std::fs::write(&doc_path, "").unwrap();
880
881 let config = SourceDrivenConfig::prepare(&defs, &doc_path)
882 .unwrap()
883 .unwrap();
884
885 assert_eq!(config.primary_name, "users");
886 assert!(config.dimensions.is_empty());
887
888 let vars = config.next_row_variables().unwrap().unwrap();
889 assert_eq!(vars.get("users.id"), Some(&Value::String("1".into())));
890 assert_eq!(vars.get("users.name"), Some(&Value::String("Alice".into())));
891
892 let vars2 = config.next_row_variables().unwrap().unwrap();
893 assert_eq!(vars2.get("users.name"), Some(&Value::String("Bob".into())));
894
895 let vars3 = config.next_row_variables().unwrap();
896 assert!(vars3.is_none());
897
898 std::fs::remove_dir_all(&dir).ok();
899 }
900
901 #[test]
902 #[cfg(not(miri))]
903 fn primary_with_dimension_join() {
904 let dir = std::env::temp_dir().join("gctf_driven_join_test");
905 std::fs::create_dir_all(&dir).unwrap();
906
907 create_temp_csv(
908 &dir,
909 "pvz.csv",
910 "pvz_id,region_id,name\n1,R01,PVZ Alpha\n2,R02,PVZ Beta\n",
911 );
912 create_temp_csv(
913 &dir,
914 "regions.csv",
915 "region_id,region_name\nR01,Moscow\nR02,Saint Petersburg\n",
916 );
917
918 let defs: Vec<SourceDefinition> = serde_yaml_ng::from_str(
919 "- file: pvz.csv\n name: pvz\n- file: regions.csv\n name: regions\n indexed_by: [region_id]\n"
920 ).unwrap();
921
922 let doc_path = dir.join("test.gctf");
923 std::fs::write(&doc_path, "").unwrap();
924
925 let config = SourceDrivenConfig::prepare(&defs, &doc_path)
926 .unwrap()
927 .unwrap();
928
929 assert_eq!(config.dimensions.len(), 1);
930
931 let vars = config.next_row_variables().unwrap().unwrap();
932
933 assert_eq!(vars.get("pvz.pvz_id"), Some(&Value::String("1".into())));
934 assert_eq!(
935 vars.get("pvz.region_id"),
936 Some(&Value::String("R01".into()))
937 );
938 assert_eq!(
939 vars.get("pvz.name"),
940 Some(&Value::String("PVZ Alpha".into()))
941 );
942
943 assert_eq!(
944 vars.get("regions.region_id"),
945 Some(&Value::String("R01".into()))
946 );
947 assert_eq!(
948 vars.get("regions.region_name"),
949 Some(&Value::String("Moscow".into()))
950 );
951
952 let vars2 = config.next_row_variables().unwrap().unwrap();
953 assert_eq!(
954 vars2.get("pvz.name"),
955 Some(&Value::String("PVZ Beta".into()))
956 );
957 assert_eq!(
958 vars2.get("regions.region_name"),
959 Some(&Value::String("Saint Petersburg".into()))
960 );
961
962 std::fs::remove_dir_all(&dir).ok();
963 }
964
965 #[test]
966 #[cfg(not(miri))]
967 fn dimension_missing_fk_still_injects_primary() {
968 let dir = std::env::temp_dir().join("gctf_driven_fk_test");
969 std::fs::create_dir_all(&dir).unwrap();
970
971 create_temp_csv(&dir, "data.csv", "id,ref_id,val\n1,MISSING,hello\n");
972 create_temp_csv(&dir, "ref.csv", "ref_id,label\nOK,Found\n");
973
974 let defs: Vec<SourceDefinition> = serde_yaml_ng::from_str(
975 "- file: data.csv\n name: data\n- file: ref.csv\n name: ref\n indexed_by: [ref_id]\n",
976 )
977 .unwrap();
978
979 let doc_path = dir.join("test.gctf");
980 std::fs::write(&doc_path, "").unwrap();
981
982 let config = SourceDrivenConfig::prepare(&defs, &doc_path)
983 .unwrap()
984 .unwrap();
985
986 let vars = config.next_row_variables().unwrap().unwrap();
987 assert_eq!(vars.get("data.val"), Some(&Value::String("hello".into())));
988 assert!(!vars.contains_key("ref.label"));
989
990 std::fs::remove_dir_all(&dir).ok();
991 }
992
993 #[test]
994 #[cfg(not(miri))]
995 fn primary_filter_skips_non_matching_rows() {
996 let dir = std::env::temp_dir().join("gctf_driven_filter_test");
997 std::fs::create_dir_all(&dir).unwrap();
998
999 create_temp_csv(
1000 &dir,
1001 "pvz.csv",
1002 "pvz_id,status,name\n1,inactive,Old\n2,active,New\n",
1003 );
1004
1005 let defs: Vec<SourceDefinition> = serde_yaml_ng::from_str(
1006 "- file: pvz.csv\n name: pvz\n filter:\n - field: status\n equals: active\n",
1007 )
1008 .unwrap();
1009
1010 let doc_path = dir.join("test.gctf");
1011 std::fs::write(&doc_path, "").unwrap();
1012
1013 let config = SourceDrivenConfig::prepare(&defs, &doc_path)
1014 .unwrap()
1015 .unwrap();
1016 let vars = config.next_row_variables().unwrap().unwrap();
1017 assert_eq!(vars.get("pvz.pvz_id"), Some(&Value::String("2".into())));
1018 assert_eq!(vars.get("pvz.name"), Some(&Value::String("New".into())));
1019 assert!(config.next_row_variables().unwrap().is_none());
1020
1021 std::fs::remove_dir_all(&dir).ok();
1022 }
1023}