1use anyhow::{Context, Result};
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::fs;
5use std::path::{Path, PathBuf};
6use walkdir::WalkDir;
7
8use super::dag::{Materialization, ModelDag, ModelLayer, OutputFormat};
9use super::paths::{resolve_configured_path, resolve_project_path};
10
11pub const DEFAULT_MEMTABLE_MAX_ROWS: usize = 50_000;
13
14#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
19#[serde(rename_all = "snake_case")]
20pub enum RefStrategy {
21 #[default]
23 #[serde(alias = "parquet_reread", alias = "lake", alias = "file")]
24 Parquet,
25 #[serde(alias = "mem_table", alias = "memory", alias = "arc")]
27 Memtable,
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum RefBackend {
33 MemTable,
35 LakeFile,
37}
38
39#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
44#[serde(rename_all = "snake_case")]
45pub enum MaterializeMode {
46 #[default]
48 #[serde(alias = "streaming")]
49 Stream,
50 #[serde(alias = "batch", alias = "legacy")]
52 Collect,
53}
54
55pub const DEFAULT_MAX_ROW_GROUP_ROWS: usize = 1_000_000;
57pub const DEFAULT_MAX_ROW_GROUP_BYTES: usize = 128 * 1024 * 1024;
59
60#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
65pub struct MaterializeConfig {
66 #[serde(default)]
68 pub mode: MaterializeMode,
69 #[serde(default)]
71 pub ref_strategy: RefStrategy,
72 #[serde(default = "default_memtable_max_rows")]
74 pub memtable_max_rows: usize,
75 #[serde(default = "default_max_row_group_rows")]
77 pub max_row_group_rows: usize,
78 #[serde(default = "default_max_row_group_bytes")]
80 pub max_row_group_bytes: usize,
81}
82
83fn default_memtable_max_rows() -> usize {
84 DEFAULT_MEMTABLE_MAX_ROWS
85}
86
87fn default_max_row_group_rows() -> usize {
88 DEFAULT_MAX_ROW_GROUP_ROWS
89}
90
91fn default_max_row_group_bytes() -> usize {
92 DEFAULT_MAX_ROW_GROUP_BYTES
93}
94
95impl Default for MaterializeConfig {
96 fn default() -> Self {
97 Self {
98 mode: MaterializeMode::Stream,
99 ref_strategy: RefStrategy::Parquet,
100 memtable_max_rows: DEFAULT_MEMTABLE_MAX_ROWS,
101 max_row_group_rows: DEFAULT_MAX_ROW_GROUP_ROWS,
102 max_row_group_bytes: DEFAULT_MAX_ROW_GROUP_BYTES,
103 }
104 }
105}
106
107pub const DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES: u64 = 1024 * 1024 * 1024;
109
110#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
112pub struct ScanConfig {
113 #[serde(default = "default_protobuf_max_payload_bytes")]
117 pub protobuf_max_payload_bytes: u64,
118}
119
120fn default_protobuf_max_payload_bytes() -> u64 {
121 DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES
122}
123
124impl Default for ScanConfig {
125 fn default() -> Self {
126 Self {
127 protobuf_max_payload_bytes: DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES,
128 }
129 }
130}
131
132impl MaterializeConfig {
133 pub fn effective_mode(&self) -> MaterializeMode {
135 match std::env::var("RBT_MATERIALIZE_MODE")
136 .or_else(|_| std::env::var("RBT_STREAM_MATERIALIZE"))
137 .ok()
138 .as_deref()
139 {
140 Some("1") | Some("true") | Some("TRUE") | Some("yes") | Some("on") => {
142 MaterializeMode::Stream
143 }
144 Some("0") | Some("false") | Some("FALSE") | Some("no") | Some("off") => {
145 MaterializeMode::Collect
146 }
147 Some(s) if s.eq_ignore_ascii_case("stream") || s.eq_ignore_ascii_case("streaming") => {
148 MaterializeMode::Stream
149 }
150 Some(s) if s.eq_ignore_ascii_case("collect") || s.eq_ignore_ascii_case("batch") => {
151 MaterializeMode::Collect
152 }
153 _ => self.mode,
154 }
155 }
156
157 pub fn choose_ref_backend(&self, row_count: usize) -> RefBackend {
159 match self.ref_strategy {
160 RefStrategy::Parquet => RefBackend::LakeFile,
161 RefStrategy::Memtable if row_count < self.memtable_max_rows => RefBackend::MemTable,
162 RefStrategy::Memtable => RefBackend::LakeFile,
163 }
164 }
165}
166
167#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
169pub struct LayerConfig {
170 pub path: PathBuf,
171 pub target_path: PathBuf,
172 pub default_format: Option<String>,
173}
174
175#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
177pub struct RbtProjectConfig {
178 pub name: String,
179 pub version: String,
180 pub models_dir: PathBuf,
181 pub target_path: PathBuf,
182 #[serde(default)]
183 pub layers: HashMap<String, LayerConfig>,
184 #[serde(default)]
186 pub materialize: MaterializeConfig,
187 #[serde(default)]
189 pub scan: ScanConfig,
190 #[serde(default)]
194 pub roots: HashMap<String, String>,
195}
196
197impl Default for RbtProjectConfig {
198 fn default() -> Self {
199 let mut layers = HashMap::new();
200 layers.insert(
201 "staging".to_string(),
202 LayerConfig {
203 path: PathBuf::from("models/staging"),
204 target_path: PathBuf::from("lake/silver"),
205 default_format: Some("parquet".to_string()),
206 },
207 );
208 layers.insert(
209 "transforms".to_string(),
210 LayerConfig {
211 path: PathBuf::from("models/transforms"),
212 target_path: PathBuf::from("lake/gold"),
213 default_format: Some("parquet".to_string()),
214 },
215 );
216 layers.insert(
217 "marts".to_string(),
218 LayerConfig {
219 path: PathBuf::from("models/marts"),
220 target_path: PathBuf::from("lake/gold"),
221 default_format: Some("parquet_and_iceberg".to_string()),
222 },
223 );
224
225 Self {
226 name: "rbt_project".to_string(),
227 version: "1.0.0".to_string(),
228 models_dir: PathBuf::from("models"),
229 target_path: PathBuf::from("lake/gold"),
230 layers,
231 materialize: MaterializeConfig::default(),
232 scan: ScanConfig::default(),
233 roots: HashMap::new(),
234 }
235 }
236}
237
238impl RbtProjectConfig {
239 pub fn load(project_dir: &Path) -> Result<Self> {
241 let project_file = project_dir.join("rbt_project.yml");
242 if project_file.exists() {
243 let content = fs::read_to_string(&project_file).with_context(|| {
244 format!(
245 "E_RBT_PROJECT_LOAD: cannot read project file {}",
246 project_file.display()
247 )
248 })?;
249 let mut config: RbtProjectConfig = serde_yaml::from_str(&content).with_context(|| {
250 format!(
251 "E_RBT_PROJECT_LOAD: failed to parse {}. \
252 Check required keys (name, version, models_dir, target_path) and \
253 optional materialize:/scan:/roots:/layers blocks.",
254 project_file.display()
255 )
256 })?;
257
258 let defaults = Self::default();
259 for (key, val) in defaults.layers {
260 config.layers.entry(key).or_insert(val);
261 }
262 Ok(config)
263 } else {
264 Ok(Self::default())
265 }
266 }
267
268 pub fn resolve_path(&self, project_dir: &Path, configured: &str) -> Result<PathBuf> {
270 resolve_project_path(project_dir, configured, &self.roots)
271 }
272
273 pub fn resolve_layer_target_dir(
275 &self,
276 project_dir: &Path,
277 layer: ModelLayer,
278 ) -> Result<PathBuf> {
279 let layer_key = match layer {
280 ModelLayer::Staging => "staging",
281 ModelLayer::Transform => "transforms",
282 ModelLayer::Mart => "marts",
283 };
284 if let Some(layer_cfg) = self.layers.get(layer_key) {
285 resolve_configured_path(project_dir, &layer_cfg.target_path, &self.roots)
286 } else {
287 resolve_configured_path(project_dir, &self.target_path, &self.roots)
288 }
289 }
290
291 pub fn resolve_model_target_path(
296 &self,
297 project_dir: &Path,
298 model_name: &str,
299 layer: ModelLayer,
300 ext: &str,
301 ) -> Result<PathBuf> {
302 let dir = self
303 .resolve_layer_target_dir(project_dir, layer)
304 .with_context(|| {
305 format!(
306 "E_RBT_MODEL_TARGET: cannot resolve output directory for model '{model_name}' \
307 (layer={layer:?}). Check `layers.*.target_path`, top-level `target_path`, and \
308 `roots:` in rbt_project.yml."
309 )
310 })?;
311 Ok(dir.join(format!("{model_name}.{ext}")))
312 }
313
314 pub fn resolve_model_target_dir(
316 &self,
317 project_dir: &Path,
318 model_name: &str,
319 layer: ModelLayer,
320 ) -> Result<PathBuf> {
321 let dir = self
322 .resolve_layer_target_dir(project_dir, layer)
323 .with_context(|| {
324 format!(
325 "E_RBT_MODEL_TARGET: cannot resolve table directory for model '{model_name}' \
326 (layer={layer:?}). Check layer target_path and roots:."
327 )
328 })?;
329 Ok(dir.join(model_name))
330 }
331
332 pub fn build_dag(
334 &self,
335 project_dir: &Path,
336 cli_format_override: Option<OutputFormat>,
337 ) -> Result<ModelDag> {
338 let models_dir = project_dir.join(&self.models_dir);
339 let mut dag = ModelDag::new();
340
341 if !models_dir.exists() {
342 let default_fmt = cli_format_override.unwrap_or(OutputFormat::Parquet);
343 dag.add_model_with_format(
344 "stg_users",
345 "SELECT 1 AS id, 'Alice' AS name, 'admin' AS role",
346 Materialization::Table,
347 default_fmt,
348 None,
349 "",
350 )?;
351 dag.build_graph()?;
352 return Ok(dag);
353 }
354
355 let mut model_count = 0;
356 for entry in WalkDir::new(&models_dir).into_iter().filter_map(|e| e.ok()) {
357 let path = entry.path();
358 if path.is_file() && path.extension().is_some_and(|ext| ext == "sql") {
359 let stem = path.file_stem().and_then(|s| s.to_str()).with_context(|| {
360 format!(
361 "E_RBT_MODEL_NAME: invalid file stem for model path {}",
362 path.display()
363 )
364 })?;
365 let raw_sql = fs::read_to_string(path).with_context(|| {
366 format!(
367 "E_RBT_MODEL_IO: failed reading model SQL {}",
368 path.display()
369 )
370 })?;
371
372 let layer = ModelLayer::from_name(stem);
373 let format = cli_format_override.clone().unwrap_or_else(|| {
374 let layer_key = match layer {
375 ModelLayer::Staging => "staging",
376 ModelLayer::Transform => "transforms",
377 ModelLayer::Mart => "marts",
378 };
379 if let Some(l_cfg) = self.layers.get(layer_key) {
380 match l_cfg.default_format.as_deref() {
381 Some("parquet") => OutputFormat::Parquet,
382 Some("jsonl") => OutputFormat::Jsonl,
383 Some("csv") => OutputFormat::Csv,
384 Some("iceberg") => OutputFormat::Iceberg,
385 Some("parquet_and_iceberg") => OutputFormat::ParquetAndIceberg,
386 _ => OutputFormat::Parquet,
387 }
388 } else {
389 OutputFormat::Parquet
390 }
391 });
392
393 let target_file_path = match format {
394 OutputFormat::Iceberg => self
395 .resolve_model_target_dir(project_dir, stem, layer)
396 .with_context(|| {
397 format!(
398 "E_RBT_MODEL_TARGET: model '{stem}' (Iceberg) — \
399 failed resolving layer target. \
400 layer={layer:?}; project={}",
401 project_dir.display()
402 )
403 })?,
404 OutputFormat::Parquet
405 | OutputFormat::ParquetAndIceberg
406 | OutputFormat::ZeroCopyClone => self
407 .resolve_model_target_path(project_dir, stem, layer, "parquet")
408 .with_context(|| {
409 format!(
410 "E_RBT_MODEL_TARGET: model '{stem}' (parquet) — \
411 failed resolving layer target. \
412 layer={layer:?}; project={}",
413 project_dir.display()
414 )
415 })?,
416 OutputFormat::Jsonl => self
417 .resolve_model_target_path(project_dir, stem, layer, "jsonl")
418 .with_context(|| {
419 format!(
420 "E_RBT_MODEL_TARGET: model '{stem}' (jsonl) — \
421 failed resolving layer target"
422 )
423 })?,
424 OutputFormat::Csv => self
425 .resolve_model_target_path(project_dir, stem, layer, "csv")
426 .with_context(|| {
427 format!(
428 "E_RBT_MODEL_TARGET: model '{stem}' (csv) — \
429 failed resolving layer target"
430 )
431 })?,
432 };
433
434 dag.add_model_with_format(
435 stem,
436 &raw_sql,
437 Materialization::Table,
438 format,
439 Some(target_file_path.to_string_lossy().to_string()),
440 "",
441 )?;
442 model_count += 1;
443 }
444 }
445
446 if model_count == 0 {
447 let default_fmt = cli_format_override.unwrap_or(OutputFormat::Parquet);
448 dag.add_model_with_format(
449 "stg_users",
450 "SELECT 1 AS id, 'Alice' AS name, 'admin' AS role",
451 Materialization::Table,
452 default_fmt,
453 None,
454 "",
455 )?;
456 }
457
458 dag.build_graph()?;
459 Ok(dag)
460 }
461}
462
463#[cfg(test)]
464mod tests {
465 use super::*;
466
467 #[test]
468 fn test_layer_target_path_resolution() -> Result<()> {
469 let config = RbtProjectConfig::default();
470 let project_dir = Path::new("/tmp/test_project");
471
472 let stg_path = config.resolve_model_target_path(
473 project_dir,
474 "stg_trades",
475 ModelLayer::Staging,
476 "parquet",
477 )?;
478 assert_eq!(stg_path, project_dir.join("lake/silver/stg_trades.parquet"));
479
480 let tf_path = config.resolve_model_target_path(
481 project_dir,
482 "tf_1m_bars",
483 ModelLayer::Transform,
484 "parquet",
485 )?;
486 assert_eq!(tf_path, project_dir.join("lake/gold/tf_1m_bars.parquet"));
487
488 let mart_path = config.resolve_model_target_path(
489 project_dir,
490 "fact_1d_bars",
491 ModelLayer::Mart,
492 "parquet",
493 )?;
494 assert_eq!(
495 mart_path,
496 project_dir.join("lake/gold/fact_1d_bars.parquet")
497 );
498
499 Ok(())
500 }
501
502 #[test]
503 fn materialize_defaults_to_parquet_reread() {
504 let cfg = MaterializeConfig::default();
505 assert_eq!(cfg.ref_strategy, RefStrategy::Parquet);
506 assert_eq!(cfg.mode, MaterializeMode::Stream);
507 assert_eq!(cfg.memtable_max_rows, DEFAULT_MEMTABLE_MAX_ROWS);
508 assert_eq!(cfg.max_row_group_rows, DEFAULT_MAX_ROW_GROUP_ROWS);
509 assert_eq!(cfg.choose_ref_backend(0), RefBackend::LakeFile);
510 assert_eq!(cfg.choose_ref_backend(1_000_000), RefBackend::LakeFile);
511 }
512
513 #[test]
514 fn materialize_mode_from_yaml() -> Result<()> {
515 let yml = r#"
516name: t
517version: "1"
518models_dir: models
519target_path: lake/gold
520materialize:
521 mode: collect
522 max_row_group_rows: 1000
523"#;
524 let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
525 assert_eq!(cfg.materialize.mode, MaterializeMode::Collect);
526 assert_eq!(cfg.materialize.max_row_group_rows, 1000);
527 Ok(())
528 }
529
530 #[test]
531 fn materialize_memtable_respects_cutoff() {
532 let cfg = MaterializeConfig {
533 ref_strategy: RefStrategy::Memtable,
534 memtable_max_rows: 50_000,
535 ..Default::default()
536 };
537 assert_eq!(cfg.choose_ref_backend(49_999), RefBackend::MemTable);
538 assert_eq!(cfg.choose_ref_backend(50_000), RefBackend::LakeFile);
539 assert_eq!(cfg.choose_ref_backend(50_001), RefBackend::LakeFile);
540 }
541
542 #[test]
543 fn parse_materialize_block_from_yaml() -> Result<()> {
544 let yml = r#"
545name: t
546version: "1"
547models_dir: models
548target_path: lake/gold
549materialize:
550 ref_strategy: memtable
551 memtable_max_rows: 10000
552"#;
553 let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
554 assert_eq!(cfg.materialize.ref_strategy, RefStrategy::Memtable);
555 assert_eq!(cfg.materialize.memtable_max_rows, 10_000);
556 assert_eq!(
557 cfg.materialize.choose_ref_backend(9_999),
558 RefBackend::MemTable
559 );
560 assert_eq!(
561 cfg.materialize.choose_ref_backend(10_000),
562 RefBackend::LakeFile
563 );
564 Ok(())
565 }
566
567 #[test]
568 fn parse_project_without_materialize_uses_defaults() -> Result<()> {
569 let yml = r#"
570name: t
571version: "1"
572models_dir: models
573target_path: lake/gold
574"#;
575 let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
576 assert_eq!(cfg.materialize, MaterializeConfig::default());
577 Ok(())
578 }
579
580 #[test]
581 fn parse_memtable_without_max_rows_defaults_cutoff() -> Result<()> {
582 let yml = r#"
583name: t
584version: "1"
585models_dir: models
586target_path: lake/gold
587materialize:
588 ref_strategy: memtable
589"#;
590 let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
591 assert_eq!(cfg.materialize.ref_strategy, RefStrategy::Memtable);
592 assert_eq!(cfg.materialize.memtable_max_rows, DEFAULT_MEMTABLE_MAX_ROWS);
593 Ok(())
594 }
595
596 #[test]
597 fn absolute_layer_target_not_nested_under_project() -> Result<()> {
598 let yml = r#"
599name: multi_root_demo
600version: "1"
601models_dir: models
602target_path: /mnt/datalake/acme/nonprod/lake_us/lake/gold
603layers:
604 staging:
605 path: models/staging
606 target_path: /mnt/datalake/acme/nonprod/lake_us/lake/silver/stage
607 default_format: parquet
608"#;
609 let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
610 let project = Path::new("/home/dev/rbt_projects/demo");
611 let stg =
612 cfg.resolve_model_target_path(project, "stg_events", ModelLayer::Staging, "parquet")?;
613 assert_eq!(
614 stg,
615 PathBuf::from(
616 "/mnt/datalake/acme/nonprod/lake_us/lake/silver/stage/stg_events.parquet"
617 )
618 );
619 assert!(!stg.starts_with(project));
620 Ok(())
621 }
622
623 #[test]
624 fn multi_root_template_in_layer_target() -> Result<()> {
625 let yml = r#"
626name: multi_root_demo
627version: "1"
628models_dir: models
629target_path: $nonprod_lake/gold
630roots:
631 nonprod_lake: /mnt/datalake/acme/nonprod/lake_us/lake
632layers:
633 staging:
634 path: models/staging
635 target_path: $nonprod_lake/silver/stage
636 default_format: parquet
637"#;
638 let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
639 let project = Path::new("/home/dev/proj");
640 let dir = cfg.resolve_layer_target_dir(project, ModelLayer::Staging)?;
641 assert_eq!(
642 dir,
643 PathBuf::from("/mnt/datalake/acme/nonprod/lake_us/lake/silver/stage")
644 );
645 Ok(())
646 }
647
648 #[test]
649 fn bad_root_in_layer_target_is_error() {
650 let yml = r#"
651name: t
652version: "1"
653models_dir: models
654target_path: lake/gold
655layers:
656 staging:
657 path: models/staging
658 target_path: $missing_root/silver
659 default_format: parquet
660"#;
661 let cfg: RbtProjectConfig = serde_yaml::from_str(yml).unwrap();
662 let err = cfg
663 .resolve_layer_target_dir(Path::new("/proj"), ModelLayer::Staging)
664 .unwrap_err()
665 .to_string();
666 assert!(err.contains("E_RBT_ROOT_UNKNOWN") || err.contains("E_RBT_LAYER_PATH"));
667 }
668
669 #[test]
670 fn scan_config_defaults_protobuf_cap() {
671 let cfg = ScanConfig::default();
672 assert_eq!(
673 cfg.protobuf_max_payload_bytes,
674 DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES
675 );
676 assert_eq!(DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES, 1024 * 1024 * 1024);
677 }
678
679 #[test]
680 fn scan_config_override_from_yml() -> Result<()> {
681 let yml = r#"
682name: t
683version: "1"
684models_dir: models
685target_path: lake/gold
686scan:
687 protobuf_max_payload_bytes: 4096
688"#;
689 let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
690 assert_eq!(cfg.scan.protobuf_max_payload_bytes, 4096);
691 let yml2 = r#"
693name: t
694version: "1"
695models_dir: models
696target_path: lake/gold
697"#;
698 let cfg2: RbtProjectConfig = serde_yaml::from_str(yml2)?;
699 assert_eq!(
700 cfg2.scan.protobuf_max_payload_bytes,
701 DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES
702 );
703 Ok(())
704 }
705
706 #[test]
708 fn load_workspace_example_projects() -> Result<()> {
709 let manifest = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
710 let repo = manifest.join("../..");
711 for (rel, name, expect_root) in [
712 ("examples/smoke_fixture", "smoke_fixture", "lake"),
713 ("examples/full_e2e_rbt_example", "market_bars", "lake"),
714 ] {
715 let dir = repo.join(rel);
716 if !dir.join("rbt_project.yml").is_file() {
717 continue;
719 }
720 let cfg = RbtProjectConfig::load(&dir)?;
721 assert_eq!(cfg.name, name, "example {rel}");
722 assert_eq!(
723 cfg.roots.get("lake").map(String::as_str),
724 Some(expect_root),
725 "example {rel} should declare roots.lake"
726 );
727 assert_eq!(cfg.materialize.ref_strategy, RefStrategy::Parquet);
728 assert_eq!(
729 cfg.scan.protobuf_max_payload_bytes,
730 DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES
731 );
732 let silver = cfg.resolve_layer_target_dir(&dir, ModelLayer::Staging)?;
733 assert!(
734 silver.ends_with("lake/silver") || silver.ends_with("lake\\silver"),
735 "staging target for {rel}: {}",
736 silver.display()
737 );
738 if dir.join("models").is_dir() {
740 let dag = cfg.build_dag(&dir, None)?;
741 assert!(
742 dag.graph.node_count() >= 3,
743 "example {rel} expected ≥3 models, got {}",
744 dag.graph.node_count()
745 );
746 }
747 }
748 Ok(())
749 }
750}