sqlglot.dialects.athena

 15class Athena(Trino):
 16    """
 17    Over the years, it looks like AWS has taken various execution engines, bolted on AWS-specific
 18    modifications and then built the Athena service around them.
 19
 20    Thus, Athena is not simply hosted Trino, it's more like a router that routes SQL queries to an
 21    execution engine depending on the query type.
 22
 23    As at 2024-09-10, assuming your Athena workgroup is configured to use "Athena engine version 3",
 24    the following engines exist:
 25
 26    Hive:
 27     - Accepts mostly the same syntax as Hadoop / Hive
 28     - Uses backticks to quote identifiers
 29     - Has a distinctive DDL syntax (around things like setting table properties, storage locations etc)
 30       that is different from Trino
 31     - Used for *most* DDL, with some exceptions that get routed to the Trino engine instead:
 32        - CREATE [EXTERNAL] TABLE (without AS SELECT)
 33        - ALTER
 34        - DROP
 35
 36    Trino:
 37      - Uses double quotes to quote identifiers
 38      - Used for DDL operations that involve SELECT queries, eg:
 39        - CREATE VIEW / DROP VIEW
 40        - CREATE TABLE... AS SELECT
 41      - Used for DML operations
 42        - SELECT, INSERT, UPDATE, DELETE, MERGE
 43
 44    The SQLGlot Athena dialect tries to identify which engine a query would be routed to and then uses the
 45    tokenizer / parser / generator for that engine. This is unfortunately necessary, as there are certain
 46    incompatibilities between the engines' dialects and thus can't be handled by a single, unifying dialect.
 47
 48    References:
 49    - https://docs.aws.amazon.com/athena/latest/ug/ddl-reference.html
 50    - https://docs.aws.amazon.com/athena/latest/ug/dml-queries-functions-operators.html
 51    """
 52
 53    # This Tokenizer consumes a combination of HiveQL and Trino SQL and then processes the tokens
 54    # to disambiguate which dialect needs to be actually used in order to tokenize correctly.
 55    class Tokenizer(tokens.Tokenizer):
 56        IDENTIFIERS = Trino.Tokenizer.IDENTIFIERS + Hive.Tokenizer.IDENTIFIERS
 57        STRING_ESCAPES = Trino.Tokenizer.STRING_ESCAPES + Hive.Tokenizer.STRING_ESCAPES
 58        HEX_STRINGS = Trino.Tokenizer.HEX_STRINGS + Hive.Tokenizer.HEX_STRINGS
 59        UNICODE_STRINGS = Trino.Tokenizer.UNICODE_STRINGS + Hive.Tokenizer.UNICODE_STRINGS
 60
 61        NUMERIC_LITERALS = {
 62            **Trino.Tokenizer.NUMERIC_LITERALS,
 63            **Hive.Tokenizer.NUMERIC_LITERALS,
 64        }
 65
 66        KEYWORDS = {
 67            **Hive.Tokenizer.KEYWORDS,
 68            **Trino.Tokenizer.KEYWORDS,
 69            "UNLOAD": TokenType.COMMAND,
 70        }
 71
 72        def __init__(self, dialect: DialectType = None) -> None:
 73            super().__init__(dialect=dialect)
 74
 75            self._hive_tokenizer = Hive().tokenizer()
 76            self._trino_tokenizer = _TrinoTokenizer(Trino())
 77
 78        def tokenize(self, sql: str) -> list[Token]:
 79            # DML only needs Trino's tokenizer; possible Hive statements need disambiguation.
 80            try:
 81                trino_tokens: list[Token] | None = self._trino_tokenizer.tokenize(sql)
 82            except TokenError:
 83                trino_tokens = None
 84
 85            if trino_tokens is not None and not _may_be_hive(trino_tokens):
 86                return trino_tokens
 87
 88            try:
 89                tokens = super().tokenize(sql)
 90            except TokenError:
 91                # Combined rules can reject literals valid in only one engine.
 92                if trino_tokens is not None and not _tokenize_as_hive(trino_tokens):
 93                    return trino_tokens
 94
 95                try:
 96                    tokens = self._hive_tokenizer.tokenize(sql)
 97                except TokenError:
 98                    pass
 99                else:
100                    if _tokenize_as_hive(tokens):
101                        return [Token(TokenType.HIVE_TOKEN_STREAM, "")] + tokens
102                raise
103
104            if _tokenize_as_hive(tokens):
105                return [Token(TokenType.HIVE_TOKEN_STREAM, "")] + self._hive_tokenizer.tokenize(sql)
106
107            return trino_tokens if trino_tokens is not None else self._trino_tokenizer.tokenize(sql)
108
109    Parser = AthenaParser
110
111    Generator = AthenaGenerator

 55    class Tokenizer(tokens.Tokenizer):
 56        IDENTIFIERS = Trino.Tokenizer.IDENTIFIERS + Hive.Tokenizer.IDENTIFIERS
 57        STRING_ESCAPES = Trino.Tokenizer.STRING_ESCAPES + Hive.Tokenizer.STRING_ESCAPES
 58        HEX_STRINGS = Trino.Tokenizer.HEX_STRINGS + Hive.Tokenizer.HEX_STRINGS
 59        UNICODE_STRINGS = Trino.Tokenizer.UNICODE_STRINGS + Hive.Tokenizer.UNICODE_STRINGS
 60
 61        NUMERIC_LITERALS = {
 62            **Trino.Tokenizer.NUMERIC_LITERALS,
 63            **Hive.Tokenizer.NUMERIC_LITERALS,
 64        }
 65
 66        KEYWORDS = {
 67            **Hive.Tokenizer.KEYWORDS,
 68            **Trino.Tokenizer.KEYWORDS,
 69            "UNLOAD": TokenType.COMMAND,
 70        }
 71
 72        def __init__(self, dialect: DialectType = None) -> None:
 73            super().__init__(dialect=dialect)
 74
 75            self._hive_tokenizer = Hive().tokenizer()
 76            self._trino_tokenizer = _TrinoTokenizer(Trino())
 77
 78        def tokenize(self, sql: str) -> list[Token]:
 79            # DML only needs Trino's tokenizer; possible Hive statements need disambiguation.
 80            try:
 81                trino_tokens: list[Token] | None = self._trino_tokenizer.tokenize(sql)
 82            except TokenError:
 83                trino_tokens = None
 84
 85            if trino_tokens is not None and not _may_be_hive(trino_tokens):
 86                return trino_tokens
 87
 88            try:
 89                tokens = super().tokenize(sql)
 90            except TokenError:
 91                # Combined rules can reject literals valid in only one engine.
 92                if trino_tokens is not None and not _tokenize_as_hive(trino_tokens):
 93                    return trino_tokens
 94
 95                try:
 96                    tokens = self._hive_tokenizer.tokenize(sql)
 97                except TokenError:
 98                    pass
 99                else:
100                    if _tokenize_as_hive(tokens):
101                        return [Token(TokenType.HIVE_TOKEN_STREAM, "")] + tokens
102                raise
103
104            if _tokenize_as_hive(tokens):
105                return [Token(TokenType.HIVE_TOKEN_STREAM, "")] + self._hive_tokenizer.tokenize(sql)
106
107            return trino_tokens if trino_tokens is not None else self._trino_tokenizer.tokenize(sql)