sqlglot.tokenizer_core

   1from __future__ import annotations
   2
   3import typing as t
   4from enum import IntEnum, auto
   5
   6from sqlglot.errors import TokenError
   7
   8# dict lookup is faster than .upper() and .isdigit()
   9_CHAR_UPPER: dict[str, str] = {chr(i): chr(i).upper() for i in range(97, 123)}
  10_DIGIT_CHARS: frozenset[str] = frozenset("0123456789")
  11_DIGIT_VALUES: dict[str, int] = {c: int(c, 16) for c in "0123456789abcdefABCDEF"}
  12_OCTAL_CHARS: frozenset[str] = frozenset("01234567")
  13
  14
  15class TokenType(IntEnum):
  16    L_PAREN = auto()
  17    R_PAREN = auto()
  18    L_BRACKET = auto()
  19    R_BRACKET = auto()
  20    L_BRACE = auto()
  21    R_BRACE = auto()
  22    COMMA = auto()
  23    DOT = auto()
  24    DASH = auto()
  25    PLUS = auto()
  26    COLON = auto()
  27    DOTCOLON = auto()
  28    DOTCARET = auto()
  29    DCOLON = auto()
  30    DCOLONDOLLAR = auto()
  31    DCOLONPERCENT = auto()
  32    DCOLONQMARK = auto()
  33    DQMARK = auto()
  34    SEMICOLON = auto()
  35    STAR = auto()
  36    BACKSLASH = auto()
  37    SLASH = auto()
  38    LT = auto()
  39    LTE = auto()
  40    GT = auto()
  41    GTE = auto()
  42    NOT = auto()
  43    EQ = auto()
  44    NEQ = auto()
  45    NULLSAFE_EQ = auto()
  46    COLON_EQ = auto()
  47    COLON_GT = auto()
  48    NCOLON_GT = auto()
  49    AND = auto()
  50    OR = auto()
  51    AMP = auto()
  52    DPIPE = auto()
  53    PIPE_GT = auto()
  54    PIPE = auto()
  55    PIPE_SLASH = auto()
  56    DPIPE_SLASH = auto()
  57    CARET = auto()
  58    CARET_AT = auto()
  59    TILDE = auto()
  60    ARROW = auto()
  61    DARROW = auto()
  62    FARROW = auto()
  63    HASH = auto()
  64    HASH_ARROW = auto()
  65    DHASH_ARROW = auto()
  66    LR_ARROW = auto()
  67    LLRR_ARROW = auto()
  68    DAT = auto()
  69    AT_QMARK = auto()
  70    LT_AT = auto()
  71    AT_GT = auto()
  72    DOLLAR = auto()
  73    PARAMETER = auto()
  74    SESSION = auto()
  75    SESSION_PARAMETER = auto()
  76    SESSION_USER = auto()
  77    DAMP = auto()
  78    AMP_LT = auto()
  79    AMP_GT = auto()
  80    ADJACENT = auto()
  81    XOR = auto()
  82    DSTAR = auto()
  83    QMARK_AMP = auto()
  84    QMARK_PIPE = auto()
  85    HASH_DASH = auto()
  86    EXCLAMATION = auto()
  87
  88    URI_START = auto()
  89
  90    BLOCK_START = auto()
  91    BLOCK_END = auto()
  92
  93    SPACE = auto()
  94    BREAK = auto()
  95
  96    STRING = auto()
  97    NUMBER = auto()
  98    IDENTIFIER = auto()
  99    DATABASE = auto()
 100    COLUMN = auto()
 101    COLUMN_DEF = auto()
 102    SCHEMA = auto()
 103    TABLE = auto()
 104    WAREHOUSE = auto()
 105    STAGE = auto()
 106    STREAM = auto()
 107    STREAMLIT = auto()
 108    VAR = auto()
 109    BIT_STRING = auto()
 110    HEX_STRING = auto()
 111    BYTE_STRING = auto()
 112    NATIONAL_RAW_STRING = auto()
 113    NATIONAL_STRING = auto()
 114    RAW_STRING = auto()
 115    HEREDOC_STRING = auto()
 116    UNICODE_STRING = auto()
 117
 118    # types
 119    BIT = auto()
 120    BOOLEAN = auto()
 121    TINYINT = auto()
 122    UTINYINT = auto()
 123    SMALLINT = auto()
 124    USMALLINT = auto()
 125    MEDIUMINT = auto()
 126    UMEDIUMINT = auto()
 127    INT = auto()
 128    UINT = auto()
 129    BIGINT = auto()
 130    UBIGINT = auto()
 131    BIGNUM = auto()
 132    INT128 = auto()
 133    UINT128 = auto()
 134    INT256 = auto()
 135    UINT256 = auto()
 136    FLOAT = auto()
 137    DOUBLE = auto()
 138    UDOUBLE = auto()
 139    DECIMAL = auto()
 140    DECIMAL32 = auto()
 141    DECIMAL64 = auto()
 142    DECIMAL128 = auto()
 143    DECIMAL256 = auto()
 144    DECFLOAT = auto()
 145    UDECIMAL = auto()
 146    BIGDECIMAL = auto()
 147    CHAR = auto()
 148    NCHAR = auto()
 149    VARCHAR = auto()
 150    NVARCHAR = auto()
 151    BPCHAR = auto()
 152    TEXT = auto()
 153    MEDIUMTEXT = auto()
 154    LONGTEXT = auto()
 155    BLOB = auto()
 156    MEDIUMBLOB = auto()
 157    LONGBLOB = auto()
 158    TINYBLOB = auto()
 159    TINYTEXT = auto()
 160    NAME = auto()
 161    BINARY = auto()
 162    VARBINARY = auto()
 163    JSON = auto()
 164    JSONB = auto()
 165    TIME = auto()
 166    TIMETZ = auto()
 167    TIME_NS = auto()
 168    TIMESTAMP = auto()
 169    TIMESTAMPTZ = auto()
 170    TIMESTAMPLTZ = auto()
 171    TIMESTAMPNTZ = auto()
 172    TIMESTAMP_S = auto()
 173    TIMESTAMP_MS = auto()
 174    TIMESTAMP_NS = auto()
 175    DATETIME = auto()
 176    DATETIME2 = auto()
 177    DATETIME64 = auto()
 178    SMALLDATETIME = auto()
 179    DATE = auto()
 180    DATE32 = auto()
 181    INT4RANGE = auto()
 182    INT4MULTIRANGE = auto()
 183    INT8RANGE = auto()
 184    INT8MULTIRANGE = auto()
 185    NUMRANGE = auto()
 186    NUMMULTIRANGE = auto()
 187    TSRANGE = auto()
 188    TSMULTIRANGE = auto()
 189    TSTZRANGE = auto()
 190    TSTZMULTIRANGE = auto()
 191    DATERANGE = auto()
 192    DATEMULTIRANGE = auto()
 193    UUID = auto()
 194    GEOGRAPHY = auto()
 195    GEOGRAPHYPOINT = auto()
 196    NULLABLE = auto()
 197    GEOMETRY = auto()
 198    POINT = auto()
 199    RING = auto()
 200    LINESTRING = auto()
 201    LOCALTIME = auto()
 202    LOCALTIMESTAMP = auto()
 203    SYSTIMESTAMP = auto()
 204    MULTILINESTRING = auto()
 205    POLYGON = auto()
 206    MULTIPOLYGON = auto()
 207    HLLSKETCH = auto()
 208    HSTORE = auto()
 209    SUPER = auto()
 210    SERIAL = auto()
 211    SMALLSERIAL = auto()
 212    BIGSERIAL = auto()
 213    XML = auto()
 214    YEAR = auto()
 215    USERDEFINED = auto()
 216    MONEY = auto()
 217    SMALLMONEY = auto()
 218    ROWVERSION = auto()
 219    IMAGE = auto()
 220    VARIANT = auto()
 221    OBJECT = auto()
 222    INET = auto()
 223    IPADDRESS = auto()
 224    IPPREFIX = auto()
 225    IPV4 = auto()
 226    IPV6 = auto()
 227    ENUM = auto()
 228    ENUM8 = auto()
 229    ENUM16 = auto()
 230    FIXEDSTRING = auto()
 231    LOWCARDINALITY = auto()
 232    NESTED = auto()
 233    AGGREGATEFUNCTION = auto()
 234    SIMPLEAGGREGATEFUNCTION = auto()
 235    TDIGEST = auto()
 236    UNKNOWN = auto()
 237    VECTOR = auto()
 238    DYNAMIC = auto()
 239    VOID = auto()
 240
 241    # keywords
 242    ALIAS = auto()
 243    ALTER = auto()
 244    ALL = auto()
 245    ANTI = auto()
 246    ANY = auto()
 247    APPLY = auto()
 248    ARRAY = auto()
 249    ASC = auto()
 250    ASOF = auto()
 251    ATTACH = auto()
 252    AUTO_INCREMENT = auto()
 253    BEGIN = auto()
 254    BETWEEN = auto()
 255    BULK_COLLECT_INTO = auto()
 256    CACHE = auto()
 257    CASE = auto()
 258    CHARACTER_SET = auto()
 259    CLUSTER_BY = auto()
 260    COLLATE = auto()
 261    COMMAND = auto()
 262    COMMENT = auto()
 263    COMMIT = auto()
 264    CONNECT_BY = auto()
 265    CONSTRAINT = auto()
 266    COPY = auto()
 267    CREATE = auto()
 268    CROSS = auto()
 269    CUBE = auto()
 270    CURRENT_DATE = auto()
 271    CURRENT_DATETIME = auto()
 272    CURRENT_SCHEMA = auto()
 273    CURRENT_TIME = auto()
 274    CURRENT_TIMESTAMP = auto()
 275    CURRENT_USER = auto()
 276    CURRENT_USER_ID = auto()
 277    CURRENT_ROLE = auto()
 278    CURRENT_CATALOG = auto()
 279    DECLARE = auto()
 280    DEFAULT = auto()
 281    DELETE = auto()
 282    DESC = auto()
 283    DESCRIBE = auto()
 284    DETACH = auto()
 285    DICTIONARY = auto()
 286    DISTINCT = auto()
 287    DISTRIBUTE_BY = auto()
 288    DIV = auto()
 289    DROP = auto()
 290    ELSE = auto()
 291    END = auto()
 292    ESCAPE = auto()
 293    EXCEPT = auto()
 294    EXECUTE = auto()
 295    EXISTS = auto()
 296    FALSE = auto()
 297    FETCH = auto()
 298    FILE = auto()
 299    FILE_FORMAT = auto()
 300    FILTER = auto()
 301    FINAL = auto()
 302    FIRST = auto()
 303    FOR = auto()
 304    FORCE = auto()
 305    FOREIGN_KEY = auto()
 306    FORMAT = auto()
 307    FROM = auto()
 308    FULL = auto()
 309    FUNCTION = auto()
 310    GET = auto()
 311    GLOB = auto()
 312    GLOBAL = auto()
 313    GRANT = auto()
 314    GROUP_BY = auto()
 315    GROUPING_SETS = auto()
 316    HAVING = auto()
 317    HINT = auto()
 318    IGNORE = auto()
 319    ILIKE = auto()
 320    IN = auto()
 321    INDEX = auto()
 322    INDEXED_BY = auto()
 323    INNER = auto()
 324    INSERT = auto()
 325    INSTALL = auto()
 326    INTEGRATION = auto()
 327    INTERSECT = auto()
 328    INTERVAL = auto()
 329    INTO = auto()
 330    INTRODUCER = auto()
 331    IRLIKE = auto()
 332    IS = auto()
 333    ISNULL = auto()
 334    JOIN = auto()
 335    JOIN_MARKER = auto()
 336    KEEP = auto()
 337    KEY = auto()
 338    KILL = auto()
 339    LANGUAGE = auto()
 340    LATERAL = auto()
 341    LEFT = auto()
 342    LIKE = auto()
 343    LIMIT = auto()
 344    LIST = auto()
 345    LOAD = auto()
 346    LOCK = auto()
 347    MAP = auto()
 348    MATCH = auto()
 349    MATCH_CONDITION = auto()
 350    MATCH_RECOGNIZE = auto()
 351    MEMBER_OF = auto()
 352    MERGE = auto()
 353    MOD = auto()
 354    MODEL = auto()
 355    NATURAL = auto()
 356    NEXT = auto()
 357    NOTHING = auto()
 358    NOTNULL = auto()
 359    NULL = auto()
 360    OBJECT_IDENTIFIER = auto()
 361    OFFSET = auto()
 362    ON = auto()
 363    ONLY = auto()
 364    OPERATOR = auto()
 365    ORDER_BY = auto()
 366    ORDER_SIBLINGS_BY = auto()
 367    ORDERED = auto()
 368    ORDINALITY = auto()
 369    OUT = auto()
 370    INOUT = auto()
 371    OUTER = auto()
 372    OVER = auto()
 373    OVERLAPS = auto()
 374    OVERWRITE = auto()
 375    PACKAGE = auto()
 376    PARTITION = auto()
 377    PARTITION_BY = auto()
 378    PASTE = auto()
 379    PERCENT = auto()
 380    PIVOT = auto()
 381    PLACEHOLDER = auto()
 382    POLICY = auto()
 383    POOL = auto()
 384    POSITIONAL = auto()
 385    PRAGMA = auto()
 386    PREWHERE = auto()
 387    PRIMARY_KEY = auto()
 388    PROCEDURE = auto()
 389    PROPERTIES = auto()
 390    PROJECTION = auto()
 391    PSEUDO_TYPE = auto()
 392    PUT = auto()
 393    QUALIFY = auto()
 394    QUOTE = auto()
 395    QDCOLON = auto()
 396    RANGE = auto()
 397    RECURSIVE = auto()
 398    REFRESH = auto()
 399    RENAME = auto()
 400    REPLACE = auto()
 401    RETURNING = auto()
 402    REVOKE = auto()
 403    REFERENCES = auto()
 404    RIGHT = auto()
 405    RLIKE = auto()
 406    ROLE = auto()
 407    ROLLBACK = auto()
 408    ROLLUP = auto()
 409    ROW = auto()
 410    ROWS = auto()
 411    RULE = auto()
 412    SELECT = auto()
 413    SEMI = auto()
 414    SEPARATOR = auto()
 415    SEQUENCE = auto()
 416    SERDE_PROPERTIES = auto()
 417    SET = auto()
 418    SETTINGS = auto()
 419    SHOW = auto()
 420    SIMILAR_TO = auto()
 421    SOME = auto()
 422    SORT_BY = auto()
 423    SOUNDS_LIKE = auto()
 424    SQL_SECURITY = auto()
 425    START_WITH = auto()
 426    STORAGE_INTEGRATION = auto()
 427    STRAIGHT_JOIN = auto()
 428    STRUCT = auto()
 429    SUMMARIZE = auto()
 430    TABLE_SAMPLE = auto()
 431    TAG = auto()
 432    TEMPORARY = auto()
 433    TOP = auto()
 434    THEN = auto()
 435    TRUE = auto()
 436    TRUNCATE = auto()
 437    TRIGGER = auto()
 438    TYPE = auto()
 439    UNCACHE = auto()
 440    UNDROP = auto()
 441    UNION = auto()
 442    UNNEST = auto()
 443    UNPIVOT = auto()
 444    UPDATE = auto()
 445    USE = auto()
 446    USING = auto()
 447    VALUES = auto()
 448    VARIADIC = auto()
 449    VIEW = auto()
 450    SEMANTIC_VIEW = auto()
 451    VOLATILE = auto()
 452    VOLUME = auto()
 453    WHEN = auto()
 454    WHERE = auto()
 455    WINDOW = auto()
 456    WITH = auto()
 457    UNIQUE = auto()
 458    UTC_DATE = auto()
 459    UTC_TIME = auto()
 460    UTC_TIMESTAMP = auto()
 461    OPTION = auto()
 462    SINK = auto()
 463    SOURCE = auto()
 464    ANALYZE = auto()
 465    NAMESPACE = auto()
 466    EXPORT = auto()
 467
 468    # sentinels
 469    HIVE_TOKEN_STREAM = auto()
 470    SENTINEL = auto()
 471
 472    def __str__(self) -> str:
 473        return f"TokenType.{self.name}"
 474
 475
 476class Token:
 477    # mypyc doesn't expose slots
 478    _attrs: t.ClassVar[tuple[str, ...]] = (
 479        "token_type",
 480        "text",
 481        "line",
 482        "col",
 483        "start",
 484        "end",
 485        "comments",
 486    )
 487    __slots__ = _attrs
 488
 489    @classmethod
 490    def number(cls, number: int) -> Token:
 491        """Returns a NUMBER token with `number` as its text."""
 492        return cls(TokenType.NUMBER, str(number))
 493
 494    @classmethod
 495    def string(cls, string: str) -> Token:
 496        """Returns a STRING token with `string` as its text."""
 497        return cls(TokenType.STRING, string)
 498
 499    @classmethod
 500    def identifier(cls, identifier: str) -> Token:
 501        """Returns an IDENTIFIER token with `identifier` as its text."""
 502        return cls(TokenType.IDENTIFIER, identifier)
 503
 504    @classmethod
 505    def var(cls, var: str) -> Token:
 506        """Returns an VAR token with `var` as its text."""
 507        return cls(TokenType.VAR, var)
 508
 509    def __init__(
 510        self,
 511        token_type: TokenType,
 512        text: str,
 513        line: int = 1,
 514        col: int = 1,
 515        start: int = 0,
 516        end: int = 0,
 517        comments: list[str] | None = None,
 518    ) -> None:
 519        self.token_type = token_type
 520        self.text = text
 521        self.line = line
 522        self.col = col
 523        self.start = start
 524        self.end = end
 525        self.comments = [] if comments is None else comments
 526
 527    def __bool__(self) -> bool:
 528        return self.token_type != TokenType.SENTINEL
 529
 530    def __repr__(self) -> str:
 531        attributes = ", ".join(
 532            f"{k}: TokenType.{self.token_type.name}"
 533            if k == "token_type"
 534            else f"{k}: {getattr(self, k)}"
 535            for k in self._attrs
 536        )
 537        return f"<Token {attributes}>"
 538
 539
 540class TokenizerCore:
 541    __slots__ = (
 542        "sql",
 543        "size",
 544        "tokens",
 545        "_start",
 546        "_current",
 547        "_line",
 548        "_col",
 549        "_comments",
 550        "_char",
 551        "_end",
 552        "_peek",
 553        "_prev_token_line",
 554        "single_tokens",
 555        "keywords",
 556        "quotes",
 557        "format_strings",
 558        "identifiers",
 559        "comments",
 560        "string_escapes",
 561        "byte_string_escapes",
 562        "identifier_escapes",
 563        "commands",
 564        "command_prefix_tokens",
 565        "nested_comments",
 566        "dash_comment_requires_boundary",
 567        "comments_terminate_at_newline_only",
 568        "hint_start",
 569        "tokens_preceding_hint",
 570        "has_bit_strings",
 571        "has_hex_strings",
 572        "numeric_literals",
 573        "var_single_tokens",
 574        "string_escapes_allowed_in_raw_strings",
 575        "heredoc_tag_is_identifier",
 576        "heredoc_string_alternative",
 577        "keyword_trie",
 578        "numbers_can_be_underscore_separated",
 579        "numbers_can_have_decimals",
 580        "identifiers_can_start_with_digit",
 581        "unescaped_sequences",
 582        "numeric_escapes",
 583        "drop_unknown_escapes",
 584        "numeric_escapes_are_bytes",
 585        "byte_strings_are_bytes",
 586        "lone_surrogate_replacement",
 587    )
 588
 589    def __init__(
 590        self,
 591        single_tokens: dict[str, TokenType],
 592        keywords: dict[str, TokenType],
 593        quotes: dict[str, str],
 594        format_strings: dict[str, tuple[str, TokenType]],
 595        identifiers: dict[str, str],
 596        comments: dict[str, str | None],
 597        string_escapes: set[str],
 598        byte_string_escapes: set[str],
 599        identifier_escapes: set[str],
 600        commands: set[TokenType],
 601        command_prefix_tokens: set[TokenType],
 602        nested_comments: bool,
 603        dash_comment_requires_boundary: bool,
 604        comments_terminate_at_newline_only: bool,
 605        hint_start: str,
 606        tokens_preceding_hint: set[TokenType],
 607        has_bit_strings: bool,
 608        has_hex_strings: bool,
 609        numeric_literals: dict[str, str],
 610        var_single_tokens: set[str],
 611        string_escapes_allowed_in_raw_strings: bool,
 612        heredoc_tag_is_identifier: bool,
 613        heredoc_string_alternative: TokenType,
 614        keyword_trie: dict,
 615        numbers_can_be_underscore_separated: bool,
 616        numbers_can_have_decimals: bool,
 617        identifiers_can_start_with_digit: bool,
 618        unescaped_sequences: dict[str, str],
 619        numeric_escapes: dict[str, tuple[int, int, int, int]],
 620        drop_unknown_escapes: bool,
 621        numeric_escapes_are_bytes: bool,
 622        byte_strings_are_bytes: bool,
 623        lone_surrogate_replacement: str,
 624    ) -> None:
 625        self.single_tokens = single_tokens
 626        self.keywords = keywords
 627        self.quotes = quotes
 628        self.format_strings = format_strings
 629        self.identifiers = identifiers
 630        self.comments = comments
 631        self.string_escapes = string_escapes
 632        self.byte_string_escapes = byte_string_escapes
 633        self.identifier_escapes = identifier_escapes
 634        self.commands = commands
 635        self.command_prefix_tokens = command_prefix_tokens
 636        self.nested_comments = nested_comments
 637        self.dash_comment_requires_boundary = dash_comment_requires_boundary
 638        self.comments_terminate_at_newline_only = comments_terminate_at_newline_only
 639        self.hint_start = hint_start
 640        self.tokens_preceding_hint = tokens_preceding_hint
 641        self.has_bit_strings = has_bit_strings
 642        self.has_hex_strings = has_hex_strings
 643        self.numeric_literals = numeric_literals
 644        self.var_single_tokens = var_single_tokens
 645        self.string_escapes_allowed_in_raw_strings = string_escapes_allowed_in_raw_strings
 646        self.heredoc_tag_is_identifier = heredoc_tag_is_identifier
 647        self.heredoc_string_alternative = heredoc_string_alternative
 648        self.keyword_trie = keyword_trie
 649        self.numbers_can_be_underscore_separated = numbers_can_be_underscore_separated
 650        self.numbers_can_have_decimals = numbers_can_have_decimals
 651        self.identifiers_can_start_with_digit = identifiers_can_start_with_digit
 652        self.unescaped_sequences = unescaped_sequences
 653        self.numeric_escapes = numeric_escapes
 654        self.drop_unknown_escapes = drop_unknown_escapes
 655        self.numeric_escapes_are_bytes = numeric_escapes_are_bytes
 656        self.byte_strings_are_bytes = byte_strings_are_bytes
 657        self.lone_surrogate_replacement = lone_surrogate_replacement
 658        self.sql = ""
 659        self.size = 0
 660        self.tokens: list[Token] = []
 661        self._start = 0
 662        self._current = 0
 663        self._line = 1
 664        self._col = 0
 665        self._comments: list[str] = []
 666        self._char = ""
 667        self._end = False
 668        self._peek = ""
 669        self._prev_token_line = -1
 670
 671    def reset(self) -> None:
 672        self.sql = ""
 673        self.size = 0
 674        self.tokens = []
 675        self._start = 0
 676        self._current = 0
 677        self._line = 1
 678        self._col = 0
 679        self._comments = []
 680        self._char = ""
 681        self._end = False
 682        self._peek = ""
 683        self._prev_token_line = -1
 684
 685    def tokenize(self, sql: str) -> list[Token]:
 686        """Returns a list of tokens corresponding to the SQL string `sql`."""
 687        self.reset()
 688        self.sql = sql
 689        self.size = len(sql)
 690
 691        try:
 692            self._scan()
 693        except Exception as e:
 694            start = max(self._current - 50, 0)
 695            end = min(self._current + 50, self.size - 1)
 696            context = self.sql[start:end]
 697            raise TokenError(f"Error tokenizing '{context}'", start=start, end=end) from e
 698
 699        return self.tokens
 700
 701    def _scan(self, check_semicolon: bool = False) -> None:
 702        identifiers = self.identifiers
 703        digit_chars = _DIGIT_CHARS
 704
 705        while self.size and not self._end:
 706            current = self._current
 707
 708            # Skip spaces here rather than iteratively calling advance() for performance reasons
 709            while current < self.size:
 710                char = self.sql[current]
 711
 712                if char == " " or char == "\t":
 713                    current += 1
 714                else:
 715                    break
 716
 717            offset = current - self._current if current > self._current else 1
 718
 719            self._start = current
 720            self._advance(offset)
 721
 722            if not self._char.isspace():
 723                if self._char in digit_chars:
 724                    self._scan_number()
 725                elif self._char in identifiers:
 726                    self._scan_identifier(identifiers[self._char])
 727                else:
 728                    self._scan_keywords()
 729
 730            if check_semicolon and self._peek == ";":
 731                break
 732
 733        if self.tokens and self._comments:
 734            self.tokens[-1].comments.extend(self._comments)
 735
 736    def _chars(self, size: int) -> str:
 737        if size == 1:
 738            return self._char
 739
 740        start = self._current - 1
 741        end = start + size
 742
 743        return self.sql[start:end] if end <= self.size else ""
 744
 745    def _advance(self, i: int = 1, alnum: bool = False) -> None:
 746        char = self._char
 747
 748        if char == "\n" or char == "\r":
 749            # Ensures we don't count an extra line if we get a \r\n line break sequence
 750            if not (char == "\r" and self._peek == "\n"):
 751                self._col = i
 752                self._line += 1
 753        else:
 754            self._col += i
 755
 756        self._current += i
 757        sql = self.sql
 758        size = self.size
 759        self._end = self._current >= size
 760        self._char = sql[self._current - 1]
 761        self._peek = "" if self._end else sql[self._current]
 762
 763        if alnum and self._char.isalnum():
 764            # Cache to local variables instead of attributes for better performance
 765            _col = self._col
 766            _current = self._current
 767            _end = self._end
 768            _peek = self._peek
 769
 770            while _peek.isalnum():
 771                _col += 1
 772                _current += 1
 773                _end = _current >= size
 774                _peek = "" if _end else sql[_current]
 775
 776            self._col = _col
 777            self._current = _current
 778            self._end = _end
 779            self._peek = _peek
 780            self._char = sql[_current - 1]
 781
 782    @property
 783    def _text(self) -> str:
 784        return self.sql[self._start : self._current]
 785
 786    def _add(self, token_type: TokenType, text: str | None = None) -> None:
 787        self._prev_token_line = self._line
 788
 789        if self._comments and token_type == TokenType.SEMICOLON and self.tokens:
 790            self.tokens[-1].comments.extend(self._comments)
 791            self._comments = []
 792
 793        if text is None:
 794            text = self.sql[self._start : self._current]
 795
 796        self.tokens.append(
 797            Token(
 798                token_type,
 799                text=text,
 800                line=self._line,
 801                col=self._col,
 802                start=self._start,
 803                end=self._current - 1,
 804                comments=self._comments,
 805            )
 806        )
 807        self._comments = []
 808
 809        # If we have either a semicolon or a begin token before the command's token, we'll parse
 810        # whatever follows the command's token as a string
 811        if (
 812            token_type in self.commands
 813            and self._peek != ";"
 814            and (len(self.tokens) == 1 or self.tokens[-2].token_type in self.command_prefix_tokens)
 815        ):
 816            start = self._current
 817            tokens = len(self.tokens)
 818            self._scan(check_semicolon=True)
 819            self.tokens = self.tokens[:tokens]
 820            text = self.sql[start : self._current].strip()
 821            if text:
 822                self._add(TokenType.STRING, text)
 823
 824    def _scan_keywords(self) -> None:
 825        sql = self.sql
 826        sql_size = self.size
 827        single_tokens = self.single_tokens
 828        char_upper = _CHAR_UPPER
 829        size = 0
 830        word = None
 831        chars = self._char
 832        char = chars
 833        prev_space = False
 834        skip = False
 835        trie = self.keyword_trie
 836        single_token = char in single_tokens
 837
 838        while chars:
 839            if not skip:
 840                sub = trie.get(char_upper.get(char, char))
 841                if sub is None:
 842                    break
 843                trie = sub
 844                if 0 in trie:
 845                    word = chars
 846
 847            end = self._current + size
 848            size += 1
 849
 850            if end < sql_size:
 851                char = sql[end]
 852                single_token = single_token or char in single_tokens
 853                is_space = char.isspace()
 854
 855                if not is_space or not prev_space:
 856                    if is_space:
 857                        char = " "
 858                    chars += char
 859                    prev_space = is_space
 860                    skip = False
 861                else:
 862                    skip = True
 863            else:
 864                char = ""
 865                break
 866
 867        if word:
 868            if self._scan_string(word):
 869                return
 870            if self._scan_comment(word):
 871                return
 872            if prev_space or single_token or not char:
 873                word = word.upper()
 874                if word in self.keywords:
 875                    self._advance(size - 1)
 876                    self._add(self.keywords[word], text=word)
 877                    return
 878
 879        if self._char in single_tokens:
 880            self._add(single_tokens[self._char], text=self._char)
 881            return
 882
 883        self._scan_var()
 884
 885    def _scan_comment(self, comment_start: str) -> bool:
 886        if comment_start not in self.comments:
 887            return False
 888
 889        if comment_start == "--" and self.dash_comment_requires_boundary:
 890            boundary_index = self._current + 1
 891            if boundary_index < self.size:
 892                boundary_char = self.sql[boundary_index]
 893                if not (
 894                    boundary_char.isspace() or ord(boundary_char) < 32 or ord(boundary_char) == 127
 895                ):
 896                    return False
 897
 898        comment_start_line = self._line
 899        comment_start_size = len(comment_start)
 900        comment_end = self.comments[comment_start]
 901
 902        if comment_end:
 903            # Skip the comment's start delimiter
 904            self._advance(comment_start_size)
 905
 906            comment_count = 1
 907            comment_end_size = len(comment_end)
 908            nested_comments = self.nested_comments
 909
 910            while not self._end:
 911                if self._chars(comment_end_size) == comment_end:
 912                    comment_count -= 1
 913                    if not comment_count:
 914                        break
 915
 916                self._advance(alnum=True)
 917
 918                # Nested comments are allowed by some dialects, e.g. databricks, duckdb, postgres
 919                if (
 920                    nested_comments
 921                    and not self._end
 922                    and self._chars(comment_end_size) == comment_start
 923                ):
 924                    self._advance(comment_start_size)
 925                    comment_count += 1
 926
 927            self._comments.append(self._text[comment_start_size : -comment_end_size + 1])
 928            self._advance(comment_end_size - 1)
 929        else:
 930            _peek = self._peek
 931            terminates_at_carriage_return = not self.comments_terminate_at_newline_only
 932            while (
 933                not self._end
 934                and _peek != "\n"
 935                and (not terminates_at_carriage_return or _peek != "\r")
 936            ):
 937                self._advance(alnum=True)
 938                _peek = self._peek
 939            self._comments.append(self._text[comment_start_size:])
 940
 941        if (
 942            comment_start == self.hint_start
 943            and self.tokens
 944            and self.tokens[-1].token_type in self.tokens_preceding_hint
 945        ):
 946            self._add(TokenType.HINT)
 947
 948        # Leading comment is attached to the succeeding token, whilst trailing comment to the preceding.
 949        # Multiple consecutive comments are preserved by appending them to the current comments list.
 950        if comment_start_line == self._prev_token_line:
 951            self.tokens[-1].comments.extend(self._comments)
 952            self._comments = []
 953            self._prev_token_line = self._line
 954
 955        return True
 956
 957    def _scan_number(self) -> None:
 958        if self._char == "0":
 959            peek = _CHAR_UPPER.get(self._peek, self._peek)
 960            if peek == "B" and self.has_bit_strings:
 961                return self._scan_bits()
 962            elif peek == "X":
 963                return self._scan_hex() if self.has_hex_strings else self._add(TokenType.NUMBER)
 964
 965        decimal = False
 966        scientific = 0
 967        numbers_can_be_underscore_separated = self.numbers_can_be_underscore_separated
 968        single_tokens = self.single_tokens
 969        keywords = self.keywords
 970        numeric_literals = self.numeric_literals
 971        identifiers_can_start_with_digit = self.identifiers_can_start_with_digit
 972
 973        is_underscore_separated: bool = False
 974        number_text: str = ""
 975        numeric_literal: str = ""
 976        numeric_type: TokenType | None = None
 977
 978        while True:
 979            if self._peek in _DIGIT_CHARS:
 980                # Batch consecutive digits: scan ahead to find how many
 981                sql = self.sql
 982                end = self._current + 1
 983                size = self.size
 984                while end < size and sql[end] in _DIGIT_CHARS:
 985                    end += 1
 986                self._advance(end - self._current)
 987            elif self._peek == "." and not decimal:
 988                if (
 989                    self.tokens and self.tokens[-1].token_type == TokenType.PARAMETER
 990                ) or not self.numbers_can_have_decimals:
 991                    break
 992                decimal = True
 993                self._advance()
 994            elif self._peek in ("-", "+") and scientific == 1:
 995                # Only consume +/- if followed by a digit
 996                if self._current + 1 < self.size and self.sql[self._current + 1] in _DIGIT_CHARS:
 997                    scientific += 1
 998                    self._advance()
 999                else:
1000                    break
1001            elif _CHAR_UPPER.get(self._peek, self._peek) == "E" and not scientific:
1002                scientific += 1
1003                self._advance()
1004            elif self._peek == "_" and numbers_can_be_underscore_separated:
1005                is_underscore_separated = True
1006                self._advance()
1007            elif self._peek.isidentifier():
1008                number_text = self._text
1009
1010                while self._peek and not self._peek.isspace() and self._peek not in single_tokens:
1011                    numeric_literal += self._peek
1012                    self._advance()
1013
1014                numeric_type = keywords.get(numeric_literals.get(numeric_literal.upper(), ""))
1015
1016                if numeric_type:
1017                    break
1018                elif identifiers_can_start_with_digit:
1019                    return self._add(TokenType.VAR)
1020
1021                self._advance(-len(numeric_literal))
1022                break
1023            else:
1024                break
1025
1026        number_text = number_text or self.sql[self._start : self._current]
1027
1028        # Normalize inputs such as 100_000 to 100000
1029        if is_underscore_separated:
1030            number_text = number_text.replace("_", "")
1031
1032        self._add(TokenType.NUMBER, number_text)
1033
1034        # Normalize inputs such as 123L to 123::BIGINT so that they're parsed as casts
1035        if numeric_type:
1036            self._add(TokenType.DCOLON, "::")
1037            self._add(numeric_type, numeric_literal)
1038
1039    def _scan_bits(self) -> None:
1040        self._advance()
1041        value = self._extract_value()
1042        try:
1043            # If `value` can't be converted to a binary, fallback to tokenizing it as an identifier
1044            int(value, 2)
1045            self._add(TokenType.BIT_STRING, value[2:])  # Drop the 0b
1046        except ValueError:
1047            self._add(TokenType.IDENTIFIER)
1048
1049    def _scan_hex(self) -> None:
1050        self._advance()
1051        value = self._extract_value()
1052        try:
1053            # If `value` can't be converted to a hex, fallback to tokenizing it as an identifier
1054            int(value, 16)
1055            self._add(TokenType.HEX_STRING, value[2:])  # Drop the 0x
1056        except ValueError:
1057            self._add(TokenType.IDENTIFIER)
1058
1059    def _extract_value(self) -> str:
1060        single_tokens = self.single_tokens
1061
1062        while True:
1063            char = self._peek.strip()
1064            if char and char not in single_tokens:
1065                self._advance(alnum=True)
1066            else:
1067                break
1068
1069        return self._text
1070
1071    def _scan_string(self, start: str) -> bool:
1072        base = None
1073        token_type = TokenType.STRING
1074
1075        if start in self.quotes:
1076            end = self.quotes[start]
1077        elif start in self.format_strings:
1078            end, token_type = self.format_strings[start]
1079
1080            if token_type == TokenType.HEX_STRING:
1081                base = 16
1082            elif token_type == TokenType.BIT_STRING:
1083                base = 2
1084            elif token_type == TokenType.HEREDOC_STRING:
1085                self._advance()
1086
1087                if self._char == end:
1088                    tag = ""
1089                else:
1090                    tag = self._extract_string(
1091                        end,
1092                        raw_string=True,
1093                        raise_unmatched=not self.heredoc_tag_is_identifier,
1094                    )
1095
1096                if (
1097                    tag
1098                    and self.heredoc_tag_is_identifier
1099                    and (self._end or tag.isdigit() or any(c.isspace() for c in tag))
1100                ):
1101                    if not self._end:
1102                        self._advance(-1)
1103
1104                    self._advance(-len(tag))
1105                    self._add(self.heredoc_string_alternative)
1106                    return True
1107
1108                end = f"{start}{tag}{end}"
1109        else:
1110            return False
1111
1112        self._advance(len(start))
1113        text = self._extract_string(
1114            end,
1115            escapes=(
1116                self.byte_string_escapes
1117                if token_type == TokenType.BYTE_STRING
1118                else self.string_escapes
1119            ),
1120            raw_string=token_type
1121            in (TokenType.RAW_STRING, TokenType.NATIONAL_RAW_STRING, TokenType.HEREDOC_STRING),
1122            bytes_literal=token_type == TokenType.BYTE_STRING and self.byte_strings_are_bytes,
1123        )
1124
1125        if base and text:
1126            try:
1127                int(text, base)
1128            except Exception:
1129                raise TokenError(
1130                    f"Numeric string contains invalid characters from {self._line}:{self._start}"
1131                )
1132
1133        self._add(token_type, text)
1134        return True
1135
1136    def _scan_identifier(self, identifier_end: str) -> None:
1137        self._advance()
1138        text = self._extract_string(
1139            identifier_end, escapes=self.identifier_escapes | {identifier_end}
1140        )
1141        self._add(TokenType.IDENTIFIER, text)
1142
1143    def _scan_var(self) -> None:
1144        var_single_tokens = self.var_single_tokens
1145        single_tokens = self.single_tokens
1146
1147        while True:
1148            peek = self._peek
1149            if not peek or peek.isspace():
1150                break
1151            if peek not in var_single_tokens and peek in single_tokens:
1152                break
1153            self._advance(alnum=True)
1154
1155        self._add(
1156            TokenType.VAR
1157            if self.tokens and self.tokens[-1].token_type == TokenType.PARAMETER
1158            else self.keywords.get(self.sql[self._start : self._current].upper(), TokenType.VAR)
1159        )
1160
1161    def _read_numeric_escape(
1162        self, start: int, base: int, min_digits: int, max_digits: int, max_value: int
1163    ) -> tuple[int, int]:
1164        """
1165        Reads up to `max_digits` digits from `start`, stopping before the value exceeds `max_value`.
1166        Returns the value and the end index, or -1 as the value if fewer than `min_digits` were read.
1167        """
1168        sql = self.sql
1169        value = 0
1170        end = start
1171        limit = min(start + max_digits, self.size)
1172
1173        while end < limit:
1174            digit = _DIGIT_VALUES.get(sql[end], 16)
1175            if digit >= base or value * base + digit > max_value:
1176                break
1177            value = value * base + digit
1178            end += 1
1179
1180        return (value, end) if end - start >= min_digits else (-1, end)
1181
1182    def _read_byte_escape(self, pos: int) -> tuple[int, int]:
1183        """Reads the numeric escape at `pos` (a backslash) if it encodes a byte, else returns -1."""
1184        peek = self.sql[pos + 1 : pos + 2]
1185        is_octal = peek in _OCTAL_CHARS
1186        spec = self.numeric_escapes.get("0" if is_octal else peek)
1187        if not spec or (spec[0] != 8 and spec[3] > 0xFF):
1188            return -1, pos
1189
1190        base, min_digits, max_digits, max_value = spec
1191        value, end = self._read_numeric_escape(
1192            pos + 1 if is_octal else pos + 2, base, min_digits, max_digits, max_value
1193        )
1194        return (value & 0xFF if value >= 0 else value), end
1195
1196    def _scan_numeric_escape(self) -> str:
1197        """
1198        Decodes the numeric escape sequence starting at the current char, which is a backslash,
1199        according to `numeric_escapes`. Returns an empty string if nothing is decoded.
1200        """
1201        sql = self.sql
1202        start = self._current
1203        peek = self._peek
1204
1205        # Octal escapes start with a digit (e.g. \101), the others with a letter (e.g. \x41)
1206        is_octal = peek in _OCTAL_CHARS
1207        spec = self.numeric_escapes.get("0" if is_octal else peek)
1208
1209        if spec:
1210            base, min_digits, max_digits, max_value = spec
1211            value, end = self._read_numeric_escape(
1212                start if is_octal else start + 1, base, min_digits, max_digits, max_value
1213            )
1214
1215            # With byte semantics, octal and hex escapes encode a byte, so octal values wrap around,
1216            # e.g. \541 is \x61
1217            is_byte = self.numeric_escapes_are_bytes and (base == 8 or max_value <= 0xFF)
1218            if is_byte and value >= 0:
1219                value &= 0xFF
1220
1221            # A UTF-16 surrogate is only valid as a high half (D800-DBFF) followed by an escaped
1222            # low half (DC00-DFFF), e.g. \uD83D\uDE00 is U+1F600
1223            if 0xD800 <= value <= 0xDFFF:
1224                low, low_end = -1, end
1225                if value <= 0xDBFF and sql[end : end + 2] == "\\" + peek:
1226                    low, low_end = self._read_numeric_escape(
1227                        end + 2, base, min_digits, max_digits, max_value
1228                    )
1229
1230                if 0xDC00 <= low <= 0xDFFF:
1231                    # Combine the two halves into a single code point (standard formula)
1232                    value = 0x10000 + ((value - 0xD800) << 10) + (low - 0xDC00)
1233                    end = low_end
1234                elif self.lone_surrogate_replacement and end < self.size:
1235                    # A lone surrogate isn't a character, so some engines replace it (e.g. with ?)
1236                    self._advance(end - start + 1)
1237                    return self.lone_surrogate_replacement
1238                else:
1239                    # A lone surrogate isn't a character, so the escape isn't decoded
1240                    value = -1
1241
1242            # Decode one UTF-8 character at a time to avoid rescanning long invalid byte runs.
1243            # Invalid UTF-8 falls back to decoding each byte as a code point. This first check
1244            # detects whether the byte can start a multi-byte UTF-8 character.
1245            if is_byte and 0xC2 <= value <= 0xF4:
1246                run = bytearray([value])
1247                run_end = end
1248                byte_count = 2 if value < 0xE0 else 3 if value < 0xF0 else 4
1249                while len(run) < byte_count and sql[run_end : run_end + 1] == "\\":
1250                    byte, byte_end = self._read_byte_escape(run_end)
1251                    if not 0x80 <= byte <= 0xBF:
1252                        break
1253                    run.append(byte)
1254                    run_end = byte_end
1255
1256                try:
1257                    decoded = run.decode("utf-8")
1258                except UnicodeDecodeError:
1259                    decoded = ""
1260
1261                if decoded and run_end < self.size:
1262                    self._advance(run_end - start + 1)
1263                    return decoded
1264
1265            # -1 means the escape isn't decoded, e.g. because it has too few digits. An escape at the
1266            # end of the input is left to the caller, which reports the unterminated string
1267            if value >= 0 and end < self.size:
1268                # Move past the whole escape, from the backslash through its last digit
1269                self._advance(end - start + 1)
1270                return chr(value)
1271
1272        return ""
1273
1274    def _extract_string(
1275        self,
1276        delimiter: str,
1277        escapes: set[str] | None = None,
1278        raw_string: bool = False,
1279        raise_unmatched: bool = True,
1280        bytes_literal: bool = False,
1281    ) -> str:
1282        """
1283        Extracts a string up to the delimiter, decoding its escape sequences. For byte literals,
1284        each char of the result is a byte, so non-ASCII chars are stored as their UTF-8 bytes.
1285        """
1286        text = ""
1287        delim_size = len(delimiter)
1288        escapes = self.string_escapes if escapes is None else escapes
1289        unescaped_sequences = self.unescaped_sequences
1290        numeric_escapes = self.numeric_escapes
1291        drop_unknown_escapes = self.drop_unknown_escapes
1292        string_escapes_allowed_in_raw_strings = self.string_escapes_allowed_in_raw_strings
1293        quotes = self.quotes
1294        sql = self.sql
1295
1296        # use str.find() when the string is simple... no \ or other escapes
1297        if delim_size == 1:
1298            pos = self._current - 1
1299            end = sql.find(delimiter, pos)
1300
1301            if (
1302                # the closing delimiter was found
1303                end != -1
1304                # there's no doubled delimiter (e.g. '' escape), or the delimiter isn't an escape char
1305                and (end + 1 >= self.size or sql[end + 1] != delimiter or delimiter not in escapes)
1306                # no backslash in the string that would need escape processing
1307                and (not (unescaped_sequences or "\\" in escapes) or sql.find("\\", pos, end) == -1)
1308            ):
1309                newlines = sql.count("\n", pos, end)
1310                if newlines:
1311                    self._line += newlines
1312                    self._col = end - sql.rfind("\n", pos, end)
1313                else:
1314                    self._col += end - pos
1315
1316                self._current = end + 1
1317                self._end = self._current >= self.size
1318                self._char = sql[end]
1319                self._peek = "" if self._end else sql[self._current]
1320                text = sql[pos:end]
1321                if bytes_literal and not text.isascii():
1322                    text = text.encode("utf-8").decode("latin-1")
1323                return text
1324
1325        while True:
1326            backslash_escape = not raw_string and self._char == "\\" and "\\" in escapes
1327
1328            # Numeric escapes come first, since they may overlap fixed sequences (e.g. \000 vs \0)
1329            if backslash_escape and numeric_escapes:
1330                decoded = self._scan_numeric_escape()
1331                if decoded:
1332                    text += decoded
1333                    continue
1334
1335            # An escape at the end of the input is left to the checks below, which report the
1336            # unterminated string
1337            if (
1338                not raw_string
1339                and unescaped_sequences
1340                and self._char in escapes
1341                and self._current + 1 < self.size
1342            ):
1343                unescaped_sequence = unescaped_sequences.get(self._char + self._peek)
1344                if unescaped_sequence is not None:
1345                    # Advance one char at a time so that line numbers are updated for a newline
1346                    if self._peek in "\n\r":
1347                        self._advance()
1348                        self._advance()
1349                    else:
1350                        self._advance(2)
1351                    text += unescaped_sequence
1352                    continue
1353
1354            # Not decoded: drop the backslash if configured, unless the string is unterminated
1355            if backslash_escape and drop_unknown_escapes and self._current + 1 < self.size:
1356                # Advance one char at a time so that line numbers are updated if `peek` is a newline
1357                peek = self._peek
1358                self._advance()
1359                self._advance()
1360                text += peek.encode("utf-8").decode("latin-1") if bytes_literal else peek
1361                continue
1362
1363            # An escaped quote before the closing delimiter (e.g. \" in """a\"""") must be
1364            # consumed here, otherwise it'd be picked up by the delimiter check below and
1365            # terminate the string early. This is only relevant for multi-char delimiters
1366            # made up of quote chars, e.g. it shouldn't apply to Snowflake's $$ strings
1367            escaped_delimiter = self._peek == delimiter or (
1368                delim_size > 1 and self._peek == delimiter[0] and self._peek in quotes
1369            )
1370
1371            if (
1372                (string_escapes_allowed_in_raw_strings or not raw_string)
1373                and self._char in escapes
1374                and (escaped_delimiter or self._peek in escapes)
1375                and (self._char not in quotes or self._char == self._peek)
1376            ):
1377                if escaped_delimiter:
1378                    text += self._peek if not raw_string else self._char + self._peek
1379                else:
1380                    text += self._char + self._peek
1381
1382                if self._current + 1 < self.size:
1383                    self._advance(2)
1384                else:
1385                    raise TokenError(f"Missing {delimiter} from {self._line}:{self._current}")
1386            else:
1387                if self._chars(delim_size) == delimiter:
1388                    if delim_size > 1:
1389                        self._advance(delim_size - 1)
1390                    break
1391
1392                if self._end:
1393                    if not raise_unmatched:
1394                        return text + self._char
1395
1396                    raise TokenError(f"Missing {delimiter} from {self._line}:{self._start}")
1397
1398                current = self._current - 1
1399                self._advance(alnum=True)
1400                chunk = sql[current : self._current - 1]
1401                if bytes_literal and not chunk.isascii():
1402                    chunk = chunk.encode("utf-8").decode("latin-1")
1403                text += chunk
1404
1405        return text