From f65c7dd38ce209773ebf5a9a868810bc2f810d7f Mon Sep 17 00:00:00 2001 From: "cathleen.yan" <58714163+cathleeny@users.noreply.github.com> Date: Fri, 6 Nov 2026 00:11:15 +0000 Subject: [PATCH 1/7] Fix inline sequence rendering after IN (#629) Since 4.0.4, inline parameters render every Python sequence as an ARRAY literal, so the documented `x IN %(ids)s` pattern produces `x IN ARRAY(1,2,3)`, which Spark rejects with PARSE_SYNTAX_ERROR. A sequence whose marker directly follows IN (or NOT IN) now renders as a parenthesised list again. Everywhere else, including `IN (%(p)s)` where the caller wrote the parentheses, it stays an ARRAY literal, so inline writes to ARRAY columns keep working. Because `x IN ARRAY(...)` never parses, no query that runs today renders differently. Co-authored-by: Isaac Signed-off-by: cathleen.yan <58714163+cathleeny@users.noreply.github.com> --- docs/parameters.md | 2 + src/databricks/sql/client.py | 2 +- src/databricks/sql/utils.py | 95 ++++++++++++++++++++++-- tests/e2e/test_parameterized_queries.py | 24 ++++++ tests/unit/test_client.py | 25 +++++++ tests/unit/test_param_escaper.py | 98 +++++++++++++++++++++++++ 6 files changed, 239 insertions(+), 7 deletions(-) diff --git a/docs/parameters.md b/docs/parameters.md index f9f4c5ff9..283da8e88 100644 --- a/docs/parameters.md +++ b/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 diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 940b12775..157edb032 100755 --- a/src/databricks/sql/client.py +++ b/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 diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 4e82370b9..02f490cf6 100644 --- a/src/databricks/sql/utils.py +++ b/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 diff --git a/tests/e2e/test_parameterized_queries.py b/tests/e2e/test_parameterized_queries.py index 7370eea93..03406e024 100644 --- a/tests/e2e/test_parameterized_queries.py +++ b/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 diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 66a3722ec..047252488 100644 --- a/tests/unit/test_client.py +++ b/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()] diff --git a/tests/unit/test_param_escaper.py b/tests/unit/test_param_escaper.py index 9b6b9c246..0e1fbac1e 100644 --- a/tests/unit/test_param_escaper.py +++ b/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"), From fe0dfb6dcd5745c45fd3a87dc43cf37fe5a3d9a1 Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Fri, 6 Nov 2026 00:53:44 +0000 Subject: [PATCH 2/7] Choose the IN-list rendering per marker and probe only sequences Follow-ups from review of the inline IN fix: - Decide per marker instead of per parameter name. A name used both directly after IN and elsewhere now keeps its ARRAY rendering at the other markers. Before, every occurrence became `(...)`, which could break a query that runs on main, e.g. one with a commented-out filter: `arr = %(tags)s -- AND tag IN %(tags)s`. - Probe only sequence values and give every other marker its real escaped value, so a `%d` or `%(n)d` marker elsewhere in the query no longer turns the fix off. Probe tokens carry a random nonce so that a string value cannot pass for one. - Leave escape_args exactly as it is on main and move the logic into ParamEscaper.render_inline, which Cursor._prepare_inline_parameters now calls. Co-authored-by: Isaac Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- src/databricks/sql/client.py | 4 +- src/databricks/sql/utils.py | 182 +++++++++++++++--------- tests/e2e/test_parameterized_queries.py | 12 ++ tests/unit/test_param_escaper.py | 100 +++++++++++-- 4 files changed, 223 insertions(+), 75 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 157edb032..6e03ed227 100755 --- a/src/databricks/sql/client.py +++ b/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, operation=stmt) - rendered_statement = inject_parameters(stmt, escaped_values) + rendered_statement = self.escaper.render_inline(stmt, params) return rendered_statement, NO_NATIVE_PARAMS diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 02f490cf6..4870ef856 100644 --- a/src/databricks/sql/utils.py +++ b/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,17 +569,24 @@ 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) +# 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 _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 unique token so a marker's position in the rendered operation can be found.""" + """Formats as a token that identifies the marker it was substituted into.""" - def __init__(self, index: int) -> None: - self._token = f"\x00{index}\x00" + def __init__(self, nonce: str, index: int) -> None: + self._token = f"\x00{nonce}:{index}\x00" def __str__(self) -> str: return self._token @@ -575,45 +594,46 @@ def __str__(self) -> str: __repr__ = __str__ -class _MarkerProbeMapping: - """Mapping stand-in that records the key looked up by each %(name)s marker.""" +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 __init__(self) -> None: - self.keys: List[str] = [] + def __str__(self) -> str: + # A bare %s marker formats the whole mapping; keep rendering the escaped dict there. + return str(self._escaped) - def __getitem__(self, key: str) -> _MarkerProbe: - self.keys.append(key) - return _MarkerProbe(len(self.keys) - 1) + __repr__ = __str__ -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`. +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)`. - The operation is formatted with probes rather than regex-parsed so marker positions match - exactly what the %-formatting in `inject_parameters` will substitute. + `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. """ - if isinstance(parameters, dict): - probe_mapping = _MarkerProbeMapping() - probe_args: Any = probe_mapping - keys: List[Any] = probe_mapping.keys + values: Any + if isinstance(escaped, dict): + values = _NumberedLookups(value_for, escaped) 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) - } + values = tuple(value_for(i, i) for i in range(len(escaped))) + return inject_parameters(operation, values) # Taken from PyHive @@ -622,34 +642,68 @@ class ParamEscaper: _TIME_FORMAT = "%H:%M:%S.%f %z" _DATETIME_FORMAT = "{} {}".format(_DATE_FORMAT, _TIME_FORMAT) - 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)): + 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: 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 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. + """ + escaped = self.escape_args(parameters) + in_list_markers = self._sequence_markers_following_in( + operation, parameters, escaped ) - 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 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 and + every other value keeps its escaped value, so scalar markers such as `%d` still format. + """ + # The nonce stops a string value that happens to contain a probe-shaped token from being + # mistaken for a probe, since real values are substituted alongside the probes. + 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) + 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 + } def escape_in_list(self, item): return "(" + ",".join(map(str, map(self.escape_item, item))) + ")" diff --git a/tests/e2e/test_parameterized_queries.py b/tests/e2e/test_parameterized_queries.py index 03406e024..f8c64bc7e 100644 --- a/tests/e2e/test_parameterized_queries.py +++ b/tests/e2e/test_parameterized_queries.py @@ -557,6 +557,18 @@ def test_params_as_sequence(self): {"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], + ), ], ) def test_inline_sequence_in_clause(self, query, params, expected_ids): diff --git a/tests/unit/test_param_escaper.py b/tests/unit/test_param_escaper.py index 0e1fbac1e..5fd0120b6 100644 --- a/tests/unit/test_param_escaper.py +++ b/tests/unit/test_param_escaper.py @@ -256,11 +256,28 @@ class TestInClauseSequenceRendering(object): (["a"], [1, 2]), "SELECT '100%' AS pct, ARRAY('a') AS tags FROM t WHERE x IN (1,2)", ), + ( + "named %d marker elsewhere in the query", + "SELECT * FROM t WHERE x IN %(ids)s LIMIT %(n)d", + {"ids": [1, 2], "n": 5}, + "SELECT * FROM t WHERE x IN (1,2) LIMIT 5", + ), + ( + "positional %d marker elsewhere in the query", + "SELECT * FROM t WHERE x IN %s LIMIT %d", + ([1, 2], 5), + "SELECT * FROM t WHERE x IN (1,2) LIMIT 5", + ), + ( + "PIVOT value list", + "SELECT * FROM t PIVOT (sum(v) FOR k IN %(ks)s)", + {"ks": ["a", "b"]}, + "SELECT * FROM t PIVOT (sum(v) FOR k IN ('a','b'))", + ), ), ) 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 + assert pe.render_inline(query, params) == expected @pytest.mark.parametrize( ("label", "query", "params", "expected"), @@ -286,13 +303,80 @@ def test_sequence_after_in_renders_as_list(self, label, query, params, expected) ), ) def test_sequence_elsewhere_stays_array(self, label, query, params, expected): - rendered = inject_parameters(query, pe.escape_args(params, operation=query)) - assert rendered == expected + assert pe.render_inline(query, params) == expected + + @pytest.mark.parametrize( + ("label", "query", "params", "expected"), + ( + ( + "same name after IN and as a function argument", + "SELECT * FROM t WHERE x IN %(ids)s AND array_contains(%(ids)s, x)", + {"ids": [1, 2]}, + "SELECT * FROM t WHERE x IN (1,2) AND array_contains(ARRAY(1,2), x)", + ), + ( + "commented-out IN does not change the live marker", + "SELECT * FROM t WHERE arr = %(tags)s\n-- AND tag IN %(tags)s\n", + {"tags": ["a", "b"]}, + "SELECT * FROM t WHERE arr = ARRAY('a','b')\n-- AND tag IN ('a','b')\n", + ), + ), + ) + def test_rendering_is_chosen_per_marker(self, label, query, params, expected): + assert pe.render_inline(query, params) == expected + + @pytest.mark.parametrize( + ("label", "query", "params"), + ( + ( + "named markers", + "INSERT INTO t VALUES (%(s)s, %(n)s, %(arr)s, %(m)s, %(none)s)", + {"s": "it's", "n": 1.5, "arr": [1, [2]], "m": {"k": 1}, "none": None}, + ), + ( + "bare %s with dict parameters formats the whole dict", + "SELECT * FROM t WHERE x IN %s", + {"ids": [1]}, + ), + ), + ) + def test_matches_escape_args_without_in_list(self, label, query, params): + assert pe.render_inline(query, params) == inject_parameters( + query, pe.escape_args(params) + ) - 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)) + def test_probe_shaped_string_value_is_not_a_probe(self): + """Real values are substituted while probing, so a string value that contains a probe-like + token must not make another marker render as a list.""" + query = "INSERT INTO t VALUES (%(note)s, %(tags)s)" + params = {"note": "x IN \x000000000000000000:1\x00", "tags": ["a"]} + assert pe.render_inline(query, params) == ( + "INSERT INTO t VALUES ('x IN \x000000000000000000:1\x00', ARRAY('a'))" + ) + + @pytest.mark.parametrize( + ("query", "params", "error", "message"), + ( + ( + "SELECT * FROM t WHERE x IN %s AND name LIKE '%sonite'", + ([1, 2],), + TypeError, + "not enough arguments for format string", + ), + ( + "SELECT * FROM t WHERE x IN %d", + ([1, 2],), + TypeError, + "a real number is required, not str", + ), + ("SELECT * FROM t WHERE x IN %(nope)s", {"ids": [1]}, KeyError, "nope"), + ), + ) + def test_errors_match_escape_args(self, query, params, error, message): + with pytest.raises(error, match=message): + inject_parameters(query, pe.escape_args(params)) + with pytest.raises(error, match=message): + pe.render_inline(query, params) class TestInlineToNativeTransformer(object): From 26c671a3c35260f00278140b129cdbeb7662b23b Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Fri, 6 Nov 2026 19:39:48 +0000 Subject: [PATCH 3/7] Render tuples in an IN list as rows; skip probing without sequences Follow-ups from review: - Inside an IN list a tuple now renders as a row, so a multi-column `(a, b) IN %(pairs)s` with `[(1, 2), (3, 4)]` renders as `(a, b) IN ((1,2),(3,4))`, as it did in 3.x. Lists still render as ARRAY literals so ARRAY columns can be matched. Spark does not widen the field types of a row, so each value must have its column's type; docs/parameters.md says so. - When no parameter is a sequence, render_inline skips the probe pass and renders exactly as escape_args and inject_parameters do on main. Co-authored-by: Isaac Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- docs/parameters.md | 13 ++++++ src/databricks/sql/utils.py | 17 +++++++- tests/e2e/test_parameterized_queries.py | 8 ++++ tests/unit/test_param_escaper.py | 53 ++++++++++++++++++++++++- 4 files changed, 88 insertions(+), 3 deletions(-) diff --git a/docs/parameters.md b/docs/parameters.md index 283da8e88..8047cf53c 100644 --- a/docs/parameters.md +++ b/docs/parameters.md @@ -237,6 +237,19 @@ 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 diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 4870ef856..d7f43be78 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -657,9 +657,15 @@ def render_inline(self, operation: str, parameters) -> str: 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. + 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 ) @@ -706,7 +712,14 @@ def probe_marker(index: int, key: Any) -> Any: } def escape_in_list(self, item): - return "(" + ",".join(map(str, map(self.escape_item, 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 diff --git a/tests/e2e/test_parameterized_queries.py b/tests/e2e/test_parameterized_queries.py index f8c64bc7e..d690601b1 100644 --- a/tests/e2e/test_parameterized_queries.py +++ b/tests/e2e/test_parameterized_queries.py @@ -569,6 +569,14 @@ def test_params_as_sequence(self): {"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): diff --git a/tests/unit/test_param_escaper.py b/tests/unit/test_param_escaper.py index 5fd0120b6..78f53371c 100644 --- a/tests/unit/test_param_escaper.py +++ b/tests/unit/test_param_escaper.py @@ -1,7 +1,9 @@ from datetime import date, datetime import unittest, pytest, decimal from typing import Any, Dict +from unittest.mock import patch from databricks.sql.parameters.native import dbsql_parameter_from_primitive +from databricks.sql.types import Row from databricks.sql.utils import ( ParamEscaper, @@ -239,11 +241,35 @@ class TestInClauseSequenceRendering(object): "SELECT * FROM t WHERE x in\n (1,2)", ), ( - "nested sequences stay arrays", + "nested lists 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))", ), + ( + "nested tuples render as rows", + "SELECT * FROM t WHERE (a, b) IN %(pairs)s", + {"pairs": [(1, "x"), (2, "O'Connell")]}, + r"SELECT * FROM t WHERE (a, b) IN ((1,'x'),(2,'O\'Connell'))", + ), + ( + "tuple of tuples with NOT IN, positional", + "SELECT * FROM t WHERE (a, b) NOT IN %s", + (((1, 2), (3, 4)),), + "SELECT * FROM t WHERE (a, b) NOT IN ((1,2),(3,4))", + ), + ( + "Row objects render as rows", + "SELECT * FROM t WHERE (a, b) IN %(rows)s", + {"rows": [Row(a=1, b=2), Row(a=3, b=4)]}, + "SELECT * FROM t WHERE (a, b) IN ((1,2),(3,4))", + ), + ( + "a row can hold an array and a nested row", + "SELECT * FROM t WHERE (a, arr, s) IN %(rows)s", + {"rows": [(1, [2, 3], (4, 5))]}, + "SELECT * FROM t WHERE (a, arr, s) IN ((1,ARRAY(2,3),(4,5)))", + ), ( "only the positional marker after IN", "SELECT %s AS tags FROM t WHERE x IN %s AND y = %s", @@ -300,6 +326,12 @@ def test_sequence_after_in_renders_as_list(self, label, query, params, expected) {"p": [1, 2]}, "SELECT * FROM t WHERE arr IN (ARRAY(1,2))", ), + ( + "a tuple outside an IN list is still an array", + "INSERT INTO t (arr) VALUES (%(p)s)", + {"p": (1, 2)}, + "INSERT INTO t (arr) VALUES (ARRAY(1,2))", + ), ), ) def test_sequence_elsewhere_stays_array(self, label, query, params, expected): @@ -345,6 +377,25 @@ def test_matches_escape_args_without_in_list(self, label, query, params): query, pe.escape_args(params) ) + @pytest.mark.parametrize( + ("query", "params"), + ( + ("SELECT * FROM t WHERE a = %(a)s AND b IN (%(b)s)", {"a": "x", "b": 2}), + ("SELECT * FROM t WHERE a = %s LIMIT %d", ("x", 5)), + ("SELECT * FROM t WHERE a = %s", {"a": "x"}), + ), + ) + def test_without_sequences_renders_as_escape_args_without_probing( + self, query, params + ): + with patch.object( + ParamEscaper, "_sequence_markers_following_in", side_effect=AssertionError + ) as probe: + assert pe.render_inline(query, params) == inject_parameters( + query, pe.escape_args(params) + ) + probe.assert_not_called() + def test_probe_shaped_string_value_is_not_a_probe(self): """Real values are substituted while probing, so a string value that contains a probe-like token must not make another marker render as a list.""" From dd02cbdbd3864e708dc1661bbb0a20db830a795f Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Fri, 6 Nov 2026 22:49:40 +0000 Subject: [PATCH 4/7] Ignore an IN that ends a line comment; keep values out of probing Review bot follow-up: a `--` comment line that ends with the word IN made the marker on the next line render as a list. For example WHERE arr = -- tags we filter IN %(tags)s rendered as `arr = ('a','b')` instead of `arr = ARRAY('a','b')`, which breaks a query that runs on main. A match is now ignored when a `--` comes before its IN on the same line and the marker is on a later line. The probe pass also replaces escaped strings with '' (numbers stay, so %d still formats), so a value containing `--` cannot hide a real IN. Co-authored-by: Isaac Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- src/databricks/sql/utils.py | 29 ++++++++++++++++--- tests/unit/test_param_escaper.py | 48 ++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 4 deletions(-) diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index d7f43be78..ded2e49c7 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -577,6 +577,21 @@ def user_friendly_error_message(self, no_retry_reason, attempt, elapsed): _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) @@ -684,11 +699,11 @@ def _sequence_markers_following_in( 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 and - every other value keeps its escaped value, so scalar markers such as `%d` still format. + 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 a string value that happens to contain a probe-shaped token from being - # mistaken for a probe, since real values are substituted alongside the probes. + # 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: @@ -696,6 +711,11 @@ def probe_marker(index: int, key: Any) -> Any: 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: @@ -709,6 +729,7 @@ def probe_marker(index: int, key: Any) -> Any: 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): diff --git a/tests/unit/test_param_escaper.py b/tests/unit/test_param_escaper.py index 78f53371c..f357a8662 100644 --- a/tests/unit/test_param_escaper.py +++ b/tests/unit/test_param_escaper.py @@ -357,6 +357,54 @@ def test_sequence_elsewhere_stays_array(self, label, query, params, expected): def test_rendering_is_chosen_per_marker(self, label, query, params, expected): assert pe.render_inline(query, params) == expected + @pytest.mark.parametrize( + ("label", "query", "params", "expected"), + ( + ( + "IN ending a line comment does not change the marker on the next line", + "SELECT * FROM t WHERE arr = -- tags we filter IN\n %(tags)s", + {"tags": ["a", "b"]}, + "SELECT * FROM t WHERE arr = -- tags we filter IN\n ARRAY('a','b')", + ), + ( + "the same with CRLF line endings", + "SELECT * FROM t WHERE arr = -- filter IN\r\n %(tags)s", + {"tags": ["a"]}, + "SELECT * FROM t WHERE arr = -- filter IN\r\n ARRAY('a')", + ), + ( + "a value containing -- does not hide an IN that ends the line", + "SELECT * FROM t WHERE name = %(name)s AND id IN\n %(ids)s", + {"name": "a--b", "ids": [1, 2]}, + "SELECT * FROM t WHERE name = 'a--b' AND id IN\n (1,2)", + ), + ( + "a -- string literal on the line of IN and the marker", + "SELECT * FROM t WHERE sep = '--' AND id IN %(ids)s", + {"ids": [1, 2]}, + "SELECT * FROM t WHERE sep = '--' AND id IN (1,2)", + ), + ( + # Detection reads the SQL text without parsing string literals, so this falls + # back to the ARRAY rendering, as on main. + "a -- string literal before an IN that ends its line keeps ARRAY", + "SELECT * FROM t WHERE sep = '--' AND id IN\n %(ids)s", + {"ids": [1, 2]}, + "SELECT * FROM t WHERE sep = '--' AND id IN\n ARRAY(1,2)", + ), + ), + ) + def test_line_comments(self, label, query, params, expected): + assert pe.render_inline(query, params) == expected + + def test_single_digit_decimal_at_a_c_marker_still_renders_the_in_list(self): + query = "SELECT %(d)c AS digit FROM t WHERE x IN %(ids)s" + params = {"d": decimal.Decimal("5"), "ids": [1, 2]} + assert ( + pe.render_inline(query, params) + == "SELECT 5 AS digit FROM t WHERE x IN (1,2)" + ) + @pytest.mark.parametrize( ("label", "query", "params"), ( From a710fa845806ab02a7bf36acbaad60b65305b410 Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Fri, 6 Nov 2026 23:13:32 +0000 Subject: [PATCH 5/7] Accept a set directly after IN in inline parameters Up to 4.0.3 the inline escaper rendered any iterable as `(...)`, so a set after IN worked. 4.0.4 switched to Sequence for ARRAY rendering and sets started raising ProgrammingError. A set, frozenset or dict key view is accepted again when every marker it fills directly follows IN. Its items are sorted so the SQL text is stable, and a set of tuples renders as rows. Anywhere else a set raises the same ProgrammingError as before, since a set has no order to keep in an ARRAY. Also adds a CHANGELOG entry for the inline IN fix. Co-authored-by: Isaac Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- CHANGELOG.md | 1 + docs/parameters.md | 2 + src/databricks/sql/utils.py | 111 ++++++++++++++++++++---- tests/e2e/test_parameterized_queries.py | 6 ++ tests/unit/test_param_escaper.py | 84 +++++++++++++++++- 5 files changed, 184 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b4ae65801..4fbb64b68 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # Release History # Unreleased +- Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) # 4.9.0 (2026-11-04) - Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes. diff --git a/docs/parameters.md b/docs/parameters.md index 8047cf53c..a46155985 100644 --- a/docs/parameters.md +++ b/docs/parameters.md @@ -250,6 +250,8 @@ 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. +A `set` or `frozenset` can also be passed directly after `IN`. Its items are sorted, so the query text is the same every time. Anywhere else a set raises `ProgrammingError`, because a set has no order to keep in 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 diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index ded2e49c7..52af95913 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -17,7 +17,7 @@ import decimal from abc import ABC, abstractmethod from collections import OrderedDict, namedtuple -from collections.abc import Mapping +from collections.abc import Mapping, Set as AbstractSet from decimal import Decimal from enum import Enum import re @@ -597,6 +597,40 @@ def _renders_as_array(value: Any) -> bool: return isinstance(value, Sequence) and not isinstance(value, str) +def _is_set(value: Any) -> bool: + """Whether `value` is a set, frozenset or dict key view. A set has no order, so it can only + render as an IN list, never as an ARRAY.""" + return isinstance(value, AbstractSet) + + +def _in_list_items(value: Any) -> Any: + """Return the items to render for an IN list, with a set's items sorted so the SQL text is + the same on every run.""" + if not _is_set(value): + return value + try: + return sorted(value) + except TypeError: + # The items do not compare with each other, for example None and an int. + return sorted(value, key=repr) + + +class _InListOnly: + """Stands in for an escaped set, which can only be rendered as an IN list. + + `render_inline` rejects a set at any other marker before formatting. Raising here as well + keeps an unescaped set from ever reaching the SQL text. + """ + + def __init__(self, value: Any) -> None: + self.value = value + + def __str__(self) -> str: + raise ProgrammingError("Unsupported object {}".format(self.value)) + + __repr__ = __str__ + + class _MarkerProbe: """Formats as a token that identifies the marker it was substituted into.""" @@ -635,13 +669,29 @@ def __str__(self) -> str: __repr__ = __str__ +def _reject_sets_outside_in_lists( + parameters: Any, in_list_markers: Set[int], marker_keys: Dict[int, Any] +) -> None: + """Raise the error `escape_args` raises for a set, unless the set is used and every marker + it fills renders as an IN list.""" + items = ( + parameters.items() if isinstance(parameters, dict) else enumerate(parameters) + ) + for key, value in items: + if not _is_set(value): + continue + markers = {index for index, k in marker_keys.items() if k == key} + if not markers or not markers <= in_list_markers: + raise ProgrammingError("Unsupported object {}".format(value)) + + 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. + `escaped` holds the escaped parameter values. `key` is the marker's name when it is a dict and + the marker's position otherwise. """ values: Any if isinstance(escaped, dict): @@ -674,42 +724,67 @@ def render_inline(self, operation: str, parameters) -> str: 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. + + A set (or frozenset, or dict key view) is accepted only where every marker it fills + directly follows IN, and its items are sorted. Anywhere else it raises the same + `ProgrammingError` as `escape_args`, because a set has no order to put into an ARRAY. """ - escaped = self.escape_args(parameters) + escaped = self._escape_values(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. + if not any(_renders_as_array(value) or _is_set(value) for value in values): + # Only a sequence or a set 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( + in_list_markers, marker_keys = self._markers_following_in( operation, parameters, escaped ) + _reject_sets_outside_in_lists(parameters, in_list_markers, marker_keys) def render_marker(index: int, key: Any) -> Any: if index in in_list_markers: - return self.escape_in_list(parameters[key]) + return self.escape_in_list(_in_list_items(parameters[key])) return escaped[key] return _format_markers(operation, escaped, render_marker) - def _sequence_markers_following_in( + def _escape_values(self, parameters): + """Escape like `escape_args`, except that a set becomes an `_InListOnly` placeholder + instead of raising, so `render_inline` can decide where it is allowed.""" + + def escape(value: Any) -> Any: + return _InListOnly(value) if _is_set(value) else self.escape_item(value) + + if isinstance(parameters, dict): + return {k: escape(v) for k, v in parameters.items()} + elif isinstance(parameters, (list, tuple)): + return tuple(escape(x) for x in parameters) + else: + raise ProgrammingError("Unsupported param format: {}".format(parameters)) + + def _markers_following_in( self, operation: str, parameters, escaped - ) -> Set[int]: - """Return the numbers of the markers that directly follow IN and receive a sequence. + ) -> Tuple[Set[int], Dict[int, Any]]: + """Return the numbers of the markers that directly follow IN and receive a sequence or a + set, and the parameter key that each marker number refers to. 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. + positional markers line up exactly. Each sequence or set 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. + If the operation does not format, both results are empty. """ # The nonce stops text that only looks like a probe token from being taken for one. nonce = secrets.token_hex(8) + marker_keys: Dict[int, Any] = {} 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]): + marker_keys[index] = key + value = parameters[key] + if _renders_as_array(value) or _is_set(value): 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 @@ -723,14 +798,16 @@ def probe_marker(index: int, key: Any) -> Any: 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() + # A set gets the ProgrammingError escape_args raises, since none of its markers count. + return set(), {} - return { + in_list_markers = { 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) } + return in_list_markers, marker_keys def escape_in_list(self, item): return "(" + ",".join(map(str, map(self._escape_in_list_element, item))) + ")" diff --git a/tests/e2e/test_parameterized_queries.py b/tests/e2e/test_parameterized_queries.py index d690601b1..6ea06af2e 100644 --- a/tests/e2e/test_parameterized_queries.py +++ b/tests/e2e/test_parameterized_queries.py @@ -577,6 +577,12 @@ def test_params_as_sequence(self): {"pairs": [(1, "a"), (2, "x"), (3, "c")]}, [1, 3], ), + ( + # A set is accepted directly after IN and rendered sorted. + "SELECT id FROM range(5) WHERE id NOT IN %(ids)s ORDER BY id", + {"ids": {4, 0, 2}}, + [1, 3], + ), ], ) def test_inline_sequence_in_clause(self, query, params, expected_ids): diff --git a/tests/unit/test_param_escaper.py b/tests/unit/test_param_escaper.py index f357a8662..56b61cc65 100644 --- a/tests/unit/test_param_escaper.py +++ b/tests/unit/test_param_escaper.py @@ -3,6 +3,7 @@ from typing import Any, Dict from unittest.mock import patch from databricks.sql.parameters.native import dbsql_parameter_from_primitive +from databricks.sql.exc import ProgrammingError from databricks.sql.types import Row from databricks.sql.utils import ( @@ -437,7 +438,7 @@ def test_without_sequences_renders_as_escape_args_without_probing( self, query, params ): with patch.object( - ParamEscaper, "_sequence_markers_following_in", side_effect=AssertionError + ParamEscaper, "_markers_following_in", side_effect=AssertionError ) as probe: assert pe.render_inline(query, params) == inject_parameters( query, pe.escape_args(params) @@ -445,8 +446,8 @@ def test_without_sequences_renders_as_escape_args_without_probing( probe.assert_not_called() def test_probe_shaped_string_value_is_not_a_probe(self): - """Real values are substituted while probing, so a string value that contains a probe-like - token must not make another marker render as a list.""" + """A string value that contains a probe-like token must not make another marker render as + a list.""" query = "INSERT INTO t VALUES (%(note)s, %(tags)s)" params = {"note": "x IN \x000000000000000000:1\x00", "tags": ["a"]} assert pe.render_inline(query, params) == ( @@ -477,6 +478,83 @@ def test_errors_match_escape_args(self, query, params, error, message): with pytest.raises(error, match=message): pe.render_inline(query, params) + @pytest.mark.parametrize( + ("label", "query", "params", "expected"), + ( + ( + "set, sorted", + "SELECT * FROM t WHERE x IN %(ids)s", + {"ids": {3, 1, 2}}, + "SELECT * FROM t WHERE x IN (1,2,3)", + ), + ( + "frozenset, positional, NOT IN", + "SELECT * FROM t WHERE x NOT IN %s", + (frozenset({2, 1}),), + "SELECT * FROM t WHERE x NOT IN (1,2)", + ), + ( + "dict keys", + "SELECT * FROM t WHERE name IN %(names)s", + {"names": {"b": 1, "a": 2}.keys()}, + "SELECT * FROM t WHERE name IN ('a','b')", + ), + ( + "items that do not compare are sorted by repr", + "SELECT * FROM t WHERE x IN %(ids)s", + {"ids": {None, 1}}, + "SELECT * FROM t WHERE x IN (1,NULL)", + ), + ( + "set of tuples renders sorted rows", + "SELECT * FROM t WHERE (a, b) IN %(pairs)s", + {"pairs": {(2, "b"), (1, "a")}}, + "SELECT * FROM t WHERE (a, b) IN ((1,'a'),(2,'b'))", + ), + ( + "same set after IN twice", + "SELECT * FROM t WHERE x IN %(ids)s OR y IN %(ids)s", + {"ids": {2, 1}}, + "SELECT * FROM t WHERE x IN (1,2) OR y IN (1,2)", + ), + ( + "set next to an array elsewhere", + "SELECT %(tags)s AS tags FROM t WHERE x IN %(ids)s", + {"tags": ["a"], "ids": {1}}, + "SELECT ARRAY('a') AS tags FROM t WHERE x IN (1)", + ), + ), + ) + def test_set_after_in_renders_as_sorted_list(self, label, query, params, expected): + assert pe.render_inline(query, params) == expected + + @pytest.mark.parametrize( + ("label", "query", "params"), + ( + ( + "set at a marker not after IN", + "INSERT INTO t VALUES (%(s)s)", + {"s": {1}}, + ), + ( + "set used after IN and elsewhere", + "SELECT %(ids)s FROM t WHERE x IN %(ids)s", + {"ids": {1}}, + ), + ("unused set", "SELECT * FROM t WHERE x = %(a)s", {"a": 1, "s": {1}}), + ("set after IN in caller-written parentheses", "SELECT 1 IN (%s)", ({1},)), + ("set inside a list", "SELECT * FROM t WHERE x IN %(v)s", {"v": [{1}]}), + ("set at a %d marker after IN", "SELECT * FROM t WHERE x IN %d", ({1},)), + ("set in an operation that does not format", "SELECT %s, %s", ({1},)), + ("set in an IN that ends a line comment", "SELECT -- IN\n %s", ({1},)), + ), + ) + def test_set_elsewhere_raises_as_before(self, label, query, params): + with pytest.raises(ProgrammingError, match="Unsupported object"): + inject_parameters(query, pe.escape_args(params)) + with pytest.raises(ProgrammingError, match="Unsupported object"): + pe.render_inline(query, params) + class TestInlineToNativeTransformer(object): @pytest.mark.parametrize( From 7689b3733ad52151841408a27bbaefcad4d6f27e Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Fri, 6 Nov 2026 23:39:37 +0000 Subject: [PATCH 6/7] Describe IN detection as a text heuristic Review bot follow-up: the comments and docs described comment handling as if it were settled. They now say that IN is found by matching the SQL text rather than parsing it, which case is handled (an IN that ends a `--` comment line) and which is not (an IN inside a string literal or block comment directly before a marker). Co-authored-by: Isaac Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- docs/parameters.md | 2 ++ src/databricks/sql/utils.py | 18 ++++++++++++++---- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/docs/parameters.md b/docs/parameters.md index a46155985..cd70876b1 100644 --- a/docs/parameters.md +++ b/docs/parameters.md @@ -237,6 +237,8 @@ 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. +The connector finds `IN` by reading the query text, not by parsing SQL. It ignores an `IN` at the end of a `--` comment line, but an `IN` inside a string literal, directly before a marker, is still treated as the keyword. + 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 diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 52af95913..8e858502e 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -574,15 +574,21 @@ def user_friendly_error_message(self, no_retry_reason, attempt, elapsed): # `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. +# +# This matches the SQL text rather than parsing it, so it is a heuristic. It skips an IN that ends +# a `--` comment line (see `_in_ends_line_comment`), but an IN inside a string literal or a block +# comment can still count as the keyword when a marker follows it. _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. + """Whether the IN that `match` found looks like it is in a `--` comment, with the marker 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. + This is a check on the text, not a parse. In that case the marker is most likely live code, so + it keeps its ARRAY rendering. A marker on the same line as a commented IN is usually inside the + comment too, so it is not checked. A `--` inside a string literal earlier on the line also + counts, which falls back to the ARRAY rendering. """ if not re.search(r"[\r\n]", match.group(0)): return False @@ -725,6 +731,10 @@ def render_inline(self, operation: str, parameters) -> str: ARRAY rendering at the other markers. Inside the list, tuples render as rows and lists as ARRAY literals. + IN is found by matching the SQL text, not by parsing it (see `_IN_MARKER_PROBE`), so an + unusual query, such as one with IN inside a string literal just before a marker, can + still get a list where an ARRAY was meant. + A set (or frozenset, or dict key view) is accepted only where every marker it fills directly follows IN, and its items are sorted. Anywhere else it raises the same `ProgrammingError` as `escape_args`, because a set has no order to put into an ARRAY. From 7256b6b63a34d963ef40263ef52b9b11ab16a4b4 Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Fri, 6 Nov 2026 23:41:43 +0000 Subject: [PATCH 7/7] Don't promise identical SQL text for sets Review bot follow-up: sorting a set's items does not give a stable order when items compare inconsistently, such as float("nan") among other numbers. The SQL still matches the same rows, but the docstring and docs/parameters.md no longer say the query text is always the same. Co-authored-by: Isaac Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- docs/parameters.md | 2 +- src/databricks/sql/utils.py | 8 ++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/docs/parameters.md b/docs/parameters.md index cd70876b1..578e75cfc 100644 --- a/docs/parameters.md +++ b/docs/parameters.md @@ -252,7 +252,7 @@ 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. -A `set` or `frozenset` can also be passed directly after `IN`. Its items are sorted, so the query text is the same every time. Anywhere else a set raises `ProgrammingError`, because a set has no order to keep in an `ARRAY(...)` literal. +A `set` or `frozenset` can also be passed directly after `IN`. Its items are sorted, which usually keeps the query text the same from run to run. Items without a consistent order, such as `float("nan")`, can still change position. Anywhere else a set raises `ProgrammingError`, because a set has no order to keep in an `ARRAY(...)` literal. **Note**: this behavior is not specified by PEP-249 and only works due to the specific implementation of our inline renderer. diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 8e858502e..80e90005e 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -610,8 +610,12 @@ def _is_set(value: Any) -> bool: def _in_list_items(value: Any) -> Any: - """Return the items to render for an IN list, with a set's items sorted so the SQL text is - the same on every run.""" + """Return the items to render for an IN list, with a set's items sorted. + + Sorting usually makes the SQL text the same from run to run, but not always: items without a + consistent order, such as `float("nan")` among other numbers, can still come out in a + different order. The IN list matches the same rows either way. + """ if not _is_set(value): return value try: