sqlglot.tokens

   1from __future__ import annotations
   2
   3import typing as t
   4from enum import auto
   5
   6from sqlglot.helper import AutoName
   7from sqlglot.trie import in_trie, new_trie
   8
   9
  10class TokenType(AutoName):
  11    L_PAREN = auto()
  12    R_PAREN = auto()
  13    L_BRACKET = auto()
  14    R_BRACKET = auto()
  15    L_BRACE = auto()
  16    R_BRACE = auto()
  17    COMMA = auto()
  18    DOT = auto()
  19    DASH = auto()
  20    PLUS = auto()
  21    COLON = auto()
  22    DCOLON = auto()
  23    SEMICOLON = auto()
  24    STAR = auto()
  25    BACKSLASH = auto()
  26    SLASH = auto()
  27    LT = auto()
  28    LTE = auto()
  29    GT = auto()
  30    GTE = auto()
  31    NOT = auto()
  32    EQ = auto()
  33    NEQ = auto()
  34    NULLSAFE_EQ = auto()
  35    AND = auto()
  36    OR = auto()
  37    AMP = auto()
  38    DPIPE = auto()
  39    PIPE = auto()
  40    CARET = auto()
  41    TILDA = auto()
  42    ARROW = auto()
  43    DARROW = auto()
  44    FARROW = auto()
  45    HASH = auto()
  46    HASH_ARROW = auto()
  47    DHASH_ARROW = auto()
  48    LR_ARROW = auto()
  49    LT_AT = auto()
  50    AT_GT = auto()
  51    DOLLAR = auto()
  52    PARAMETER = auto()
  53    SESSION_PARAMETER = auto()
  54    DAMP = auto()
  55
  56    BLOCK_START = auto()
  57    BLOCK_END = auto()
  58
  59    SPACE = auto()
  60    BREAK = auto()
  61
  62    STRING = auto()
  63    NUMBER = auto()
  64    IDENTIFIER = auto()
  65    DATABASE = auto()
  66    COLUMN = auto()
  67    COLUMN_DEF = auto()
  68    SCHEMA = auto()
  69    TABLE = auto()
  70    VAR = auto()
  71    BIT_STRING = auto()
  72    HEX_STRING = auto()
  73    BYTE_STRING = auto()
  74    NATIONAL_STRING = auto()
  75    RAW_STRING = auto()
  76
  77    # types
  78    BIT = auto()
  79    BOOLEAN = auto()
  80    TINYINT = auto()
  81    UTINYINT = auto()
  82    SMALLINT = auto()
  83    USMALLINT = auto()
  84    INT = auto()
  85    UINT = auto()
  86    BIGINT = auto()
  87    UBIGINT = auto()
  88    INT128 = auto()
  89    UINT128 = auto()
  90    INT256 = auto()
  91    UINT256 = auto()
  92    FLOAT = auto()
  93    DOUBLE = auto()
  94    DECIMAL = auto()
  95    BIGDECIMAL = auto()
  96    CHAR = auto()
  97    NCHAR = auto()
  98    VARCHAR = auto()
  99    NVARCHAR = auto()
 100    TEXT = auto()
 101    MEDIUMTEXT = auto()
 102    LONGTEXT = auto()
 103    MEDIUMBLOB = auto()
 104    LONGBLOB = auto()
 105    BINARY = auto()
 106    VARBINARY = auto()
 107    JSON = auto()
 108    JSONB = auto()
 109    TIME = auto()
 110    TIMESTAMP = auto()
 111    TIMESTAMPTZ = auto()
 112    TIMESTAMPLTZ = auto()
 113    DATETIME = auto()
 114    DATETIME64 = auto()
 115    DATE = auto()
 116    INT4RANGE = auto()
 117    INT4MULTIRANGE = auto()
 118    INT8RANGE = auto()
 119    INT8MULTIRANGE = auto()
 120    NUMRANGE = auto()
 121    NUMMULTIRANGE = auto()
 122    TSRANGE = auto()
 123    TSMULTIRANGE = auto()
 124    TSTZRANGE = auto()
 125    TSTZMULTIRANGE = auto()
 126    DATERANGE = auto()
 127    DATEMULTIRANGE = auto()
 128    UUID = auto()
 129    GEOGRAPHY = auto()
 130    NULLABLE = auto()
 131    GEOMETRY = auto()
 132    HLLSKETCH = auto()
 133    HSTORE = auto()
 134    SUPER = auto()
 135    SERIAL = auto()
 136    SMALLSERIAL = auto()
 137    BIGSERIAL = auto()
 138    XML = auto()
 139    UNIQUEIDENTIFIER = auto()
 140    MONEY = auto()
 141    SMALLMONEY = auto()
 142    ROWVERSION = auto()
 143    IMAGE = auto()
 144    VARIANT = auto()
 145    OBJECT = auto()
 146    INET = auto()
 147
 148    # keywords
 149    ALIAS = auto()
 150    ALTER = auto()
 151    ALWAYS = auto()
 152    ALL = auto()
 153    ANTI = auto()
 154    ANY = auto()
 155    APPLY = auto()
 156    ARRAY = auto()
 157    ASC = auto()
 158    ASOF = auto()
 159    AUTO_INCREMENT = auto()
 160    BEGIN = auto()
 161    BETWEEN = auto()
 162    CACHE = auto()
 163    CASE = auto()
 164    CHARACTER_SET = auto()
 165    COLLATE = auto()
 166    COMMAND = auto()
 167    COMMENT = auto()
 168    COMMIT = auto()
 169    CONSTRAINT = auto()
 170    CREATE = auto()
 171    CROSS = auto()
 172    CUBE = auto()
 173    CURRENT_DATE = auto()
 174    CURRENT_DATETIME = auto()
 175    CURRENT_TIME = auto()
 176    CURRENT_TIMESTAMP = auto()
 177    CURRENT_USER = auto()
 178    DEFAULT = auto()
 179    DELETE = auto()
 180    DESC = auto()
 181    DESCRIBE = auto()
 182    DICTIONARY = auto()
 183    DISTINCT = auto()
 184    DIV = auto()
 185    DROP = auto()
 186    ELSE = auto()
 187    END = auto()
 188    ESCAPE = auto()
 189    EXCEPT = auto()
 190    EXECUTE = auto()
 191    EXISTS = auto()
 192    FALSE = auto()
 193    FETCH = auto()
 194    FILTER = auto()
 195    FINAL = auto()
 196    FIRST = auto()
 197    FOR = auto()
 198    FOREIGN_KEY = auto()
 199    FORMAT = auto()
 200    FROM = auto()
 201    FULL = auto()
 202    FUNCTION = auto()
 203    GLOB = auto()
 204    GLOBAL = auto()
 205    GROUP_BY = auto()
 206    GROUPING_SETS = auto()
 207    HAVING = auto()
 208    HINT = auto()
 209    IF = auto()
 210    ILIKE = auto()
 211    ILIKE_ANY = auto()
 212    IN = auto()
 213    INDEX = auto()
 214    INNER = auto()
 215    INSERT = auto()
 216    INTERSECT = auto()
 217    INTERVAL = auto()
 218    INTO = auto()
 219    INTRODUCER = auto()
 220    IRLIKE = auto()
 221    IS = auto()
 222    ISNULL = auto()
 223    JOIN = auto()
 224    JOIN_MARKER = auto()
 225    KEEP = auto()
 226    LANGUAGE = auto()
 227    LATERAL = auto()
 228    LEFT = auto()
 229    LIKE = auto()
 230    LIKE_ANY = auto()
 231    LIMIT = auto()
 232    LOAD = auto()
 233    LOCK = auto()
 234    MAP = auto()
 235    MATCH_RECOGNIZE = auto()
 236    MERGE = auto()
 237    MOD = auto()
 238    NATURAL = auto()
 239    NEXT = auto()
 240    NEXT_VALUE_FOR = auto()
 241    NOTNULL = auto()
 242    NULL = auto()
 243    OFFSET = auto()
 244    ON = auto()
 245    ORDER_BY = auto()
 246    ORDERED = auto()
 247    ORDINALITY = auto()
 248    OUTER = auto()
 249    OVER = auto()
 250    OVERLAPS = auto()
 251    OVERWRITE = auto()
 252    PARTITION = auto()
 253    PARTITION_BY = auto()
 254    PERCENT = auto()
 255    PIVOT = auto()
 256    PLACEHOLDER = auto()
 257    PRAGMA = auto()
 258    PRIMARY_KEY = auto()
 259    PROCEDURE = auto()
 260    PROPERTIES = auto()
 261    PSEUDO_TYPE = auto()
 262    QUALIFY = auto()
 263    QUOTE = auto()
 264    RANGE = auto()
 265    RECURSIVE = auto()
 266    REPLACE = auto()
 267    RETURNING = auto()
 268    REFERENCES = auto()
 269    RIGHT = auto()
 270    RLIKE = auto()
 271    ROLLBACK = auto()
 272    ROLLUP = auto()
 273    ROW = auto()
 274    ROWS = auto()
 275    SELECT = auto()
 276    SEMI = auto()
 277    SEPARATOR = auto()
 278    SERDE_PROPERTIES = auto()
 279    SET = auto()
 280    SETTINGS = auto()
 281    SHOW = auto()
 282    SIMILAR_TO = auto()
 283    SOME = auto()
 284    STRUCT = auto()
 285    TABLE_SAMPLE = auto()
 286    TEMPORARY = auto()
 287    TOP = auto()
 288    THEN = auto()
 289    TRUE = auto()
 290    UNCACHE = auto()
 291    UNION = auto()
 292    UNNEST = auto()
 293    UNPIVOT = auto()
 294    UPDATE = auto()
 295    USE = auto()
 296    USING = auto()
 297    VALUES = auto()
 298    VIEW = auto()
 299    VOLATILE = auto()
 300    WHEN = auto()
 301    WHERE = auto()
 302    WINDOW = auto()
 303    WITH = auto()
 304    UNIQUE = auto()
 305
 306
 307class Token:
 308    __slots__ = ("token_type", "text", "line", "col", "start", "end", "comments")
 309
 310    @classmethod
 311    def number(cls, number: int) -> Token:
 312        """Returns a NUMBER token with `number` as its text."""
 313        return cls(TokenType.NUMBER, str(number))
 314
 315    @classmethod
 316    def string(cls, string: str) -> Token:
 317        """Returns a STRING token with `string` as its text."""
 318        return cls(TokenType.STRING, string)
 319
 320    @classmethod
 321    def identifier(cls, identifier: str) -> Token:
 322        """Returns an IDENTIFIER token with `identifier` as its text."""
 323        return cls(TokenType.IDENTIFIER, identifier)
 324
 325    @classmethod
 326    def var(cls, var: str) -> Token:
 327        """Returns an VAR token with `var` as its text."""
 328        return cls(TokenType.VAR, var)
 329
 330    def __init__(
 331        self,
 332        token_type: TokenType,
 333        text: str,
 334        line: int = 1,
 335        col: int = 1,
 336        start: int = 0,
 337        end: int = 0,
 338        comments: t.List[str] = [],
 339    ) -> None:
 340        """Token initializer.
 341
 342        Args:
 343            token_type: The TokenType Enum.
 344            text: The text of the token.
 345            line: The line that the token ends on.
 346            col: The column that the token ends on.
 347            start: The start index of the token.
 348            end: The ending index of the token.
 349        """
 350        self.token_type = token_type
 351        self.text = text
 352        self.line = line
 353        self.col = col
 354        self.start = start
 355        self.end = end
 356        self.comments = comments
 357
 358    def __repr__(self) -> str:
 359        attributes = ", ".join(f"{k}: {getattr(self, k)}" for k in self.__slots__)
 360        return f"<Token {attributes}>"
 361
 362
 363class _Tokenizer(type):
 364    def __new__(cls, clsname, bases, attrs):
 365        klass = super().__new__(cls, clsname, bases, attrs)
 366
 367        def _convert_quotes(arr: t.List[str | t.Tuple[str, str]]) -> t.Dict[str, str]:
 368            return dict(
 369                (item, item) if isinstance(item, str) else (item[0], item[1]) for item in arr
 370            )
 371
 372        def _quotes_to_format(
 373            token_type: TokenType, arr: t.List[str | t.Tuple[str, str]]
 374        ) -> t.Dict[str, t.Tuple[str, TokenType]]:
 375            return {k: (v, token_type) for k, v in _convert_quotes(arr).items()}
 376
 377        klass._QUOTES = _convert_quotes(klass.QUOTES)
 378        klass._IDENTIFIERS = _convert_quotes(klass.IDENTIFIERS)
 379
 380        klass._FORMAT_STRINGS = {
 381            **{
 382                p + s: (e, TokenType.NATIONAL_STRING)
 383                for s, e in klass._QUOTES.items()
 384                for p in ("n", "N")
 385            },
 386            **_quotes_to_format(TokenType.BIT_STRING, klass.BIT_STRINGS),
 387            **_quotes_to_format(TokenType.BYTE_STRING, klass.BYTE_STRINGS),
 388            **_quotes_to_format(TokenType.HEX_STRING, klass.HEX_STRINGS),
 389            **_quotes_to_format(TokenType.RAW_STRING, klass.RAW_STRINGS),
 390        }
 391
 392        klass._STRING_ESCAPES = set(klass.STRING_ESCAPES)
 393        klass._IDENTIFIER_ESCAPES = set(klass.IDENTIFIER_ESCAPES)
 394        klass._COMMENTS = dict(
 395            (comment, None) if isinstance(comment, str) else (comment[0], comment[1])
 396            for comment in klass.COMMENTS
 397        )
 398
 399        klass.KEYWORD_TRIE = new_trie(
 400            key.upper()
 401            for key in (
 402                *klass.KEYWORDS,
 403                *klass._COMMENTS,
 404                *klass._QUOTES,
 405                *klass._FORMAT_STRINGS,
 406            )
 407            if " " in key or any(single in key for single in klass.SINGLE_TOKENS)
 408        )
 409
 410        return klass
 411
 412
 413class Tokenizer(metaclass=_Tokenizer):
 414    SINGLE_TOKENS = {
 415        "(": TokenType.L_PAREN,
 416        ")": TokenType.R_PAREN,
 417        "[": TokenType.L_BRACKET,
 418        "]": TokenType.R_BRACKET,
 419        "{": TokenType.L_BRACE,
 420        "}": TokenType.R_BRACE,
 421        "&": TokenType.AMP,
 422        "^": TokenType.CARET,
 423        ":": TokenType.COLON,
 424        ",": TokenType.COMMA,
 425        ".": TokenType.DOT,
 426        "-": TokenType.DASH,
 427        "=": TokenType.EQ,
 428        ">": TokenType.GT,
 429        "<": TokenType.LT,
 430        "%": TokenType.MOD,
 431        "!": TokenType.NOT,
 432        "|": TokenType.PIPE,
 433        "+": TokenType.PLUS,
 434        ";": TokenType.SEMICOLON,
 435        "/": TokenType.SLASH,
 436        "\\": TokenType.BACKSLASH,
 437        "*": TokenType.STAR,
 438        "~": TokenType.TILDA,
 439        "?": TokenType.PLACEHOLDER,
 440        "@": TokenType.PARAMETER,
 441        # used for breaking a var like x'y' but nothing else
 442        # the token type doesn't matter
 443        "'": TokenType.QUOTE,
 444        "`": TokenType.IDENTIFIER,
 445        '"': TokenType.IDENTIFIER,
 446        "#": TokenType.HASH,
 447    }
 448
 449    BIT_STRINGS: t.List[str | t.Tuple[str, str]] = []
 450    BYTE_STRINGS: t.List[str | t.Tuple[str, str]] = []
 451    HEX_STRINGS: t.List[str | t.Tuple[str, str]] = []
 452    RAW_STRINGS: t.List[str | t.Tuple[str, str]] = []
 453    IDENTIFIERS: t.List[str | t.Tuple[str, str]] = ['"']
 454    IDENTIFIER_ESCAPES = ['"']
 455    QUOTES: t.List[t.Tuple[str, str] | str] = ["'"]
 456    STRING_ESCAPES = ["'"]
 457    VAR_SINGLE_TOKENS: t.Set[str] = set()
 458
 459    _COMMENTS: t.Dict[str, str] = {}
 460    _FORMAT_STRINGS: t.Dict[str, t.Tuple[str, TokenType]] = {}
 461    _IDENTIFIERS: t.Dict[str, str] = {}
 462    _IDENTIFIER_ESCAPES: t.Set[str] = set()
 463    _QUOTES: t.Dict[str, str] = {}
 464    _STRING_ESCAPES: t.Set[str] = set()
 465
 466    KEYWORDS: t.Dict[t.Optional[str], TokenType] = {
 467        **{f"{{%{postfix}": TokenType.BLOCK_START for postfix in ("", "+", "-")},
 468        **{f"{prefix}%}}": TokenType.BLOCK_END for prefix in ("", "+", "-")},
 469        "{{+": TokenType.BLOCK_START,
 470        "{{-": TokenType.BLOCK_START,
 471        "+}}": TokenType.BLOCK_END,
 472        "-}}": TokenType.BLOCK_END,
 473        "/*+": TokenType.HINT,
 474        "==": TokenType.EQ,
 475        "::": TokenType.DCOLON,
 476        "||": TokenType.DPIPE,
 477        ">=": TokenType.GTE,
 478        "<=": TokenType.LTE,
 479        "<>": TokenType.NEQ,
 480        "!=": TokenType.NEQ,
 481        "<=>": TokenType.NULLSAFE_EQ,
 482        "->": TokenType.ARROW,
 483        "->>": TokenType.DARROW,
 484        "=>": TokenType.FARROW,
 485        "#>": TokenType.HASH_ARROW,
 486        "#>>": TokenType.DHASH_ARROW,
 487        "<->": TokenType.LR_ARROW,
 488        "&&": TokenType.DAMP,
 489        "ALL": TokenType.ALL,
 490        "ALWAYS": TokenType.ALWAYS,
 491        "AND": TokenType.AND,
 492        "ANTI": TokenType.ANTI,
 493        "ANY": TokenType.ANY,
 494        "ASC": TokenType.ASC,
 495        "AS": TokenType.ALIAS,
 496        "ASOF": TokenType.ASOF,
 497        "AUTOINCREMENT": TokenType.AUTO_INCREMENT,
 498        "AUTO_INCREMENT": TokenType.AUTO_INCREMENT,
 499        "BEGIN": TokenType.BEGIN,
 500        "BETWEEN": TokenType.BETWEEN,
 501        "CACHE": TokenType.CACHE,
 502        "UNCACHE": TokenType.UNCACHE,
 503        "CASE": TokenType.CASE,
 504        "CHARACTER SET": TokenType.CHARACTER_SET,
 505        "COLLATE": TokenType.COLLATE,
 506        "COLUMN": TokenType.COLUMN,
 507        "COMMIT": TokenType.COMMIT,
 508        "CONSTRAINT": TokenType.CONSTRAINT,
 509        "CREATE": TokenType.CREATE,
 510        "CROSS": TokenType.CROSS,
 511        "CUBE": TokenType.CUBE,
 512        "CURRENT_DATE": TokenType.CURRENT_DATE,
 513        "CURRENT_TIME": TokenType.CURRENT_TIME,
 514        "CURRENT_TIMESTAMP": TokenType.CURRENT_TIMESTAMP,
 515        "CURRENT_USER": TokenType.CURRENT_USER,
 516        "DATABASE": TokenType.DATABASE,
 517        "DEFAULT": TokenType.DEFAULT,
 518        "DELETE": TokenType.DELETE,
 519        "DESC": TokenType.DESC,
 520        "DESCRIBE": TokenType.DESCRIBE,
 521        "DISTINCT": TokenType.DISTINCT,
 522        "DIV": TokenType.DIV,
 523        "DROP": TokenType.DROP,
 524        "ELSE": TokenType.ELSE,
 525        "END": TokenType.END,
 526        "ESCAPE": TokenType.ESCAPE,
 527        "EXCEPT": TokenType.EXCEPT,
 528        "EXECUTE": TokenType.EXECUTE,
 529        "EXISTS": TokenType.EXISTS,
 530        "FALSE": TokenType.FALSE,
 531        "FETCH": TokenType.FETCH,
 532        "FILTER": TokenType.FILTER,
 533        "FIRST": TokenType.FIRST,
 534        "FULL": TokenType.FULL,
 535        "FUNCTION": TokenType.FUNCTION,
 536        "FOR": TokenType.FOR,
 537        "FOREIGN KEY": TokenType.FOREIGN_KEY,
 538        "FORMAT": TokenType.FORMAT,
 539        "FROM": TokenType.FROM,
 540        "GEOGRAPHY": TokenType.GEOGRAPHY,
 541        "GEOMETRY": TokenType.GEOMETRY,
 542        "GLOB": TokenType.GLOB,
 543        "GROUP BY": TokenType.GROUP_BY,
 544        "GROUPING SETS": TokenType.GROUPING_SETS,
 545        "HAVING": TokenType.HAVING,
 546        "IF": TokenType.IF,
 547        "ILIKE": TokenType.ILIKE,
 548        "IN": TokenType.IN,
 549        "INDEX": TokenType.INDEX,
 550        "INET": TokenType.INET,
 551        "INNER": TokenType.INNER,
 552        "INSERT": TokenType.INSERT,
 553        "INTERVAL": TokenType.INTERVAL,
 554        "INTERSECT": TokenType.INTERSECT,
 555        "INTO": TokenType.INTO,
 556        "IS": TokenType.IS,
 557        "ISNULL": TokenType.ISNULL,
 558        "JOIN": TokenType.JOIN,
 559        "KEEP": TokenType.KEEP,
 560        "LATERAL": TokenType.LATERAL,
 561        "LEFT": TokenType.LEFT,
 562        "LIKE": TokenType.LIKE,
 563        "LIMIT": TokenType.LIMIT,
 564        "LOAD": TokenType.LOAD,
 565        "LOCK": TokenType.LOCK,
 566        "MERGE": TokenType.MERGE,
 567        "NATURAL": TokenType.NATURAL,
 568        "NEXT": TokenType.NEXT,
 569        "NEXT VALUE FOR": TokenType.NEXT_VALUE_FOR,
 570        "NOT": TokenType.NOT,
 571        "NOTNULL": TokenType.NOTNULL,
 572        "NULL": TokenType.NULL,
 573        "OBJECT": TokenType.OBJECT,
 574        "OFFSET": TokenType.OFFSET,
 575        "ON": TokenType.ON,
 576        "OR": TokenType.OR,
 577        "ORDER BY": TokenType.ORDER_BY,
 578        "ORDINALITY": TokenType.ORDINALITY,
 579        "OUTER": TokenType.OUTER,
 580        "OVER": TokenType.OVER,
 581        "OVERLAPS": TokenType.OVERLAPS,
 582        "OVERWRITE": TokenType.OVERWRITE,
 583        "PARTITION": TokenType.PARTITION,
 584        "PARTITION BY": TokenType.PARTITION_BY,
 585        "PARTITIONED BY": TokenType.PARTITION_BY,
 586        "PARTITIONED_BY": TokenType.PARTITION_BY,
 587        "PERCENT": TokenType.PERCENT,
 588        "PIVOT": TokenType.PIVOT,
 589        "PRAGMA": TokenType.PRAGMA,
 590        "PRIMARY KEY": TokenType.PRIMARY_KEY,
 591        "PROCEDURE": TokenType.PROCEDURE,
 592        "QUALIFY": TokenType.QUALIFY,
 593        "RANGE": TokenType.RANGE,
 594        "RECURSIVE": TokenType.RECURSIVE,
 595        "REGEXP": TokenType.RLIKE,
 596        "REPLACE": TokenType.REPLACE,
 597        "REFERENCES": TokenType.REFERENCES,
 598        "RIGHT": TokenType.RIGHT,
 599        "RLIKE": TokenType.RLIKE,
 600        "ROLLBACK": TokenType.ROLLBACK,
 601        "ROLLUP": TokenType.ROLLUP,
 602        "ROW": TokenType.ROW,
 603        "ROWS": TokenType.ROWS,
 604        "SCHEMA": TokenType.SCHEMA,
 605        "SELECT": TokenType.SELECT,
 606        "SEMI": TokenType.SEMI,
 607        "SET": TokenType.SET,
 608        "SETTINGS": TokenType.SETTINGS,
 609        "SHOW": TokenType.SHOW,
 610        "SIMILAR TO": TokenType.SIMILAR_TO,
 611        "SOME": TokenType.SOME,
 612        "TABLE": TokenType.TABLE,
 613        "TABLESAMPLE": TokenType.TABLE_SAMPLE,
 614        "TEMP": TokenType.TEMPORARY,
 615        "TEMPORARY": TokenType.TEMPORARY,
 616        "THEN": TokenType.THEN,
 617        "TRUE": TokenType.TRUE,
 618        "UNION": TokenType.UNION,
 619        "UNNEST": TokenType.UNNEST,
 620        "UNPIVOT": TokenType.UNPIVOT,
 621        "UPDATE": TokenType.UPDATE,
 622        "USE": TokenType.USE,
 623        "USING": TokenType.USING,
 624        "UUID": TokenType.UUID,
 625        "VALUES": TokenType.VALUES,
 626        "VIEW": TokenType.VIEW,
 627        "VOLATILE": TokenType.VOLATILE,
 628        "WHEN": TokenType.WHEN,
 629        "WHERE": TokenType.WHERE,
 630        "WINDOW": TokenType.WINDOW,
 631        "WITH": TokenType.WITH,
 632        "APPLY": TokenType.APPLY,
 633        "ARRAY": TokenType.ARRAY,
 634        "BIT": TokenType.BIT,
 635        "BOOL": TokenType.BOOLEAN,
 636        "BOOLEAN": TokenType.BOOLEAN,
 637        "BYTE": TokenType.TINYINT,
 638        "TINYINT": TokenType.TINYINT,
 639        "SHORT": TokenType.SMALLINT,
 640        "SMALLINT": TokenType.SMALLINT,
 641        "INT2": TokenType.SMALLINT,
 642        "INTEGER": TokenType.INT,
 643        "INT": TokenType.INT,
 644        "INT4": TokenType.INT,
 645        "LONG": TokenType.BIGINT,
 646        "BIGINT": TokenType.BIGINT,
 647        "INT8": TokenType.BIGINT,
 648        "DEC": TokenType.DECIMAL,
 649        "DECIMAL": TokenType.DECIMAL,
 650        "BIGDECIMAL": TokenType.BIGDECIMAL,
 651        "BIGNUMERIC": TokenType.BIGDECIMAL,
 652        "MAP": TokenType.MAP,
 653        "NULLABLE": TokenType.NULLABLE,
 654        "NUMBER": TokenType.DECIMAL,
 655        "NUMERIC": TokenType.DECIMAL,
 656        "FIXED": TokenType.DECIMAL,
 657        "REAL": TokenType.FLOAT,
 658        "FLOAT": TokenType.FLOAT,
 659        "FLOAT4": TokenType.FLOAT,
 660        "FLOAT8": TokenType.DOUBLE,
 661        "DOUBLE": TokenType.DOUBLE,
 662        "DOUBLE PRECISION": TokenType.DOUBLE,
 663        "JSON": TokenType.JSON,
 664        "CHAR": TokenType.CHAR,
 665        "CHARACTER": TokenType.CHAR,
 666        "NCHAR": TokenType.NCHAR,
 667        "VARCHAR": TokenType.VARCHAR,
 668        "VARCHAR2": TokenType.VARCHAR,
 669        "NVARCHAR": TokenType.NVARCHAR,
 670        "NVARCHAR2": TokenType.NVARCHAR,
 671        "STR": TokenType.TEXT,
 672        "STRING": TokenType.TEXT,
 673        "TEXT": TokenType.TEXT,
 674        "CLOB": TokenType.TEXT,
 675        "LONGVARCHAR": TokenType.TEXT,
 676        "BINARY": TokenType.BINARY,
 677        "BLOB": TokenType.VARBINARY,
 678        "BYTEA": TokenType.VARBINARY,
 679        "VARBINARY": TokenType.VARBINARY,
 680        "TIME": TokenType.TIME,
 681        "TIMESTAMP": TokenType.TIMESTAMP,
 682        "TIMESTAMPTZ": TokenType.TIMESTAMPTZ,
 683        "TIMESTAMPLTZ": TokenType.TIMESTAMPLTZ,
 684        "DATE": TokenType.DATE,
 685        "DATETIME": TokenType.DATETIME,
 686        "INT4RANGE": TokenType.INT4RANGE,
 687        "INT4MULTIRANGE": TokenType.INT4MULTIRANGE,
 688        "INT8RANGE": TokenType.INT8RANGE,
 689        "INT8MULTIRANGE": TokenType.INT8MULTIRANGE,
 690        "NUMRANGE": TokenType.NUMRANGE,
 691        "NUMMULTIRANGE": TokenType.NUMMULTIRANGE,
 692        "TSRANGE": TokenType.TSRANGE,
 693        "TSMULTIRANGE": TokenType.TSMULTIRANGE,
 694        "TSTZRANGE": TokenType.TSTZRANGE,
 695        "TSTZMULTIRANGE": TokenType.TSTZMULTIRANGE,
 696        "DATERANGE": TokenType.DATERANGE,
 697        "DATEMULTIRANGE": TokenType.DATEMULTIRANGE,
 698        "UNIQUE": TokenType.UNIQUE,
 699        "STRUCT": TokenType.STRUCT,
 700        "VARIANT": TokenType.VARIANT,
 701        "ALTER": TokenType.ALTER,
 702        "ANALYZE": TokenType.COMMAND,
 703        "CALL": TokenType.COMMAND,
 704        "COMMENT": TokenType.COMMENT,
 705        "COPY": TokenType.COMMAND,
 706        "EXPLAIN": TokenType.COMMAND,
 707        "GRANT": TokenType.COMMAND,
 708        "OPTIMIZE": TokenType.COMMAND,
 709        "PREPARE": TokenType.COMMAND,
 710        "TRUNCATE": TokenType.COMMAND,
 711        "VACUUM": TokenType.COMMAND,
 712    }
 713
 714    WHITE_SPACE: t.Dict[t.Optional[str], TokenType] = {
 715        " ": TokenType.SPACE,
 716        "\t": TokenType.SPACE,
 717        "\n": TokenType.BREAK,
 718        "\r": TokenType.BREAK,
 719        "\r\n": TokenType.BREAK,
 720    }
 721
 722    COMMANDS = {
 723        TokenType.COMMAND,
 724        TokenType.EXECUTE,
 725        TokenType.FETCH,
 726        TokenType.SHOW,
 727    }
 728
 729    COMMAND_PREFIX_TOKENS = {TokenType.SEMICOLON, TokenType.BEGIN}
 730
 731    # handle numeric literals like in hive (3L = BIGINT)
 732    NUMERIC_LITERALS: t.Dict[str, str] = {}
 733    ENCODE: t.Optional[str] = None
 734
 735    COMMENTS = ["--", ("/*", "*/"), ("{#", "#}")]
 736    KEYWORD_TRIE: t.Dict = {}  # autofilled
 737
 738    __slots__ = (
 739        "sql",
 740        "size",
 741        "tokens",
 742        "_start",
 743        "_current",
 744        "_line",
 745        "_col",
 746        "_comments",
 747        "_char",
 748        "_end",
 749        "_peek",
 750        "_prev_token_line",
 751        "identifiers_can_start_with_digit",
 752    )
 753
 754    def __init__(self) -> None:
 755        self.reset()
 756
 757    def reset(self) -> None:
 758        self.sql = ""
 759        self.size = 0
 760        self.tokens: t.List[Token] = []
 761        self._start = 0
 762        self._current = 0
 763        self._line = 1
 764        self._col = 0
 765        self._comments: t.List[str] = []
 766
 767        self._char = ""
 768        self._end = False
 769        self._peek = ""
 770        self._prev_token_line = -1
 771
 772    def tokenize(self, sql: str) -> t.List[Token]:
 773        """Returns a list of tokens corresponding to the SQL string `sql`."""
 774        self.reset()
 775        self.sql = sql
 776        self.size = len(sql)
 777
 778        try:
 779            self._scan()
 780        except Exception as e:
 781            start = max(self._current - 50, 0)
 782            end = min(self._current + 50, self.size - 1)
 783            context = self.sql[start:end]
 784            raise ValueError(f"Error tokenizing '{context}'") from e
 785
 786        return self.tokens
 787
 788    def _scan(self, until: t.Optional[t.Callable] = None) -> None:
 789        while self.size and not self._end:
 790            self._start = self._current
 791            self._advance()
 792
 793            if self._char is None:
 794                break
 795
 796            if self._char not in self.WHITE_SPACE:
 797                if self._char.isdigit():
 798                    self._scan_number()
 799                elif self._char in self._IDENTIFIERS:
 800                    self._scan_identifier(self._IDENTIFIERS[self._char])
 801                else:
 802                    self._scan_keywords()
 803
 804            if until and until():
 805                break
 806
 807        if self.tokens and self._comments:
 808            self.tokens[-1].comments.extend(self._comments)
 809
 810    def _chars(self, size: int) -> str:
 811        if size == 1:
 812            return self._char
 813
 814        start = self._current - 1
 815        end = start + size
 816
 817        return self.sql[start:end] if end <= self.size else ""
 818
 819    def _advance(self, i: int = 1, alnum: bool = False) -> None:
 820        if self.WHITE_SPACE.get(self._char) is TokenType.BREAK:
 821            self._col = 1
 822            self._line += 1
 823        else:
 824            self._col += i
 825
 826        self._current += i
 827        self._end = self._current >= self.size
 828        self._char = self.sql[self._current - 1]
 829        self._peek = "" if self._end else self.sql[self._current]
 830
 831        if alnum and self._char.isalnum():
 832            # Here we use local variables instead of attributes for better performance
 833            _col = self._col
 834            _current = self._current
 835            _end = self._end
 836            _peek = self._peek
 837
 838            while _peek.isalnum():
 839                _col += 1
 840                _current += 1
 841                _end = _current >= self.size
 842                _peek = "" if _end else self.sql[_current]
 843
 844            self._col = _col
 845            self._current = _current
 846            self._end = _end
 847            self._peek = _peek
 848            self._char = self.sql[_current - 1]
 849
 850    @property
 851    def _text(self) -> str:
 852        return self.sql[self._start : self._current]
 853
 854    def _add(self, token_type: TokenType, text: t.Optional[str] = None) -> None:
 855        self._prev_token_line = self._line
 856        self.tokens.append(
 857            Token(
 858                token_type,
 859                text=self._text if text is None else text,
 860                line=self._line,
 861                col=self._col,
 862                start=self._start,
 863                end=self._current - 1,
 864                comments=self._comments,
 865            )
 866        )
 867        self._comments = []
 868
 869        # If we have either a semicolon or a begin token before the command's token, we'll parse
 870        # whatever follows the command's token as a string
 871        if (
 872            token_type in self.COMMANDS
 873            and self._peek != ";"
 874            and (len(self.tokens) == 1 or self.tokens[-2].token_type in self.COMMAND_PREFIX_TOKENS)
 875        ):
 876            start = self._current
 877            tokens = len(self.tokens)
 878            self._scan(lambda: self._peek == ";")
 879            self.tokens = self.tokens[:tokens]
 880            text = self.sql[start : self._current].strip()
 881            if text:
 882                self._add(TokenType.STRING, text)
 883
 884    def _scan_keywords(self) -> None:
 885        size = 0
 886        word = None
 887        chars = self._text
 888        char = chars
 889        prev_space = False
 890        skip = False
 891        trie = self.KEYWORD_TRIE
 892        single_token = char in self.SINGLE_TOKENS
 893
 894        while chars:
 895            if skip:
 896                result = 1
 897            else:
 898                result, trie = in_trie(trie, char.upper())
 899
 900            if result == 0:
 901                break
 902            if result == 2:
 903                word = chars
 904
 905            size += 1
 906            end = self._current - 1 + size
 907
 908            if end < self.size:
 909                char = self.sql[end]
 910                single_token = single_token or char in self.SINGLE_TOKENS
 911                is_space = char in self.WHITE_SPACE
 912
 913                if not is_space or not prev_space:
 914                    if is_space:
 915                        char = " "
 916                    chars += char
 917                    prev_space = is_space
 918                    skip = False
 919                else:
 920                    skip = True
 921            else:
 922                char = ""
 923                chars = " "
 924
 925        word = None if not single_token and chars[-1] not in self.WHITE_SPACE else word
 926
 927        if not word:
 928            if self._char in self.SINGLE_TOKENS:
 929                self._add(self.SINGLE_TOKENS[self._char], text=self._char)
 930                return
 931            self._scan_var()
 932            return
 933
 934        if self._scan_string(word):
 935            return
 936        if self._scan_comment(word):
 937            return
 938
 939        self._advance(size - 1)
 940        word = word.upper()
 941        self._add(self.KEYWORDS[word], text=word)
 942
 943    def _scan_comment(self, comment_start: str) -> bool:
 944        if comment_start not in self._COMMENTS:
 945            return False
 946
 947        comment_start_line = self._line
 948        comment_start_size = len(comment_start)
 949        comment_end = self._COMMENTS[comment_start]
 950
 951        if comment_end:
 952            # Skip the comment's start delimiter
 953            self._advance(comment_start_size)
 954
 955            comment_end_size = len(comment_end)
 956            while not self._end and self._chars(comment_end_size) != comment_end:
 957                self._advance(alnum=True)
 958
 959            self._comments.append(self._text[comment_start_size : -comment_end_size + 1])
 960            self._advance(comment_end_size - 1)
 961        else:
 962            while not self._end and not self.WHITE_SPACE.get(self._peek) is TokenType.BREAK:
 963                self._advance(alnum=True)
 964            self._comments.append(self._text[comment_start_size:])
 965
 966        # Leading comment is attached to the succeeding token, whilst trailing comment to the preceding.
 967        # Multiple consecutive comments are preserved by appending them to the current comments list.
 968        if comment_start_line == self._prev_token_line:
 969            self.tokens[-1].comments.extend(self._comments)
 970            self._comments = []
 971            self._prev_token_line = self._line
 972
 973        return True
 974
 975    def _scan_number(self) -> None:
 976        if self._char == "0":
 977            peek = self._peek.upper()
 978            if peek == "B":
 979                return self._scan_bits() if self.BIT_STRINGS else self._add(TokenType.NUMBER)
 980            elif peek == "X":
 981                return self._scan_hex() if self.HEX_STRINGS else self._add(TokenType.NUMBER)
 982
 983        decimal = False
 984        scientific = 0
 985
 986        while True:
 987            if self._peek.isdigit():
 988                self._advance()
 989            elif self._peek == "." and not decimal:
 990                decimal = True
 991                self._advance()
 992            elif self._peek in ("-", "+") and scientific == 1:
 993                scientific += 1
 994                self._advance()
 995            elif self._peek.upper() == "E" and not scientific:
 996                scientific += 1
 997                self._advance()
 998            elif self._peek.isidentifier():
 999                number_text = self._text
