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