Skip to main content

polyglot_sql/dialects/
athena.rs

1//! Athena Dialect
2//!
3//! AWS Athena-specific transformations based on sqlglot patterns.
4//! Athena routes between Hive (DDL) and Trino (DML) engines:
5//!
6//! - **Hive** (backticks): CREATE EXTERNAL TABLE, CREATE TABLE (no AS SELECT),
7//!   ALTER, DROP (except VIEW), DESCRIBE, SHOW
8//! - **Trino** (double quotes): CREATE VIEW, CREATE TABLE AS SELECT, DROP VIEW,
9//!   SELECT, INSERT, UPDATE, DELETE, MERGE
10
11use super::{DialectImpl, DialectType};
12use crate::error::Result;
13use crate::expressions::{
14    AggFunc, Case, Cast, DataType, Expression, Function, LikeOp, UnaryFunc, VarArgFunc,
15};
16#[cfg(feature = "generate")]
17use crate::generator::{GeneratorConfig, IdentifierQuoteStyle};
18use crate::tokens::TokenizerConfig;
19
20/// Athena dialect (based on Trino for DML operations)
21pub struct AthenaDialect;
22
23impl DialectImpl for AthenaDialect {
24    fn dialect_type(&self) -> DialectType {
25        DialectType::Athena
26    }
27
28    fn tokenizer_config(&self) -> TokenizerConfig {
29        let mut config = TokenizerConfig::default();
30        // Athena uses double quotes for identifiers (Trino-style for DML)
31        config.identifiers.insert('"', '"');
32        // Also supports backticks (Hive-style for DDL)
33        config.identifiers.insert('`', '`');
34        config.nested_comments = false;
35        // Athena/Hive supports backslash escapes in string literals (e.g., \' for escaped quote)
36        config.string_escapes.push('\\');
37        config
38    }
39
40    #[cfg(feature = "generate")]
41
42    fn generator_config(&self) -> GeneratorConfig {
43        // Default config uses Trino style (double quotes)
44        GeneratorConfig {
45            identifier_quote: '"',
46            identifier_quote_style: IdentifierQuoteStyle::DOUBLE_QUOTE,
47            dialect: Some(DialectType::Athena),
48            schema_comment_with_eq: false,
49            ..Default::default()
50        }
51    }
52
53    #[cfg(feature = "generate")]
54
55    fn generator_config_for_expr(&self, expr: &Expression) -> GeneratorConfig {
56        if should_use_hive_engine(expr) {
57            // Hive mode: backticks for identifiers
58            GeneratorConfig {
59                identifier_quote: '`',
60                identifier_quote_style: IdentifierQuoteStyle::BACKTICK,
61                dialect: Some(DialectType::Athena),
62                schema_comment_with_eq: false,
63                ..Default::default()
64            }
65        } else {
66            // Trino mode: double quotes for identifiers
67            GeneratorConfig {
68                identifier_quote: '"',
69                identifier_quote_style: IdentifierQuoteStyle::DOUBLE_QUOTE,
70                dialect: Some(DialectType::Athena),
71                schema_comment_with_eq: false,
72                ..Default::default()
73            }
74        }
75    }
76
77    #[cfg(feature = "transpile")]
78
79    fn transform_expr(&self, expr: Expression) -> Result<Expression> {
80        match expr {
81            // IFNULL -> COALESCE in Athena
82            Expression::IfNull(f) => Ok(Expression::Coalesce(Box::new(VarArgFunc {
83                original_name: None,
84                expressions: vec![f.this, f.expression],
85                inferred_type: None,
86            }))),
87
88            // NVL -> COALESCE in Athena
89            Expression::Nvl(f) => Ok(Expression::Coalesce(Box::new(VarArgFunc {
90                original_name: None,
91                expressions: vec![f.this, f.expression],
92                inferred_type: None,
93            }))),
94
95            // Coalesce with original_name (e.g., IFNULL parsed as Coalesce) -> clear original_name
96            Expression::Coalesce(mut f) => {
97                f.original_name = None;
98                Ok(Expression::Coalesce(f))
99            }
100
101            // TryCast stays as TryCast (Athena/Trino supports TRY_CAST)
102            Expression::TryCast(c) => Ok(Expression::TryCast(c)),
103
104            // SafeCast -> TRY_CAST in Athena
105            Expression::SafeCast(c) => Ok(Expression::TryCast(c)),
106
107            // ILike -> LOWER() LIKE LOWER() (Trino doesn't support ILIKE)
108            Expression::ILike(op) => {
109                let lower_left = Expression::Lower(Box::new(UnaryFunc::new(op.left.clone())));
110                let lower_right = Expression::Lower(Box::new(UnaryFunc::new(op.right.clone())));
111                Ok(Expression::Like(Box::new(LikeOp {
112                    left: lower_left,
113                    right: lower_right,
114                    escape: op.escape,
115                    quantifier: op.quantifier.clone(),
116                    inferred_type: None,
117                })))
118            }
119
120            // CountIf -> SUM(CASE WHEN condition THEN 1 ELSE 0 END)
121            Expression::CountIf(f) => {
122                let case_expr = Expression::Case(Box::new(Case {
123                    operand: None,
124                    whens: vec![(f.this.clone(), Expression::number(1))],
125                    else_: Some(Expression::number(0)),
126                    comments: Vec::new(),
127                    inferred_type: None,
128                }));
129                Ok(Expression::Sum(Box::new(AggFunc {
130                    ignore_nulls: None,
131                    having_max: None,
132                    this: case_expr,
133                    distinct: f.distinct,
134                    filter: f.filter,
135                    order_by: Vec::new(),
136                    name: None,
137                    limit: None,
138                    inferred_type: None,
139                })))
140            }
141
142            // EXPLODE -> UNNEST in Athena
143            Expression::Explode(f) => Ok(Expression::Unnest(Box::new(
144                crate::expressions::UnnestFunc {
145                    this: f.this,
146                    expressions: Vec::new(),
147                    with_ordinality: false,
148                    alias: None,
149                    offset_alias: None,
150                    inferred_type: None,
151                },
152            ))),
153
154            // ExplodeOuter -> UNNEST in Athena
155            Expression::ExplodeOuter(f) => Ok(Expression::Unnest(Box::new(
156                crate::expressions::UnnestFunc {
157                    this: f.this,
158                    expressions: Vec::new(),
159                    with_ordinality: false,
160                    alias: None,
161                    offset_alias: None,
162                    inferred_type: None,
163                },
164            ))),
165
166            // Generic function transformations
167            Expression::Function(f) => self.transform_function(*f),
168
169            // Generic aggregate function transformations
170            Expression::AggregateFunction(f) => self.transform_aggregate_function(f),
171
172            // Cast transformations
173            Expression::Cast(c) => self.transform_cast(*c),
174
175            // Pass through everything else
176            _ => Ok(expr),
177        }
178    }
179}
180
181#[cfg(feature = "transpile")]
182impl AthenaDialect {
183    fn transform_function(&self, f: Function) -> Result<Expression> {
184        let name_upper = f.name.to_uppercase();
185        match name_upper.as_str() {
186            // IFNULL -> COALESCE
187            "IFNULL" if f.args.len() == 2 => Ok(Expression::Coalesce(Box::new(VarArgFunc {
188                original_name: None,
189                expressions: f.args,
190                inferred_type: None,
191            }))),
192
193            // NVL -> COALESCE
194            "NVL" if f.args.len() == 2 => Ok(Expression::Coalesce(Box::new(VarArgFunc {
195                original_name: None,
196                expressions: f.args,
197                inferred_type: None,
198            }))),
199
200            // ISNULL -> COALESCE
201            "ISNULL" if f.args.len() == 2 => Ok(Expression::Coalesce(Box::new(VarArgFunc {
202                original_name: None,
203                expressions: f.args,
204                inferred_type: None,
205            }))),
206
207            // GETDATE -> CURRENT_TIMESTAMP
208            "GETDATE" => Ok(Expression::CurrentTimestamp(
209                crate::expressions::CurrentTimestamp {
210                    precision: None,
211                    sysdate: false,
212                },
213            )),
214
215            // NOW -> CURRENT_TIMESTAMP
216            "NOW" => Ok(Expression::CurrentTimestamp(
217                crate::expressions::CurrentTimestamp {
218                    precision: None,
219                    sysdate: false,
220                },
221            )),
222
223            // RAND -> RANDOM in Athena
224            "RAND" => Ok(Expression::Function(Box::new(Function::new(
225                "RANDOM".to_string(),
226                vec![],
227            )))),
228
229            // GROUP_CONCAT -> LISTAGG in Athena (Trino-style)
230            "GROUP_CONCAT" if !f.args.is_empty() => Ok(Expression::Function(Box::new(
231                Function::new("LISTAGG".to_string(), f.args),
232            ))),
233
234            // STRING_AGG -> LISTAGG in Athena
235            "STRING_AGG" if !f.args.is_empty() => Ok(Expression::Function(Box::new(
236                Function::new("LISTAGG".to_string(), f.args),
237            ))),
238
239            // SUBSTR -> SUBSTRING
240            "SUBSTR" => Ok(Expression::Function(Box::new(Function::new(
241                "SUBSTRING".to_string(),
242                f.args,
243            )))),
244
245            // LEN -> LENGTH
246            "LEN" if f.args.len() == 1 => Ok(Expression::Length(Box::new(UnaryFunc::new(
247                f.args.into_iter().next().unwrap(),
248            )))),
249
250            // CHARINDEX -> STRPOS in Athena (with swapped args)
251            "CHARINDEX" if f.args.len() >= 2 => {
252                let mut args = f.args;
253                let substring = args.remove(0);
254                let string = args.remove(0);
255                Ok(Expression::Function(Box::new(Function::new(
256                    "STRPOS".to_string(),
257                    vec![string, substring],
258                ))))
259            }
260
261            // INSTR -> STRPOS
262            "INSTR" if f.args.len() >= 2 => Ok(Expression::Function(Box::new(Function::new(
263                "STRPOS".to_string(),
264                f.args,
265            )))),
266
267            // LOCATE -> STRPOS in Athena (with swapped args)
268            "LOCATE" if f.args.len() >= 2 => {
269                let mut args = f.args;
270                let substring = args.remove(0);
271                let string = args.remove(0);
272                Ok(Expression::Function(Box::new(Function::new(
273                    "STRPOS".to_string(),
274                    vec![string, substring],
275                ))))
276            }
277
278            // ARRAY_LENGTH -> CARDINALITY in Athena
279            "ARRAY_LENGTH" if f.args.len() == 1 => Ok(Expression::Function(Box::new(
280                Function::new("CARDINALITY".to_string(), f.args),
281            ))),
282
283            // SIZE -> CARDINALITY in Athena
284            "SIZE" if f.args.len() == 1 => Ok(Expression::Function(Box::new(Function::new(
285                "CARDINALITY".to_string(),
286                f.args,
287            )))),
288
289            // TO_DATE -> CAST to DATE or DATE_PARSE
290            "TO_DATE" if !f.args.is_empty() => {
291                if f.args.len() == 1 {
292                    Ok(Expression::Cast(Box::new(Cast {
293                        this: f.args.into_iter().next().unwrap(),
294                        to: DataType::Date,
295                        trailing_comments: Vec::new(),
296                        double_colon_syntax: false,
297                        format: None,
298                        default: None,
299                        inferred_type: None,
300                    })))
301                } else {
302                    Ok(Expression::Function(Box::new(Function::new(
303                        "DATE_PARSE".to_string(),
304                        f.args,
305                    ))))
306                }
307            }
308
309            // TO_TIMESTAMP -> CAST or DATE_PARSE
310            "TO_TIMESTAMP" if !f.args.is_empty() => {
311                if f.args.len() == 1 {
312                    Ok(Expression::Cast(Box::new(Cast {
313                        this: f.args.into_iter().next().unwrap(),
314                        to: DataType::Timestamp {
315                            precision: None,
316                            timezone: false,
317                        },
318                        trailing_comments: Vec::new(),
319                        double_colon_syntax: false,
320                        format: None,
321                        default: None,
322                        inferred_type: None,
323                    })))
324                } else {
325                    Ok(Expression::Function(Box::new(Function::new(
326                        "DATE_PARSE".to_string(),
327                        f.args,
328                    ))))
329                }
330            }
331
332            // strftime -> DATE_FORMAT in Athena
333            "STRFTIME" if f.args.len() >= 2 => {
334                let mut args = f.args;
335                let format = args.remove(0);
336                let date = args.remove(0);
337                Ok(Expression::Function(Box::new(Function::new(
338                    "DATE_FORMAT".to_string(),
339                    vec![date, format],
340                ))))
341            }
342
343            // TO_CHAR -> DATE_FORMAT in Athena
344            "TO_CHAR" if f.args.len() >= 2 => Ok(Expression::Function(Box::new(Function::new(
345                "DATE_FORMAT".to_string(),
346                f.args,
347            )))),
348
349            // GET_JSON_OBJECT -> JSON_EXTRACT_SCALAR in Athena
350            "GET_JSON_OBJECT" if f.args.len() == 2 => Ok(Expression::Function(Box::new(
351                Function::new("JSON_EXTRACT_SCALAR".to_string(), f.args),
352            ))),
353
354            // COLLECT_LIST -> ARRAY_AGG
355            "COLLECT_LIST" if !f.args.is_empty() => Ok(Expression::Function(Box::new(
356                Function::new("ARRAY_AGG".to_string(), f.args),
357            ))),
358
359            // Pass through everything else
360            _ => Ok(Expression::Function(Box::new(f))),
361        }
362    }
363
364    fn transform_aggregate_function(
365        &self,
366        f: Box<crate::expressions::AggregateFunction>,
367    ) -> Result<Expression> {
368        let name_upper = f.name.to_uppercase();
369        match name_upper.as_str() {
370            // COUNT_IF -> SUM(CASE WHEN...)
371            "COUNT_IF" if !f.args.is_empty() => {
372                let condition = f.args.into_iter().next().unwrap();
373                let case_expr = Expression::Case(Box::new(Case {
374                    operand: None,
375                    whens: vec![(condition, Expression::number(1))],
376                    else_: Some(Expression::number(0)),
377                    comments: Vec::new(),
378                    inferred_type: None,
379                }));
380                Ok(Expression::Sum(Box::new(AggFunc {
381                    ignore_nulls: None,
382                    having_max: None,
383                    this: case_expr,
384                    distinct: f.distinct,
385                    filter: f.filter,
386                    order_by: Vec::new(),
387                    name: None,
388                    limit: None,
389                    inferred_type: None,
390                })))
391            }
392
393            // ANY_VALUE -> ARBITRARY in Athena (Trino)
394            "ANY_VALUE" if !f.args.is_empty() => Ok(Expression::Function(Box::new(Function::new(
395                "ARBITRARY".to_string(),
396                f.args,
397            )))),
398
399            // GROUP_CONCAT -> LISTAGG in Athena
400            "GROUP_CONCAT" if !f.args.is_empty() => Ok(Expression::Function(Box::new(
401                Function::new("LISTAGG".to_string(), f.args),
402            ))),
403
404            // STRING_AGG -> LISTAGG in Athena
405            "STRING_AGG" if !f.args.is_empty() => Ok(Expression::Function(Box::new(
406                Function::new("LISTAGG".to_string(), f.args),
407            ))),
408
409            // Pass through everything else
410            _ => Ok(Expression::AggregateFunction(f)),
411        }
412    }
413
414    fn transform_cast(&self, c: Cast) -> Result<Expression> {
415        // Athena type mappings are handled in the generator
416        Ok(Expression::Cast(Box::new(c)))
417    }
418}
419
420/// Determine if an expression should be generated using Hive engine (backticks)
421/// or Trino engine (double quotes).
422///
423/// Hive is used for:
424/// - CREATE EXTERNAL TABLE
425/// - CREATE TABLE (without AS SELECT)
426/// - CREATE SCHEMA / CREATE DATABASE
427/// - ALTER statements
428/// - DROP statements (except DROP VIEW)
429/// - DESCRIBE / SHOW statements
430///
431/// Trino is used for everything else (DML, CREATE VIEW, etc.)
432fn should_use_hive_engine(expr: &Expression) -> bool {
433    match expr {
434        // CREATE TABLE: Hive if EXTERNAL or no AS SELECT
435        Expression::CreateTable(ct) => {
436            // CREATE EXTERNAL TABLE → Hive
437            if let Some(ref modifier) = ct.table_modifier {
438                if modifier.to_uppercase() == "EXTERNAL" {
439                    return true;
440                }
441            }
442            // CREATE TABLE ... AS SELECT → Trino
443            // CREATE TABLE (without query) → Hive
444            ct.as_select.is_none()
445        }
446
447        // CREATE VIEW → Trino
448        Expression::CreateView(_) => false,
449
450        // CREATE SCHEMA / DATABASE → Hive
451        Expression::CreateSchema(_) => true,
452        Expression::CreateDatabase(_) => true,
453
454        // ALTER statements → Hive
455        Expression::AlterTable(_) => true,
456        Expression::AlterView(_) => true,
457        Expression::AlterIndex(_) => true,
458        Expression::AlterSequence(_) => true,
459
460        // DROP VIEW → Trino (because CREATE VIEW is Trino)
461        Expression::DropView(_) => false,
462
463        // Other DROP statements → Hive
464        Expression::DropTable(_) => true,
465        Expression::DropSchema(_) => true,
466        Expression::DropDatabase(_) => true,
467        Expression::DropIndex(_) => true,
468        Expression::DropFunction(_) => true,
469        Expression::DropProcedure(_) => true,
470        Expression::DropSequence(_) => true,
471
472        // DESCRIBE / SHOW → Hive
473        Expression::Describe(_) => true,
474        Expression::Show(_) => true,
475
476        // Everything else (SELECT, INSERT, UPDATE, DELETE, MERGE, etc.) → Trino
477        _ => false,
478    }
479}