Skip to main content

rbt/core/
dag.rs

1use super::frontmatter::{
2    scan_path_exists, BronzeCheckMode, BronzeDiagnostic, BronzeValidationReport,
3    DiagnosticSeverity, StagingFrontmatter,
4};
5use super::parser::{DependencyRef, SqlModelParser};
6use anyhow::{bail, Result};
7use petgraph::algo::{is_cyclic_directed, toposort};
8use petgraph::graph::{DiGraph, NodeIndex};
9use serde::{Deserialize, Serialize};
10use std::collections::{HashMap, HashSet};
11use std::path::Path;
12
13/// Model materialization strategy in Apache Iceberg & data lakes.
14#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
15pub enum Materialization {
16    View,
17    Table,
18    IncrementalAppend,
19    IncrementalMerge,
20    /// Native zero-copy metadata table clone (zero disk byte duplication)
21    ZeroCopyClone,
22}
23
24/// Configurable model output format supported by the engine.
25#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
26pub enum OutputFormat {
27    /// Output directly as Parquet file(s)
28    #[default]
29    Parquet,
30    /// Output as JSONL (jshift zero-copy compatible)
31    Jsonl,
32    /// Output as CSV
33    Csv,
34    /// Output directly into Apache Iceberg catalog table
35    Iceberg,
36    /// Dual-write output: local Parquet files AND Apache Iceberg catalog registration
37    ParquetAndIceberg,
38    /// Native zero-copy metadata pointer clone
39    ZeroCopyClone,
40}
41
42/// DAG Layer classification following project architecture conventions.
43#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
44pub enum ModelLayer {
45    Staging,
46    Transform,
47    Mart,
48}
49
50impl ModelLayer {
51    pub fn from_name(name: &str) -> Self {
52        if name.starts_with("stg_") {
53            Self::Staging
54        } else if name.starts_with("tf_") || name.starts_with("int_") {
55            Self::Transform
56        } else if name.starts_with("dim_")
57            || name.starts_with("fact_")
58            || name.starts_with("obt_")
59            || name.starts_with("fct_")
60        {
61            Self::Mart
62        } else {
63            Self::Transform
64        }
65    }
66}
67
68/// Pipeline Model Node definition.
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct ModelNode {
71    pub name: String,
72    pub description: Option<String>,
73    pub raw_sql: String,
74    pub compiled_sql: String,
75    pub materialization: Materialization,
76    pub output_format: OutputFormat,
77    pub output_path: Option<String>,
78    pub dependencies: Vec<DependencyRef>,
79    pub layer: ModelLayer,
80    pub frontmatter: Option<StagingFrontmatter>,
81}
82
83/// Complete pipeline DAG with topological execution tiering and cycle detection.
84#[derive(Debug, Default)]
85pub struct ModelDag {
86    pub graph: DiGraph<ModelNode, ()>,
87    pub node_map: HashMap<String, NodeIndex>,
88}
89
90impl ModelDag {
91    pub fn new() -> Self {
92        Self::default()
93    }
94
95    /// Registers a raw model SQL definition into the DAG with a default Parquet format.
96    pub fn add_model(
97        &mut self,
98        name: impl Into<String>,
99        raw_sql: &str,
100        materialization: Materialization,
101        catalog_prefix: &str,
102    ) -> Result<NodeIndex> {
103        self.add_model_with_format(
104            name,
105            raw_sql,
106            materialization,
107            OutputFormat::Parquet,
108            None,
109            catalog_prefix,
110        )
111    }
112
113    /// Registers a raw model SQL definition with explicit `OutputFormat` and optional output path.
114    pub fn add_model_with_format(
115        &mut self,
116        name: impl Into<String>,
117        raw_sql: &str,
118        materialization: Materialization,
119        output_format: OutputFormat,
120        output_path: Option<String>,
121        catalog_prefix: &str,
122    ) -> Result<NodeIndex> {
123        let name = name.into();
124        let (frontmatter, pure_sql) = SqlModelParser::parse_frontmatter(raw_sql)
125            .map_err(|e| anyhow::anyhow!("model '{}': {}", name, e))?;
126        let dependencies = SqlModelParser::extract_dependencies(&pure_sql)?;
127        let compiled_sql = SqlModelParser::compile_sql(&pure_sql, catalog_prefix)?;
128        let layer = ModelLayer::from_name(&name);
129        let description = frontmatter.as_ref().and_then(|f| f.description.clone());
130
131        let node = ModelNode {
132            name: name.clone(),
133            description,
134            raw_sql: raw_sql.to_string(),
135            compiled_sql,
136            materialization,
137            output_format,
138            output_path,
139            dependencies,
140            layer,
141            frontmatter,
142        };
143
144        let idx = self.graph.add_node(node);
145        self.node_map.insert(name, idx);
146        Ok(idx)
147    }
148
149    /// Resolves dependency edges between models and validates that there are no circular dependencies or layer boundary violations.
150    pub fn build_graph(&mut self) -> Result<()> {
151        let name_map = self.node_map.clone();
152        let mut edges = Vec::new();
153
154        for (idx, node) in self.node_map.values().map(|&i| (i, &self.graph[i])) {
155            for dep in &node.dependencies {
156                if let DependencyRef::Model(dep_name) = dep {
157                    if let Some(&dep_idx) = name_map.get(dep_name) {
158                        let dep_node = &self.graph[dep_idx];
159
160                        // Enforce Layer Boundary Rule: Transform models cannot depend on Mart models
161                        if node.layer == ModelLayer::Transform && dep_node.layer == ModelLayer::Mart
162                        {
163                            bail!(
164                                "Illegal DAG Layer Boundary: Transform model '{}' (tf_) cannot depend on Mart model '{}' (dim_/fact_/obt_)",
165                                node.name,
166                                dep_node.name
167                            );
168                        }
169
170                        // Enforce Layer Boundary Rule: Staging models cannot depend on downstream models
171                        if node.layer == ModelLayer::Staging {
172                            bail!(
173                                "Illegal DAG Layer Boundary: Staging model '{}' (stg_) cannot depend on downstream model '{}'",
174                                node.name,
175                                dep_node.name
176                            );
177                        }
178
179                        edges.push((dep_idx, idx));
180                    } else {
181                        bail!(
182                            "Missing model dependency '{}' referenced by model '{}'",
183                            dep_name,
184                            node.name
185                        );
186                    }
187                }
188            }
189        }
190
191        for (from, to) in edges {
192            self.graph.add_edge(from, to, ());
193        }
194
195        if is_cyclic_directed(&self.graph) {
196            bail!("Circular dependency detected in model pipeline DAG!");
197        }
198
199        Ok(())
200    }
201
202    /// Computes a flat topologically sorted execution sequence.
203    pub fn topological_sequence(&self) -> Result<Vec<ModelNode>> {
204        let indices = toposort(&self.graph, None)
205            .map_err(|_| anyhow::anyhow!("Circular dependency found during topological sort"))?;
206
207        Ok(indices.into_iter().map(|i| self.graph[i].clone()).collect())
208    }
209
210    /// Computes parallel execution tiers where all models in a tier can be run concurrently.
211    pub fn execution_tiers(&self) -> Result<Vec<Vec<ModelNode>>> {
212        let _sorted = self.topological_sequence()?;
213        let mut in_degrees: HashMap<NodeIndex, usize> = HashMap::new();
214
215        for idx in self.graph.node_indices() {
216            in_degrees.insert(
217                idx,
218                self.graph
219                    .neighbors_directed(idx, petgraph::Direction::Incoming)
220                    .count(),
221            );
222        }
223
224        let mut current_tier: Vec<NodeIndex> = in_degrees
225            .iter()
226            .filter(|(_, &deg)| deg == 0)
227            .map(|(&idx, _)| idx)
228            .collect();
229
230        let mut tiers = Vec::new();
231        let mut visited = HashSet::new();
232
233        while !current_tier.is_empty() {
234            let tier_nodes: Vec<ModelNode> = current_tier
235                .iter()
236                .map(|&i| self.graph[i].clone())
237                .collect();
238            tiers.push(tier_nodes);
239
240            for &node_idx in &current_tier {
241                visited.insert(node_idx);
242            }
243
244            let mut next_tier = Vec::new();
245            for idx in self.graph.node_indices() {
246                if visited.contains(&idx) {
247                    continue;
248                }
249                let incoming_unvisited = self
250                    .graph
251                    .neighbors_directed(idx, petgraph::Direction::Incoming)
252                    .filter(|n| !visited.contains(n))
253                    .count();
254
255                if incoming_unvisited == 0 {
256                    next_tier.push(idx);
257                }
258            }
259
260            current_tier = next_tier;
261        }
262
263        Ok(tiers)
264    }
265
266    /// First `source('ns','table')` dependency, if any.
267    pub fn primary_source(node: &ModelNode) -> Option<(&str, &str)> {
268        node.dependencies.iter().find_map(|d| match d {
269            DependencyRef::Source {
270                source_name,
271                table_name,
272            } => Some((source_name.as_str(), table_name.as_str())),
273            _ => None,
274        })
275    }
276
277    /// Resolve registration identity for a bronze-backed model.
278    pub fn bronze_source_ident(node: &ModelNode) -> Option<(String, String)> {
279        if let Some(fm) = &node.frontmatter {
280            if let (Some(s), Some(t)) = (&fm.source_name, &fm.source_table) {
281                return Some((s.clone(), t.clone()));
282            }
283            if let Some((s, t)) = Self::primary_source(node) {
284                let source = fm.source_name.clone().unwrap_or_else(|| s.to_string());
285                let table = fm.source_table.clone().unwrap_or_else(|| t.to_string());
286                return Some((source, table));
287            }
288            if fm.has_scan_contract() {
289                // Fall back to model name under schema "bronze"
290                let table = fm.source_table.clone().unwrap_or_else(|| node.name.clone());
291                let source = fm
292                    .source_name
293                    .clone()
294                    .unwrap_or_else(|| "bronze".to_string());
295                return Some((source, table));
296            }
297        }
298        Self::primary_source(node).map(|(s, t)| (s.to_string(), t.to_string()))
299    }
300
301    /// Compile-time bronze frontmatter / scan_path validation.
302    ///
303    /// * `Off` — no filesystem checks
304    /// * `Warn` — missing paths become warnings
305    /// * `Fail` — missing paths become errors (`report.has_errors()`)
306    pub fn validate_bronze_sources(
307        &self,
308        project_dir: &Path,
309        mode: BronzeCheckMode,
310    ) -> Result<BronzeValidationReport> {
311        if mode == BronzeCheckMode::Off {
312            return Ok(BronzeValidationReport::default());
313        }
314
315        let severity = match mode {
316            BronzeCheckMode::Fail => DiagnosticSeverity::Error,
317            BronzeCheckMode::Warn => DiagnosticSeverity::Warning,
318            BronzeCheckMode::Off => unreachable!(),
319        };
320
321        let mut report = BronzeValidationReport::default();
322
323        for idx in self.graph.node_indices() {
324            let node = &self.graph[idx];
325            let has_source_dep = node
326                .dependencies
327                .iter()
328                .any(|d| matches!(d, DependencyRef::Source { .. }));
329            let fm = node.frontmatter.as_ref();
330
331            // Staging models with source() should declare a scan contract.
332            if has_source_dep && node.layer == ModelLayer::Staging {
333                let missing_scan = fm.map(|f| !f.has_scan_contract()).unwrap_or(true);
334                if missing_scan {
335                    report.diagnostics.push(BronzeDiagnostic {
336                        model: node.name.clone(),
337                        severity,
338                        code: "E_RBT_BRONZE_SCAN_PATH_MISSING",
339                        message:
340                            "staging model references source() but has no frontmatter scan_path; \
341                             add YAML frontmatter with scan_path (and source_format)"
342                                .to_string(),
343                    });
344                    continue;
345                }
346            }
347
348            let Some(fm) = fm else {
349                continue;
350            };
351
352            if !fm.has_scan_contract() {
353                continue;
354            }
355
356            let scan_path = fm.scan_path.as_deref().unwrap();
357
358            // Format resolvable?
359            if let Err(e) = fm.resolve_format() {
360                report.diagnostics.push(BronzeDiagnostic {
361                    model: node.name.clone(),
362                    severity,
363                    code: "E_RBT_BRONZE_FORMAT_UNKNOWN",
364                    message: e.to_string(),
365                });
366            }
367
368            if !scan_path_exists(project_dir, scan_path) {
369                let resolved = super::frontmatter::resolve_scan_path(project_dir, scan_path);
370                report.diagnostics.push(BronzeDiagnostic {
371                    model: node.name.clone(),
372                    severity,
373                    code: "E_RBT_BRONZE_SCAN_PATH_NOT_FOUND",
374                    message: format!(
375                        "scan_path '{}' does not exist (resolved: {})",
376                        scan_path,
377                        resolved.display()
378                    ),
379                });
380            }
381
382            // source() identity vs frontmatter overrides — soft check
383            if let (Some((dep_s, dep_t)), Some(ident)) =
384                (Self::primary_source(node), Self::bronze_source_ident(node))
385            {
386                if let Some(fm_s) = &fm.source_name {
387                    if fm_s != dep_s {
388                        report.diagnostics.push(BronzeDiagnostic {
389                            model: node.name.clone(),
390                            severity: DiagnosticSeverity::Warning,
391                            code: "W_RBT_BRONZE_SOURCE_NAME_MISMATCH",
392                            message: format!(
393                                "frontmatter source_name='{}' differs from source('{}', ...); \
394                                 registration will use '{}'.{} ",
395                                fm_s, dep_s, ident.0, ident.1
396                            ),
397                        });
398                    }
399                }
400                if let Some(fm_t) = &fm.source_table {
401                    if fm_t != dep_t {
402                        report.diagnostics.push(BronzeDiagnostic {
403                            model: node.name.clone(),
404                            severity: DiagnosticSeverity::Warning,
405                            code: "W_RBT_BRONZE_SOURCE_TABLE_MISMATCH",
406                            message: format!(
407                                "frontmatter source_table='{}' differs from source(..., '{}')",
408                                fm_t, dep_t
409                            ),
410                        });
411                    }
412                }
413            }
414        }
415
416        Ok(report)
417    }
418}
419
420#[cfg(test)]
421mod tests {
422    use super::*;
423
424    #[test]
425    fn test_dag_building_and_tiering() -> Result<()> {
426        let mut dag = ModelDag::new();
427
428        // Tier 0: Staging models (no dependencies)
429        dag.add_model_with_format(
430            "stg_users",
431            "SELECT * FROM {{ source('raw', 'users') }}",
432            Materialization::View,
433            OutputFormat::Jsonl,
434            None,
435            "db",
436        )?;
437        dag.add_model_with_format(
438            "stg_orders",
439            "SELECT * FROM {{ source('raw', 'orders') }}",
440            Materialization::View,
441            OutputFormat::ParquetAndIceberg,
442            None,
443            "db",
444        )?;
445
446        // Tier 1: Intermediate model depending on stg_users and stg_orders
447        dag.add_model(
448            "int_user_orders",
449            "SELECT * FROM {{ ref('stg_users') }} u JOIN {{ ref('stg_orders') }} o ON u.id = o.user_id",
450            Materialization::Table,
451            "db",
452        )?;
453
454        // Tier 2: Mart model depending on int_user_orders
455        dag.add_model(
456            "fct_revenue",
457            "SELECT user_id, SUM(amount) FROM {{ ref('int_user_orders') }} GROUP BY user_id",
458            Materialization::IncrementalAppend,
459            "db",
460        )?;
461
462        dag.build_graph()?;
463
464        let tiers = dag.execution_tiers()?;
465        assert_eq!(tiers.len(), 3);
466
467        // Tier 0 has 2 parallel models with formats Jsonl and ParquetAndIceberg
468        assert_eq!(tiers[0].len(), 2);
469        let tier0_formats: Vec<OutputFormat> =
470            tiers[0].iter().map(|m| m.output_format.clone()).collect();
471        assert!(tier0_formats.contains(&OutputFormat::Jsonl));
472        assert!(tier0_formats.contains(&OutputFormat::ParquetAndIceberg));
473
474        Ok(())
475    }
476
477    #[test]
478    fn test_cycle_detection() -> Result<()> {
479        let mut dag = ModelDag::new();
480
481        dag.add_model(
482            "model_a",
483            "SELECT * FROM {{ ref('model_b') }}",
484            Materialization::View,
485            "db",
486        )?;
487        dag.add_model(
488            "model_b",
489            "SELECT * FROM {{ ref('model_a') }}",
490            Materialization::View,
491            "db",
492        )?;
493
494        let res = dag.build_graph();
495        assert!(res.is_err());
496        assert!(res.unwrap_err().to_string().contains("Circular dependency"));
497
498        Ok(())
499    }
500
501    #[test]
502    fn test_layer_boundary_enforcement() -> Result<()> {
503        let mut dag = ModelDag::new();
504
505        // Mart model
506        dag.add_model(
507            "fact_orders",
508            "SELECT 1 AS order_id",
509            Materialization::Table,
510            "db",
511        )?;
512
513        // Illegal transform model attempting to pull from a mart model!
514        dag.add_model(
515            "tf_illegal_transform",
516            "SELECT * FROM {{ ref('fact_orders') }}",
517            Materialization::Table,
518            "db",
519        )?;
520
521        let res = dag.build_graph();
522        assert!(res.is_err());
523        assert!(res
524            .unwrap_err()
525            .to_string()
526            .contains("Illegal DAG Layer Boundary"));
527
528        Ok(())
529    }
530
531    #[test]
532    fn test_bronze_scan_path_validation_warn_and_fail() -> Result<()> {
533        let mut dag = ModelDag::new();
534        dag.add_model_with_format(
535            "stg_events",
536            r#"---
537source_format: jsonl
538scan_path: "lake/bronze/missing.jsonl"
539---
540SELECT * FROM {{ source('bronze', 'events') }}
541"#,
542            Materialization::Table,
543            OutputFormat::Parquet,
544            None,
545            "",
546        )?;
547        dag.build_graph()?;
548
549        let project = Path::new("/tmp/rbt_nonexistent_project_root");
550        let warn = dag.validate_bronze_sources(project, BronzeCheckMode::Warn)?;
551        assert_eq!(warn.error_count(), 0);
552        assert!(warn.warning_count() >= 1);
553        assert!(warn
554            .diagnostics
555            .iter()
556            .any(|d| d.code == "E_RBT_BRONZE_SCAN_PATH_NOT_FOUND"));
557
558        let fail = dag.validate_bronze_sources(project, BronzeCheckMode::Fail)?;
559        assert!(fail.has_errors());
560        Ok(())
561    }
562}