fix(FIR-43230): string escape in query parameters (#418)

Co-authored-by: Luciano Scarpulla <[email protected]>

Author: ptiurin

src/firebolt/async_db/cursor.py

@@ -27,12 +27,7 @@
 )
 
 from firebolt.client.client import AsyncClient, AsyncClientV1, AsyncClientV2
-from firebolt.common._types import (
-    ColType,
-    ParameterType,
-    SetParameter,
-    split_format_sql,
-)
+from firebolt.common._types import ColType, ParameterType, SetParameter
 from firebolt.common.base_cursor import (
     JSON_OUTPUT_FORMAT,
     RESET_SESSION_HEADER,
@@ -47,6 +42,7 @@
     check_not_closed,
     check_query_executed,
 )
+from firebolt.common.statement_formatter import create_statement_formatter
 from firebolt.utils.exception import (
     EngineNotRunningError,
     FireboltDatabaseError,
@@ -209,7 +205,9 @@ async def _do_execute(
     ) -> None:
         self._reset()
         queries: List[Union[SetParameter, str]] = (
-            [raw_query] if skip_parsing else split_format_sql(raw_query, parameters)
+            [raw_query]
+            if skip_parsing
+            else self._formatter.split_format_sql(raw_query, parameters)
         )
         timeout_controller = TimeoutController(timeout)
 
@@ -435,7 +433,13 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         assert isinstance(client, AsyncClientV2)
-        super().__init__(*args, client=client, connection=connection, **kwargs)
+        super().__init__(
+            *args,
+            client=client,
+            connection=connection,
+            formatter=create_statement_formatter(version=2),
+            **kwargs,
+        )
 
     @check_not_closed
     async def execute_async(
@@ -512,7 +516,13 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         assert isinstance(client, AsyncClientV1)
-        super().__init__(*args, client=client, connection=connection, **kwargs)
+        super().__init__(
+            *args,
+            client=client,
+            connection=connection,
+            formatter=create_statement_formatter(version=1),
+            **kwargs,
+        )
 
     async def is_db_available(self, database_name: str) -> bool:
         """

src/firebolt/common/_types.py

@@ -2,24 +2,11 @@
 
 import re
 from collections import namedtuple
-from datetime import date, datetime, timezone
+from datetime import date, datetime
 from decimal import Decimal
 from enum import Enum
 from io import StringIO
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
-
-from sqlparse import parse as parse_sql  # type: ignore
-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 typing import Any, Dict, List, Sequence, Tuple, Union
 
 try:
     from ciso8601 import parse_datetime  # type: ignore
@@ -46,11 +33,7 @@ def parse_datetime(datetime_string: str) -> datetime:
         return datetime.fromisoformat(_fix_timezone(_fix_milliseconds(datetime_string)))
 
 
-from firebolt.utils.exception import (
-    DataError,
-    InterfaceError,
-    NotSupportedError,
-)
+from firebolt.utils.exception import DataError, NotSupportedError
 from firebolt.utils.util import cached_property
 
 _NoneType = type(None)
@@ -372,158 +355,4 @@ def parse_value(
     raise DataError(f"Unsupported data type returned: {ctype.__name__}")
 
 
-escape_chars = {
-    "\0": "\\0",
-    "\\": "\\\\",
-    "'": "\\'",
-}
-
-
-def format_value(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(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(format_value(it) for it in value)}]"
-
-    raise DataError(f"unsupported parameter type {type(value)}")
-
-
-def format_statement(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 = 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 = statement_to_sql(process_token(statement))
-
-    if idx < len(parameters):
-        raise DataError(
-            f"too many parameters provided for substitution: given {len(parameters)}, "
-            f"used only {idx}"
-        )
-
-    return formatted_sql
-
-
 SetParameter = namedtuple("SetParameter", ["name", "value"])
-
-
-def statement_to_set(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(
-                statement_to_sql(tokens[1].left),
-                statement_to_sql(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(statement_to_sql(t) for t in left_tokens),
-                    "".join(statement_to_sql(t) for t in right_tokens).strip("'"),
-                )
-
-        raise InterfaceError(
-            f"Invalid set statement format: {statement_to_sql(statement)},"
-            " expected SET <param> = <value>"
-        )
-    return None
-
-
-def statement_to_sql(statement: Statement) -> str:
-    return str(statement).strip().rstrip(";")
-
-
-def split_format_sql(
-    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 statement_to_set(statements[0]):
-            raise NotSupportedError("Formatting set statements is not supported.")
-        return [format_statement(statements[0], paramset) for paramset in parameters]
-
-    # Try parsing each statement as a SET, otherwise return as a plain sql string
-    return [statement_to_set(st) or statement_to_sql(st) for st in statements]

src/firebolt/common/base_cursor.py

@@ -18,6 +18,7 @@
     parse_type,
     parse_value,
 )
+from firebolt.common.statement_formatter import StatementFormatter
 from firebolt.utils.exception import (
     ConfigurationError,
     CursorClosedError,
@@ -187,6 +188,7 @@ class BaseCursor:
         "_rows",
         "_idx",
         "_row_sets",
+        "_formatter",
         "_next_set_idx",
         "_set_parameters",
         "_query_id",
@@ -196,13 +198,16 @@ class BaseCursor:
 
     default_arraysize = 1
 
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
+    def __init__(
+        self, *args: Any, formatter: StatementFormatter, **kwargs: Any
+    ) -> None:
         self._arraysize = self.default_arraysize
         # These fields initialized here for type annotations purpose
         self._rows: Optional[List[List[RawColType]]] = None
         self._descriptions: Optional[List[Column]] = None
         self._statistics: Optional[Statistics] = None
         self._row_sets: List[RowSet] = []
+        self._formatter = formatter
         # User-defined set parameters
         self._set_parameters: Dict[str, Any] = dict()
         # Server-side parameters (user can't change them)