Skip to main content

rbt_core/
parser.rs

1use anyhow::{bail, Context, Result};
2use minijinja::value::Kwargs;
3use minijinja::Environment;
4use regex::Regex;
5use serde::{Deserialize, Serialize};
6use std::collections::HashSet;
7use std::sync::Arc;
8
9pub use crate::frontmatter::{
10    resolve_scan_path, scan_path_exists, BronzeCheckMode, BronzeDiagnostic, BronzeValidationReport,
11    DiagnosticSeverity, SourceFormat, StagingFrontmatter,
12};
13
14/// Parsed dependency reference from a Jinja-style model query.
15#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
16pub enum DependencyRef {
17    Model(String),
18    Source { source_name: String, table_name: String },
19}
20
21/// Fast-path SQL Model template parser extracting `{{ ref(...) }}` and `{{ source(...) }}` references.
22pub struct SqlModelParser;
23
24impl SqlModelParser {
25    /// Extracts YAML frontmatter and pure SQL content from raw model text.
26    ///
27    /// If a `---` block is present, YAML must parse successfully or this returns an error
28    /// (no silent fallback that leaves `---` in the SQL body).
29    pub fn parse_frontmatter(raw: &str) -> Result<(Option<StagingFrontmatter>, String)> {
30        let trimmed = raw.trim_start();
31        if !trimmed.starts_with("---") {
32            return Ok((None, raw.to_string()));
33        }
34
35        // Opening fence: line starting with --- (optional trailing spaces).
36        let after_open = match trimmed.strip_prefix("---") {
37            Some(rest) => rest.strip_prefix('\r').unwrap_or(rest),
38            None => return Ok((None, raw.to_string())),
39        };
40        let after_open = after_open
41            .strip_prefix('\n')
42            .or_else(|| after_open.strip_prefix("\r\n"))
43            .unwrap_or(after_open);
44
45        // Closing fence must be its own line (`---`), not a substring of comments
46        // like `# --------------------------------`.
47        let Some((yaml_str, sql_content)) = split_closing_frontmatter_fence(after_open) else {
48            bail!("Unclosed frontmatter block: found opening '---' but no closing '---' line");
49        };
50
51        // Empty YAML between fences is allowed → default frontmatter
52        if yaml_str.trim().is_empty() {
53            return Ok((Some(StagingFrontmatter::default()), sql_content));
54        }
55
56        let frontmatter: StagingFrontmatter = serde_yaml::from_str(yaml_str).with_context(|| {
57            format!(
58                "Invalid frontmatter YAML (between --- delimiters):\n{}",
59                yaml_str.trim()
60            )
61        })?;
62        Ok((Some(frontmatter), sql_content))
63    }
64
65    /// Fast-path extraction of all model and source dependencies from raw model SQL.
66    pub fn extract_dependencies(sql: &str) -> Result<Vec<DependencyRef>> {
67        let mut deps = Vec::new();
68        let mut seen = HashSet::new();
69
70        let ref_re = Regex::new(r#"\{\{\s*ref\s*\(\s*['"]([^'"]+)['"]\s*\)\s*\}\}"#)?;
71        for cap in ref_re.captures_iter(sql) {
72            let model_name = cap[1].trim().to_string();
73            if seen.insert(DependencyRef::Model(model_name.clone())) {
74                deps.push(DependencyRef::Model(model_name));
75            }
76        }
77
78        let source_re = Regex::new(
79            r#"\{\{\s*source\s*\(\s*['"]([^'"]+)['"]\s*,\s*['"]([^'"]+)['"]\s*\)\s*\}\}"#,
80        )?;
81        for cap in source_re.captures_iter(sql) {
82            let source_name = cap[1].trim().to_string();
83            let table_name = cap[2].trim().to_string();
84            let dep = DependencyRef::Source {
85                source_name,
86                table_name,
87            };
88            if seen.insert(dep.clone()) {
89                deps.push(dep);
90            }
91        }
92
93        Ok(deps)
94    }
95
96    /// Compiles raw SQL model by resolving `{{ ref(...) }}` and `{{ source(...) }}`.
97    ///
98    /// - `ref('m')` → `m` when `catalog_prefix` is empty, else `{prefix}.m`
99    /// - `source('s','t')` → `s.t` when prefix empty, else `{prefix}.s.t`
100    pub fn compile_sql(sql: &str, catalog_prefix: &str) -> Result<String> {
101        let ref_re = Regex::new(r#"\{\{\s*ref\s*\(\s*['"]([^'"]+)['"]\s*\)\s*\}\}"#)?;
102        let compiled_refs = ref_re.replace_all(sql, |caps: &regex::Captures| {
103            if catalog_prefix.is_empty() {
104                caps[1].to_string()
105            } else {
106                format!("{}.{}", catalog_prefix, &caps[1])
107            }
108        });
109
110        let source_re = Regex::new(
111            r#"\{\{\s*source\s*\(\s*['"]([^'"]+)['"]\s*,\s*['"]([^'"]+)['"]\s*\)\s*\}\}"#,
112        )?;
113        let compiled_sources = source_re.replace_all(&compiled_refs, |caps: &regex::Captures| {
114            if catalog_prefix.is_empty() {
115                format!("{}.{}", &caps[1], &caps[2])
116            } else {
117                format!("{}.{}.{}", catalog_prefix, &caps[1], &caps[2])
118            }
119        });
120
121        Ok(compiled_sources.to_string())
122    }
123}
124
125/// Split body after the opening `---` into (yaml, sql) at a closing fence line.
126fn split_closing_frontmatter_fence(after_open: &str) -> Option<(&str, String)> {
127    let mut offset = 0usize;
128    for line in after_open.split_inclusive('\n') {
129        let line_body = line.trim_end_matches(['\n', '\r']);
130        if line_body.trim() == "---" {
131            let yaml_str = &after_open[..offset];
132            let sql = after_open[offset + line.len()..].trim_start().to_string();
133            return Some((yaml_str, sql));
134        }
135        offset += line.len();
136    }
137    None
138}
139
140/// Advanced Jinja-compatible templating engine backed by `minijinja`.
141pub struct RbtTemplateEngine {
142    catalog_prefix: String,
143}
144
145impl RbtTemplateEngine {
146    pub fn new(catalog_prefix: impl Into<String>) -> Self {
147        Self {
148            catalog_prefix: catalog_prefix.into(),
149        }
150    }
151
152    /// Compiles a SQL model template using full Jinja2 evaluation with dbt-compatible macros.
153    pub fn render(&self, template_name: &str, template_source: &str) -> Result<String> {
154        let mut env = Environment::new();
155        let prefix = Arc::new(self.catalog_prefix.clone());
156
157        let prefix_ref = prefix.clone();
158        env.add_function("ref", move |model_name: &str| -> String {
159            if prefix_ref.is_empty() {
160                model_name.to_string()
161            } else {
162                format!("{}.{}", prefix_ref, model_name)
163            }
164        });
165
166        let prefix_src = prefix.clone();
167        env.add_function(
168            "source",
169            move |source_name: &str, table_name: &str| -> String {
170                if prefix_src.is_empty() {
171                    format!("{}.{}", source_name, table_name)
172                } else {
173                    format!("{}.{}.{}", prefix_src, source_name, table_name)
174                }
175            },
176        );
177
178        env.add_function("config", |_kwargs: Kwargs| -> String { String::new() });
179
180        env.add_template(template_name, template_source)?;
181        let tmpl = env.get_template(template_name)?;
182        let rendered = tmpl.render(())?;
183        Ok(rendered)
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    #[test]
192    fn test_extract_dependencies() -> Result<()> {
193        let sql = r#"
194            SELECT
195                o.order_id,
196                u.user_name,
197                p.product_name
198            FROM {{ ref('stg_orders') }} o
199            JOIN {{ ref('stg_users') }} u ON o.user_id = u.user_id
200            JOIN {{ source('raw_store', 'products') }} p ON o.product_id = p.id
201        "#;
202
203        let deps = SqlModelParser::extract_dependencies(sql)?;
204        assert_eq!(deps.len(), 3);
205        assert!(deps.contains(&DependencyRef::Model("stg_orders".to_string())));
206        assert!(deps.contains(&DependencyRef::Model("stg_users".to_string())));
207        assert!(deps.contains(&DependencyRef::Source {
208            source_name: "raw_store".to_string(),
209            table_name: "products".to_string()
210        }));
211
212        Ok(())
213    }
214
215    #[test]
216    fn test_compile_sql() -> Result<()> {
217        let sql = "SELECT * FROM {{ ref('stg_orders') }} JOIN {{ source('raw', 'users') }}";
218        let compiled = SqlModelParser::compile_sql(sql, "iceberg_db")?;
219        assert_eq!(
220            compiled,
221            "SELECT * FROM iceberg_db.stg_orders JOIN iceberg_db.raw.users"
222        );
223
224        let compiled_local = SqlModelParser::compile_sql(sql, "")?;
225        assert_eq!(
226            compiled_local,
227            "SELECT * FROM stg_orders JOIN raw.users"
228        );
229        Ok(())
230    }
231
232    #[test]
233    fn test_minijinja_template_engine() -> Result<()> {
234        let engine = RbtTemplateEngine::new("prod_lake");
235        let template = r#"
236            {{ config(materialized="table") }}
237            SELECT * FROM {{ ref('stg_events') }}
238            WHERE event_type = 'click'
239            {% if true %}
240                AND source = '{{ source("telemetry", "clicks") }}'
241            {% endif %}
242        "#;
243
244        let rendered = engine.render("my_model.sql", template)?;
245        assert!(rendered.contains("FROM prod_lake.stg_events"));
246        assert!(rendered.contains("AND source = 'prod_lake.telemetry.clicks'"));
247        Ok(())
248    }
249
250    #[test]
251    fn test_parse_frontmatter() -> Result<()> {
252        let raw_sql = r#"---
253source_format: parquet
254scan_path: "s3://lake/events/*/*.parquet"
255partition_by: ["year", "month"]
256paths: [id, tenant_id]
257---
258SELECT * FROM {{ source('raw', 'events') }}
259"#;
260        let (frontmatter, sql) = SqlModelParser::parse_frontmatter(raw_sql)?;
261        assert!(frontmatter.is_some());
262        let fm = frontmatter.unwrap();
263        assert_eq!(fm.source_format, Some(SourceFormat::Parquet));
264        assert_eq!(
265            fm.scan_path.as_deref(),
266            Some("s3://lake/events/*/*.parquet")
267        );
268        assert_eq!(
269            fm.partition_by,
270            Some(vec!["year".to_string(), "month".to_string()])
271        );
272        assert_eq!(
273            fm.paths,
274            Some(vec!["id".to_string(), "tenant_id".to_string()])
275        );
276        assert!(sql.contains("SELECT * FROM {{ source('raw', 'events') }}"));
277        Ok(())
278    }
279
280    #[test]
281    fn test_invalid_frontmatter_is_error() {
282        let raw = "---\nsource_format: [not, valid, for, enum\n---\nSELECT 1";
283        let err = SqlModelParser::parse_frontmatter(raw).unwrap_err();
284        assert!(err.to_string().contains("Invalid frontmatter") || err.to_string().contains("YAML"));
285    }
286
287    #[test]
288    fn test_unclosed_frontmatter_is_error() {
289        let raw = "---\nsource_format: jsonl\nSELECT 1";
290        let err = SqlModelParser::parse_frontmatter(raw).unwrap_err();
291        assert!(err.to_string().contains("Unclosed frontmatter"));
292    }
293
294    #[test]
295    fn test_frontmatter_ignores_triple_dash_inside_comments() -> Result<()> {
296        let raw = r#"---
297# --------------------------------
298# decorative banner with many dashes
299# --------------------------------
300source_format: arrow_ipc
301scan_path: "lake/bronze"
302---
303SELECT 1
304"#;
305        let (fm, sql) = SqlModelParser::parse_frontmatter(raw)?;
306        let fm = fm.expect("frontmatter");
307        assert_eq!(fm.scan_path.as_deref(), Some("lake/bronze"));
308        assert!(sql.contains("SELECT 1"));
309        Ok(())
310    }
311}