use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use walkdir::WalkDir;
use super::dag::{Materialization, ModelDag, ModelLayer, OutputFormat};
use super::paths::{resolve_configured_path, resolve_project_path};
pub const DEFAULT_MEMTABLE_MAX_ROWS: usize = 50_000;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RefStrategy {
#[default]
#[serde(alias = "parquet_reread", alias = "lake", alias = "file")]
Parquet,
#[serde(alias = "mem_table", alias = "memory", alias = "arc")]
Memtable,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RefBackend {
MemTable,
LakeFile,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MaterializeMode {
#[default]
#[serde(alias = "streaming")]
Stream,
#[serde(alias = "batch", alias = "legacy")]
Collect,
}
pub const DEFAULT_MAX_ROW_GROUP_ROWS: usize = 1_000_000;
pub const DEFAULT_MAX_ROW_GROUP_BYTES: usize = 128 * 1024 * 1024;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum IcebergWriteMode {
#[default]
#[serde(alias = "catalog_commit", alias = "sor")]
Catalog,
#[serde(alias = "fs", alias = "layout")]
Filesystem,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct IcebergConfig {
#[serde(default)]
pub mode: IcebergWriteMode,
#[serde(default = "default_iceberg_namespace")]
pub namespace: String,
}
fn default_iceberg_namespace() -> String {
"rbt".into()
}
impl Default for IcebergConfig {
fn default() -> Self {
Self {
mode: IcebergWriteMode::Catalog,
namespace: default_iceberg_namespace(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MaterializeConfig {
#[serde(default)]
pub mode: MaterializeMode,
#[serde(default)]
pub ref_strategy: RefStrategy,
#[serde(default = "default_memtable_max_rows")]
pub memtable_max_rows: usize,
#[serde(default = "default_max_row_group_rows")]
pub max_row_group_rows: usize,
#[serde(default = "default_max_row_group_bytes")]
pub max_row_group_bytes: usize,
#[serde(default)]
pub iceberg: IcebergConfig,
#[serde(default)]
pub wap: bool,
}
fn default_memtable_max_rows() -> usize {
DEFAULT_MEMTABLE_MAX_ROWS
}
fn default_max_row_group_rows() -> usize {
DEFAULT_MAX_ROW_GROUP_ROWS
}
fn default_max_row_group_bytes() -> usize {
DEFAULT_MAX_ROW_GROUP_BYTES
}
impl Default for MaterializeConfig {
fn default() -> Self {
Self {
mode: MaterializeMode::Stream,
ref_strategy: RefStrategy::Parquet,
memtable_max_rows: DEFAULT_MEMTABLE_MAX_ROWS,
max_row_group_rows: DEFAULT_MAX_ROW_GROUP_ROWS,
max_row_group_bytes: DEFAULT_MAX_ROW_GROUP_BYTES,
iceberg: IcebergConfig::default(),
wap: false,
}
}
}
pub const DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES: u64 = 1024 * 1024 * 1024;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ScanConfig {
#[serde(default = "default_protobuf_max_payload_bytes")]
pub protobuf_max_payload_bytes: u64,
#[serde(default = "default_spill_arrow_ipc")]
pub spill_arrow_ipc: bool,
#[serde(default = "default_spill_dir")]
pub spill_dir: String,
}
fn default_protobuf_max_payload_bytes() -> u64 {
DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES
}
fn default_spill_arrow_ipc() -> bool {
true
}
fn default_spill_dir() -> String {
".rbt/bronze_spill".into()
}
impl Default for ScanConfig {
fn default() -> Self {
Self {
protobuf_max_payload_bytes: DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES,
spill_arrow_ipc: true,
spill_dir: default_spill_dir(),
}
}
}
impl MaterializeConfig {
pub fn effective_mode(&self) -> MaterializeMode {
match std::env::var("RBT_MATERIALIZE_MODE")
.or_else(|_| std::env::var("RBT_STREAM_MATERIALIZE"))
.ok()
.as_deref()
{
Some("1") | Some("true") | Some("TRUE") | Some("yes") | Some("on") => {
MaterializeMode::Stream
}
Some("0") | Some("false") | Some("FALSE") | Some("no") | Some("off") => {
MaterializeMode::Collect
}
Some(s) if s.eq_ignore_ascii_case("stream") || s.eq_ignore_ascii_case("streaming") => {
MaterializeMode::Stream
}
Some(s) if s.eq_ignore_ascii_case("collect") || s.eq_ignore_ascii_case("batch") => {
MaterializeMode::Collect
}
_ => self.mode,
}
}
pub fn choose_ref_backend(&self, row_count: usize) -> RefBackend {
match self.ref_strategy {
RefStrategy::Parquet => RefBackend::LakeFile,
RefStrategy::Memtable if row_count < self.memtable_max_rows => RefBackend::MemTable,
RefStrategy::Memtable => RefBackend::LakeFile,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct LayerConfig {
pub path: PathBuf,
pub target_path: PathBuf,
pub default_format: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct RbtProjectConfig {
pub name: String,
pub version: String,
pub models_dir: PathBuf,
pub target_path: PathBuf,
#[serde(default)]
pub layers: HashMap<String, LayerConfig>,
#[serde(default)]
pub materialize: MaterializeConfig,
#[serde(default)]
pub scan: ScanConfig,
#[serde(default)]
pub roots: HashMap<String, String>,
}
impl Default for RbtProjectConfig {
fn default() -> Self {
let mut layers = HashMap::new();
layers.insert(
"staging".to_string(),
LayerConfig {
path: PathBuf::from("models/staging"),
target_path: PathBuf::from("lake/silver"),
default_format: Some("parquet".to_string()),
},
);
layers.insert(
"transforms".to_string(),
LayerConfig {
path: PathBuf::from("models/transforms"),
target_path: PathBuf::from("lake/gold"),
default_format: Some("parquet".to_string()),
},
);
layers.insert(
"marts".to_string(),
LayerConfig {
path: PathBuf::from("models/marts"),
target_path: PathBuf::from("lake/gold"),
default_format: Some("parquet_and_iceberg".to_string()),
},
);
Self {
name: "rbt_project".to_string(),
version: "1.0.0".to_string(),
models_dir: PathBuf::from("models"),
target_path: PathBuf::from("lake/gold"),
layers,
materialize: MaterializeConfig::default(),
scan: ScanConfig::default(),
roots: HashMap::new(),
}
}
}
impl RbtProjectConfig {
pub fn load(project_dir: &Path) -> Result<Self> {
let project_file = project_dir.join("rbt_project.yml");
if project_file.exists() {
let content = fs::read_to_string(&project_file).with_context(|| {
format!(
"E_RBT_PROJECT_LOAD: cannot read project file {}",
project_file.display()
)
})?;
let mut config: RbtProjectConfig = serde_yaml::from_str(&content).with_context(|| {
format!(
"E_RBT_PROJECT_LOAD: failed to parse {}. \
Check required keys (name, version, models_dir, target_path) and \
optional materialize:/scan:/roots:/layers blocks.",
project_file.display()
)
})?;
let defaults = Self::default();
for (key, val) in defaults.layers {
config.layers.entry(key).or_insert(val);
}
Ok(config)
} else {
Ok(Self::default())
}
}
pub fn resolve_path(&self, project_dir: &Path, configured: &str) -> Result<PathBuf> {
resolve_project_path(project_dir, configured, &self.roots)
}
pub fn resolve_layer_target_dir(
&self,
project_dir: &Path,
layer: ModelLayer,
) -> Result<PathBuf> {
let layer_key = match layer {
ModelLayer::Staging => "staging",
ModelLayer::Transform => "transforms",
ModelLayer::Mart => "marts",
};
if let Some(layer_cfg) = self.layers.get(layer_key) {
resolve_configured_path(project_dir, &layer_cfg.target_path, &self.roots)
} else {
resolve_configured_path(project_dir, &self.target_path, &self.roots)
}
}
pub fn resolve_model_target_path(
&self,
project_dir: &Path,
model_name: &str,
layer: ModelLayer,
ext: &str,
) -> Result<PathBuf> {
let dir = self
.resolve_layer_target_dir(project_dir, layer)
.with_context(|| {
format!(
"E_RBT_MODEL_TARGET: cannot resolve output directory for model '{model_name}' \
(layer={layer:?}). Check `layers.*.target_path`, top-level `target_path`, and \
`roots:` in rbt_project.yml."
)
})?;
Ok(dir.join(format!("{model_name}.{ext}")))
}
pub fn resolve_model_target_dir(
&self,
project_dir: &Path,
model_name: &str,
layer: ModelLayer,
) -> Result<PathBuf> {
let dir = self
.resolve_layer_target_dir(project_dir, layer)
.with_context(|| {
format!(
"E_RBT_MODEL_TARGET: cannot resolve table directory for model '{model_name}' \
(layer={layer:?}). Check layer target_path and roots:."
)
})?;
Ok(dir.join(model_name))
}
pub fn build_dag(
&self,
project_dir: &Path,
cli_format_override: Option<OutputFormat>,
) -> Result<ModelDag> {
let models_dir = project_dir.join(&self.models_dir);
let mut dag = ModelDag::new();
if !models_dir.exists() {
let default_fmt = cli_format_override.unwrap_or(OutputFormat::Parquet);
dag.add_model_with_format(
"stg_users",
"SELECT 1 AS id, 'Alice' AS name, 'admin' AS role",
Materialization::Table,
default_fmt,
None,
"",
)?;
dag.build_graph()?;
return Ok(dag);
}
let mut model_count = 0;
for entry in WalkDir::new(&models_dir).into_iter().filter_map(|e| e.ok()) {
let path = entry.path();
if path.is_file() && path.extension().is_some_and(|ext| ext == "sql") {
let stem = path.file_stem().and_then(|s| s.to_str()).with_context(|| {
format!(
"E_RBT_MODEL_NAME: invalid file stem for model path {}",
path.display()
)
})?;
let raw_sql = fs::read_to_string(path).with_context(|| {
format!(
"E_RBT_MODEL_IO: failed reading model SQL {}",
path.display()
)
})?;
let layer = ModelLayer::from_name(stem);
let format = cli_format_override.clone().unwrap_or_else(|| {
let layer_key = match layer {
ModelLayer::Staging => "staging",
ModelLayer::Transform => "transforms",
ModelLayer::Mart => "marts",
};
if let Some(l_cfg) = self.layers.get(layer_key) {
match l_cfg.default_format.as_deref() {
Some("parquet") => OutputFormat::Parquet,
Some("jsonl") => OutputFormat::Jsonl,
Some("csv") => OutputFormat::Csv,
Some("iceberg") => OutputFormat::Iceberg,
Some("parquet_and_iceberg") => OutputFormat::ParquetAndIceberg,
_ => OutputFormat::Parquet,
}
} else {
OutputFormat::Parquet
}
});
let target_file_path = match format {
OutputFormat::Iceberg => self
.resolve_model_target_dir(project_dir, stem, layer)
.with_context(|| {
format!(
"E_RBT_MODEL_TARGET: model '{stem}' (Iceberg) — \
failed resolving layer target. \
layer={layer:?}; project={}",
project_dir.display()
)
})?,
OutputFormat::Parquet
| OutputFormat::ParquetAndIceberg
| OutputFormat::ZeroCopyClone => self
.resolve_model_target_path(project_dir, stem, layer, "parquet")
.with_context(|| {
format!(
"E_RBT_MODEL_TARGET: model '{stem}' (parquet) — \
failed resolving layer target. \
layer={layer:?}; project={}",
project_dir.display()
)
})?,
OutputFormat::Jsonl => self
.resolve_model_target_path(project_dir, stem, layer, "jsonl")
.with_context(|| {
format!(
"E_RBT_MODEL_TARGET: model '{stem}' (jsonl) — \
failed resolving layer target"
)
})?,
OutputFormat::Csv => self
.resolve_model_target_path(project_dir, stem, layer, "csv")
.with_context(|| {
format!(
"E_RBT_MODEL_TARGET: model '{stem}' (csv) — \
failed resolving layer target"
)
})?,
};
dag.add_model_with_format(
stem,
&raw_sql,
Materialization::Table,
format,
Some(target_file_path.to_string_lossy().to_string()),
"",
)?;
model_count += 1;
}
}
if model_count == 0 {
let default_fmt = cli_format_override.unwrap_or(OutputFormat::Parquet);
dag.add_model_with_format(
"stg_users",
"SELECT 1 AS id, 'Alice' AS name, 'admin' AS role",
Materialization::Table,
default_fmt,
None,
"",
)?;
}
dag.build_graph()?;
Ok(dag)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_layer_target_path_resolution() -> Result<()> {
let config = RbtProjectConfig::default();
let project_dir = Path::new("/tmp/test_project");
let stg_path = config.resolve_model_target_path(
project_dir,
"stg_trades",
ModelLayer::Staging,
"parquet",
)?;
assert_eq!(stg_path, project_dir.join("lake/silver/stg_trades.parquet"));
let tf_path = config.resolve_model_target_path(
project_dir,
"tf_1m_bars",
ModelLayer::Transform,
"parquet",
)?;
assert_eq!(tf_path, project_dir.join("lake/gold/tf_1m_bars.parquet"));
let mart_path = config.resolve_model_target_path(
project_dir,
"fact_1d_bars",
ModelLayer::Mart,
"parquet",
)?;
assert_eq!(
mart_path,
project_dir.join("lake/gold/fact_1d_bars.parquet")
);
Ok(())
}
#[test]
fn materialize_defaults_to_parquet_reread() {
let cfg = MaterializeConfig::default();
assert_eq!(cfg.ref_strategy, RefStrategy::Parquet);
assert_eq!(cfg.mode, MaterializeMode::Stream);
assert_eq!(cfg.memtable_max_rows, DEFAULT_MEMTABLE_MAX_ROWS);
assert_eq!(cfg.max_row_group_rows, DEFAULT_MAX_ROW_GROUP_ROWS);
assert_eq!(cfg.choose_ref_backend(0), RefBackend::LakeFile);
assert_eq!(cfg.choose_ref_backend(1_000_000), RefBackend::LakeFile);
}
#[test]
fn materialize_mode_from_yaml() -> Result<()> {
let yml = r#"
name: t
version: "1"
models_dir: models
target_path: lake/gold
materialize:
mode: collect
max_row_group_rows: 1000
"#;
let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
assert_eq!(cfg.materialize.mode, MaterializeMode::Collect);
assert_eq!(cfg.materialize.max_row_group_rows, 1000);
Ok(())
}
#[test]
fn materialize_memtable_respects_cutoff() {
let cfg = MaterializeConfig {
ref_strategy: RefStrategy::Memtable,
memtable_max_rows: 50_000,
..Default::default()
};
assert_eq!(cfg.choose_ref_backend(49_999), RefBackend::MemTable);
assert_eq!(cfg.choose_ref_backend(50_000), RefBackend::LakeFile);
assert_eq!(cfg.choose_ref_backend(50_001), RefBackend::LakeFile);
}
#[test]
fn parse_materialize_block_from_yaml() -> Result<()> {
let yml = r#"
name: t
version: "1"
models_dir: models
target_path: lake/gold
materialize:
ref_strategy: memtable
memtable_max_rows: 10000
"#;
let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
assert_eq!(cfg.materialize.ref_strategy, RefStrategy::Memtable);
assert_eq!(cfg.materialize.memtable_max_rows, 10_000);
assert_eq!(
cfg.materialize.choose_ref_backend(9_999),
RefBackend::MemTable
);
assert_eq!(
cfg.materialize.choose_ref_backend(10_000),
RefBackend::LakeFile
);
Ok(())
}
#[test]
fn parse_project_without_materialize_uses_defaults() -> Result<()> {
let yml = r#"
name: t
version: "1"
models_dir: models
target_path: lake/gold
"#;
let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
assert_eq!(cfg.materialize, MaterializeConfig::default());
Ok(())
}
#[test]
fn parse_memtable_without_max_rows_defaults_cutoff() -> Result<()> {
let yml = r#"
name: t
version: "1"
models_dir: models
target_path: lake/gold
materialize:
ref_strategy: memtable
"#;
let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
assert_eq!(cfg.materialize.ref_strategy, RefStrategy::Memtable);
assert_eq!(cfg.materialize.memtable_max_rows, DEFAULT_MEMTABLE_MAX_ROWS);
Ok(())
}
#[test]
fn absolute_layer_target_not_nested_under_project() -> Result<()> {
let yml = r#"
name: multi_root_demo
version: "1"
models_dir: models
target_path: /mnt/datalake/acme/nonprod/lake_us/lake/gold
layers:
staging:
path: models/staging
target_path: /mnt/datalake/acme/nonprod/lake_us/lake/silver/stage
default_format: parquet
"#;
let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
let project = Path::new("/home/dev/rbt_projects/demo");
let stg =
cfg.resolve_model_target_path(project, "stg_events", ModelLayer::Staging, "parquet")?;
assert_eq!(
stg,
PathBuf::from(
"/mnt/datalake/acme/nonprod/lake_us/lake/silver/stage/stg_events.parquet"
)
);
assert!(!stg.starts_with(project));
Ok(())
}
#[test]
fn multi_root_template_in_layer_target() -> Result<()> {
let yml = r#"
name: multi_root_demo
version: "1"
models_dir: models
target_path: $nonprod_lake/gold
roots:
nonprod_lake: /mnt/datalake/acme/nonprod/lake_us/lake
layers:
staging:
path: models/staging
target_path: $nonprod_lake/silver/stage
default_format: parquet
"#;
let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
let project = Path::new("/home/dev/proj");
let dir = cfg.resolve_layer_target_dir(project, ModelLayer::Staging)?;
assert_eq!(
dir,
PathBuf::from("/mnt/datalake/acme/nonprod/lake_us/lake/silver/stage")
);
Ok(())
}
#[test]
fn bad_root_in_layer_target_is_error() {
let yml = r#"
name: t
version: "1"
models_dir: models
target_path: lake/gold
layers:
staging:
path: models/staging
target_path: $missing_root/silver
default_format: parquet
"#;
let cfg: RbtProjectConfig = serde_yaml::from_str(yml).unwrap();
let err = cfg
.resolve_layer_target_dir(Path::new("/proj"), ModelLayer::Staging)
.unwrap_err()
.to_string();
assert!(err.contains("E_RBT_ROOT_UNKNOWN") || err.contains("E_RBT_LAYER_PATH"));
}
#[test]
fn scan_config_defaults_protobuf_cap() {
let cfg = ScanConfig::default();
assert_eq!(
cfg.protobuf_max_payload_bytes,
DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES
);
assert_eq!(DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES, 1024 * 1024 * 1024);
}
#[test]
fn scan_config_override_from_yml() -> Result<()> {
let yml = r#"
name: t
version: "1"
models_dir: models
target_path: lake/gold
scan:
protobuf_max_payload_bytes: 4096
"#;
let cfg: RbtProjectConfig = serde_yaml::from_str(yml)?;
assert_eq!(cfg.scan.protobuf_max_payload_bytes, 4096);
let yml2 = r#"
name: t
version: "1"
models_dir: models
target_path: lake/gold
"#;
let cfg2: RbtProjectConfig = serde_yaml::from_str(yml2)?;
assert_eq!(
cfg2.scan.protobuf_max_payload_bytes,
DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES
);
Ok(())
}
#[test]
fn load_workspace_example_projects() -> Result<()> {
let manifest = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
let repo = manifest.join("../..");
for (rel, name, expect_root) in [
("examples/smoke_fixture", "smoke_fixture", "lake"),
("examples/full_e2e_rbt_example", "market_bars", "lake"),
] {
let dir = repo.join(rel);
if !dir.join("rbt_project.yml").is_file() {
continue;
}
let cfg = RbtProjectConfig::load(&dir)?;
assert_eq!(cfg.name, name, "example {rel}");
assert_eq!(
cfg.roots.get("lake").map(String::as_str),
Some(expect_root),
"example {rel} should declare roots.lake"
);
assert_eq!(cfg.materialize.ref_strategy, RefStrategy::Parquet);
assert_eq!(
cfg.scan.protobuf_max_payload_bytes,
DEFAULT_PROTOBUF_MAX_PAYLOAD_BYTES
);
let silver = cfg.resolve_layer_target_dir(&dir, ModelLayer::Staging)?;
assert!(
silver.ends_with("lake/silver") || silver.ends_with("lake\\silver"),
"staging target for {rel}: {}",
silver.display()
);
if dir.join("models").is_dir() {
let dag = cfg.build_dag(&dir, None)?;
assert!(
dag.graph.node_count() >= 3,
"example {rel} expected ≥3 models, got {}",
dag.graph.node_count()
);
}
}
Ok(())
}
}