_diffctx/edges/semantic/
dbt.rs

1use std::path::{Path, PathBuf};
2
3use once_cell::sync::Lazy;
4use regex::Regex;
5use rustc_hash::{FxHashMap, FxHashSet};
6
7use crate::config::weights::EDGE_WEIGHTS;
8use crate::types::Fragment;
9
10use super::super::EdgeDict;
11use super::super::base::{self, EdgeBuilder, add_edge, discover_files_by_refs};
12
13fn is_dbt_file(content: &str) -> bool {
14    content.contains("{{ ref(") || content.contains("{{ source(") || content.contains("{{ config(")
15}
16
17fn is_sql_file(path: &Path) -> bool {
18    base::file_ext(path) == ".sql"
19}
20
21static REF_RE: Lazy<Regex> =
22    Lazy::new(|| Regex::new(r#"\{\{\s*ref\s*\(\s*['"](\w+)['"]\s*\)\s*\}\}"#).unwrap());
23static SOURCE_RE: Lazy<Regex> = Lazy::new(|| {
24    Regex::new(r#"\{\{\s*source\s*\(\s*['"](\w+)['"]\s*,\s*['"](\w+)['"]\s*\)\s*\}\}"#).unwrap()
25});
26static MACRO_CALL_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\{\{\s*(\w+)\s*\(").unwrap());
27static MACRO_DEF_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\{%-?\s*macro\s+(\w+)").unwrap());
28
29static DBT_BUILTINS: Lazy<FxHashSet<&str>> = Lazy::new(|| {
30    [
31        "ref", "source", "config", "set", "if", "for", "endif", "endfor", "else", "elif", "block",
32        "endblock", "macro", "endmacro", "do", "call", "filter",
33    ]
34    .iter()
35    .copied()
36    .collect()
37});
38
39fn extract_refs(content: &str) -> FxHashSet<String> {
40    REF_RE
41        .captures_iter(content)
42        .map(|c| c[1].to_string())
43        .collect()
44}
45
46fn extract_sources(content: &str) -> FxHashSet<String> {
47    SOURCE_RE
48        .captures_iter(content)
49        .map(|c| c[2].to_string())
50        .collect()
51}
52
53fn extract_macro_calls(content: &str) -> FxHashSet<String> {
54    MACRO_CALL_RE
55        .captures_iter(content)
56        .map(|c| c[1].to_string())
57        .filter(|n| !DBT_BUILTINS.contains(n.as_str()))
58        .collect()
59}
60
61fn extract_macro_defs(content: &str) -> FxHashSet<String> {
62    MACRO_DEF_RE
63        .captures_iter(content)
64        .map(|c| c[1].to_string())
65        .collect()
66}
67
68pub struct DbtEdgeBuilder;
69
70impl EdgeBuilder for DbtEdgeBuilder {
71    fn build(&self, fragments: &[Fragment], repo_root: Option<&Path>) -> EdgeDict {
72        let frags: Vec<&Fragment> = fragments
73            .iter()
74            .filter(|f| is_sql_file(Path::new(f.path())) && is_dbt_file(&f.content))
75            .collect();
76        if frags.is_empty() {
77            return FxHashMap::default();
78        }
79
80        let ref_w = EDGE_WEIGHTS["dbt_ref"].forward;
81        let source_w = EDGE_WEIGHTS["dbt_source"].forward;
82        let macro_w = EDGE_WEIGHTS["dbt_macro"].forward;
83        let ref_rev = EDGE_WEIGHTS["dbt_ref"].reverse_factor;
84        let source_rev = EDGE_WEIGHTS["dbt_source"].reverse_factor;
85        let macro_rev = EDGE_WEIGHTS["dbt_macro"].reverse_factor;
86
87        let idx = base::FragmentIndex::new(fragments, repo_root);
88        let mut macro_to_frags: FxHashMap<String, Vec<_>> = FxHashMap::default();
89        for f in &frags {
90            for name in extract_macro_defs(&f.content) {
91                macro_to_frags
92                    .entry(name.to_lowercase())
93                    .or_default()
94                    .push(f.id.clone());
95            }
96        }
97
98        let mut edges: EdgeDict = FxHashMap::default();
99
100        for f in &frags {
101            for r in extract_refs(&f.content) {
102                base::link_by_name(&f.id, &r, &idx, &mut edges, ref_w, ref_rev);
103            }
104            for s in extract_sources(&f.content) {
105                base::link_by_name(&f.id, &s, &idx, &mut edges, source_w, source_rev);
106            }
107            for mc in extract_macro_calls(&f.content) {
108                if let Some(targets) = macro_to_frags.get(&mc.to_lowercase()) {
109                    for t in targets {
110                        if t != &f.id {
111                            add_edge(&mut edges, &f.id, t, macro_w, macro_rev);
112                        }
113                    }
114                }
115            }
116        }
117        edges
118    }
119
120    fn discover_related_files(
121        &self,
122        changed: &[PathBuf],
123        candidates: &[PathBuf],
124        repo_root: Option<&Path>,
125        file_cache: Option<&FxHashMap<PathBuf, String>>,
126    ) -> Vec<PathBuf> {
127        let mut refs = FxHashSet::default();
128        for f in changed {
129            if !is_sql_file(f) {
130                continue;
131            }
132            if let Some(content) = base::read_file_cached(f, file_cache) {
133                if !is_dbt_file(&content) {
134                    continue;
135                }
136                refs.extend(extract_refs(&content));
137                refs.extend(extract_sources(&content));
138                refs.extend(extract_macro_calls(&content));
139            }
140        }
141        if refs.is_empty() {
142            return vec![];
143        }
144        discover_files_by_refs(&refs, changed, candidates, repo_root)
145    }
146}
_diffctx/edges/semantic/dbt.rs

_diffctx/edges/semantic/
dbt.rs
