statement_formatter.py

from datetime import date, datetime, timezone
from decimal import Decimal
from typing import Dict, List, Optional, Sequence, Union

from sqlparse import parse as parse_sql  # type: ignore
from sqlparse import tokens as _T
from sqlparse.engine.statement_splitter import (
    StatementSplitter as _StatementSplitter,
)
from sqlparse.sql import (  # type: ignore
    Comment,
    Comparison,
    Statement,
    Token,
    TokenList,
)
from sqlparse.tokens import Comparison as ComparisonType  # type: ignore
from sqlparse.tokens import Newline  # type: ignore
from sqlparse.tokens import Whitespace  # type: ignore
from sqlparse.tokens import Token as TokenType  # type: ignore

from firebolt.common._types import ParameterType, SetParameter
from firebolt.utils.exception import (
    DataError,
    InterfaceError,
    NotSupportedError,
)


def _patched_change_splitlevel(self, ttype, value):  # type: ignore[no-untyped-def]
    """Patched version of StatementSplitter._change_splitlevel.

    Fixes CASE...END level tracking outside of CREATE blocks.
    See: https://github.com/andialbrecht/sqlparse/pull/839
    """
    if ttype is _T.Punctuation and value == "(":
        return 1
    elif ttype is _T.Punctuation and value == ")":
        return -1
    elif ttype not in _T.Keyword:
        return 0

    unified = value.upper()

    if ttype is _T.Keyword.DDL and unified.startswith("CREATE"):
        self._is_create = True
        return 0

    if unified == "DECLARE" and self._is_create and self._begin_depth == 0:
        self._in_declare = True
        return 1

    if unified == "BEGIN":
        self._begin_depth += 1
        self._seen_begin = True
        if self._is_create:
            return 1
        return 0

    if (
        self._seen_begin
        and (ttype is _T.Keyword or ttype is _T.Name)
        and unified
        in (
            "TRANSACTION",
            "WORK",
            "TRAN",
            "DISTRIBUTED",
            "DEFERRED",
            "IMMEDIATE",
            "EXCLUSIVE",
        )
    ):
        self._begin_depth = max(0, self._begin_depth - 1)
        self._seen_begin = False
        return 0

    if unified == "END":
        if not self._in_case:
            self._begin_depth = max(0, self._begin_depth - 1)
        else:
            self._in_case = False
        return -1

    if unified == "CASE":
        self._in_case = True
        return 1

    if unified in ("IF", "FOR", "WHILE") and self._is_create and self._begin_depth > 0:
        return 1

    if unified in ("END IF", "END FOR", "END WHILE"):
        return -1

    return 0


setattr(_StatementSplitter, "_change_splitlevel", _patched_change_splitlevel)


escape_chars_v2 = {
    "\0": "\\0",
    "'": "''",
}

escape_chars_v1 = {
    "\0": "\\0",
    "'": "''",
    "\\": "\\\\",
}


