rbt_engine/

lib.rs

//! `rbt-engine`: Apache DataFusion query engine integration, bronze registration, and DAG execution.

pub mod bronze;

use anyhow::{Context, Result};
use datafusion::datasource::MemTable;
use datafusion::execution::context::SessionContext;
use datafusion::physical_plan::SendableRecordBatchStream;
use iceberg::Catalog;
use iceberg_datafusion::IcebergCatalogProvider;
use rbt_core::dag::{ModelDag, OutputFormat};
use rbt_materializer::MultiFormatWriter;
use rbt_testing::{assertions_from_model_tests, RecordBatchValidator};
use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::Arc;

pub use bronze::{
    register_bronze_for_model, register_bronze_sources_for_dag, BronzeRegistrationMode,
    BronzeSourceMeta, BronzeTableProvider,
};

/// Execution metric summary for a executed model DAG.
#[derive(Debug, Clone)]
pub struct DagExecutionSummary {
    pub models_executed: usize,
    pub total_rows_produced: usize,
    pub bronze_sources_registered: usize,
}

/// Fluent Builder for configuring and launching `TransformationEngine` instances.
#[derive(Default)]
pub struct RbtEngineBuilder {
    catalogs: Vec<(String, Arc<dyn Catalog>)>,
}

impl RbtEngineBuilder {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_catalog(mut self, name: impl Into<String>, catalog: Arc<dyn Catalog>) -> Self {
        self.catalogs.push((name.into(), catalog));
        self
    }

    pub async fn build(self) -> Result<TransformationEngine> {
        let engine = TransformationEngine::new();
        for (name, cat) in self.catalogs {
            engine.register_iceberg_catalog(&name, cat).await?;
        }
        Ok(engine)
    }
}

pub struct TransformationEngine {
    pub ctx: SessionContext,
}

impl Default for TransformationEngine {
    fn default() -> Self {
        Self::new()
    }
}

impl TransformationEngine {
    pub fn new() -> Self {
        Self {
            ctx: SessionContext::new(),
        }
    }

    /// Registers an Apache Iceberg catalog directly into the DataFusion query context.
    pub async fn register_iceberg_catalog(
        &self,
        catalog_name: &str,
        catalog: Arc<dyn Catalog>,
    ) -> Result<()> {
        tracing::info!(
            "Registering Iceberg catalog '{}' into DataFusion SessionContext",
            catalog_name
        );
        let provider = IcebergCatalogProvider::try_new(catalog).await?;
        self.ctx.register_catalog(catalog_name, Arc::new(provider));
        Ok(())
    }

    /// Executes a SQL transform query against registered tables.
    pub async fn execute_sql(&self, sql: &str) -> Result<SendableRecordBatchStream> {
        tracing::info!(
            "Executing SQL transform via Apache DataFusion engine: {}",
            sql
        );
        let df = self.ctx.sql(sql).await?;
        let stream = df.execute_stream().await?;
        Ok(stream)
    }

