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