Fix inline sequence rendering after IN (#629)

docs/parameters.md

@@ -235,6 +235,21 @@ Output:
 SELECT field FROM table WHERE field IN (1,2,3,4,5)
 ```
 
+The marker must directly follow `IN` (or `NOT IN`), as in `IN %(value_list)s`. Anywhere else, including inside parentheses you write yourself such as `IN (%(value_list)s)`, a sequence is rendered as an `ARRAY(...)` literal.
+
+Inside the list, a tuple is rendered as a row and a list as an `ARRAY(...)` literal. For a multi-column `IN`, pass tuples:
+
+```python
+query = "SELECT * FROM table WHERE (id, tag) IN %(pairs)s"
+parameters = {"pairs": [(1, "a"), (2, "b")]}
+```
+
+```sql
+SELECT * FROM table WHERE (id, tag) IN ((1,'a'),(2,'b'))
+```
+
+Spark does not convert the types of the values inside a row, so each value must have the same type as its column. For example, the literal `1` is an `INT` and does not match a `BIGINT` column.
+
 **Note**: this behavior is not specified by PEP-249 and only works due to the specific implementation of our inline renderer.
 
 ### Migrating to native parameters

src/databricks/sql/client.py

@@ -46,7 +46,6 @@
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.utils import (
     ParamEscaper,
-    inject_parameters,
     transform_paramstyle,
     ColumnTable,
     ColumnQueue,
@@ -1211,8 +1210,7 @@ def _prepare_inline_parameters(
                 The list is always empty because native parameters are never used under the inline approach
         """
 
-        escaped_values = self.escaper.escape_args(params)
-        rendered_statement = inject_parameters(stmt, escaped_values)
+        rendered_statement = self.escaper.render_inline(stmt, params)
 
         return rendered_statement, NO_NATIVE_PARAMS
 

src/databricks/sql/utils.py

@@ -1,5 +1,16 @@
 from __future__ import annotations
-from typing import Any, Dict, List, Optional, Tuple, Union, Sequence, TYPE_CHECKING
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Union,
+    Sequence,
+    TYPE_CHECKING,
+)
 
 from dateutil import parser
 import datetime
@@ -10,6 +21,7 @@
 from decimal import Decimal
 from enum import Enum
 import re
+import secrets
 
 import lz4.frame
 
@@ -557,6 +569,88 @@ def user_friendly_error_message(self, no_retry_reason, attempt, elapsed):
         return user_friendly_error_message + ". " + error_context
 
 
+# A probe token directly after the IN keyword (also matches NOT IN). Where Spark's grammar accepts a
+# value directly after IN, in the IN predicate and in PIVOT, it requires a parenthesised list, so
+# `IN ARRAY(...)` is a syntax error. (The one other place, `position(a IN b)`, takes a string and
+# rejects an ARRAY too.) Requiring the marker to follow IN directly leaves `IN (%(p)s)` alone: there
+# the caller's parentheses already form the list and an ARRAY value is a valid list element.
+_IN_MARKER_PROBE = re.compile(r"\bIN\s*\x00([0-9a-f]+):(\d+)\x00", re.IGNORECASE)
+
+
+def _in_ends_line_comment(text: str, match: re.Match[str]) -> bool:
+    """Whether the IN that `match` found is inside a `--` comment and the marker is on a later line.
+
+    The marker is then live code and keeps its ARRAY rendering. A marker on the same line as a
+    commented IN is inside the comment too, so its rendering does not matter. A `--` inside a
+    string literal earlier on the line also counts; that only falls back to the ARRAY rendering.
+    """
+    if not re.search(r"[\r\n]", match.group(0)):
+        return False
+    line_start = max(
+        text.rfind("\n", 0, match.start()), text.rfind("\r", 0, match.start())
+    )
+    return "--" in text[line_start + 1 : match.start()]
+
+
+def _renders_as_array(value: Any) -> bool:
+    """Whether ParamEscaper.escape_item renders `value` as an ARRAY literal."""
+    return isinstance(value, Sequence) and not isinstance(value, str)
+
+
+class _MarkerProbe:
+    """Formats as a token that identifies the marker it was substituted into."""
+
+    def __init__(self, nonce: str, index: int) -> None:
+        self._token = f"\x00{nonce}:{index}\x00"
+
+    def __str__(self) -> str:
+        return self._token
+
+    __repr__ = __str__
+
+
+class _NumberedLookups:
+    """Mapping stand-in for %-formatting that numbers its lookups.
+
+    %-formatting looks up each %(name)s marker once, left to right, so lookup n is always the n-th
+    named marker of the operation. `value_for(n, name)` supplies the value substituted there.
+    """
+
+    def __init__(
+        self, value_for: Callable[[int, str], Any], escaped: Dict[str, Any]
+    ) -> None:
+        self._value_for = value_for
+        self._escaped = escaped
+        self._count = 0
+
+    def __getitem__(self, name: str) -> Any:
+        index = self._count
+        self._count += 1
+        return self._value_for(index, name)
+
+    def __str__(self) -> str:
+        # A bare %s marker formats the whole mapping; keep rendering the escaped dict there.
+        return str(self._escaped)
+
+    __repr__ = __str__
+
+
+def _format_markers(
+    operation: str, escaped: Any, value_for: Callable[[int, Any], Any]
+) -> str:
+    """%-format `operation`, taking the value of marker n from `value_for(n, key)`.
+
+    `escaped` is the output of `ParamEscaper.escape_args`. `key` is the marker's name when it is a
+    dict and the marker's position otherwise.
+    """
+    values: Any
+    if isinstance(escaped, dict):
+        values = _NumberedLookups(value_for, escaped)
+    else:
+        values = tuple(value_for(i, i) for i in range(len(escaped)))
+    return inject_parameters(operation, values)
+
+
 # Taken from PyHive
 class ParamEscaper:
     _DATE_FORMAT = "%Y-%m-%d"
@@ -571,6 +665,83 @@ def escape_args(self, parameters):
         else:
             raise ProgrammingError("Unsupported param format: {}".format(parameters))
 
+    def render_inline(self, operation: str, parameters) -> str:
+        """Render `parameters` into the markers of `operation` as inline SQL literals.
+
+        Values are escaped as `escape_args` escapes them, with one exception: a sequence at a
+        marker that directly follows IN (e.g. `x IN %(ids)s`) renders as the parenthesised list
+        Spark requires there, `x IN (1,2,3)`, instead of `x IN ARRAY(1,2,3)`, which Spark rejects.
+        The choice is made per marker, so a parameter used both after IN and elsewhere keeps its
+        ARRAY rendering at the other markers. Inside the list, tuples render as rows and lists as
+        ARRAY literals.
+        """
+        escaped = self.escape_args(parameters)
+        values = parameters.values() if isinstance(parameters, dict) else parameters
+        if not any(_renders_as_array(value) for value in values):
+            # Only a sequence can render as an IN list, so without one render exactly as before.
+            return inject_parameters(operation, escaped)
+
+        in_list_markers = self._sequence_markers_following_in(
+            operation, parameters, escaped
+        )
+
+        def render_marker(index: int, key: Any) -> Any:
+            if index in in_list_markers:
+                return self.escape_in_list(parameters[key])
+            return escaped[key]
+
+        return _format_markers(operation, escaped, render_marker)
+
+    def _sequence_markers_following_in(
+        self, operation: str, parameters, escaped
+    ) -> Set[int]:
+        """Return the numbers of the markers that directly follow IN and receive a sequence.
+
+        Markers are numbered as `_format_markers` numbers them. To find them, the operation is
+        formatted once through the same %-formatting as the real substitution, so `%%` and
+        positional markers line up exactly. Each sequence value is replaced by a probe token.
+        Numbers keep their value, so markers such as `%d` still format, and escaped strings are
+        replaced by an empty string literal, so the text of a value cannot change the result.
+        """
+        # The nonce stops text that only looks like a probe token from being taken for one.
+        nonce = secrets.token_hex(8)
+
+        def probe_marker(index: int, key: Any) -> Any:
+            # A missing name raises KeyError here, as it does when rendering.
+            escaped_value = escaped[key]
+            if _renders_as_array(parameters[key]):
+                return _MarkerProbe(nonce, index)
+            # Replace text with '' so the content of a value cannot change the result. A string
+            # of two or more characters formats under exactly the markers '' does. A one-character
+            # string can only be a single-digit Decimal, which %c accepts, so it is kept.
+            if isinstance(escaped_value, str) and len(escaped_value) > 1:
+                return "''"
+            return escaped_value
+
+        try:
+            probed_operation = _format_markers(operation, escaped, probe_marker)
+        except (KeyError, TypeError, ValueError):
+            # Rendering raises the same error: a probe only replaces a sequence, and a marker that
+            # cannot format a probe (such as %d) cannot format the sequence's ARRAY string either.
+            return set()
+
+        return {
+            int(match.group(2))
+            for match in _IN_MARKER_PROBE.finditer(probed_operation)
+            if match.group(1) == nonce
+            and not _in_ends_line_comment(probed_operation, match)
+        }
+
+    def escape_in_list(self, item):
+        return "(" + ",".join(map(str, map(self._escape_in_list_element, item))) + ")"
+
+    def _escape_in_list_element(self, item):
+        # Inside an IN list a tuple is a row, as in `(a, b) IN ((1, 2), (3, 4))`. A list stays an
+        # ARRAY literal so ARRAY columns can still be matched, as in `arr IN (ARRAY(1, 2))`.
+        if isinstance(item, tuple):
+            return self.escape_in_list(item)
+        return self.escape_item(item)
+
     def escape_number(self, item):
         return item
 

tests/e2e/test_parameterized_queries.py

@@ -543,6 +543,50 @@ def test_params_as_sequence(self):
                 assert result.bar == 2
                 assert result.baz == 3
 
+    @pytest.mark.parametrize(
+        "query,params,expected_ids",
+        [
+            (
+                "SELECT id FROM range(5) WHERE id IN %(ids)s ORDER BY id",
+                {"ids": [1, 3]},
+                [1, 3],
+            ),
+            ("SELECT id FROM range(5) WHERE id IN %s ORDER BY id", [[1, 3]], [1, 3]),
+            (
+                "SELECT id FROM range(5) WHERE id NOT IN %(ids)s ORDER BY id",
+                {"ids": (1, 3)},
+                [0, 2, 4],
+            ),
+            (
+                # The same parameter renders as a list after IN and as an ARRAY elsewhere.
+                "SELECT id FROM range(5) WHERE id IN %(ids)s"
+                " AND array_contains(%(ids)s, id) ORDER BY id",
+                {"ids": [1, 3]},
+                [1, 3],
+            ),
+            (
+                "SELECT id FROM range(5) WHERE id IN %(ids)s ORDER BY id LIMIT %(n)d",
+                {"ids": [1, 3, 4], "n": 2},
+                [1, 3],
+            ),
+            (
+                # Tuples render as rows for a multi-column IN. Spark does not widen the field
+                # types of a row, so the INT and STRING columns match the rendered literals.
+                "SELECT id FROM VALUES (1, 'a'), (2, 'b'), (3, 'c') AS t(id, tag)"
+                " WHERE (id, tag) IN %(pairs)s ORDER BY id",
+                {"pairs": [(1, "a"), (2, "x"), (3, "c")]},
+                [1, 3],
+            ),
+        ],
+    )
+    def test_inline_sequence_in_clause(self, query, params, expected_ids):
+        """Regression test for #629: a sequence bound directly after IN must render as a
+        parenthesised list, not an ARRAY literal, or the server rejects the query."""
+        with self.cursor(extra_params={"use_inline_params": True}) as cursor:
+            result = cursor.execute(query, parameters=params).fetchall()
+
+        assert [row.id for row in result] == expected_ids
+
     def test_inline_ordinals_can_break_sql(self):
         """With inline mode, ordinal parameters can break the SQL syntax
         because `%` symbols are used to wildcard match within LIKE statements. This test

tests/unit/test_client.py

@@ -494,6 +494,31 @@ def test_execute_parameter_passthrough(self):
                 expected_query,
             )
 
+    def test_execute_inline_sequence_after_in_renders_as_list(self):
+        """Regression test for databricks-sql-python#629: `x IN ARRAY(...)` is a syntax error."""
+        mock_thrift_backend = ThriftDatabricksClientMockFactory.new()
+        cursor = client.Cursor(Mock(), mock_thrift_backend)
+
+        tests = [
+            (
+                "SELECT * FROM t WHERE x IN %(ids)s",
+                "SELECT * FROM t WHERE x IN (1,2,3)",
+                {"ids": [1, 2, 3]},
+            ),
+            (
+                "SELECT * FROM t WHERE x IN %s",
+                "SELECT * FROM t WHERE x IN (1,2,3)",
+                [[1, 2, 3]],
+            ),
+        ]
+
+        for query, expected_query, params in tests:
+            cursor.execute(query, parameters=params)
+            self.assertEqual(
+                mock_thrift_backend.execute_command.call_args[1]["operation"],
+                expected_query,
+            )
+
     def test_executemany_parameter_passhthrough_and_uses_last_result_set(self):
         # Create a new mock result set each time the class is instantiated
         mock_result_set_instances = [Mock(), Mock(), Mock()]