comments.py

"""Extract and strip SQL comments using the sqlglot tokenizer.

sqlglot's tokenizer skips comments during tokenization, which means
comments live in the *gaps* between consecutive token positions.  This
module exploits that property: it tokenizes the SQL, then scans each gap
for comment delimiters (``--``, ``/* */``, ``#``).

Two public entry points exist:

* :func:`extract_comments` — returns the raw comment texts (delimiters
  included) for inspection or logging.
* :func:`strip_comments` — returns the SQL with all comments removed and
  whitespace normalised, used by :class:`Parser` for the ``without_comments``
  property.

A third, internal variant :func:`strip_comments_for_parsing` is consumed
by :mod:`_ast` before handing SQL to ``sqlglot.parse()``; it always uses
the MySQL tokenizer so that ``#``-style comments are reliably stripped.
"""

import re
from typing import Any

from sqlglot.dialects.mysql import MySQL
from sqlglot.errors import TokenError
from sqlglot.tokens import Tokenizer


class _MySQLCommentTokenizer(MySQL.Tokenizer):
    """MySQL tokenizer that treats any ``--`` as a line comment.

    sqlglot >=30.22 follows MySQL semantics where ``--`` only starts a
    comment when followed by whitespace, so ``--foo`` would be tokenized
    as two dashes and an identifier.  We use this tokenizer purely to
    strip comments, and ``--foo`` is meant as a comment in practice.
    """

    DASH_COMMENT_REQUIRES_BOUNDARY = False


def _choose_tokenizer(sql: str) -> Tokenizer:
    """Select the appropriate sqlglot tokenizer for *sql*.

    The default sqlglot tokenizer does **not** treat ``#`` as a comment
    delimiter, but MySQL does.  When ``#`` appears in the SQL and is used
    as a comment (not as a variable/template prefix), we switch to the
    MySQL tokenizer so that ``#``-style comments are properly skipped.

    :param sql: Raw SQL string to inspect.
    :type sql: str
    :returns: An instantiated tokenizer (MySQL or default).
    :rtype: sqlglot.tokens.Tokenizer
    """
    if "#" in sql and not _has_hash_variables(sql):
        # Pass dialect="mysql" so the cached TokenizerCore (sqlglot >=30.7.0
        # caches per-class) is built with MySQL identifier semantics — without
        # it the cache is primed from the default dialect and later mysql
        # parses misclassify e.g. ``0020_big_table`` as NUMBER + VAR.
        return _MySQLCommentTokenizer(dialect="mysql")
    return Tokenizer()


#: Matches T-SQL temporary table references (#tmp, ##tmp) following keywords
#: like FROM, JOIN, INTO, TABLE, EXISTS, UPDATE, or after a comma.
_TEMP_TABLE_RE = re.compile(
    r"(?i)(?:\b(?:FROM|JOIN|INTO|TABLE|EXISTS|UPDATE)\s+|,\s*)#+\w"
)


def _has_hash_variables(sql: str) -> bool:
    """Determine whether ``#`` characters in *sql* are variable references.

    MSSQL uses ``#table`` for temporary tables and some template engines
    use ``#VAR#`` placeholders.  This function distinguishes those from
    MySQL-style ``# comment`` lines so that :func:`_choose_tokenizer`
    picks the right dialect.

    Heuristics (checked via regex):

    * ``#WORD#`` — bracketed template variable.
    * ``= #WORD`` or ``(#WORD`` — assignment / parameter context.
    * ``FROM #table``, ``JOIN #table``, etc. — T-SQL temporary tables.

    :param sql: Raw SQL string.
    :type sql: str
    :returns: ``True`` if at least one ``#`` looks like a variable prefix.
    :rtype: bool
    """
    # #WORD# template variable (e.g. #VAR#)
    if re.search(r"#\w+#", sql):
        return True
    # = #WORD or (#WORD with optional whitespace before #
    if re.search(r"[=(]\s*#\w", sql):
        return True
    # T-SQL temporary table references (e.g. #tmp, ##tmp, #HistRaw)
    if _TEMP_TABLE_RE.search(sql):
        return True
    return False


