Skip to main content

rbt/core/
frontmatter.rs

1//! Staging SQL frontmatter: bronze scan contract and compile-time path checks.
2
3use crate::core::run_scope::OnMissing;
4use anyhow::{bail, Context, Result};
5use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
6use serde::{Deserialize, Serialize};
7use std::fmt;
8use std::path::{Path, PathBuf};
9use std::sync::Arc;
10
11/// How `rbt compile` treats missing/unresolvable bronze `scan_path` entries.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
13#[serde(rename_all = "snake_case")]
14pub enum BronzeCheckMode {
15    /// Skip filesystem checks (DAG structure only).
16    Off,
17    /// Emit warnings; compile still succeeds (default for `compile`).
18    #[default]
19    Warn,
20    /// Missing or invalid bronze sources fail compile.
21    Fail,
22}
23
24impl BronzeCheckMode {
25    pub fn as_str(self) -> &'static str {
26        match self {
27            Self::Off => "off",
28            Self::Warn => "warn",
29            Self::Fail => "fail",
30        }
31    }
32}
33
34impl fmt::Display for BronzeCheckMode {
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        f.write_str(self.as_str())
37    }
38}
39
40/// Supported bronze file formats for staging lake scans.
41#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
42#[serde(rename_all = "snake_case")]
43pub enum SourceFormat {
44    /// Newline-delimited JSON (also accepts alias `ndjson`).
45    #[serde(alias = "ndjson")]
46    Jsonl,
47    /// Single JSON document or JSON array of objects.
48    Json,
49    Parquet,
50    Csv,
51    /// Arrow IPC file (random-access footer).
52    #[serde(alias = "arrow", alias = "arrow_file", alias = "ipc")]
53    ArrowIpc,
54    /// Arrow IPC stream (append-friendly / WAL-style).
55    #[serde(alias = "arrow_stream", alias = "ipc_stream")]
56    ArrowIpcStream,
57    /// Line-oriented application / server logs.
58    Log,
59    /// Line-oriented text (llms.txt, docs dumps, structured line files).
60    Txt,
61    /// TOML tables / array-of-tables as rows.
62    Toml,
63    /// Length-delimited or whole-file protobuf blobs (opaque bronze).
64    ///
65    /// Each file becomes one row: `_source_path` (Utf8) + `payload` (Binary).
66    /// Typed decode of domain messages is a later step (Rust models / schema registry).
67    #[serde(alias = "pb", alias = "proto")]
68    Protobuf,
69}
70
71impl SourceFormat {
72    pub fn as_str(self) -> &'static str {
73        match self {
74            Self::Jsonl => "jsonl",
75            Self::Json => "json",
76            Self::Parquet => "parquet",
77            Self::Csv => "csv",
78            Self::ArrowIpc => "arrow_ipc",
79            Self::ArrowIpcStream => "arrow_ipc_stream",
80            Self::Log => "log",
81            Self::Txt => "txt",
82            Self::Toml => "toml",
83            Self::Protobuf => "protobuf",
84        }
85    }
86
87    /// Prefer DataFusion listing / external table registration when true.
88    pub fn prefers_datafusion_listing(self) -> bool {
89        matches!(self, Self::Parquet | Self::Csv | Self::Json | Self::Jsonl)
90    }
91
92    /// Infer format from a file extension (without the dot).
93    pub fn from_extension(ext: &str) -> Option<Self> {
94        match ext.to_ascii_lowercase().as_str() {
95            "jsonl" | "ndjson" => Some(Self::Jsonl),
96            "json" => Some(Self::Json),
97            "parquet" | "pq" => Some(Self::Parquet),
98            "csv" | "tsv" => Some(Self::Csv),
99            "arrow" | "arrows" | "ipc" | "feather" => Some(Self::ArrowIpc),
100            "arrows_stream" | "ipc_stream" => Some(Self::ArrowIpcStream),
101            "log" => Some(Self::Log),
102            "txt" | "text" | "md" => Some(Self::Txt),
103            "toml" => Some(Self::Toml),
104            "pb" | "protobuf" | "protobin" => Some(Self::Protobuf),
105            _ => None,
106        }
107    }
108
109    /// Parse free-form frontmatter / CLI format strings.
110    pub fn parse(s: &str) -> Result<Self> {
111        let key = s.trim().to_ascii_lowercase().replace('-', "_");
112        match key.as_str() {
113            "jsonl" | "ndjson" | "json_lines" => Ok(Self::Jsonl),
114            "json" => Ok(Self::Json),
115            "parquet" | "pq" => Ok(Self::Parquet),
116            "csv" | "tsv" => Ok(Self::Csv),
117            "arrow_ipc" | "arrow" | "arrow_file" | "ipc" | "feather" => Ok(Self::ArrowIpc),
118            "arrow_ipc_stream" | "arrow_stream" | "ipc_stream" => Ok(Self::ArrowIpcStream),
119            "log" => Ok(Self::Log),
120            "txt" | "text" => Ok(Self::Txt),
121            "toml" => Ok(Self::Toml),
122            "protobuf" | "pb" | "proto" | "protobin" => Ok(Self::Protobuf),
123            other => bail!(
124                "Unknown source_format '{}'. Expected one of: jsonl, json, parquet, csv, arrow_ipc, arrow_ipc_stream, log, txt, toml, protobuf",
125                other
126            ),
127        }
128    }
129}
130
131impl fmt::Display for SourceFormat {
132    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
133        f.write_str(self.as_str())
134    }
135}
136
137/// Declared data-quality tests for a model (run after materialization when present).
138#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
139pub struct ModelTests {
140    /// Columns that must have zero nulls.
141    #[serde(default)]
142    pub not_null: Option<Vec<String>>,
143    /// Single column unique, or multi-column composite unique when len > 1.
144    #[serde(default)]
145    pub unique: Option<Vec<String>>,
146    /// Map of column → allowed string values.
147    #[serde(default)]
148    pub accepted_values: Option<std::collections::HashMap<String, Vec<String>>>,
149    /// When true (default), failed tests abort `rbt run` for that model.
150    #[serde(default)]
151    pub fail_on_error: Option<bool>,
152}
153
154impl ModelTests {
155    pub fn is_empty(&self) -> bool {
156        self.not_null.as_ref().map(|v| v.is_empty()).unwrap_or(true)
157            && self.unique.as_ref().map(|v| v.is_empty()).unwrap_or(true)
158            && self
159                .accepted_values
160                .as_ref()
161                .map(|m| m.is_empty())
162                .unwrap_or(true)
163    }
164
165    pub fn should_fail_on_error(&self) -> bool {
166        self.fail_on_error.unwrap_or(true)
167    }
168}
169
170/// Per-column documentation for humans and AI agents.
171///
172/// * `description` — short label (1–2 lines)
173/// * `context` — longer intent, units, lineage, caveats (agent-oriented)
174#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
175pub struct ColumnMeta {
176    #[serde(default)]
177    pub description: Option<String>,
178    #[serde(default)]
179    pub context: Option<String>,
180    /// Optional logical type hint (`utf8`, `int64`, `float64`, `timestamp`, …).
181    #[serde(default)]
182    pub dtype: Option<String>,
183    /// Optional unit (`USD`, `shares`, `ratio`, `ns_epoch`, …).
184    #[serde(default)]
185    pub unit: Option<String>,
186}
187
188/// YAML frontmatter embedded in model SQL files (`---` … `---`).
189///
190/// Used on staging, transforms, and marts. Scan-related fields only apply when
191/// a bronze scan contract (`scan_path`) is present.
192#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
193pub struct StagingFrontmatter {
194    /// Human-readable model purpose (docs + future catalog).
195    #[serde(default)]
196    pub description: Option<String>,
197    /// Longer model-level context for AI agents (why it exists, consumers, caveats).
198    #[serde(default)]
199    pub context: Option<String>,
200    /// Column-level description + context map (name → meta).
201    #[serde(default)]
202    pub columns: Option<std::collections::BTreeMap<String, ColumnMeta>>,
203
204    /// Explicit bronze format. If omitted, inferred from `scan_path` extension.
205    #[serde(default)]
206    pub source_format: Option<SourceFormat>,
207    /// File or directory to scan (project-relative, absolute, or `$root/...` template).
208    #[serde(default)]
209    pub scan_path: Option<String>,
210    /// Filename / relative-path glob(s) under `scan_path` (OR semantics).
211    ///
212    /// Examples: `crawlplan.parquet`, `**/raw_snoop/crawlplan.parquet`, `*.jsonl`.
213    /// Empty / omitted = all files matching `source_format`.
214    /// Accepts a single string or a YAML list.
215    ///
216    /// **Pushdown note:** any non-empty `path_glob` forces the scan→MemTable bronze path
217    /// (DataFusion directory listing / predicate pushdown is **not** used for that source),
218    /// because listing providers cannot apply rbt's filename globs or hive path injection.
219    #[serde(default, deserialize_with = "deserialize_string_or_vec")]
220    pub path_glob: Option<Vec<String>>,
221    /// Optional hive-style partition keys (path injection + future pruning).
222    #[serde(default)]
223    pub partition_by: Option<Vec<String>>,
224    /// jshift / projection field paths for selective JSON(L) extract.
225    #[serde(default)]
226    pub paths: Option<Vec<String>>,
227    /// Override catalog/schema name for registration (default: first `source()` name).
228    #[serde(default)]
229    pub source_name: Option<String>,
230    /// Override table name for registration (default: first `source()` table).
231    #[serde(default)]
232    pub source_table: Option<String>,
233    /// TOML: key of the array-of-tables to expand into rows (default: auto-detect).
234    #[serde(default)]
235    pub toml_rows_key: Option<String>,
236    /// When true, use scan→MemTable path even for formats that support DF listing.
237    #[serde(default)]
238    pub force_scan: Option<bool>,
239    /// Only scan hive-partitioned files whose path segments match these values.
240    /// Example: `{ timeframe: "1m" }` keeps `.../timeframe=1m/...` and skips `timeframe=1d`.
241    #[serde(default)]
242    pub require_partitions: Option<std::collections::HashMap<String, String>>,
243    /// Inject `_source_path` (Utf8) with the absolute file path for each row.
244    /// Enables "latest chunk wins" dedupe via `ORDER BY _source_path DESC`.
245    #[serde(default)]
246    pub inject_source_path: Option<bool>,
247    /// When scan root is missing or filters match no files: `error` (default) | `empty`.
248    ///
249    /// `empty` registers a zero-row table with a declared schema from `columns.*.dtype`
250    /// (plus `partition_by` keys as Utf8). Required for partial multi-artifact bronze.
251    #[serde(default)]
252    pub on_missing: Option<OnMissing>,
253    /// Silver stage policy hint (docs + future engine): `full_refresh` | `latest_only` |
254    /// `append` | `mirror_bronze`. Does not change SQL by itself — authors implement
255    /// semantics in the model; rbt may use this for materialization defaults later.
256    #[serde(default)]
257    pub stage_mode: Option<String>,
258
259    /// Logical grain of the model (e.g. `[symbol, timestamp_ns]`).
260    #[serde(default)]
261    pub grain: Option<Vec<String>>,
262    /// Primary uniqueness contract (usually same as grain for staging facts).
263    #[serde(default)]
264    pub unique_key: Option<Vec<String>>,
265    /// Free-form tags for selection / docs.
266    #[serde(default)]
267    pub tags: Option<Vec<String>>,
268    /// Materialization hint: `table` | `view` | `incremental_append` (engine may ignore for now).
269    #[serde(default)]
270    pub materialization: Option<String>,
271    /// Post-materialization assertions.
272    #[serde(default)]
273    pub tests: Option<ModelTests>,
274    /// Opaque metadata for tools and agents (strings, lists, nested maps OK).
275    #[serde(default)]
276    pub meta: Option<std::collections::BTreeMap<String, serde_yaml::Value>>,
277}
278
279impl StagingFrontmatter {
280    /// Resolve format from explicit field or path extension.
281    pub fn resolve_format(&self) -> Result<SourceFormat> {
282        if let Some(fmt) = self.source_format {
283            return Ok(fmt);
284        }
285        let path = self
286            .scan_path
287            .as_deref()
288            .context("frontmatter missing both source_format and scan_path")?;
289        // Strip globs for extension sniffing: `foo/*.jsonl` → look at last segment
290        let candidate = path.rsplit('/').next().unwrap_or(path);
291        let candidate = candidate.trim_matches(|c| c == '*' || c == '?');
292        if let Some(ext) = Path::new(candidate).extension().and_then(|e| e.to_str()) {
293            if let Some(fmt) = SourceFormat::from_extension(ext) {
294                return Ok(fmt);
295            }
296        }
297        // Directory paths: no extension — require explicit format
298        bail!(
299            "Cannot infer source_format from scan_path '{}'; set source_format explicitly",
300            path
301        );
302    }
303
304    pub fn has_scan_contract(&self) -> bool {
305        self.scan_path
306            .as_ref()
307            .map(|s| !s.trim().is_empty())
308            .unwrap_or(false)
309    }
310
311    pub fn on_missing_policy(&self) -> OnMissing {
312        self.on_missing.unwrap_or(OnMissing::Error)
313    }
314
315    /// Build Arrow schema for empty bronze frames (`on_missing: empty`).
316    ///
317    /// Fields: declared `columns` with `dtype`, then any `partition_by` keys not already
318    /// present (Utf8), then optional `_source_path`.
319    pub fn empty_frame_schema(&self) -> Result<SchemaRef> {
320        let mut fields: Vec<Field> = Vec::new();
321        let mut seen = std::collections::HashSet::new();
322
323        if let Some(cols) = &self.columns {
324            for (name, meta) in cols {
325                let dtype = meta
326                    .dtype
327                    .as_deref()
328                    .with_context(|| {
329                        format!(
330                            "E_RBT_EMPTY_SCHEMA: column '{name}' needs dtype: for on_missing: empty \
331                             (e.g. utf8, int64, float64, bool, binary, timestamp)"
332                        )
333                    })?;
334                let dt = parse_logical_dtype(dtype).with_context(|| {
335                    format!("E_RBT_EMPTY_SCHEMA: column '{name}' dtype '{dtype}'")
336                })?;
337                fields.push(Field::new(name, dt, true));
338                seen.insert(name.clone());
339            }
340        }
341
342        if let Some(parts) = &self.partition_by {
343            for p in parts {
344                if seen.insert(p.clone()) {
345                    fields.push(Field::new(p, DataType::Utf8, true));
346                }
347            }
348        }
349
350        if self.inject_source_path.unwrap_or(false) && seen.insert("_source_path".into()) {
351            fields.push(Field::new("_source_path", DataType::Utf8, true));
352        }
353
354        if fields.is_empty() {
355            bail!(
356                "E_RBT_EMPTY_SCHEMA: on_missing: empty requires columns with dtype \
357                 and/or partition_by (model scan contract has no schema fields)"
358            );
359        }
360        Ok(Arc::new(Schema::new(fields)))
361    }
362}
363
364/// Parse logical dtype strings used in frontmatter `columns.*.dtype`.
365pub fn parse_logical_dtype(s: &str) -> Result<DataType> {
366    let key = s.trim().to_ascii_lowercase().replace('-', "_");
367    Ok(match key.as_str() {
368        "utf8" | "string" | "str" | "varchar" | "text" => DataType::Utf8,
369        "int64" | "long" | "bigint" | "i64" => DataType::Int64,
370        "int32" | "int" | "i32" => DataType::Int32,
371        "int16" | "smallint" | "i16" => DataType::Int16,
372        "int8" | "tinyint" | "i8" => DataType::Int8,
373        "uint64" | "u64" => DataType::UInt64,
374        "uint32" | "u32" => DataType::UInt32,
375        "float64" | "double" | "f64" => DataType::Float64,
376        "float32" | "float" | "f32" => DataType::Float32,
377        "bool" | "boolean" => DataType::Boolean,
378        "binary" | "bytes" | "blob" => DataType::Binary,
379        "date" | "date32" => DataType::Date32,
380        "timestamp" | "timestamp_us" | "timestamptz" => {
381            DataType::Timestamp(TimeUnit::Microsecond, None)
382        }
383        "timestamp_ms" => DataType::Timestamp(TimeUnit::Millisecond, None),
384        "timestamp_ns" => DataType::Timestamp(TimeUnit::Nanosecond, None),
385        "timestamp_s" => DataType::Timestamp(TimeUnit::Second, None),
386        other => bail!(
387            "unknown dtype '{other}' (expected utf8|int64|int32|float64|bool|binary|date|timestamp…)"
388        ),
389    })
390}
391
392/// Severity of a bronze compile diagnostic.
393#[derive(Debug, Clone, Copy, PartialEq, Eq)]
394pub enum DiagnosticSeverity {
395    Warning,
396    Error,
397}
398
399/// One compile-time finding about bronze frontmatter / scan paths.
400#[derive(Debug, Clone, PartialEq, Eq)]
401pub struct BronzeDiagnostic {
402    pub model: String,
403    pub severity: DiagnosticSeverity,
404    pub code: &'static str,
405    pub message: String,
406}
407
408impl fmt::Display for BronzeDiagnostic {
409    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
410        let level = match self.severity {
411            DiagnosticSeverity::Warning => "warning",
412            DiagnosticSeverity::Error => "error",
413        };
414        write!(
415            f,
416            "{}[{}] model={}: {}",
417            level, self.code, self.model, self.message
418        )
419    }
420}
421
422/// Result of bronze path validation during compile.
423#[derive(Debug, Clone, Default)]
424pub struct BronzeValidationReport {
425    pub diagnostics: Vec<BronzeDiagnostic>,
426}
427
428impl BronzeValidationReport {
429    pub fn warning_count(&self) -> usize {
430        self.diagnostics
431            .iter()
432            .filter(|d| d.severity == DiagnosticSeverity::Warning)
433            .count()
434    }
435
436    pub fn error_count(&self) -> usize {
437        self.diagnostics
438            .iter()
439            .filter(|d| d.severity == DiagnosticSeverity::Error)
440            .count()
441    }
442
443    pub fn has_errors(&self) -> bool {
444        self.error_count() > 0
445    }
446}
447
448/// Resolve `scan_path` against the project root (no named roots). Prefer
449/// [`crate::core::paths::resolve_project_path`] when `roots:` are in play.
450pub fn resolve_scan_path(project_dir: &Path, scan_path: &str) -> PathBuf {
451    crate::core::paths::resolve_project_path(project_dir, scan_path, &Default::default())
452        .unwrap_or_else(|_| project_dir.to_path_buf())
453}
454
455pub use crate::core::paths::is_remote_uri;
456
457/// Deserialize either a single string or a sequence into `Option<Vec<String>>`.
458fn deserialize_string_or_vec<'de, D>(
459    deserializer: D,
460) -> std::result::Result<Option<Vec<String>>, D::Error>
461where
462    D: serde::Deserializer<'de>,
463{
464    use serde::de::{self, SeqAccess, Visitor};
465    use std::fmt;
466
467    struct StringOrVec;
468    impl<'de> Visitor<'de> for StringOrVec {
469        type Value = Option<Vec<String>>;
470
471        fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
472            f.write_str("a string or list of strings")
473        }
474
475        fn visit_none<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
476            Ok(None)
477        }
478
479        fn visit_unit<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
480            Ok(None)
481        }
482
483        fn visit_str<E: de::Error>(self, v: &str) -> std::result::Result<Self::Value, E> {
484            Ok(Some(vec![v.to_string()]))
485        }
486
487        fn visit_string<E: de::Error>(self, v: String) -> std::result::Result<Self::Value, E> {
488            Ok(Some(vec![v]))
489        }
490
491        fn visit_seq<A: SeqAccess<'de>>(
492            self,
493            mut seq: A,
494        ) -> std::result::Result<Self::Value, A::Error> {
495            let mut out = Vec::new();
496            while let Some(s) = seq.next_element::<String>()? {
497                out.push(s);
498            }
499            Ok(Some(out))
500        }
501    }
502
503    deserializer.deserialize_any(StringOrVec)
504}
505
506/// Whether a resolved local scan path currently exists (file or directory).
507/// Remote URIs are treated as "exists" for compile (runtime / object-store later).
508pub fn scan_path_exists(project_dir: &Path, scan_path: &str) -> bool {
509    scan_path_exists_with_roots(project_dir, scan_path, &std::collections::HashMap::new())
510}
511
512/// Like [`scan_path_exists`] but expands `$root` templates from project config.
513pub fn scan_path_exists_with_roots(
514    project_dir: &Path,
515    scan_path: &str,
516    roots: &std::collections::HashMap<String, String>,
517) -> bool {
518    if is_remote_uri(scan_path.trim()) {
519        return true;
520    }
521    let Ok(resolved) = crate::core::paths::resolve_project_path(project_dir, scan_path, roots)
522    else {
523        return false;
524    };
525    // Support simple trailing globs: `dir/*.jsonl` → check parent dir
526    let check = strip_simple_glob(&resolved);
527    check.exists()
528}
529
530fn strip_simple_glob(path: &Path) -> PathBuf {
531    let s = path.to_string_lossy();
532    if s.contains('*') || s.contains('?') {
533        if let Some(parent) = path.parent() {
534            return parent.to_path_buf();
535        }
536    }
537    path.to_path_buf()
538}
539
540#[cfg(test)]
541mod tests {
542    use super::*;
543
544    #[test]
545    fn format_from_extension_and_parse() {
546        assert_eq!(
547            SourceFormat::from_extension("jsonl"),
548            Some(SourceFormat::Jsonl)
549        );
550        assert_eq!(
551            SourceFormat::from_extension("toml"),
552            Some(SourceFormat::Toml)
553        );
554        assert_eq!(SourceFormat::from_extension("log"), Some(SourceFormat::Log));
555        assert_eq!(
556            SourceFormat::parse("arrow-ipc").unwrap(),
557            SourceFormat::ArrowIpc
558        );
559        assert_eq!(SourceFormat::parse("ndjson").unwrap(), SourceFormat::Jsonl);
560    }
561
562    #[test]
563    fn resolve_format_from_path() {
564        let fm = StagingFrontmatter {
565            scan_path: Some("lake/bronze/raw.jsonl".into()),
566            ..Default::default()
567        };
568        assert_eq!(fm.resolve_format().unwrap(), SourceFormat::Jsonl);
569    }
570
571    #[test]
572    fn remote_uri_exists_for_compile() {
573        assert!(scan_path_exists(
574            Path::new("/tmp"),
575            "s3://bucket/bronze/x.jsonl"
576        ));
577    }
578}