class StatementFormatter:
    def __init__(self, escape_chars: Dict[str, str]):
        self.escape_chars = escape_chars

    def format_value(self, value: ParameterType) -> str:
        """For Python value to be used in a SQL query."""
        if isinstance(value, bool):
            return "true" if value else "false"
        if isinstance(value, (int, float, Decimal)):
            return str(value)
        elif isinstance(value, str):
            return f"'{''.join(self.escape_chars.get(c, c) for c in value)}'"
        elif isinstance(value, datetime):
            if value.tzinfo is not None:
                value = value.astimezone(timezone.utc)
            return f"'{value.strftime('%Y-%m-%d %H:%M:%S')}'"
        elif isinstance(value, date):
            return f"'{value.isoformat()}'"
        elif isinstance(value, bytes):
            # Encode each byte into hex
            return "E'" + "".join(f"\\x{b:02x}" for b in value) + "'"
        if value is None:
            return "NULL"
        elif isinstance(value, Sequence):
            return f"[{', '.join(self.format_value(it) for it in value)}]"

        raise DataError(f"unsupported parameter type {type(value)}")

    def convert_parameter_for_serialization(
        self, value: ParameterType
    ) -> Union[int, float, bool, None, str, List]:
        """
        Convert parameter values for fb_numeric paramstyle to JSON-serializable
        format. This is used for server-side parameter substitution.

        Basic types (int, float, bool, None) are preserved as-is.
        All other types are converted to strings for JSON serialization.

        Args:
            value: The parameter value to convert

        Returns:
            JSON-serializable value (int, float, bool, None, or string)
        """
        if isinstance(value, (int, float, bool)) or value is None:
            return value

        if isinstance(value, Decimal):
            return str(value)
        elif isinstance(value, bytes):
            return value.decode("utf-8")
        elif isinstance(value, list):
            return [self.convert_parameter_for_serialization(item) for item in value]
        else:
            return str(value)

    def format_statement(
        self, statement: Statement, parameters: Sequence[ParameterType]
    ) -> str:
        """
        Substitute placeholders in a `sqlparse` statement with provided values.
        """
        idx = 0

        def process_token(token: Token) -> Token:
            nonlocal idx
            if token.ttype == TokenType.Name.Placeholder:
                # Replace placeholder with formatted parameter
                if idx >= len(parameters):
                    raise DataError(
                        "not enough parameters provided for substitution: given "
                        f"{len(parameters)}, found one more"
                    )
                formatted = self.format_value(parameters[idx])
                idx += 1
                return Token(TokenType.Text, formatted)
            if isinstance(token, TokenList):
                # Process all children tokens

                return TokenList([process_token(t) for t in token.tokens])
            return token

        formatted_sql = self._clean_sql_string(process_token(statement))

        if idx < len(parameters):
            raise DataError(
                "too many parameters provided for substitution:"
                f" given {len(parameters)}, "
                f"used only {idx}"
            )

        return formatted_sql

    def statement_to_set(self, statement: Statement) -> Optional[SetParameter]:
        """
        Try to parse `statement` as a `SET` command.
        Return `None` if it's not a `SET` command.
        """
        # Filter out meaningless tokens like Punctuation and Whitespaces
        skip_types = [Whitespace, Newline]
        tokens = [
            token
            for token in statement.tokens
            if token.ttype not in skip_types and not isinstance(token, Comment)
        ]
        # Trim tail punctuation
        right_idx = len(tokens) - 1
        while str(tokens[right_idx]) == ";":
            right_idx -= 1

        tokens = tokens[: right_idx + 1]

        # Check if it's a SET statement by checking if it starts with set
        if (
            len(tokens) > 0
            and tokens[0].ttype == TokenType.Keyword
            and tokens[0].value.lower() == "set"
        ):
            # Check if set statement has a valid format
            if len(tokens) == 2 and isinstance(tokens[1], Comparison):
                return SetParameter(
                    self._clean_sql_string(tokens[1].left),
                    self._clean_sql_string(tokens[1].right).strip("'"),
                )
            # Or if at least there is a comparison
            cmp_idx = next(
                (
                    i
                    for i, token in enumerate(tokens)
                    if token.ttype == ComparisonType or isinstance(token, Comparison)
                ),
                None,
            )
            if cmp_idx:
                left_tokens, right_tokens = tokens[1:cmp_idx], tokens[cmp_idx + 1 :]
                if isinstance(tokens[cmp_idx], Comparison):
                    left_tokens = left_tokens + [tokens[cmp_idx].left]
                    right_tokens = [tokens[cmp_idx].right] + right_tokens

                if left_tokens and right_tokens:
                    return SetParameter(
                        "".join(self._clean_sql_string(t) for t in left_tokens),
                        "".join(self._clean_sql_string(t) for t in right_tokens).strip(
                            "'"
                        ),
                    )

            raise InterfaceError(
                f"Invalid set statement format: {self._clean_sql_string(statement)},"
                " expected SET <param> = <value>"
            )
        return None

    def _clean_sql_string(self, token: Token) -> str:
        """Strip whitespace and trailing semicolons from a SQL token or statement."""
        return str(token).strip().rstrip(";")

    def split_format_sql(
        self, query: str, parameters: Sequence[Sequence[ParameterType]]
    ) -> List[Union[str, SetParameter]]:
        """
        Multi-statement query formatting will result in `NotSupportedError`.
        Instead, split a query into a separate statement and format with parameters.
        """
        statements = parse_sql(query)
        if not statements:
            return [query]

        if parameters:
            if len(statements) > 1:
                raise NotSupportedError(
                    "Formatting multi-statement queries is not supported."
                )
            if self.statement_to_set(statements[0]):
                raise NotSupportedError("Formatting set statements is not supported.")
            return [
                self.format_statement(statements[0], paramset)
                for paramset in parameters
            ]

        # Try parsing each statement as a SET, otherwise return as a plain sql string
        return [
            self.statement_to_set(st) or self._clean_sql_string(st) for st in statements
        ]

    def format_bulk_insert(
        self, query: str, parameters_seq: Sequence[Sequence[ParameterType]]
    ) -> str:
        """
        Format bulk insert operations by creating multiple INSERT statements.

        Args:
            query: The base INSERT query template
            parameters_seq: Sequence of parameter sets for each INSERT

        Returns:
            Combined SQL string with all INSERT statements
        """
        statements = parse_sql(query)
        if not statements:
            raise DataError("Invalid SQL query for bulk insert")

        formatted_queries = []
        for param_set in parameters_seq:
            formatted_query = self.format_statement(statements[0], param_set)
            formatted_queries.append(formatted_query)

        return "; ".join(formatted_queries)


def create_statement_formatter(version: int) -> StatementFormatter:
    if version == 1:
        return StatementFormatter(escape_chars_v1)
    elif version == 2:
        return StatementFormatter(escape_chars_v2)
    else:
        raise ValueError(f"Unsupported version: {version}")