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(), .isspace(), .isdigit()
   9_CHAR_UPPER: t.Dict[str, str] = {chr(i): chr(i).upper() for i in range(97, 123)}
  10
  11_SPACE_CHARS: t.FrozenSet[str] = frozenset(
  12    "\t\n\r \x0b\x0c\x1c\x1d\x1e\x1f\x85\xa0"
  13    "\u1680\u2000\u2001\u2002\u2003\u2004\u2005\u2006\u2007\u2008\u2009\u200a"
  14    "\u2028\u2029\u202f\u205f\u3000"
  15)
  16_DIGIT_CHARS: t.FrozenSet[str] = frozenset("0123456789")
  17
  18
  19class TokenType(IntEnum):
  20    L_PAREN = auto()
  21    R_PAREN = auto()
  22    L_BRACKET = auto()
  23    R_BRACKET = auto()
  24    L_BRACE = auto()
  25    R_BRACE = auto()
  26    COMMA = auto()
  27    DOT = auto()
  28    DASH = auto()
  29    PLUS = auto()
  30    COLON = auto()
  31    DOTCOLON = auto()
  32    DCOLON = auto()
  33    DCOLONDOLLAR = auto()
  34    DCOLONPERCENT = auto()
  35    DCOLONQMARK = auto()
  36    DQMARK = auto()
  37    SEMICOLON = auto()
  38    STAR = auto()
  39    BACKSLASH = auto()
  40    SLASH = auto()
  41    LT = auto()
  42    LTE = auto()
  43    GT = auto()
  44    GTE = auto()
  45    NOT = auto()
  46    EQ = auto()
  47    NEQ = auto()
  48    NULLSAFE_EQ = auto()
  49    COLON_EQ = auto()
  50    COLON_GT = auto()
  51    NCOLON_GT = auto()
  52    AND = auto()
  53    OR = auto()
  54    AMP = auto()
  55    DPIPE = auto()
  56    PIPE_GT = auto()
  57    PIPE = auto()
  58    PIPE_SLASH = auto()
  59    DPIPE_SLASH = auto()
  60    CARET = auto()
  61    CARET_AT = auto()
  62    TILDE = auto()
  63    ARROW = auto()
  64    DARROW = auto()
  65    FARROW = auto()
  66    HASH = auto()
  67    HASH_ARROW = auto()
  68    DHASH_ARROW = auto()
  69    LR_ARROW = auto()
  70    DAT = auto()
  71    LT_AT = auto()
  72    AT_GT = auto()
  73    DOLLAR = auto()
  74    PARAMETER = auto()
  75    SESSION = auto()
  76    SESSION_PARAMETER = auto()
  77    SESSION_USER = auto()
  78    DAMP = auto()
  79    AMP_LT = auto()
  80    AMP_GT = auto()
  81    ADJACENT = auto()
  82    XOR = auto()
  83    DSTAR = auto()
  84    QMARK_AMP = auto()
  85    QMARK_PIPE = auto()
  86    HASH_DASH = auto()
  87    EXCLAMATION = auto()
  88
  89    URI_START = auto()
  90
  91    BLOCK_START = auto()
  92    BLOCK_END = auto()
  93
  94    SPACE = auto()
  95    BREAK = auto()
  96
  97    STRING = auto()
  98    NUMBER = auto()
  99    IDENTIFIER = auto()
 100    DATABASE = auto()
 101    COLUMN = auto()
 102    COLUMN_DEF = auto()
 103    SCHEMA = auto()
 104    TABLE = auto()
 105    WAREHOUSE = auto()
 106    STAGE = auto()
 107    STREAMLIT = auto()
 108    VAR = auto()
 109    BIT_STRING = auto()
 110    HEX_STRING = auto()
 111    BYTE_STRING = auto()
 112    NATIONAL_STRING = auto()
 113    RAW_STRING = auto()
 114    HEREDOC_STRING = auto()
 115    UNICODE_STRING = auto()
 116
 117    # types
 118    BIT = auto()
 119    BOOLEAN = auto()
 120    TINYINT = auto()
 121    UTINYINT = auto()
 122    SMALLINT = auto()
 123    USMALLINT = auto()
 124    MEDIUMINT = auto()
 125    UMEDIUMINT = auto()
 126    INT = auto()
 127    UINT = auto()
 128    BIGINT = auto()
 129    UBIGINT = auto()
 130    BIGNUM = auto()
 131    INT128 = auto()
 132    UINT128 = auto()
 133    INT256 = auto()
 134    UINT256 = auto()
 135    FLOAT = auto()
 136    DOUBLE = auto()
 137    UDOUBLE = auto()
 138    DECIMAL = auto()
 139    DECIMAL32 = auto()
 140    DECIMAL64 = auto()
 141    DECIMAL128 = auto()
 142    DECIMAL256 = auto()
 143    DECFLOAT = auto()
 144    UDECIMAL = auto()
 145    BIGDECIMAL = auto()
 146    CHAR = auto()
 147    NCHAR = auto()
 148    VARCHAR = auto()
 149    NVARCHAR = auto()
 150    BPCHAR = auto()
 151    TEXT = auto()
 152    MEDIUMTEXT = auto()
 153    LONGTEXT = auto()
 154    BLOB = auto()
 155    MEDIUMBLOB = auto()
 156    LONGBLOB = auto()
 157    TINYBLOB = auto()
 158    TINYTEXT = auto()
 159    NAME = auto()
 160    BINARY = auto()
 161    VARBINARY = auto()
 162    JSON = auto()
 163    JSONB = auto()
 164    TIME = auto()
 165    TIMETZ = auto()
 166    TIME_NS = auto()
 167    TIMESTAMP = auto()
 168    TIMESTAMPTZ = auto()
 169    TIMESTAMPLTZ = auto()
 170    TIMESTAMPNTZ = auto()
 171    TIMESTAMP_S = auto()
 172    TIMESTAMP_MS = auto()
 173    TIMESTAMP_NS = auto()
 174    DATETIME = auto()
 175    DATETIME2 = auto()
 176    DATETIME64 = auto()
 177    SMALLDATETIME = auto()
 178    DATE = auto()
 179    DATE32 = auto()
 180    INT4RANGE = auto()
 181    INT4MULTIRANGE = auto()
 182    INT8RANGE = auto()
 183    INT8MULTIRANGE = auto()
 184    NUMRANGE = auto()
 185    NUMMULTIRANGE = auto()
 186    TSRANGE = auto()
 187    TSMULTIRANGE = auto()
 188    TSTZRANGE = auto()
 189    TSTZMULTIRANGE = auto()
 190    DATERANGE = auto()
 191    DATEMULTIRANGE = auto()
 192    UUID = auto()
 193    GEOGRAPHY = auto()
 194    GEOGRAPHYPOINT = auto()
 195    NULLABLE = auto()
 196    GEOMETRY = auto()
 197    POINT = auto()
 198    RING = auto()
 199    LINESTRING = auto()
 200    LOCALTIME = auto()
 201    LOCALTIMESTAMP = auto()
 202    SYSTIMESTAMP = auto()
 203    MULTILINESTRING = auto()
 204    POLYGON = auto()
 205    MULTIPOLYGON = auto()
 206    HLLSKETCH = auto()
 207    HSTORE = auto()
 208    SUPER = auto()
 209    SERIAL = auto()
 210    SMALLSERIAL = auto()
 211    BIGSERIAL = auto()
 212    XML = auto()
 213    YEAR = auto()
 214    USERDEFINED = auto()
 215    MONEY = auto()
 216    SMALLMONEY = auto()
 217    ROWVERSION = auto()
 218    IMAGE = auto()
 219    VARIANT = auto()
 220    OBJECT = auto()
 221    INET = auto()
 222    IPADDRESS = auto()
 223    IPPREFIX = auto()
 224    IPV4 = auto()
 225    IPV6 = auto()
 226    ENUM = auto()
 227    ENUM8 = auto()
 228    ENUM16 = auto()
 229    FIXEDSTRING = auto()
 230    LOWCARDINALITY = auto()
 231    NESTED = auto()
 232    AGGREGATEFUNCTION = auto()
 233    SIMPLEAGGREGATEFUNCTION = auto()
 234    TDIGEST = auto()
 235    UNKNOWN = auto()
 236    VECTOR = auto()
 237    DYNAMIC = auto()
 238    VOID = auto()
 239
 240    # keywords
 241    ALIAS = auto()
 242    ALTER = auto()
 243    ALL = auto()
 244    ANTI = auto()
 245    ANY = auto()
 246    APPLY = auto()
 247    ARRAY = auto()
 248    ASC = auto()
 249    ASOF = auto()
 250    ATTACH = auto()
 251    AUTO_INCREMENT = auto()
 252    BEGIN = auto()
 253    BETWEEN = auto()
 254    BULK_COLLECT_INTO = auto()
 255    CACHE = auto()
 256    CASE = auto()
 257    CHARACTER_SET = auto()
 258    CLUSTER_BY = auto()
 259    COLLATE = auto()
 260    COMMAND = auto()
 261    COMMENT = auto()
 262    COMMIT = auto()
 263    CONNECT_BY = auto()
 264    CONSTRAINT = auto()
 265    COPY = auto()
 266    CREATE = auto()
 267    CROSS = auto()
 268    CUBE = auto()
 269    CURRENT_DATE = auto()
 270    CURRENT_DATETIME = auto()
 271    CURRENT_SCHEMA = auto()
 272    CURRENT_TIME = auto()
 273    CURRENT_TIMESTAMP = auto()
 274    CURRENT_USER = auto()
 275    CURRENT_ROLE = auto()
 276    CURRENT_CATALOG = auto()
 277    DECLARE = auto()
 278    DEFAULT = auto()
 279    DELETE = auto()
 280    DESC = auto()
 281    DESCRIBE = auto()
 282    DETACH = auto()
 283    DICTIONARY = auto()
 284    DISTINCT = auto()
 285    DISTRIBUTE_BY = auto()
 286    DIV = auto()
 287    DROP = auto()
 288    ELSE = auto()
 289    END = auto()
 290    ESCAPE = auto()
 291    EXCEPT = auto()
 292    EXECUTE = auto()
 293    EXISTS = auto()
 294    FALSE = auto()
 295    FETCH = auto()
 296    FILE = auto()
 297    FILE_FORMAT = auto()
 298    FILTER = auto()
 299    FINAL = auto()
 300    FIRST = auto()
 301    FOR = auto()
 302    FORCE = auto()
 303    FOREIGN_KEY = auto()
 304    FORMAT = auto()
 305    FROM = auto()
 306    FULL = auto()
 307    FUNCTION = auto()
 308    GET = auto()
 309    GLOB = auto()
 310    GLOBAL = auto()
 311    GRANT = auto()
 312    GROUP_BY = auto()
 313    GROUPING_SETS = auto()
 314    HAVING = auto()
 315    HINT = auto()
 316    IGNORE = auto()
 317    ILIKE = auto()
 318    IN = auto()
 319    INDEX = auto()
 320    INDEXED_BY = auto()
 321    INNER = auto()
 322    INSERT = auto()
 323    INSTALL = auto()
 324    INTERSECT = auto()
 325    INTERVAL = auto()
 326    INTO = auto()
 327    INTRODUCER = auto()
 328    IRLIKE = auto()
 329    IS = auto()
 330    ISNULL = auto()
 331    JOIN = auto()
 332    JOIN_MARKER = auto()
 333    KEEP = auto()
 334    KEY = auto()
 335    KILL = auto()
 336    LANGUAGE = auto()
 337    LATERAL = auto()
 338    LEFT = auto()
 339    LIKE = auto()
 340    LIMIT = auto()
 341    LIST = auto()
 342    LOAD = auto()
 343    LOCK = auto()
 344    MAP = auto()
 345    MATCH = auto()
 346    MATCH_CONDITION = auto()
 347    MATCH_RECOGNIZE = auto()
 348    MEMBER_OF = auto()
 349    MERGE = auto()
 350    MOD = auto()
 351    MODEL = auto()
 352    NATURAL = auto()
 353    NEXT = auto()
 354    NOTHING = auto()
 355    NOTNULL = auto()
 356    NULL = auto()
 357    OBJECT_IDENTIFIER = auto()
 358    OFFSET = auto()
 359    ON = auto()
 360    ONLY = auto()
 361    OPERATOR = auto()
 362    ORDER_BY = auto()
 363    ORDER_SIBLINGS_BY = auto()
 364    ORDERED = auto()
 365    ORDINALITY = auto()
 366    OUT = auto()
 367    INOUT = auto()
 368    OUTER = auto()
 369    OVER = auto()
 370    OVERLAPS = auto()
 371    OVERWRITE = auto()
 372    PARTITION = auto()
 373    PARTITION_BY = auto()
 374    PERCENT = auto()
 375    PIVOT = auto()
 376    PLACEHOLDER = auto()
 377    POSITIONAL = auto()
 378    PRAGMA = auto()
 379    PREWHERE = auto()
 380    PRIMARY_KEY = auto()
 381    PROCEDURE = auto()
 382    PROPERTIES = auto()
 383    PSEUDO_TYPE = auto()
 384    PUT = auto()
 385    QUALIFY = auto()
 386    QUOTE = auto()
 387    QDCOLON = auto()
 388    RANGE = auto()
 389    RECURSIVE = auto()
 390    REFRESH = auto()
 391    RENAME = auto()
 392    REPLACE = auto()
 393    RETURNING = auto()
 394    REVOKE = auto()
 395    REFERENCES = auto()
 396    RIGHT = auto()
 397    RLIKE = auto()
 398    ROLLBACK = auto()
 399    ROLLUP = auto()
 400    ROW = auto()
 401    ROWS = auto()
 402    SELECT = auto()
 403    SEMI = auto()
 404    SEPARATOR = auto()
 405    SEQUENCE = auto()
 406    SERDE_PROPERTIES = auto()
 407    SET = auto()
 408    SETTINGS = auto()
 409    SHOW = auto()
 410    SIMILAR_TO = auto()
 411    SOME = auto()
 412    SORT_BY = auto()
 413    SOUNDS_LIKE = auto()
 414    START_WITH = auto()
 415    STORAGE_INTEGRATION = auto()
 416    STRAIGHT_JOIN = auto()
 417    STRUCT = auto()
 418    SUMMARIZE = auto()
 419    TABLE_SAMPLE = auto()
 420    TAG = auto()
 421    TEMPORARY = auto()
 422    TOP = auto()
 423    THEN = auto()
 424    TRUE = auto()
 425    TRUNCATE = auto()
 426    TRIGGER = auto()
 427    UNCACHE = auto()
 428    UNION = auto()
 429    UNNEST = auto()
 430    UNPIVOT = auto()
 431    UPDATE = auto()
 432    USE = auto()
 433    USING = auto()
 434    VALUES = auto()
 435    VARIADIC = auto()
 436    VIEW = auto()
 437    SEMANTIC_VIEW = auto()
 438    VOLATILE = auto()
 439    WHEN = auto()
 440    WHERE = auto()
 441    WINDOW = auto()
 442    WITH = auto()
 443    UNIQUE = auto()
 444    UTC_DATE = auto()
 445    UTC_TIME = auto()
 446    UTC_TIMESTAMP = auto()
 447    VERSION_SNAPSHOT = auto()
 448    TIMESTAMP_SNAPSHOT = auto()
 449    OPTION = auto()
 450    SINK = auto()
 451    SOURCE = auto()
 452    ANALYZE = auto()
 453    NAMESPACE = auto()
 454    EXPORT = auto()
 455
 456    # sentinel
 457    HIVE_TOKEN_STREAM = auto()
 458
 459    def __str__(self) -> str:
 460        return f"TokenType.{self.name}"
 461
 462
 463class Token:
 464    # mypyc doesn't expose slots
 465    _attrs: t.ClassVar[t.Tuple[str, ...]] = (
 466        "token_type",
 467        "text",
 468        "line",
 469        "col",
 470        "start",
 471        "end",
 472        "comments",
 473    )
 474    __slots__ = _attrs
 475
 476    @classmethod
 477    def number(cls, number: int) -> Token:
 478        """Returns a NUMBER token with `number` as its text."""
 479        return cls(TokenType.NUMBER, str(number))
 480
 481    @classmethod
 482    def string(cls, string: str) -> Token:
 483        """Returns a STRING token with `string` as its text."""
 484        return cls(TokenType.STRING, string)
 485
 486    @classmethod
 487    def identifier(cls, identifier: str) -> Token:
 488        """Returns an IDENTIFIER token with `identifier` as its text."""
 489        return cls(TokenType.IDENTIFIER, identifier)
 490
 491    @classmethod
 492    def var(cls, var: str) -> Token:
 493        """Returns an VAR token with `var` as its text."""
 494        return cls(TokenType.VAR, var)
 495
 496    def __init__(
 497        self,
 498        token_type: TokenType,
 499        text: str,
 500        line: int = 1,
 501        col: int = 1,
 502        start: int = 0,
 503        end: int = 0,
 504        comments: t.Optional[t.List[str]] = None,
 505    ) -> None:
 506        self.token_type = token_type
 507        self.text = text
 508        self.line = line
 509        self.col = col
 510        self.start = start
 511        self.end = end
 512        self.comments = [] if comments is None else comments
 513
 514    def __repr__(self) -> str:
 515        attributes = ", ".join(
 516            f"{k}: TokenType.{self.token_type.name}"
 517            if k == "token_type"
 518            else f"{k}: {getattr(self, k)}"
 519            for k in self._attrs
 520        )
 521        return f"<Token {attributes}>"
 522
 523
 524class TokenizerCore:
 525    __slots__ = (
 526        "sql",
 527        "size",
 528        "tokens",
 529        "_start",
 530        "_current",
 531        "_line",
 532        "_col",
 533        "_comments",
 534        "_char",
 535        "_end",
 536        "_peek",
 537        "_prev_token_line",
 538        "single_tokens",
 539        "keywords",
 540        "quotes",
 541        "format_strings",
 542        "identifiers",
 543        "comments",
 544        "string_escapes",
 545        "byte_string_escapes",
 546        "identifier_escapes",
 547        "escape_follow_chars",
 548        "commands",
 549        "command_prefix_tokens",
 550        "nested_comments",
 551        "hint_start",
 552        "tokens_preceding_hint",
 553        "bit_strings",
 554        "hex_strings",
 555        "numeric_literals",
 556        "var_single_tokens",
 557        "string_escapes_allowed_in_raw_strings",
 558        "heredoc_tag_is_identifier",
 559        "heredoc_string_alternative",
 560        "keyword_trie",
 561        "numbers_can_be_underscore_separated",
 562        "identifiers_can_start_with_digit",
 563        "unescaped_sequences",
 564    )
 565
 566    def __init__(
 567        self,
 568        single_tokens: t.Dict[str, TokenType],
 569        keywords: t.Dict[str, TokenType],
 570        quotes: t.Dict[str, str],
 571        format_strings: t.Dict[str, t.Tuple[str, TokenType]],
 572        identifiers: t.Dict[str, str],
 573        comments: t.Dict[str, t.Optional[str]],
 574        string_escapes: t.Set[str],
 575        byte_string_escapes: t.Set[str],
 576        identifier_escapes: t.Set[str],
 577        escape_follow_chars: t.Set[str],
 578        commands: t.Set[TokenType],
 579        command_prefix_tokens: t.Set[TokenType],
 580        nested_comments: bool,
 581        hint_start: str,
 582        tokens_preceding_hint: t.Set[TokenType],
 583        bit_strings: t.List[t.Union[str, t.Tuple[str, str]]],
 584        hex_strings: t.List[t.Union[str, t.Tuple[str, str]]],
 585        numeric_literals: t.Dict[str, str],
 586        var_single_tokens: t.Set[str],
 587        string_escapes_allowed_in_raw_strings: bool,
 588        heredoc_tag_is_identifier: bool,
 589        heredoc_string_alternative: TokenType,
 590        keyword_trie: t.Dict,
 591        numbers_can_be_underscore_separated: bool,
 592        identifiers_can_start_with_digit: bool,
 593        unescaped_sequences: t.Dict[str, str],
 594    ) -> None:
 595        self.single_tokens = single_tokens
 596        self.keywords = keywords
 597        self.quotes = quotes
 598        self.format_strings = format_strings
 599        self.identifiers = identifiers
 600        self.comments = comments
 601        self.string_escapes = string_escapes
 602        self.byte_string_escapes = byte_string_escapes
 603        self.identifier_escapes = identifier_escapes
 604        self.escape_follow_chars = escape_follow_chars
 605        self.commands = commands
 606        self.command_prefix_tokens = command_prefix_tokens
 607        self.nested_comments = nested_comments
 608        self.hint_start = hint_start
 609        self.tokens_preceding_hint = tokens_preceding_hint
 610        self.bit_strings = bit_strings
 611        self.hex_strings = hex_strings
 612        self.numeric_literals = numeric_literals
 613        self.var_single_tokens = var_single_tokens
 614        self.string_escapes_allowed_in_raw_strings = string_escapes_allowed_in_raw_strings
 615        self.heredoc_tag_is_identifier = heredoc_tag_is_identifier
 616        self.heredoc_string_alternative = heredoc_string_alternative
 617        self.keyword_trie = keyword_trie
 618        self.numbers_can_be_underscore_separated = numbers_can_be_underscore_separated
 619        self.identifiers_can_start_with_digit = identifiers_can_start_with_digit
 620        self.unescaped_sequences = unescaped_sequences
 621        self.reset()
 622
 623    def reset(self) -> None:
 624        self.sql = ""
 625        self.size = 0
 626        self.tokens: t.List[Token] = []
 627        self._start = 0
 628        self._current = 0
 629        self._line = 1
 630        self._col = 0
 631        self._comments: t.List[str] = []
 632        self._char = ""
 633        self._end = False
 634        self._peek = ""
 635        self._prev_token_line = -1
 636
 637    def tokenize(self, sql: str) -> t.List[Token]:
 638        """Returns a list of tokens corresponding to the SQL string `sql`."""
 639        self.reset()
 640        self.sql = sql
 641        self.size = len(sql)
 642
 643        try:
 644            self._scan()
 645        except Exception as e:
 646            start = max(self._current - 50, 0)
 647            end = min(self._current + 50, self.size - 1)
 648            context = self.sql[start:end]
 649            raise TokenError(f"Error tokenizing '{context}'") from e
 650
 651        return self.tokens
 652
 653    def _scan(self, check_semicolon: bool = False) -> None:
 654        identifiers = self.identifiers
 655        space_chars = _SPACE_CHARS
 656        digit_chars = _DIGIT_CHARS
 657
 658        while self.size and not self._end:
 659            current = self._current
 660
 661            # Skip spaces here rather than iteratively calling advance() for performance reasons
 662            while current < self.size:
 663                char = self.sql[current]
 664
 665                if char == " " or char == "\t":
 666                    current += 1
 667                else:
 668                    break
 669
 670            offset = current - self._current if current > self._current else 1
 671
 672            self._start = current
 673            self._advance(offset)
 674
 675            if self._char not in space_chars:
 676                if self._char in digit_chars:
 677                    self._scan_number()
 678                elif self._char in identifiers:
 679                    self._scan_identifier(identifiers[self._char])
 680                else:
 681                    self._scan_keywords()
 682
 683            if check_semicolon and self._peek == ";":
 684                break
 685
 686        if self.tokens and self._comments:
 687            self.tokens[-1].comments.extend(self._comments)
 688
 689    def _chars(self, size: int) -> str:
 690        if size == 1:
 691            return self._char
 692
 693        start = self._current - 1
 694        end = start + size
 695
 696        return self.sql[start:end] if end <= self.size else ""
 697
 698    def _advance(self, i: int = 1, alnum: bool = False) -> None:
 699        char = self._char
 700
 701        if char == "\n" or char == "\r":
 702            # Ensures we don't count an extra line if we get a \r\n line break sequence
 703            if not (char == "\r" and self._peek == "\n"):
 704                self._col = i
 705                self._line += 1
 706        else:
 707            self._col += i
 708
 709        self._current += i
 710        sql = self.sql
 711        size = self.size
 712        self._end = self._current >= size
 713        self._char = sql[self._current - 1]
 714        self._peek = "" if self._end else sql[self._current]
 715
 716        if alnum and self._char.isalnum():
 717            # Cache to local variables instead of attributes for better performance
 718            _col = self._col
 719            _current = self._current
 720            _end = self._end
 721            _peek = self._peek
 722
 723            while _peek.isalnum():
 724                _col += 1
 725                _current += 1
 726                _end = _current >= size
 727                _peek = "" if _end else sql[_current]
 728
 729            self._col = _col
 730            self._current = _current
 731            self._end = _end
 732            self._peek = _peek
 733            self._char = sql[_current - 1]
 734
 735    @property
 736    def _text(self) -> str:
 737        return self.sql[self._start : self._current]
 738
 739    def _add(self, token_type: TokenType, text: t.Optional[str] = None) -> None:
 740        self._prev_token_line = self._line
 741
 742        if self._comments and token_type == TokenType.SEMICOLON and self.tokens:
 743            self.tokens[-1].comments.extend(self._comments)
 744            self._comments = []
 745
 746        if text is None:
 747            text = self.sql[self._start : self._current]
 748
 749        self.tokens.append(
 750            Token(
 751                token_type,
 752                text=text,
 753                line=self._line,
 754                col=self._col,
 755                start=self._start,
 756                end=self._current - 1,
 757                comments=self._comments,
 758            )
 759        )
 760        self._comments = []
 761
 762        # If we have either a semicolon or a begin token before the command's token, we'll parse
 763        # whatever follows the command's token as a string
 764        if (
 765            token_type in self.commands
 766            and self._peek != ";"
 767            and (len(self.tokens) == 1 or self.tokens[-2].token_type in self.command_prefix_tokens)
 768        ):
 769            start = self._current
 770            tokens = len(self.tokens)
 771            self._scan(check_semicolon=True)
 772            self.tokens = self.tokens[:tokens]
 773            text = self.sql[start : self._current].strip()
 774            if text:
 775                self._add(TokenType.STRING, text)
 776
 777    def _scan_keywords(self) -> None:
 778        sql = self.sql
 779        sql_size = self.size
 780        single_tokens = self.single_tokens
 781        char_upper = _CHAR_UPPER
 782        space_chars = _SPACE_CHARS
 783        size = 0
 784        word = None
 785        chars = self._char
 786        char = chars
 787        prev_space = False
 788        skip = False
 789        trie = self.keyword_trie
 790        single_token = char in single_tokens
 791
 792        while chars:
 793            if not skip:
 794                sub = trie.get(char_upper.get(char, char))
 795                if sub is None:
 796                    break
 797                trie = sub
 798                if 0 in trie:
 799                    word = chars
 800
 801            end = self._current + size
 802            size += 1
 803
 804            if end < sql_size:
 805                char = sql[end]
 806                single_token = single_token or char in single_tokens
 807                is_space = char in space_chars
 808
 809                if not is_space or not prev_space:
 810                    if is_space:
 811                        char = " "
 812                    chars += char
 813                    prev_space = is_space
 814                    skip = False
 815                else:
 816                    skip = True
 817            else:
 818                char = ""
 819                break
 820
 821        if word:
 822            if self._scan_string(word):
 823                return
 824            if self._scan_comment(word):
 825                return
 826            if prev_space or single_token or not char:
 827                self._advance(size - 1)
 828                word = word.upper()
 829                self._add(self.keywords[word], text=word)
 830                return
 831
 832        if self._char in single_tokens:
 833            self._add(single_tokens[self._char], text=self._char)
 834            return
 835
 836        self._scan_var()
 837
 838    def _scan_comment(self, comment_start: str) -> bool:
 839        if comment_start not in self.comments:
 840            return False
 841
 842        comment_start_line = self._line
 843        comment_start_size = len(comment_start)
 844        comment_end = self.comments[comment_start]
 845
 846        if comment_end:
 847            # Skip the comment's start delimiter
 848            self._advance(comment_start_size)
 849
 850            comment_count = 1
 851            comment_end_size = len(comment_end)
 852            nested_comments = self.nested_comments
 853
 854            while not self._end:
 855                if self._chars(comment_end_size) == comment_end:
 856                    comment_count -= 1
 857                    if not comment_count:
 858                        break
 859
 860                self._advance(alnum=True)
 861
 862                # Nested comments are allowed by some dialects, e.g. databricks, duckdb, postgres
 863                if (
 864                    nested_comments
 865                    and not self._end
 866                    and self._chars(comment_end_size) == comment_start
 867                ):
 868                    self._advance(comment_start_size)
 869                    comment_count += 1
 870
 871            self._comments.append(self._text[comment_start_size : -comment_end_size + 1])
 872            self._advance(comment_end_size - 1)
 873        else:
 874            _peek = self._peek
 875            while not self._end and _peek != "\n" and _peek != "\r":
 876                self._advance(alnum=True)
 877                _peek = self._peek
 878            self._comments.append(self._text[comment_start_size:])
 879
 880        if (
 881            comment_start == self.hint_start
 882            and self.tokens
 883            and self.tokens[-1].token_type in self.tokens_preceding_hint
 884        ):
 885            self._add(TokenType.HINT)
 886
 887        # Leading comment is attached to the succeeding token, whilst trailing comment to the preceding.
 888        # Multiple consecutive comments are preserved by appending them to the current comments list.
 889        if comment_start_line == self._prev_token_line:
 890            self.tokens[-1].comments.extend(self._comments)
 891            self._comments = []
 892            self._prev_token_line = self._line
 893
 894        return True
 895
 896    def _scan_number(self) -> None:
 897        if self._char == "0":
 898            peek = _CHAR_UPPER.get(self._peek, self._peek)
 899            if peek == "B":
 900                return self._scan_bits() if self.bit_strings else self._add(TokenType.NUMBER)
 901            elif peek == "X":
 902                return self._scan_hex() if self.hex_strings else self._add(TokenType.NUMBER)
 903
 904        decimal = False
 905        scientific = 0
 906        numbers_can_be_underscore_separated = self.numbers_can_be_underscore_separated
 907        single_tokens = self.single_tokens
 908        keywords = self.keywords
 909        numeric_literals = self.numeric_literals
 910        identifiers_can_start_with_digit = self.identifiers_can_start_with_digit
 911
 912        while True:
 913            if self._peek in _DIGIT_CHARS:
 914                self._advance()
 915            elif self._peek == "." and not decimal:
 916                if self.tokens and self.tokens[-1].token_type == TokenType.PARAMETER:
 917                    return self._add(TokenType.NUMBER)
 918                decimal = True
 919                self._advance()
 920            elif self._peek in ("-", "+") and scientific == 1:
 921                # Only consume +/- if followed by a digit
 922                if self._current + 1 < self.size and self.sql[self._current + 1] in _DIGIT_CHARS:
 923                    scientific += 1
 924                    self._advance()
 925                else:
 926                    return self._add(TokenType.NUMBER)
 927            elif _CHAR_UPPER.get(self._peek, self._peek) == "E" and not scientific:
 928                scientific += 1
 929                self._advance()
 930            elif self._peek == "_" and numbers_can_be_underscore_separated:
 931                self._advance()
 932            elif self._peek.isidentifier():
 933                number_text = self._text
 934                literal = ""
 935
 936                while (
 937                    self._peek
 938                    and self._peek not in _SPACE_CHARS
 939                    and self._peek not in single_tokens
 940                ):
 941                    literal += self._peek
 942                    self._advance()
 943
 944                token_type = keywords.get(numeric_literals.get(literal.upper(), ""))
 945
 946                if token_type:
 947                    self._add(TokenType.NUMBER, number_text)
 948                    self._add(TokenType.DCOLON, "::")
 949                    return self._add(token_type, literal)
 950                elif identifiers_can_start_with_digit:
 951                    return self._add(TokenType.VAR)
 952
 953                self._advance(-len(literal))
 954                return self._add(TokenType.NUMBER, number_text)
 955            else:
 956                return self._add(TokenType.NUMBER)
 957
 958    def _scan_bits(self) -> None:
 959        self._advance()
 960        value = self._extract_value()
 961        try:
 962            # If `value` can't be converted to a binary, fallback to tokenizing it as an identifier
 963            int(value, 2)
 964            self._add(TokenType.BIT_STRING, value[2:])  # Drop the 0b
 965        except ValueError:
 966            self._add(TokenType.IDENTIFIER)
 967
 968    def _scan_hex(self) -> None:
 969        self._advance()
 970        value = self._extract_value()
 971        try:
 972            # If `value` can't be converted to a hex, fallback to tokenizing it as an identifier
 973            int(value, 16)
 974            self._add(TokenType.HEX_STRING, value[2:])  # Drop the 0x
 975        except ValueError:
 976            self._add(TokenType.IDENTIFIER)
 977
 978    def _extract_value(self) -> str:
 979        single_tokens = self.single_tokens
 980
 981        while True:
 982            char = self._peek.strip()
 983            if char and char not in single_tokens:
 984                self._advance(alnum=True)
 985            else:
 986                break
 987
 988        return self._text
 989
 990    def _scan_string(self, start: str) -> bool:
 991        base = None
 992        token_type = TokenType.STRING
 993
 994        if start in self.quotes:
 995            end = self.quotes[start]
 996        elif start in self.format_strings:
 997            end, token_type = self.format_strings[start]
 998
 999            if token_type == TokenType.HEX_STRING:
1000                base = 16
1001            elif token_type == TokenType.BIT_STRING:
1002                base = 2
1003            elif token_type == TokenType.HEREDOC_STRING:
1004                self._advance()
1005
1006                if self._char == end:
1007                    tag = ""
1008                else:
1009                    tag = self._extract_string(
1010                        end,
1011                        raw_string=True,
1012                        raise_unmatched=not self.heredoc_tag_is_identifier,
1013                    )
1014
1015                if (
1016                    tag
1017                    and self.heredoc_tag_is_identifier
1018                    and (self._end or tag.isdigit() or any(c.isspace() for c in tag))
1019                ):
1020                    if not self._end:
1021                        self._advance(-1)
1022
1023                    self._advance(-len(tag))
1024                    self._add(self.heredoc_string_alternative)
1025                    return True
1026
1027                end = f"{start}{tag}{end}"
1028        else:
1029            return False
1030
1031        self._advance(len(start))
1032        text = self._extract_string(
1033            end,
1034            escapes=(
1035                self.byte_string_escapes
1036                if token_type == TokenType.BYTE_STRING
1037                else self.string_escapes
1038            ),
1039            raw_string=token_type == TokenType.RAW_STRING,
1040        )
1041
1042        if base and text:
1043            try:
1044                int(text, base)
1045            except Exception:
1046                raise TokenError(
1047                    f"Numeric string contains invalid characters from {self._line}:{self._start}"
1048                )
1049
1050        self._add(token_type, text)
1051        return True
1052
1053    def _scan_identifier(self, identifier_end: str) -> None:
1054        self._advance()
1055        text = self._extract_string(
1056            identifier_end, escapes=self.identifier_escapes | {identifier_end}
1057        )
1058        self._add(TokenType.IDENTIFIER, text)
1059
1060    def _scan_var(self) -> None:
1061        var_single_tokens = self.var_single_tokens
1062        single_tokens = self.single_tokens
1063
1064        while True:
1065            peek = self._peek
1066            if not peek or peek in _SPACE_CHARS:
1067                break
1068            if peek not in var_single_tokens and peek in single_tokens:
1069                break
1070            self._advance(alnum=True)
1071
1072        self._add(
1073            TokenType.VAR
1074            if self.tokens and self.tokens[-1].token_type == TokenType.PARAMETER
1075            else self.keywords.get(self.sql[self._start : self._current].upper(), TokenType.VAR)
1076        )
1077
1078    def _extract_string(
1079        self,
1080        delimiter: str,
1081        escapes: t.Optional[t.Set[str]] = None,
1082        raw_string: bool = False,
1083        raise_unmatched: bool = True,
1084    ) -> str:
1085        text = ""
1086        delim_size = len(delimiter)
1087        escapes = self.string_escapes if escapes is None else escapes
1088        unescaped_sequences = self.unescaped_sequences
1089        escape_follow_chars = self.escape_follow_chars
1090        string_escapes_allowed_in_raw_strings = self.string_escapes_allowed_in_raw_strings
1091        quotes = self.quotes
1092
1093        while True:
1094            if not raw_string and unescaped_sequences and self._peek and self._char in escapes:
1095                unescaped_sequence = unescaped_sequences.get(self._char + self._peek)
1096                if unescaped_sequence:
1097                    self._advance(2)
1098                    text += unescaped_sequence
1099                    continue
1100
1101            is_valid_custom_escape = (
1102                escape_follow_chars and self._char == "\\" and self._peek not in escape_follow_chars
1103            )
1104
1105            if (
1106                (string_escapes_allowed_in_raw_strings or not raw_string)
1107                and self._char in escapes
1108                and (self._peek == delimiter or self._peek in escapes or is_valid_custom_escape)
1109                and (self._char not in quotes or self._char == self._peek)
1110            ):
1111                if self._peek == delimiter:
1112                    text += self._peek
1113                elif is_valid_custom_escape and self._char != self._peek:
1114                    text += self._peek
1115                else:
1116                    text += self._char + self._peek
1117
1118                if self._current + 1 < self.size:
1119                    self._advance(2)
1120                else:
1121                    raise TokenError(f"Missing {delimiter} from {self._line}:{self._current}")
1122            else:
1123                if self._chars(delim_size) == delimiter:
1124                    if delim_size > 1:
1125                        self._advance(delim_size - 1)
1126                    break
1127
1128                if self._end:
1129                    if not raise_unmatched:
1130                        return text + self._char
1131
1132                    raise TokenError(f"Missing {delimiter} from {self._line}:{self._start}")
1133
1134                current = self._current - 1
1135                self._advance(alnum=True)
1136                text += self.sql[current : self._current - 1]
1137
1138        return text