Skip to main content

rbt/core/
frontmatter.rs

1//! Staging SQL frontmatter: bronze scan contract and compile-time path checks.
2
3use anyhow::{bail, Context, Result};
4use serde::{Deserialize, Serialize};
5use std::fmt;
6use std::path::{Path, PathBuf};
7
8/// How `rbt compile` treats missing/unresolvable bronze `scan_path` entries.
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
10#[serde(rename_all = "snake_case")]
11pub enum BronzeCheckMode {
12    /// Skip filesystem checks (DAG structure only).
13    Off,
14    /// Emit warnings; compile still succeeds (default for `compile`).
15    #[default]
16    Warn,
17    /// Missing or invalid bronze sources fail compile.
18    Fail,
19}
20
21impl BronzeCheckMode {
22    pub fn as_str(self) -> &'static str {
23        match self {
24            Self::Off => "off",
25            Self::Warn => "warn",
26            Self::Fail => "fail",
27        }
28    }
29}
30
31impl fmt::Display for BronzeCheckMode {
32    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33        f.write_str(self.as_str())
34    }
35}
36
37/// Supported bronze file formats for staging lake scans.
38#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
39#[serde(rename_all = "snake_case")]
40pub enum SourceFormat {
41    /// Newline-delimited JSON (also accepts alias `ndjson`).
42    #[serde(alias = "ndjson")]
43    Jsonl,
44    /// Single JSON document or JSON array of objects.
45    Json,
46    Parquet,
47    Csv,
48    /// Arrow IPC file (random-access footer).
49    #[serde(alias = "arrow", alias = "arrow_file", alias = "ipc")]
50    ArrowIpc,
51    /// Arrow IPC stream (append-friendly / WAL-style).
52    #[serde(alias = "arrow_stream", alias = "ipc_stream")]
53    ArrowIpcStream,
54    /// Line-oriented application / server logs.
55    Log,
56    /// Line-oriented text (llms.txt, docs dumps, structured line files).
57    Txt,
58    /// TOML tables / array-of-tables as rows.
59    Toml,
60}
61
62impl SourceFormat {
63    pub fn as_str(self) -> &'static str {
64        match self {
65            Self::Jsonl => "jsonl",
66            Self::Json => "json",
67            Self::Parquet => "parquet",
68            Self::Csv => "csv",
69            Self::ArrowIpc => "arrow_ipc",
70            Self::ArrowIpcStream => "arrow_ipc_stream",
71            Self::Log => "log",
72            Self::Txt => "txt",
73            Self::Toml => "toml",
74        }
75    }
76
77    /// Prefer DataFusion listing / external table registration when true.
78    pub fn prefers_datafusion_listing(self) -> bool {
79        matches!(self, Self::Parquet | Self::Csv | Self::Json | Self::Jsonl)
80    }
81
82    /// Infer format from a file extension (without the dot).
83    pub fn from_extension(ext: &str) -> Option<Self> {
84        match ext.to_ascii_lowercase().as_str() {
85            "jsonl" | "ndjson" => Some(Self::Jsonl),
86            "json" => Some(Self::Json),
87            "parquet" | "pq" => Some(Self::Parquet),
88            "csv" | "tsv" => Some(Self::Csv),
89            "arrow" | "arrows" | "ipc" | "feather" => Some(Self::ArrowIpc),
90            "arrows_stream" | "ipc_stream" => Some(Self::ArrowIpcStream),
91            "log" => Some(Self::Log),
92            "txt" | "text" | "md" => Some(Self::Txt),
93            "toml" => Some(Self::Toml),
94            _ => None,
95        }
96    }
97
98    /// Parse free-form frontmatter / CLI format strings.
99    pub fn parse(s: &str) -> Result<Self> {
100        let key = s.trim().to_ascii_lowercase().replace('-', "_");
101        match key.as_str() {
102            "jsonl" | "ndjson" | "json_lines" => Ok(Self::Jsonl),
103            "json" => Ok(Self::Json),
104            "parquet" | "pq" => Ok(Self::Parquet),
105            "csv" | "tsv" => Ok(Self::Csv),
106            "arrow_ipc" | "arrow" | "arrow_file" | "ipc" | "feather" => Ok(Self::ArrowIpc),
107            "arrow_ipc_stream" | "arrow_stream" | "ipc_stream" => Ok(Self::ArrowIpcStream),
108            "log" => Ok(Self::Log),
109            "txt" | "text" => Ok(Self::Txt),
110            "toml" => Ok(Self::Toml),
111            other => bail!(
112                "Unknown source_format '{}'. Expected one of: jsonl, json, parquet, csv, arrow_ipc, arrow_ipc_stream, log, txt, toml",
113                other
114            ),
115        }
116    }
117}
118
119impl fmt::Display for SourceFormat {
120    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
121        f.write_str(self.as_str())
122    }
123}
124
125/// Declared data-quality tests for a model (run after materialization when present).
126#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
127pub struct ModelTests {
128    /// Columns that must have zero nulls.
129    #[serde(default)]
130    pub not_null: Option<Vec<String>>,
131    /// Single column unique, or multi-column composite unique when len > 1.
132    #[serde(default)]
133    pub unique: Option<Vec<String>>,
134    /// Map of column → allowed string values.
135    #[serde(default)]
136    pub accepted_values: Option<std::collections::HashMap<String, Vec<String>>>,
137    /// When true (default), failed tests abort `rbt run` for that model.
138    #[serde(default)]
139    pub fail_on_error: Option<bool>,
140}
141
142impl ModelTests {
143    pub fn is_empty(&self) -> bool {
144        self.not_null.as_ref().map(|v| v.is_empty()).unwrap_or(true)
145            && self.unique.as_ref().map(|v| v.is_empty()).unwrap_or(true)
146            && self
147                .accepted_values
148                .as_ref()
149                .map(|m| m.is_empty())
150                .unwrap_or(true)
151    }
152
153    pub fn should_fail_on_error(&self) -> bool {
154        self.fail_on_error.unwrap_or(true)
155    }
156}
157
158/// Per-column documentation for humans and AI agents.
159///
160/// * `description` — short label (1–2 lines)
161/// * `context` — longer intent, units, lineage, caveats (agent-oriented)
162#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
163pub struct ColumnMeta {
164    #[serde(default)]
165    pub description: Option<String>,
166    #[serde(default)]
167    pub context: Option<String>,
168    /// Optional logical type hint (`utf8`, `int64`, `float64`, `timestamp`, …).
169    #[serde(default)]
170    pub dtype: Option<String>,
171    /// Optional unit (`USD`, `shares`, `ratio`, `ns_epoch`, …).
172    #[serde(default)]
173    pub unit: Option<String>,
174}
175
176/// YAML frontmatter embedded in model SQL files (`---` … `---`).
177///
178/// Used on staging, transforms, and marts. Scan-related fields only apply when
179/// a bronze scan contract (`scan_path`) is present.
180#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
181pub struct StagingFrontmatter {
182    /// Human-readable model purpose (docs + future catalog).
183    #[serde(default)]
184    pub description: Option<String>,
185    /// Longer model-level context for AI agents (why it exists, consumers, caveats).
186    #[serde(default)]
187    pub context: Option<String>,
188    /// Column-level description + context map (name → meta).
189    #[serde(default)]
190    pub columns: Option<std::collections::BTreeMap<String, ColumnMeta>>,
191
192    /// Explicit bronze format. If omitted, inferred from `scan_path` extension.
193    #[serde(default)]
194    pub source_format: Option<SourceFormat>,
195    /// File or directory to scan (project-relative or absolute; `s3://` deferred).
196    #[serde(default)]
197    pub scan_path: Option<String>,
198    /// Optional hive-style partition keys (path injection + future pruning).
199    #[serde(default)]
200    pub partition_by: Option<Vec<String>>,
201    /// jshift / projection field paths for selective JSON(L) extract.
202    #[serde(default)]
203    pub paths: Option<Vec<String>>,
204    /// Override catalog/schema name for registration (default: first `source()` name).
205    #[serde(default)]
206    pub source_name: Option<String>,
207    /// Override table name for registration (default: first `source()` table).
208    #[serde(default)]
209    pub source_table: Option<String>,
210    /// TOML: key of the array-of-tables to expand into rows (default: auto-detect).
211    #[serde(default)]
212    pub toml_rows_key: Option<String>,
213    /// When true, use scan→MemTable path even for formats that support DF listing.
214    #[serde(default)]
215    pub force_scan: Option<bool>,
216    /// Only scan hive-partitioned files whose path segments match these values.
217    /// Example: `{ timeframe: "1m" }` keeps `.../timeframe=1m/...` and skips `timeframe=1d`.
218    #[serde(default)]
219    pub require_partitions: Option<std::collections::HashMap<String, String>>,
220    /// Inject `_source_path` (Utf8) with the absolute file path for each row.
221    /// Enables "latest chunk wins" dedupe via `ORDER BY _source_path DESC`.
222    #[serde(default)]
223    pub inject_source_path: Option<bool>,
224
225    /// Logical grain of the model (e.g. `[symbol, timestamp_ns]`).
226    #[serde(default)]
227    pub grain: Option<Vec<String>>,
228    /// Primary uniqueness contract (usually same as grain for staging facts).
229    #[serde(default)]
230    pub unique_key: Option<Vec<String>>,
231    /// Free-form tags for selection / docs.
232    #[serde(default)]
233    pub tags: Option<Vec<String>>,
234    /// Materialization hint: `table` | `view` | `incremental_append` (engine may ignore for now).
235    #[serde(default)]
236    pub materialization: Option<String>,
237    /// Post-materialization assertions.
238    #[serde(default)]
239    pub tests: Option<ModelTests>,
240    /// Opaque metadata for tools and agents (strings, lists, nested maps OK).
241    #[serde(default)]
242    pub meta: Option<std::collections::BTreeMap<String, serde_yaml::Value>>,
243}
244
245impl StagingFrontmatter {
246    /// Resolve format from explicit field or path extension.
247    pub fn resolve_format(&self) -> Result<SourceFormat> {
248        if let Some(fmt) = self.source_format {
249            return Ok(fmt);
250        }
251        let path = self
252            .scan_path
253            .as_deref()
254            .context("frontmatter missing both source_format and scan_path")?;
255        // Strip globs for extension sniffing: `foo/*.jsonl` → look at last segment
256        let candidate = path.rsplit('/').next().unwrap_or(path);
257        let candidate = candidate.trim_matches(|c| c == '*' || c == '?');
258        if let Some(ext) = Path::new(candidate).extension().and_then(|e| e.to_str()) {
259            if let Some(fmt) = SourceFormat::from_extension(ext) {
260                return Ok(fmt);
261            }
262        }
263        // Directory paths: no extension — require explicit format
264        bail!(
265            "Cannot infer source_format from scan_path '{}'; set source_format explicitly",
266            path
267        );
268    }
269
270    pub fn has_scan_contract(&self) -> bool {
271        self.scan_path
272            .as_ref()
273            .map(|s| !s.trim().is_empty())
274            .unwrap_or(false)
275    }
276}
277
278/// Severity of a bronze compile diagnostic.
279#[derive(Debug, Clone, Copy, PartialEq, Eq)]
280pub enum DiagnosticSeverity {
281    Warning,
282    Error,
283}
284
285/// One compile-time finding about bronze frontmatter / scan paths.
286#[derive(Debug, Clone, PartialEq, Eq)]
287pub struct BronzeDiagnostic {
288    pub model: String,
289    pub severity: DiagnosticSeverity,
290    pub code: &'static str,
291    pub message: String,
292}
293
294impl fmt::Display for BronzeDiagnostic {
295    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
296        let level = match self.severity {
297            DiagnosticSeverity::Warning => "warning",
298            DiagnosticSeverity::Error => "error",
299        };
300        write!(
301            f,
302            "{}[{}] model={}: {}",
303            level, self.code, self.model, self.message
304        )
305    }
306}
307
308/// Result of bronze path validation during compile.
309#[derive(Debug, Clone, Default)]
310pub struct BronzeValidationReport {
311    pub diagnostics: Vec<BronzeDiagnostic>,
312}
313
314impl BronzeValidationReport {
315    pub fn warning_count(&self) -> usize {
316        self.diagnostics
317            .iter()
318            .filter(|d| d.severity == DiagnosticSeverity::Warning)
319            .count()
320    }
321
322    pub fn error_count(&self) -> usize {
323        self.diagnostics
324            .iter()
325            .filter(|d| d.severity == DiagnosticSeverity::Error)
326            .count()
327    }
328
329    pub fn has_errors(&self) -> bool {
330        self.error_count() > 0
331    }
332}
333
334/// Resolve `scan_path` against the project root. Non-local URIs are returned as-is (unchecked).
335pub fn resolve_scan_path(project_dir: &Path, scan_path: &str) -> PathBuf {
336    let trimmed = scan_path.trim();
337    if trimmed.is_empty() {
338        return project_dir.to_path_buf();
339    }
340    if is_remote_uri(trimmed) {
341        return PathBuf::from(trimmed);
342    }
343    let p = Path::new(trimmed);
344    if p.is_absolute() {
345        p.to_path_buf()
346    } else {
347        project_dir.join(p)
348    }
349}
350
351pub fn is_remote_uri(path: &str) -> bool {
352    let lower = path.to_ascii_lowercase();
353    lower.starts_with("s3://")
354        || lower.starts_with("s3a://")
355        || lower.starts_with("gs://")
356        || lower.starts_with("gcs://")
357        || lower.starts_with("az://")
358        || lower.starts_with("abfs://")
359        || lower.starts_with("abfss://")
360        || lower.starts_with("http://")
361        || lower.starts_with("https://")
362        || lower.starts_with("file://")
363}
364
365/// Whether a resolved local scan path currently exists (file or directory).
366/// Remote URIs are treated as "exists" for compile (runtime / object-store later).
367pub fn scan_path_exists(project_dir: &Path, scan_path: &str) -> bool {
368    if is_remote_uri(scan_path.trim()) {
369        return true;
370    }
371    let resolved = resolve_scan_path(project_dir, scan_path);
372    // Support simple trailing globs: `dir/*.jsonl` → check parent dir
373    let check = strip_simple_glob(&resolved);
374    check.exists()
375}
376
377fn strip_simple_glob(path: &Path) -> PathBuf {
378    let s = path.to_string_lossy();
379    if s.contains('*') || s.contains('?') {
380        if let Some(parent) = path.parent() {
381            return parent.to_path_buf();
382        }
383    }
384    path.to_path_buf()
385}
386
387#[cfg(test)]
388mod tests {
389    use super::*;
390
391    #[test]
392    fn format_from_extension_and_parse() {
393        assert_eq!(
394            SourceFormat::from_extension("jsonl"),
395            Some(SourceFormat::Jsonl)
396        );
397        assert_eq!(
398            SourceFormat::from_extension("toml"),
399            Some(SourceFormat::Toml)
400        );
401        assert_eq!(SourceFormat::from_extension("log"), Some(SourceFormat::Log));
402        assert_eq!(
403            SourceFormat::parse("arrow-ipc").unwrap(),
404            SourceFormat::ArrowIpc
405        );
406        assert_eq!(SourceFormat::parse("ndjson").unwrap(), SourceFormat::Jsonl);
407    }
408
409    #[test]
410    fn resolve_format_from_path() {
411        let fm = StagingFrontmatter {
412            scan_path: Some("lake/bronze/raw.jsonl".into()),
413            ..Default::default()
414        };
415        assert_eq!(fm.resolve_format().unwrap(), SourceFormat::Jsonl);
416    }
417
418    #[test]
419    fn remote_uri_exists_for_compile() {
420        assert!(scan_path_exists(
421            Path::new("/tmp"),
422            "s3://bucket/bronze/x.jsonl"
423        ));
424    }
425}