Skip to main content

polyglot_sql/dialects/
cockroachdb.rs

1//! CockroachDB Dialect
2//!
3//! CockroachDB-specific transformations based on sqlglot patterns.
4//! CockroachDB is PostgreSQL-compatible with distributed database extensions.
5
6use super::{DialectImpl, DialectType};
7use crate::error::Result;
8use crate::expressions::{AggFunc, Case, Cast, Expression, Function, VarArgFunc};
9use crate::generator::GeneratorConfig;
10use crate::tokens::TokenizerConfig;
11
12/// CockroachDB dialect (PostgreSQL-compatible distributed database)
13pub struct CockroachDBDialect;
14
15impl DialectImpl for CockroachDBDialect {
16    fn dialect_type(&self) -> DialectType {
17        DialectType::CockroachDB
18    }
19
20    fn tokenizer_config(&self) -> TokenizerConfig {
21        let mut config = TokenizerConfig::default();
22        // CockroachDB uses double quotes for identifiers (PostgreSQL-style)
23        config.identifiers.insert('"', '"');
24        // PostgreSQL-style nested comments supported
25        config.nested_comments = true;
26        config
27    }
28
29    fn generator_config(&self) -> GeneratorConfig {
30        use crate::generator::IdentifierQuoteStyle;
31        GeneratorConfig {
32            identifier_quote: '"',
33            identifier_quote_style: IdentifierQuoteStyle::DOUBLE_QUOTE,
34            dialect: Some(DialectType::CockroachDB),
35            ..Default::default()
36        }
37    }
38
39    fn transform_expr(&self, expr: Expression) -> Result<Expression> {
40        match expr {
41            // IFNULL -> COALESCE in CockroachDB
42            Expression::IfNull(f) => Ok(Expression::Coalesce(Box::new(VarArgFunc {
43                original_name: None,
44                expressions: vec![f.this, f.expression],
45                inferred_type: None,
46            }))),
47
48            // NVL -> COALESCE in CockroachDB
49            Expression::Nvl(f) => Ok(Expression::Coalesce(Box::new(VarArgFunc {
50                original_name: None,
51                expressions: vec![f.this, f.expression],
52                inferred_type: None,
53            }))),
54
55            // Coalesce with original_name (e.g., IFNULL parsed as Coalesce) -> clear original_name
56            Expression::Coalesce(mut f) => {
57                f.original_name = None;
58                Ok(Expression::Coalesce(f))
59            }
60
61            // TryCast is not directly supported, use CAST
62            Expression::TryCast(c) => Ok(Expression::Cast(c)),
63
64            // SafeCast -> CAST in CockroachDB
65            Expression::SafeCast(c) => Ok(Expression::Cast(c)),
66
67            // ILIKE is native in CockroachDB (PostgreSQL-style)
68            Expression::ILike(op) => Ok(Expression::ILike(op)),
69
70            // CountIf -> SUM(CASE WHEN condition THEN 1 ELSE 0 END)
71            Expression::CountIf(f) => {
72                let case_expr = Expression::Case(Box::new(Case {
73                    operand: None,
74                    whens: vec![(f.this.clone(), Expression::number(1))],
75                    else_: Some(Expression::number(0)),
76                    comments: Vec::new(),
77                    inferred_type: None,
78                }));
79                Ok(Expression::Sum(Box::new(AggFunc {
80                    ignore_nulls: None,
81                    having_max: None,
82                    this: case_expr,
83                    distinct: f.distinct,
84                    filter: f.filter,
85                    order_by: Vec::new(),
86                    name: None,
87                    limit: None,
88                    inferred_type: None,
89                })))
90            }
91
92            // RAND -> RANDOM in CockroachDB (PostgreSQL-style)
93            Expression::Rand(r) => {
94                let _ = r.seed;
95                Ok(Expression::Random(crate::expressions::Random))
96            }
97
98            // Generic function transformations
99            Expression::Function(f) => self.transform_function(*f),
100
101            // Generic aggregate function transformations
102            Expression::AggregateFunction(f) => self.transform_aggregate_function(f),
103
104            // Cast transformations
105            Expression::Cast(c) => self.transform_cast(*c),
106
107            // Pass through everything else
108            _ => Ok(expr),
109        }
110    }
111}
112
113impl CockroachDBDialect {
114    fn transform_function(&self, f: Function) -> Result<Expression> {
115        let name_upper = f.name.to_uppercase();
116        match name_upper.as_str() {
117            // IFNULL -> COALESCE
118            "IFNULL" if f.args.len() == 2 => Ok(Expression::Coalesce(Box::new(VarArgFunc {
119                original_name: None,
120                expressions: f.args,
121                inferred_type: None,
122            }))),
123
124            // NVL -> COALESCE
125            "NVL" if f.args.len() == 2 => Ok(Expression::Coalesce(Box::new(VarArgFunc {
126                original_name: None,
127                expressions: f.args,
128                inferred_type: None,
129            }))),
130
131            // ISNULL -> COALESCE
132            "ISNULL" if f.args.len() == 2 => Ok(Expression::Coalesce(Box::new(VarArgFunc {
133                original_name: None,
134                expressions: f.args,
135                inferred_type: None,
136            }))),
137
138            // NOW is native in CockroachDB
139            "NOW" => Ok(Expression::CurrentTimestamp(
140                crate::expressions::CurrentTimestamp {
141                    precision: None,
142                    sysdate: false,
143                },
144            )),
145
146            // GETDATE -> NOW
147            "GETDATE" => Ok(Expression::CurrentTimestamp(
148                crate::expressions::CurrentTimestamp {
149                    precision: None,
150                    sysdate: false,
151                },
152            )),
153
154            // RAND -> RANDOM
155            "RAND" => Ok(Expression::Random(crate::expressions::Random)),
156
157            // STRING_AGG is native in CockroachDB (PostgreSQL-style)
158            "STRING_AGG" => Ok(Expression::Function(Box::new(f))),
159
160            // GROUP_CONCAT -> STRING_AGG
161            "GROUP_CONCAT" if !f.args.is_empty() => Ok(Expression::Function(Box::new(
162                Function::new("STRING_AGG".to_string(), f.args),
163            ))),
164
165            // LISTAGG -> STRING_AGG
166            "LISTAGG" if !f.args.is_empty() => Ok(Expression::Function(Box::new(Function::new(
167                "STRING_AGG".to_string(),
168                f.args,
169            )))),
170
171            // SUBSTR -> SUBSTRING
172            "SUBSTR" => Ok(Expression::Function(Box::new(Function::new(
173                "SUBSTRING".to_string(),
174                f.args,
175            )))),
176
177            // LENGTH is native in CockroachDB
178            "LENGTH" => Ok(Expression::Function(Box::new(f))),
179
180            // LEN -> LENGTH
181            "LEN" if f.args.len() == 1 => Ok(Expression::Function(Box::new(Function::new(
182                "LENGTH".to_string(),
183                f.args,
184            )))),
185
186            // CHARINDEX -> STRPOS (with swapped args)
187            "CHARINDEX" if f.args.len() >= 2 => {
188                let mut args = f.args;
189                let substring = args.remove(0);
190                let string = args.remove(0);
191                Ok(Expression::Function(Box::new(Function::new(
192                    "STRPOS".to_string(),
193                    vec![string, substring],
194                ))))
195            }
196
197            // INSTR -> STRPOS
198            "INSTR" if f.args.len() >= 2 => Ok(Expression::Function(Box::new(Function::new(
199                "STRPOS".to_string(),
200                f.args,
201            )))),
202
203            // LOCATE -> STRPOS (with swapped args)
204            "LOCATE" if f.args.len() >= 2 => {
205                let mut args = f.args;
206                let substring = args.remove(0);
207                let string = args.remove(0);
208                Ok(Expression::Function(Box::new(Function::new(
209                    "STRPOS".to_string(),
210                    vec![string, substring],
211                ))))
212            }
213
214            // STRPOS is native in CockroachDB
215            "STRPOS" => Ok(Expression::Function(Box::new(f))),
216
217            // ARRAY_LENGTH is native in CockroachDB
218            "ARRAY_LENGTH" => Ok(Expression::Function(Box::new(f))),
219
220            // SIZE -> ARRAY_LENGTH
221            "SIZE" if f.args.len() == 1 => Ok(Expression::Function(Box::new(Function::new(
222                "ARRAY_LENGTH".to_string(),
223                f.args,
224            )))),
225
226            // CARDINALITY is native in CockroachDB
227            "CARDINALITY" => Ok(Expression::Function(Box::new(f))),
228
229            // TO_CHAR is native in CockroachDB
230            "TO_CHAR" => Ok(Expression::Function(Box::new(f))),
231
232            // DATE_FORMAT -> TO_CHAR
233            "DATE_FORMAT" if f.args.len() >= 2 => Ok(Expression::Function(Box::new(
234                Function::new("TO_CHAR".to_string(), f.args),
235            ))),
236
237            // strftime -> TO_CHAR
238            "STRFTIME" if f.args.len() >= 2 => {
239                let mut args = f.args;
240                let format = args.remove(0);
241                let date = args.remove(0);
242                Ok(Expression::Function(Box::new(Function::new(
243                    "TO_CHAR".to_string(),
244                    vec![date, format],
245                ))))
246            }
247
248            // JSONB_EXTRACT_PATH_TEXT is native in CockroachDB
249            "JSONB_EXTRACT_PATH_TEXT" => Ok(Expression::Function(Box::new(f))),
250
251            // JSON_EXTRACT_PATH_TEXT is native in CockroachDB
252            "JSON_EXTRACT_PATH_TEXT" => Ok(Expression::Function(Box::new(f))),
253
254            // GET_JSON_OBJECT -> JSONB_EXTRACT_PATH_TEXT
255            "GET_JSON_OBJECT" if f.args.len() == 2 => Ok(Expression::Function(Box::new(
256                Function::new("JSONB_EXTRACT_PATH_TEXT".to_string(), f.args),
257            ))),
258
259            // JSON_EXTRACT -> JSONB_EXTRACT_PATH
260            "JSON_EXTRACT" if f.args.len() >= 2 => Ok(Expression::Function(Box::new(
261                Function::new("JSONB_EXTRACT_PATH".to_string(), f.args),
262            ))),
263
264            // REGEXP_LIKE is native in CockroachDB
265            "REGEXP_LIKE" => Ok(Expression::Function(Box::new(f))),
266
267            // RLIKE -> REGEXP_LIKE
268            "RLIKE" if f.args.len() >= 2 => Ok(Expression::Function(Box::new(Function::new(
269                "REGEXP_LIKE".to_string(),
270                f.args,
271            )))),
272
273            // Pass through everything else
274            _ => Ok(Expression::Function(Box::new(f))),
275        }
276    }
277
278    fn transform_aggregate_function(
279        &self,
280        f: Box<crate::expressions::AggregateFunction>,
281    ) -> Result<Expression> {
282        let name_upper = f.name.to_uppercase();
283        match name_upper.as_str() {
284            // COUNT_IF -> SUM(CASE WHEN...)
285            "COUNT_IF" if !f.args.is_empty() => {
286                let condition = f.args.into_iter().next().unwrap();
287                let case_expr = Expression::Case(Box::new(Case {
288                    operand: None,
289                    whens: vec![(condition, Expression::number(1))],
290                    else_: Some(Expression::number(0)),
291                    comments: Vec::new(),
292                    inferred_type: None,
293                }));
294                Ok(Expression::Sum(Box::new(AggFunc {
295                    ignore_nulls: None,
296                    having_max: None,
297                    this: case_expr,
298                    distinct: f.distinct,
299                    filter: f.filter,
300                    order_by: Vec::new(),
301                    name: None,
302                    limit: None,
303                    inferred_type: None,
304                })))
305            }
306
307            // Pass through everything else
308            _ => Ok(Expression::AggregateFunction(f)),
309        }
310    }
311
312    fn transform_cast(&self, c: Cast) -> Result<Expression> {
313        // CockroachDB type mappings are handled in the generator
314        Ok(Expression::Cast(Box::new(c)))
315    }
316}