    /// Executes a full pipeline DAG tier by tier.
    ///
    /// Before any model SQL runs, bronze sources declared in staging frontmatter are
    /// registered via [`register_bronze_sources_for_dag`].
    pub async fn execute_dag(
        &self,
        dag: &ModelDag,
        project_dir: impl AsRef<Path>,
        output_dir: impl AsRef<Path>,
    ) -> Result<DagExecutionSummary> {
        let project_dir = project_dir.as_ref();
        let output_base = output_dir.as_ref();
        tokio::fs::create_dir_all(output_base).await?;

        let mut registered = HashSet::new();
        let bronze_sources_registered =
            register_bronze_sources_for_dag(&self.ctx, dag, project_dir, &mut registered)
                .await
                .context("frontmatter-driven bronze registration failed")?;

        let tiers = dag.execution_tiers()?;
        let mut models_executed = 0;
        let mut total_rows_produced = 0;

        for (tier_idx, tier) in tiers.iter().enumerate() {
            tracing::info!(
                "Executing DAG Tier {} with {} parallel models",
                tier_idx,
                tier.len()
            );

            for model in tier {
                tracing::info!("Executing model '{}'...", model.name);

                // Late-bind: if this model carries frontmatter not registered yet
                register_bronze_for_model(&self.ctx, model, project_dir, &mut registered)
                    .await?;

                let df = self
                    .ctx
                    .sql(&model.compiled_sql)
                    .await
                    .with_context(|| {
                        format!(
                            "SQL execution failed for model '{}' (compiled: {})",
                            model.name, model.compiled_sql
                        )
                    })?;
                let batches = df.collect().await.with_context(|| {
                    format!("collect failed for model '{}'", model.name)
                })?;
                let row_count: usize = batches.iter().map(|b| b.num_rows()).sum();

                let dest_path = model.output_path.as_ref().map(PathBuf::from).unwrap_or_else(
                    || match model.output_format {
                        OutputFormat::Iceberg => output_base.join(&model.name),
                        OutputFormat::Jsonl => {
                            output_base.join(format!("{}.jsonl", model.name))
                        }
                        OutputFormat::Csv => output_base.join(format!("{}.csv", model.name)),
                        _ => output_base.join(format!("{}.parquet", model.name)),
                    },
                );

                if let Some(parent) = dest_path.parent() {
                    std::fs::create_dir_all(parent)?;
                }

                MultiFormatWriter::write_batches(&batches, &model.output_format, &dest_path)?;

                // Frontmatter-declared tests (staging grain / not_null / unique_key).
                if let Some(fm) = model.frontmatter.as_ref() {
                    if let Some(tests) = fm.tests.as_ref() {
                        if !tests.is_empty() {
                            let unique = tests
                                .unique
                                .clone()
                                .or_else(|| fm.unique_key.clone())
                                .or_else(|| fm.grain.clone());
                            let assertions = assertions_from_model_tests(
                                tests.not_null.as_deref(),
                                unique.as_deref(),
                                tests.accepted_values.as_ref(),
                            );
                            if !assertions.is_empty() {
                                let result =
                                    RecordBatchValidator::validate_batches(&batches, &assertions);
                                if result.failed_assertions > 0 {
                                    let msg = format!(
                                        "model '{}' failed {} test(s): {}",
                                        model.name,
                                        result.failed_assertions,
                                        result.errors.join("; ")
                                    );
                                    if tests.should_fail_on_error() {
                                        anyhow::bail!(msg);
                                    }
                                    tracing::warn!("{}", msg);
                                } else {
                                    tracing::info!(
                                        "model '{}': {} assertion(s) passed ({} rows)",
                                        model.name,
                                        result.passed_assertions,
                                        result.total_rows
                                    );
                                }
                            }
                        }
                    } else if let Some(uk) = fm
                        .unique_key
                        .as_ref()
                        .or(fm.grain.as_ref())
                        .filter(|v| !v.is_empty())
                    {
                        // Implicit unique_key/grain check when no tests: block declared
                        let assertions = assertions_from_model_tests(None, Some(uk.as_slice()), None);
                        let result =
                            RecordBatchValidator::validate_batches(&batches, &assertions);
                        if result.failed_assertions > 0 {
                            anyhow::bail!(
                                "model '{}' grain/unique_key violated: {}",
                                model.name,
                                result.errors.join("; ")
                            );
                        }
                    }
                }

                if !batches.is_empty() {
                    let schema = batches[0].schema();
                    let mem_table = MemTable::try_new(schema, vec![batches.clone()])?;
                    // Allow re-runs in tests
                    let _ = self.ctx.deregister_table(model.name.as_str());
                    self.ctx
                        .register_table(model.name.as_str(), Arc::new(mem_table))?;
                }

                models_executed += 1;
                total_rows_produced += row_count;
            }
        }

        Ok(DagExecutionSummary {
            models_executed,
            total_rows_produced,
            bronze_sources_registered,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rbt_core::dag::{Materialization, ModelDag, OutputFormat};

    #[tokio::test]
    async fn test_engine_initialization() -> Result<()> {
        let engine = TransformationEngine::new();
        let df = engine.ctx.sql("SELECT 1 AS col").await?;
        let batches = df.collect().await?;
        assert_eq!(batches.len(), 1);
        assert_eq!(batches[0].num_rows(), 1);
        Ok(())
    }

    #[tokio::test]
    async fn test_dag_execution_multi_format() -> Result<()> {
        let temp_dir = tempfile::tempdir()?;
        let engine = TransformationEngine::new();

        let mut dag = ModelDag::new();
        dag.add_model_with_format(
            "users",
            "SELECT 1 AS id, 'Alice' AS name",
            Materialization::Table,
            OutputFormat::Jsonl,
            None,
            "",
        )?;
        dag.add_model_with_format(
            "active_users",
            "SELECT * FROM {{ ref('users') }} WHERE id = 1",
            Materialization::Table,
            OutputFormat::Parquet,
            None,
            "",
        )?;
        dag.build_graph()?;

        let summary = engine
            .execute_dag(&dag, temp_dir.path(), temp_dir.path())
            .await?;
        assert_eq!(summary.models_executed, 2);
        assert_eq!(summary.total_rows_produced, 2);
        assert!(temp_dir.path().join("users.jsonl").exists());
        assert!(temp_dir.path().join("active_users.parquet").exists());
        Ok(())
    }

    #[tokio::test]
    async fn test_frontmatter_bronze_end_to_end() -> Result<()> {
        let temp = tempfile::tempdir()?;
        let bronze_dir = temp.path().join("lake/bronze");
        std::fs::create_dir_all(&bronze_dir)?;
        std::fs::write(
            bronze_dir.join("raw_stock_trades.jsonl"),
            r#"{"ticker":"NVDA","timestamp":"2026-07-24T09:30:01Z","price":125.5,"volume":100}
{"ticker":"AAPL","timestamp":"2026-07-24T09:30:05Z","price":190.0,"volume":50}
"#,
        )?;

        let sql = r#"---
source_format: jsonl
scan_path: "lake/bronze/raw_stock_trades.jsonl"
---
SELECT ticker, price, volume FROM {{ source('bronze', 'raw_stock_trades') }}
"#;

        let mut dag = ModelDag::new();
        dag.add_model_with_format(
            "stg_stock_trades",
            sql,
            Materialization::Table,
            OutputFormat::Parquet,
            Some(
                temp.path()
                    .join("lake/silver/stg_stock_trades.parquet")
                    .to_string_lossy()
                    .into(),
            ),
            "",
        )?;
        dag.build_graph()?;

        let engine = TransformationEngine::new();
        let summary = engine.execute_dag(&dag, temp.path(), temp.path().join("out")).await?;
        assert_eq!(summary.bronze_sources_registered, 1);
        assert_eq!(summary.models_executed, 1);
        assert_eq!(summary.total_rows_produced, 2);
        assert!(temp
            .path()
            .join("lake/silver/stg_stock_trades.parquet")
            .exists());
        Ok(())
    }
}