sqlglot.dialects.athena

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp, generator, parser, tokens
  6from sqlglot.dialects import Dialect, Hive, Trino
  7from sqlglot.tokens import TokenType, Token
  8
  9
 10class Athena(Dialect):
 11    """
 12    Over the years, it looks like AWS has taken various execution engines, bolted on AWS-specific
 13    modifications and then built the Athena service around them.
 14
 15    Thus, Athena is not simply hosted Trino, it's more like a router that routes SQL queries to an
 16    execution engine depending on the query type.
 17
 18    As at 2024-09-10, assuming your Athena workgroup is configured to use "Athena engine version 3",
 19    the following engines exist:
 20
 21    Hive:
 22     - Accepts mostly the same syntax as Hadoop / Hive
 23     - Uses backticks to quote identifiers
 24     - Has a distinctive DDL syntax (around things like setting table properties, storage locations etc)
 25       that is different from Trino
 26     - Used for *most* DDL, with some exceptions that get routed to the Trino engine instead:
 27        - CREATE [EXTERNAL] TABLE (without AS SELECT)
 28        - ALTER
 29        - DROP
 30
 31    Trino:
 32      - Uses double quotes to quote identifiers
 33      - Used for DDL operations that involve SELECT queries, eg:
 34        - CREATE VIEW / DROP VIEW
 35        - CREATE TABLE... AS SELECT
 36      - Used for DML operations
 37        - SELECT, INSERT, UPDATE, DELETE, MERGE
 38
 39    The SQLGlot Athena dialect tries to identify which engine a query would be routed to and then uses the
 40    tokenizer / parser / generator for that engine. This is unfortunately necessary, as there are certain
 41    incompatibilities between the engines' dialects and thus can't be handled by a single, unifying dialect.
 42
 43    References:
 44    - https://docs.aws.amazon.com/athena/latest/ug/ddl-reference.html
 45    - https://docs.aws.amazon.com/athena/latest/ug/dml-queries-functions-operators.html
 46    """
 47
 48    def __init__(self, **kwargs):
 49        super().__init__(**kwargs)
 50
 51        self._hive = Hive(**kwargs)
 52        self._trino = Trino(**kwargs)
 53
 54    def tokenize(self, sql: str, **opts) -> t.List[Token]:
 55        opts["hive"] = self._hive
 56        opts["trino"] = self._trino
 57        return super().tokenize(sql, **opts)
 58
 59    def parse(self, sql: str, **opts) -> t.List[t.Optional[exp.Expression]]:
 60        opts["hive"] = self._hive
 61        opts["trino"] = self._trino
 62        return super().parse(sql, **opts)
 63
 64    def parse_into(
 65        self, expression_type: exp.IntoType, sql: str, **opts
 66    ) -> t.List[t.Optional[exp.Expression]]:
 67        opts["hive"] = self._hive
 68        opts["trino"] = self._trino
 69        return super().parse_into(expression_type, sql, **opts)
 70
 71    def generate(self, expression: exp.Expression, copy: bool = True, **opts) -> str:
 72        opts["hive"] = self._hive
 73        opts["trino"] = self._trino
 74        return super().generate(expression, copy=copy, **opts)
 75
 76    # This Tokenizer consumes a combination of HiveQL and Trino SQL and then processes the tokens
 77    # to disambiguate which dialect needs to be actually used in order to tokenize correctly.
 78    class Tokenizer(tokens.Tokenizer):
 79        IDENTIFIERS = Trino.Tokenizer.IDENTIFIERS + Hive.Tokenizer.IDENTIFIERS
 80        STRING_ESCAPES = Trino.Tokenizer.STRING_ESCAPES + Hive.Tokenizer.STRING_ESCAPES
 81        HEX_STRINGS = Trino.Tokenizer.HEX_STRINGS + Hive.Tokenizer.HEX_STRINGS
 82        UNICODE_STRINGS = Trino.Tokenizer.UNICODE_STRINGS + Hive.Tokenizer.UNICODE_STRINGS
 83
 84        NUMERIC_LITERALS = {
 85            **Trino.Tokenizer.NUMERIC_LITERALS,
 86            **Hive.Tokenizer.NUMERIC_LITERALS,
 87        }
 88
 89        KEYWORDS = {
 90            **Hive.Tokenizer.KEYWORDS,
 91            **Trino.Tokenizer.KEYWORDS,
 92            "UNLOAD": TokenType.COMMAND,
 93        }
 94
 95        def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
 96            hive = kwargs.pop("hive", None) or Hive()
 97            trino = kwargs.pop("trino", None) or Trino()
 98
 99            super().__init__(*args, **kwargs)