1000                literal = ""
1001
1002                while self._peek.strip() and self._peek not in self.SINGLE_TOKENS:
1003                    literal += self._peek.upper()
1004                    self._advance()
1005
1006                token_type = self.KEYWORDS.get(self.NUMERIC_LITERALS.get(literal))
1007
1008                if token_type:
1009                    self._add(TokenType.NUMBER, number_text)
1010                    self._add(TokenType.DCOLON, "::")
1011                    return self._add(token_type, literal)
1012                elif self.identifiers_can_start_with_digit:  # type: ignore
1013                    return self._add(TokenType.VAR)
1014
1015                self._add(TokenType.NUMBER, number_text)
1016                return self._advance(-len(literal))
1017            else:
1018                return self._add(TokenType.NUMBER)
1019
1020    def _scan_bits(self) -> None:
1021        self._advance()
1022        value = self._extract_value()
1023        try:
1024            # If `value` can't be converted to a binary, fallback to tokenizing it as an identifier
1025            int(value, 2)
1026            self._add(TokenType.BIT_STRING, value[2:])  # Drop the 0b
1027        except ValueError:
1028            self._add(TokenType.IDENTIFIER)
1029
1030    def _scan_hex(self) -> None:
1031        self._advance()
1032        value = self._extract_value()
1033        try:
1034            # If `value` can't be converted to a hex, fallback to tokenizing it as an identifier
1035            int(value, 16)
1036            self._add(TokenType.HEX_STRING, value[2:])  # Drop the 0x
1037        except ValueError:
1038            self._add(TokenType.IDENTIFIER)
1039
1040    def _extract_value(self) -> str:
1041        while True:
1042            char = self._peek.strip()
1043            if char and char not in self.SINGLE_TOKENS:
1044                self._advance(alnum=True)
1045            else:
1046                break
1047
1048        return self._text
1049
1050    def _scan_string(self, start: str) -> bool:
1051        base = None
1052        token_type = TokenType.STRING
1053
1054        if start in self._QUOTES:
1055            end = self._QUOTES[start]
1056        elif start in self._FORMAT_STRINGS:
1057            end, token_type = self._FORMAT_STRINGS[start]
1058
1059            if token_type == TokenType.HEX_STRING:
1060                base = 16
1061            elif token_type == TokenType.BIT_STRING:
1062                base = 2
1063        else:
1064            return False
1065
1066        self._advance(len(start))
1067        text = self._extract_string(end)
1068
1069        if base:
1070            try:
1071                int(text, base)
1072            except:
1073                raise RuntimeError(
1074                    f"Numeric string contains invalid characters from {self._line}:{self._start}"
1075                )
1076        else:
1077            text = text.encode(self.ENCODE).decode(self.ENCODE) if self.ENCODE else text
1078
1079        self._add(token_type, text)
1080        return True
1081
1082    def _scan_identifier(self, identifier_end: str) -> None:
1083        self._advance()
1084        text = self._extract_string(identifier_end, self._IDENTIFIER_ESCAPES)
1085        self._add(TokenType.IDENTIFIER, text)
1086
1087    def _scan_var(self) -> None:
1088        while True:
1089            char = self._peek.strip()
1090            if char and (char in self.VAR_SINGLE_TOKENS or char not in self.SINGLE_TOKENS):
1091                self._advance(alnum=True)
1092            else:
1093                break
1094
1095        self._add(
1096            TokenType.VAR
1097            if self.tokens and self.tokens[-1].token_type == TokenType.PARAMETER
1098            else self.KEYWORDS.get(self._text.upper(), TokenType.VAR)
1099        )
1100
1101    def _extract_string(self, delimiter: str, escapes=None) -> str:
1102        text = ""
1103        delim_size = len(delimiter)
1104        escapes = self._STRING_ESCAPES if escapes is None else escapes
1105
1106        while True:
1107            if self._char in escapes and (self._peek == delimiter or self._peek in escapes):
1108                if self._peek == delimiter:
1109                    text += self._peek
1110                else:
1111                    text += self._char + self._peek
1112
1113                if self._current + 1 < self.size:
1114                    self._advance(2)
1115                else:
1116                    raise RuntimeError(f"Missing {delimiter} from {self._line}:{self._current}")
1117            else:
1118                if self._chars(delim_size) == delimiter:
1119                    if delim_size > 1:
1120                        self._advance(delim_size - 1)
1121                    break
1122
1123                if self._end:
1124                    raise RuntimeError(f"Missing {delimiter} from {self._line}:{self._start}")
1125
1126                current = self._current - 1
1127                self._advance(alnum=True)
1128                text += self.sql[current : self._current - 1]
1129
1130        return text