Skip to main content

dre_core/
dialect.rs

1//! What DRE knows about each source type's SQL: how to quote identifiers and string literals,
2//! how unquoted names compare, and which functions reach outside the database. One table, so
3//! `dre catalog`, `dre query` and the integrations reading the catalogue agree.
4//!
5//! A type missing here fails closed: no catalogue, and no literal an integration could build.
6
7use serde::Serialize;
8use serde_json::{Value, json};
9
10/// How a string value becomes a SQL literal.
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum LiteralStyle {
13    /// `'…'` with each `'` doubled; the dialect has no backslash escapes in literals.
14    Doubling,
15    /// The UTF-8 bytes as hex, decoded by this expression (`{hex}` stands for the digits): no
16    /// quote or backslash of the value ever reaches the dialect's lexer.
17    Hex(&'static str),
18}
19
20/// How an identifier written without quotes matches a stored name.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
22#[serde(rename_all = "snake_case")]
23pub enum IdentifierCase {
24    /// Names compare case-insensitively, quoted or not.
25    Insensitive,
26    /// Unquoted names fold to lowercase; quoted names are exact.
27    LowerUnlessQuoted,
28}
29
30#[derive(Debug, Clone, Copy)]
31pub struct Dialect {
32    pub identifier_quote: char,
33    pub literal_style: LiteralStyle,
34    pub identifier_case: IdentifierCase,
35    /// Functions that read files, run programs or reach the network; `dre query` refuses them.
36    pub denied_functions: &'static [&'static str],
37}
38
39const DUCKDB: Dialect = Dialect {
40    identifier_quote: '"',
41    literal_style: LiteralStyle::Doubling,
42    identifier_case: IdentifierCase::Insensitive,
43    denied_functions: &[
44        "read_text",
45        "read_blob",
46        "read_csv",
47        "read_csv_auto",
48        "read_json",
49        "read_json_auto",
50        "read_json_objects",
51        "read_json_objects_auto",
52        "read_ndjson",
53        "read_ndjson_auto",
54        "read_ndjson_objects",
55        "read_parquet",
56        "parquet_scan",
57        "parquet_metadata",
58        "parquet_schema",
59        "parquet_file_metadata",
60        "parquet_kv_metadata",
61        "read_xlsx",
62        "sniff_csv",
63        "glob",
64        "delta_scan",
65        "iceberg_scan",
66        "iceberg_metadata",
67        "iceberg_snapshots",
68        "sqlite_scan",
69        "sqlite_attach",
70        "postgres_scan",
71        "postgres_query",
72        "postgres_attach",
73        "mysql_query",
74        "query",
75        "query_table",
76        "getenv",
77        "duckdb_secrets",
78        "duckdb_settings",
79        "which_secret",
80    ],
81};
82
83const POSTGRES: Dialect = Dialect {
84    identifier_quote: '"',
85    literal_style: LiteralStyle::Hex("convert_from(decode('{hex}', 'hex'), 'UTF8')"),
86    identifier_case: IdentifierCase::LowerUnlessQuoted,
87    denied_functions: &[
88        "pg_read_file",
89        "pg_read_binary_file",
90        "pg_ls_dir",
91        "pg_ls_logdir",
92        "pg_ls_waldir",
93        "pg_ls_tmpdir",
94        "pg_ls_archive_statusdir",
95        "pg_stat_file",
96        "lo_import",
97        "lo_export",
98        "lo_get",
99        "lo_open",
100        "loread",
101        "dblink",
102        "dblink_exec",
103        "dblink_connect",
104        "dblink_send_query",
105        "pg_execute_server_program",
106        // Run SQL passed as a string, past this guard.
107        "query_to_xml",
108        "query_to_xml_and_xmlschema",
109        "query_to_xmlschema",
110        "cursor_to_xml",
111        "cursor_to_xmlschema",
112        "table_to_xml",
113        "table_to_xml_and_xmlschema",
114        "database_to_xml",
115        "schema_to_xml",
116        "current_setting",
117        "pg_reload_conf",
118        "pg_terminate_backend",
119        "pg_cancel_backend",
120    ],
121};
122
123const DATABRICKS: Dialect = Dialect {
124    identifier_quote: '`',
125    literal_style: LiteralStyle::Hex("decode(unhex('{hex}'), 'UTF-8')"),
126    identifier_case: IdentifierCase::Insensitive,
127    denied_functions: &[
128        "read_files",
129        "read_kafka",
130        "read_kinesis",
131        "read_pubsub",
132        "read_pulsar",
133        "read_state_metadata",
134        "read_statestore",
135        "cloud_files_state",
136        "http_request",
137        "java_method",
138        "reflect",
139        "secret",
140        "list_secrets",
141        "ai_query",
142    ],
143};
144
145/// The dialect of a source type, or `None` when DRE doesn't know its quoting rules.
146pub fn of(kind: &str) -> Option<&'static Dialect> {
147    match kind {
148        "duckdb" => Some(&DUCKDB),
149        "postgres" => Some(&POSTGRES),
150        "databricks" => Some(&DATABRICKS),
151        _ => None,
152    }
153}
154
155impl Dialect {
156    /// `value` as a SQL string literal.
157    pub fn literal(&self, value: &str) -> String {
158        match self.literal_style {
159            LiteralStyle::Hex(template) => {
160                let hex: String = value.bytes().map(|b| format!("{b:02x}")).collect();
161                template.replace("{hex}", &hex)
162            }
163            LiteralStyle::Doubling => format!("'{}'", value.replace('\'', "''")),
164        }
165    }
166
167    /// Whether an unquoted identifier `written` names the stored `actual`.
168    pub fn same_name(&self, written: &str, actual: &str) -> bool {
169        match self.identifier_case {
170            IdentifierCase::Insensitive => written.eq_ignore_ascii_case(actual),
171            IdentifierCase::LowerUnlessQuoted => written == actual,
172        }
173    }
174
175    /// Whether `function` (any case, optionally schema-qualified) reaches outside the database.
176    pub fn denies(&self, function: &str) -> bool {
177        let name = function
178            .rsplit('.')
179            .next()
180            .unwrap_or(function)
181            .trim_matches(['"', '`']);
182        self.denied_functions.iter().any(|f| f.eq_ignore_ascii_case(name))
183    }
184
185    /// The rules as `dre catalog` reports them.
186    pub fn describe(&self) -> Value {
187        json!({
188            "identifier_quote": self.identifier_quote.to_string(),
189            "identifier_case": self.identifier_case,
190            "literal_style": match self.literal_style {
191                LiteralStyle::Doubling => "doubling",
192                LiteralStyle::Hex(_) => "hex",
193            },
194            "hex_literal": match self.literal_style {
195                LiteralStyle::Hex(template) => Some(template),
196                LiteralStyle::Doubling => None,
197            },
198        })
199    }
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    #[test]
207    fn literals_never_pass_a_quote_or_backslash_to_hex_dialects() {
208        let value = "x' or '1'='1\\";
209        assert_eq!(of("duckdb").unwrap().literal(value), "'x'' or ''1''=''1\\'");
210        let pg = of("postgres").unwrap().literal(value);
211        assert!(pg.starts_with("convert_from(decode('78") && !pg.contains(" or "));
212        assert_eq!(
213            of("databricks").unwrap().literal("é"),
214            "decode(unhex('c3a9'), 'UTF-8')"
215        );
216    }
217
218    #[test]
219    fn unknown_types_have_no_dialect() {
220        assert!(of("snowflake").is_none());
221    }
222
223    #[test]
224    fn denied_functions_match_any_case_and_schema() {
225        let pg = of("postgres").unwrap();
226        assert!(pg.denies("PG_READ_FILE") && pg.denies("pg_catalog.pg_read_file"));
227        // Functions that run SQL given as a string would carry a denied call past the guard.
228        assert!(pg.denies("query_to_xml") && of("duckdb").unwrap().denies("query"));
229        assert!(!pg.denies("lower"));
230    }
231}