Skip to main content

_diffctx/edges/semantic/
sql.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_sql_file(path: &Path) -> bool {
14    base::file_ext(path) == ".sql"
15}
16
17static CREATE_RE: Lazy<Regex> = Lazy::new(|| {
18    Regex::new(r"(?mi)^\s*CREATE\s+(?:OR\s+REPLACE\s+)?(?:TABLE|VIEW|FUNCTION|PROCEDURE|TYPE|INDEX|TRIGGER|SEQUENCE|SCHEMA)\s+(?:IF\s+NOT\s+EXISTS\s+)?(?:[\w.]+\.)?(\w+)").unwrap()
19});
20static REFERENCES_RE: Lazy<Regex> =
21    Lazy::new(|| Regex::new(r"(?mi)REFERENCES\s+(?:[\w.]+\.)?(\w+)").unwrap());
22static TABLE_REF_RE: Lazy<Regex> = Lazy::new(|| {
23    Regex::new(r"(?mi)(?:FROM|JOIN|INTO|UPDATE|DELETE\s+FROM|ALTER\s+TABLE|DROP\s+TABLE|TRUNCATE\s+TABLE|INSERT\s+INTO)\s+(?:[\w.]+\.)?(\w+)").unwrap()
24});
25
26static SQL_KEYWORDS: Lazy<FxHashSet<&str>> = Lazy::new(|| {
27    [
28        "select", "from", "where", "and", "or", "not", "in", "exists", "null", "true", "false",
29        "set", "values", "as", "on", "using", "left", "right", "inner", "outer", "cross", "group",
30        "order", "by", "having", "limit", "offset", "union", "all", "distinct", "case", "when",
31        "then", "else", "end", "if", "begin", "declare", "returns", "return",
32    ]
33    .iter()
34    .copied()
35    .collect()
36});
37
38fn extract_creates(content: &str) -> FxHashSet<String> {
39    CREATE_RE
40        .captures_iter(content)
41        .map(|c| c[1].to_lowercase())
42        .filter(|n| !SQL_KEYWORDS.contains(n.as_str()))
43        .collect()
44}
45
46fn extract_table_refs(content: &str) -> FxHashSet<String> {
47    let mut refs: FxHashSet<String> = TABLE_REF_RE
48        .captures_iter(content)
49        .map(|c| c[1].to_lowercase())
50        .filter(|n| !SQL_KEYWORDS.contains(n.as_str()))
51        .collect();
52    refs.extend(
53        REFERENCES_RE
54            .captures_iter(content)
55            .map(|c| c[1].to_lowercase())
56            .filter(|n| !SQL_KEYWORDS.contains(n.as_str())),
57    );
58    refs
59}
60
61pub struct SqlEdgeBuilder;
62
63impl EdgeBuilder for SqlEdgeBuilder {
64    fn build(&self, fragments: &[Fragment], _repo_root: Option<&Path>) -> EdgeDict {
65        let frags: Vec<&Fragment> = fragments
66            .iter()
67            .filter(|f| is_sql_file(Path::new(f.path())))
68            .collect();
69        if frags.is_empty() {
70            return FxHashMap::default();
71        }
72
73        let fk_w = EDGE_WEIGHTS["sql_fk"].forward;
74        let table_ref_w = EDGE_WEIGHTS["sql_table_ref"].forward;
75        let reverse_factor = EDGE_WEIGHTS["sql_fk"].reverse_factor;
76
77        let mut table_to_frags: FxHashMap<String, Vec<_>> = FxHashMap::default();
78        for f in &frags {
79            for name in extract_creates(&f.content) {
80                table_to_frags.entry(name).or_default().push(f.id.clone());
81            }
82        }
83
84        let mut edges: EdgeDict = FxHashMap::default();
85
86        for f in &frags {
87            let self_creates = extract_creates(&f.content);
88            let has_fk = REFERENCES_RE.is_match(&f.content);
89            for tref in extract_table_refs(&f.content) {
90                if self_creates.contains(&tref) {
91                    continue;
92                }
93                let w = if has_fk
94                    && REFERENCES_RE
95                        .captures_iter(&f.content)
96                        .any(|c| c[1].to_lowercase() == tref)
97                {
98                    fk_w
99                } else {
100                    table_ref_w
101                };
102                if let Some(targets) = table_to_frags.get(&tref) {
103                    for t in targets {
104                        if t != &f.id {
105                            add_edge(&mut edges, &f.id, t, w, reverse_factor);
106                        }
107                    }
108                }
109            }
110        }
111        edges
112    }
113
114    fn discover_related_files(
115        &self,
116        changed: &[PathBuf],
117        candidates: &[PathBuf],
118        repo_root: Option<&Path>,
119        file_cache: Option<&FxHashMap<PathBuf, String>>,
120    ) -> Vec<PathBuf> {
121        let sql_changed: Vec<&PathBuf> = changed.iter().filter(|f| is_sql_file(f)).collect();
122        if sql_changed.is_empty() {
123            return vec![];
124        }
125        let mut refs = FxHashSet::default();
126        for f in &sql_changed {
127            if let Some(content) = base::read_file_cached(f, file_cache) {
128                refs.extend(extract_table_refs(&content));
129            }
130        }
131        discover_files_by_refs(&refs, changed, candidates, repo_root)
132    }
133}