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"),