Fix inline sequence rendering after IN (#629)

docs/parameters.md

@@ -235,6 +235,8 @@ 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.
+
 **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

@@ -1211,7 +1211,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)
+        escaped_values = self.escaper.escape_args(params, operation=stmt)
         rendered_statement = inject_parameters(stmt, escaped_values)
 
         return rendered_statement, NO_NATIVE_PARAMS

src/databricks/sql/utils.py

@@ -557,20 +557,103 @@ 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). 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_PREDICATE_PROBE = re.compile(r"\bIN\s*\x00(\d+)\x00", re.IGNORECASE)
+
+
+class _MarkerProbe:
+    """Formats as a unique token so a marker's position in the rendered operation can be found."""
+
+    def __init__(self, index: int) -> None:
+        self._token = f"\x00{index}\x00"
+
+    def __str__(self) -> str:
+        return self._token
+
+    __repr__ = __str__
+
+
+class _MarkerProbeMapping:
+    """Mapping stand-in that records the key looked up by each %(name)s marker."""
+
+    def __init__(self) -> None:
+        self.keys: List[str] = []
+
+    def __getitem__(self, key: str) -> _MarkerProbe:
+        self.keys.append(key)
+        return _MarkerProbe(len(self.keys) - 1)
+
+
+def _parameter_keys_following_in(
+    operation: str, parameters: Union[Dict[str, Any], Sequence]
+) -> set:
+    """Return the names (dict) or positions (sequence) of parameters whose marker directly
+    follows the IN keyword in `operation`.
+
+    The operation is formatted with probes rather than regex-parsed so marker positions match
+    exactly what the %-formatting in `inject_parameters` will substitute.
+    """
+    if isinstance(parameters, dict):
+        probe_mapping = _MarkerProbeMapping()
+        probe_args: Any = probe_mapping
+        keys: List[Any] = probe_mapping.keys
+    else:
+        keys = list(range(len(parameters)))
+        probe_args = tuple(_MarkerProbe(i) for i in keys)
+
+    try:
+        probed_operation = operation % probe_args
+    except (TypeError, ValueError):
+        # Either inject_parameters raises the same error for this operation, or a marker such as
+        # %d needs a real value; keep the default rendering in both cases.
+        return set()
+
+    return {
+        keys[int(match.group(1))]
+        for match in _IN_PREDICATE_PROBE.finditer(probed_operation)
+    }
+
+
 # Taken from PyHive
 class ParamEscaper:
     _DATE_FORMAT = "%Y-%m-%d"
     _TIME_FORMAT = "%H:%M:%S.%f %z"
     _DATETIME_FORMAT = "{} {}".format(_DATE_FORMAT, _TIME_FORMAT)
 
-    def escape_args(self, parameters):
-        if isinstance(parameters, dict):
-            return {k: self.escape_item(v) for k, v in parameters.items()}
-        elif isinstance(parameters, (list, tuple)):
-            return tuple(self.escape_item(x) for x in parameters)
-        else:
+    def escape_args(self, parameters, operation: Optional[str] = None):
+        """Escape parameter values for inline rendering into `operation`.
+
+        When `operation` is given, a sequence whose marker directly follows IN (e.g.
+        `x IN %(ids)s`) renders as the parenthesised list the IN predicate needs; anywhere else it
+        stays an ARRAY literal. Spark rejects `x IN ARRAY(...)` as a syntax error, so this cannot
+        change the rendering of a query that parses today.
+        """
+        if not isinstance(parameters, (dict, list, tuple)):
             raise ProgrammingError("Unsupported param format: {}".format(parameters))
 
+        in_list_keys = (
+            _parameter_keys_following_in(operation, parameters)
+            if operation is not None
+            else set()
+        )
+        if isinstance(parameters, dict):
+            return {
+                k: self._escape_arg(v, k in in_list_keys) for k, v in parameters.items()
+            }
+        return tuple(
+            self._escape_arg(x, i in in_list_keys) for i, x in enumerate(parameters)
+        )
+
+    def _escape_arg(self, item, follows_in: bool):
+        if follows_in and isinstance(item, Sequence) and not isinstance(item, str):
+            return self.escape_in_list(item)
+        return self.escape_item(item)
+
+    def escape_in_list(self, item):
+        return "(" + ",".join(map(str, map(self.escape_item, item))) + ")"
+
     def escape_number(self, item):
         return item
 

tests/e2e/test_parameterized_queries.py