def extract_comments(sql: str) -> list[str]:
    """Return all comments found in *sql*, with delimiters preserved.

    Tokenizes the SQL, then scans every gap between consecutive token
    positions for comment delimiters.  Returned strings include the
    opening delimiter (``--``, ``/*``, ``#``) and, for block comments,
    the closing ``*/``.

    Called by :attr:`Parser.comments`.

    :param sql: Raw SQL string.
    :type sql: str
    :returns: List of comment strings in source order.
    :rtype: List[str]
    """
    if not sql:
        return []
    try:
        tokens = list(_choose_tokenizer(sql).tokenize(sql))
    except TokenError:
        return []
    comments: list[str] = []
    prev_end = -1
    for tok in tokens:
        _scan_gap(sql, prev_end + 1, tok.start, comments)
        prev_end = tok.end
    _scan_gap(sql, prev_end + 1, len(sql), comments)
    return comments


#: Matches all three SQL comment styles in a single pass:
#: ``/* ... */`` (block, possibly unterminated), ``-- ...``, and ``# ...``.
_COMMENT_RE = re.compile(r"/\*.*?\*/|/\*.*$|--[^\n]*\n?|#[^\n]*\n?", re.DOTALL)


def _scan_gap(sql: str, start: int, end: int, out: list[str]) -> None:
    """Scan a slice of *sql* for comment delimiters and append matches.

    :param sql: The full SQL string (not just the gap).
    :param start: Start index of the gap to scan.
    :param end: End index (exclusive) of the gap.
    :param out: Mutable list to which discovered comment strings are appended.
    """
    out.extend(_COMMENT_RE.findall(sql[start:end]))


def _reconstruct_from_tokens(sql: str, tokens: list[Any]) -> str:
    """Rebuild SQL from token spans, collapsing gaps to single spaces.

    Concatenates the text of each token using its ``start`` / ``end``
    positions.  Any gap between consecutive tokens (where comments or
    extra whitespace lived) is replaced by a single space.

    :param sql: The original SQL string.
    :type sql: str
    :param tokens: Sqlglot token objects with ``start`` and ``end`` attrs.
    :type tokens: list[Any]
    :returns: Reconstructed SQL with comments removed.
    :rtype: str
    """
    if not tokens:
        return ""
    parts = [sql[tokens[0].start : tokens[0].end + 1]]
    for i in range(1, len(tokens)):
        if tokens[i].start > tokens[i - 1].end + 1:
            parts.append(" ")
        parts.append(sql[tokens[i].start : tokens[i].end + 1])
    return "".join(parts).strip()


def strip_comments_for_parsing(sql: str) -> str:
    """Strip **all** comments — including ``#`` lines — for sqlglot parsing.

    Unlike :func:`strip_comments`, this always uses the MySQL tokenizer
    (which treats ``#`` as a comment delimiter) so that hash-style
    comments are removed before ``sqlglot.parse()`` sees the SQL.  The
    only exceptions are ``CREATE FUNCTION`` bodies (which may contain
    ``#`` in procedural code) and MSSQL ``#temp`` table references.

    Called exclusively by :meth:`ASTParser._parse` in ``_ast.py``.

    :param sql: Raw SQL string.
    :type sql: str
    :returns: SQL with all comments removed and whitespace collapsed.
    :rtype: str
    """
    if not sql:
        return sql or ""
    # Skip MySQL tokenizer when # is used as variable (not comment)
    upper = sql.strip().upper()
    if upper.startswith("CREATE FUNCTION") or _has_hash_variables(sql):
        tokenizer = Tokenizer()
    else:
        # See _choose_tokenizer — the explicit dialect prevents the
        # sqlglot >=30.7.0 TokenizerCore cache from being primed with
        # default-dialect semantics.
        tokenizer = _MySQLCommentTokenizer(dialect="mysql")
    try:
        tokens = list(tokenizer.tokenize(sql))
    except TokenError:
        return sql.strip()
    return _reconstruct_from_tokens(sql, tokens)


def strip_comments(sql: str) -> str:
    """Remove comments and normalise whitespace, preserving ``#VAR`` references.

    Reconstructs the SQL from its token spans, inserting a single space
    wherever a gap (comment or extra whitespace) existed between two
    tokens.  Uses :func:`_choose_tokenizer` so that ``#VAR`` template
    variables in MSSQL queries are kept intact.

    Called by :attr:`Parser.without_comments` and
    :attr:`Generalizator.without_comments`.

    :param sql: Raw SQL string.
    :type sql: str
    :returns: SQL with comments removed and whitespace normalised.
    :rtype: str
    """
    if not sql:
        return sql or ""
    try:
        tokens = list(_choose_tokenizer(sql).tokenize(sql))
    except TokenError:
        return sql.strip()
    return _reconstruct_from_tokens(sql, tokens)