100
101            self._hive_tokenizer = hive.tokenizer(*args, **{**kwargs, "dialect": hive})
102            self._trino_tokenizer = _TrinoTokenizer(*args, **{**kwargs, "dialect": trino})
103
104        def tokenize(self, sql: str) -> t.List[Token]:
105            tokens = super().tokenize(sql)
106
107            if _tokenize_as_hive(tokens):
108                return [Token(TokenType.HIVE_TOKEN_STREAM, "")] + self._hive_tokenizer.tokenize(sql)
109
110            return self._trino_tokenizer.tokenize(sql)
111
112    class Parser(parser.Parser):
113        def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
114            hive = kwargs.pop("hive", None) or Hive()
115            trino = kwargs.pop("trino", None) or Trino()
116
117            super().__init__(*args, **kwargs)
118
119            self._hive_parser = hive.parser(*args, **{**kwargs, "dialect": hive})
120            self._trino_parser = _TrinoParser(*args, **{**kwargs, "dialect": trino})
121
122        def parse(
123            self, raw_tokens: t.List[Token], sql: t.Optional[str] = None
124        ) -> t.List[t.Optional[exp.Expression]]:
125            if raw_tokens and raw_tokens[0].token_type == TokenType.HIVE_TOKEN_STREAM:
126                return self._hive_parser.parse(raw_tokens[1:], sql)
127
128            return self._trino_parser.parse(raw_tokens, sql)
129
130        def parse_into(
131            self,
132            expression_types: exp.IntoType,
133            raw_tokens: t.List[Token],
134            sql: t.Optional[str] = None,
135        ) -> t.List[t.Optional[exp.Expression]]:
136            if raw_tokens and raw_tokens[0].token_type == TokenType.HIVE_TOKEN_STREAM:
137                return self._hive_parser.parse_into(expression_types, raw_tokens[1:], sql)
138
139            return self._trino_parser.parse_into(expression_types, raw_tokens, sql)
140
141    class Generator(generator.Generator):
142        def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
143            hive = kwargs.pop("hive", None) or Hive()
144            trino = kwargs.pop("trino", None) or Trino()
145
146            super().__init__(*args, **kwargs)
147
148            self._hive_generator = _HiveGenerator(*args, **{**kwargs, "dialect": hive})
149            self._trino_generator = _TrinoGenerator(*args, **{**kwargs, "dialect": trino})
150
151        def generate(self, expression: exp.Expression, copy: bool = True) -> str:
152            if _generate_as_hive(expression):
153                generator = self._hive_generator
154            else:
155                generator = self._trino_generator
156
157            return generator.generate(expression, copy=copy)
158
159
160def _tokenize_as_hive(tokens: t.List[Token]) -> bool:
161    if len(tokens) < 2:
162        return False
163
164    first, second, *rest = tokens
165
166    first_type = first.token_type
167    first_text = first.text.upper()
168    second_type = second.token_type
169    second_text = second.text.upper()
170
171    if first_type in (TokenType.DESCRIBE, TokenType.SHOW) or first_text == "MSCK REPAIR":
172        return True
173
174    if first_type in (TokenType.ALTER, TokenType.CREATE, TokenType.DROP):
175        if second_text in ("DATABASE", "EXTERNAL", "SCHEMA"):
176            return True
177        if second_type == TokenType.VIEW:
178            return False
179
180        return all(t.token_type != TokenType.SELECT for t in rest)
181
182    return False
183
184
185def _generate_as_hive(expression: exp.Expression) -> bool:
186    if isinstance(expression, exp.Create):
187        if expression.kind == "TABLE":
188            properties = expression.args.get("properties")
189
190            # CREATE EXTERNAL TABLE is Hive
191            if properties and properties.find(exp.ExternalProperty):
192                return True
193
194            # Any CREATE TABLE other than CREATE TABLE ... AS <query> is Hive
195            if not isinstance(expression.expression, exp.Query):
196                return True
197        else:
198            # CREATE VIEW is Trino, but CREATE SCHEMA, CREATE DATABASE, etc, is Hive
199            return expression.kind != "VIEW"
200    elif isinstance(expression, (exp.Alter, exp.Drop, exp.Describe, exp.Show)):
201        if isinstance(expression, exp.Drop) and expression.kind == "VIEW":
202            # DROP VIEW is Trino, because CREATE VIEW is as well
203            return False
204
205        # Everything else, e.g., ALTER statements, is Hive
206        return True
207
208    return False
209
210
211def _is_iceberg_table(properties: exp.Properties) -> bool:
212    for p in properties.expressions:
213        if isinstance(p, exp.Property) and p.name == "table_type":
214            return p.text("value").lower() == "iceberg"
215
216    return False
217
218
219def _location_property_sql(self: Athena.Generator, e: exp.LocationProperty):
220    # If table_type='iceberg', the LocationProperty is called 'location'
221    # Otherwise, it's called 'external_location'
222    # ref: https://docs.aws.amazon.com/athena/latest/ug/create-table-as.html
223
224    prop_name = "external_location"
225
226    if isinstance(e.parent, exp.Properties):
227        if _is_iceberg_table(e.parent):
228            prop_name = "location"
229
230    return f"{prop_name}={self.sql(e, 'this')}"
231
232
233def _partitioned_by_property_sql(self: Athena.Generator, e: exp.PartitionedByProperty) -> str:
234    # If table_type='iceberg' then the table property for partitioning is called 'partitioning'
235    # If table_type='hive' it's called 'partitioned_by'
236    # ref: https://docs.aws.amazon.com/athena/latest/ug/create-table-as.html#ctas-table-properties
237
238    prop_name = "partitioned_by"
239
240    if isinstance(e.parent, exp.Properties):
241        if _is_iceberg_table(e.parent):
242            prop_name = "partitioning"
243
244    return f"{prop_name}={self.sql(e, 'this')}"
245
246
247# Athena extensions to Hive's generator
248class _HiveGenerator(Hive.Generator):
249    def alter_sql(self, expression: exp.Alter) -> str:
250        # Package any ALTER TABLE ADD actions into a Schema object, so it gets generated as
251        # `ALTER TABLE .. ADD COLUMNS(...)`, instead of `ALTER TABLE ... ADD COLUMN`, which
252        # is invalid syntax on Athena
253        if isinstance(expression, exp.Alter) and expression.kind == "TABLE":
254            if expression.actions and isinstance(expression.actions[0], exp.ColumnDef):
255                new_actions = exp.Schema(expressions=expression.actions)
256                expression.set("actions", [new_actions])
257
258        return super().alter_sql(expression)
259
260
261# Athena extensions to Trino's tokenizer
262class _TrinoTokenizer(Trino.Tokenizer):
263    KEYWORDS = {
264        **Trino.Tokenizer.KEYWORDS,
265        "UNLOAD": TokenType.COMMAND,
266    }
267
268
269# Athena extensions to Trino's parser
270class _TrinoParser(Trino.Parser):
271    STATEMENT_PARSERS = {
272        **Trino.Parser.STATEMENT_PARSERS,
273        TokenType.USING: lambda self: self._parse_as_command(self._prev),
274    }
275
276
277# Athena extensions to Trino's generator
278class _TrinoGenerator(Trino.Generator):
279    PROPERTIES_LOCATION = {
280        **Trino.Generator.PROPERTIES_LOCATION,
281        exp.LocationProperty: exp.Properties.Location.POST_WITH,
282    }
283
284    TRANSFORMS = {
285        **Trino.Generator.TRANSFORMS,
286        exp.PartitionedByProperty: _partitioned_by_property_sql,
287        exp.LocationProperty: _location_property_sql,
288    }