@@ -543,6 +543,30 @@ 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],
+            ),
+        ],
+    )
+    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()]

tests/unit/test_param_escaper.py

@@ -197,6 +197,104 @@ def test_only_bind_in_where_clause(self):
             inject_parameters(INPUT, pe.escape_args(args))
 
 
+class TestInClauseSequenceRendering(object):
+    """Regression tests for databricks-sql-python#629.
+
+    The IN predicate needs a parenthesised list; Spark rejects `x IN ARRAY(...)` as a syntax
+    error. Everywhere else a sequence must keep rendering as an ARRAY literal so inline values can
+    still be written to ARRAY columns.
+    """
+
+    @pytest.mark.parametrize(
+        ("label", "query", "params", "expected"),
+        (
+            (
+                "named list",
+                "SELECT * FROM t WHERE x IN %(ids)s",
+                {"ids": [1, 2, 3]},
+                "SELECT * FROM t WHERE x IN (1,2,3)",
+            ),
+            (
+                "positional list",
+                "SELECT * FROM t WHERE x IN %s",
+                ([1, 2, 3],),
+                "SELECT * FROM t WHERE x IN (1,2,3)",
+            ),
+            (
+                "tuple of strings is escaped",
+                "SELECT * FROM t WHERE name IN %(names)s",
+                {"names": ("a", "O'Connell")},
+                r"SELECT * FROM t WHERE name IN ('a','O\'Connell')",
+            ),
+            (
+                "NOT IN",
+                "SELECT * FROM t WHERE x NOT IN %(ids)s",
+                {"ids": [1, 2]},
+                "SELECT * FROM t WHERE x NOT IN (1,2)",
+            ),
+            (
+                "lowercase keyword and newline before the marker",
+                "SELECT * FROM t WHERE x in\n  %(ids)s",
+                {"ids": [1, 2]},
+                "SELECT * FROM t WHERE x in\n  (1,2)",
+            ),
+            (
+                "nested sequences stay arrays",
+                "SELECT * FROM t WHERE arr IN %(arrs)s",
+                {"arrs": [[1, 2], [3]]},
+                "SELECT * FROM t WHERE arr IN (ARRAY(1,2),ARRAY(3))",
+            ),
+            (
+                "only the positional marker after IN",
+                "SELECT %s AS tags FROM t WHERE x IN %s AND y = %s",
+                (["a"], [1, 2], 5),
+                "SELECT ARRAY('a') AS tags FROM t WHERE x IN (1,2) AND y = 5",
+            ),
+            (
+                "escaped percent does not shift positional markers",
+                "SELECT '100%%' AS pct, %s AS tags FROM t WHERE x IN %s",
+                (["a"], [1, 2]),
+                "SELECT '100%' AS pct, ARRAY('a') AS tags FROM t WHERE x IN (1,2)",
+            ),
+        ),
+    )
+    def test_sequence_after_in_renders_as_list(self, label, query, params, expected):
+        rendered = inject_parameters(query, pe.escape_args(params, operation=query))
+        assert rendered == expected
+
+    @pytest.mark.parametrize(
+        ("label", "query", "params", "expected"),
+        (
+            (
+                "insert into an array column",
+                "INSERT INTO t (arr) VALUES (%(p)s)",
+                {"p": ["a", "b"]},
+                "INSERT INTO t (arr) VALUES (ARRAY('a','b'))",
+            ),
+            (
+                "function argument",
+                "SELECT array_contains(%(p)s, 2)",
+                {"p": [1, 2, 3]},
+                "SELECT array_contains(ARRAY(1,2,3), 2)",
+            ),
+            (
+                "caller-written parentheses after IN compare against an array value",
+                "SELECT * FROM t WHERE arr IN (%(p)s)",
+                {"p": [1, 2]},
+                "SELECT * FROM t WHERE arr IN (ARRAY(1,2))",
+            ),
+        ),
+    )
+    def test_sequence_elsewhere_stays_array(self, label, query, params, expected):
+        rendered = inject_parameters(query, pe.escape_args(params, operation=query))
+        assert rendered == expected
+
+    def test_malformed_operation_raises_the_same_error(self):
+        query = "SELECT * FROM t WHERE x IN %s AND name LIKE '%sonite'"
+        with pytest.raises(TypeError, match="not enough arguments for format string"):
+            inject_parameters(query, pe.escape_args(([1, 2],), operation=query))
+
+
 class TestInlineToNativeTransformer(object):
     @pytest.mark.parametrize(
         ("label", "query", "params", "expected"),