forked from databricks/databricks-sql-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_parameterized_queries.py
More file actions
450 lines (380 loc) · 17.4 KB
/
test_parameterized_queries.py
File metadata and controls
450 lines (380 loc) · 17.4 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
import datetime
from contextlib import contextmanager
from decimal import Decimal
from enum import Enum
from typing import Dict, List, Type, Union
from unittest.mock import patch
import pytest
import pytz
from databricks.sql.parameters.native import (
BigIntegerParameter,
BooleanParameter,
DateParameter,
DbsqlParameterBase,
DecimalParameter,
DoubleParameter,
FloatParameter,
IntegerParameter,
ParameterApproach,
ParameterStructure,
SmallIntParameter,
StringParameter,
TDbsqlParameter,
TimestampNTZParameter,
TimestampParameter,
TinyIntParameter,
VoidParameter,
)
from tests.e2e.test_driver import PySQLPytestTestCase
class ParamStyle(Enum):
NAMED = 1
PYFORMAT = 2
NONE = 3
class Primitive(Enum):
"""These are the inferrable types. This Enum is used for parametrized tests."""
NONE = None
BOOL = True
INT = 50
BIGINT = 2147483648
STRING = "Hello"
DECIMAL = Decimal("1234.56")
DATE = datetime.date(2023, 9, 6)
TIMESTAMP = datetime.datetime(2023, 9, 6, 3, 14, 27, 843, tzinfo=pytz.UTC)
DOUBLE = 3.14
FLOAT = 3.15
SMALLINT = 51
class PrimitiveExtra(Enum):
"""These are not inferrable types. This Enum is used for parametrized tests."""
TIMESTAMP_NTZ = datetime.datetime(2023, 9, 6, 3, 14, 27, 843)
TINYINT = 20
# We don't test inline approach with named paramstyle because it's never supported
# We don't test inline approach with positional parameters because it's never supported
# Paramstyle doesn't apply when ParameterStructure.POSITIONAL because question marks are used.
approach_paramstyle_combinations = [
(ParameterApproach.INLINE, ParamStyle.PYFORMAT, ParameterStructure.NAMED),
(ParameterApproach.NATIVE, ParamStyle.NONE, ParameterStructure.POSITIONAL),
(ParameterApproach.NATIVE, ParamStyle.PYFORMAT, ParameterStructure.NAMED),
(ParameterApproach.NATIVE, ParamStyle.NONE, ParameterStructure.POSITIONAL),
(ParameterApproach.NATIVE, ParamStyle.NAMED, ParameterStructure.NAMED),
]
class TestParameterizedQueries(PySQLPytestTestCase):
"""Namespace for tests of this connector's parameterisation behaviour.
databricks-sql-connector can approach parameterisation in two ways:
NATIVE: the connector will use server-side bound parameters implemented by DBR 14.1 and above.
INLINE: the connector will render parameter values as strings and interpolate them into the query.
Prior to connector version 3.0.0, the connector would always use the INLINE approach. This approach
is still the default but this will be changed in a subsequent release.
The INLINE and NATIVE approaches use different query syntax, which these tests verify.
There is not 1-to-1 feature parity between these approaches. Where possible, we run the same test
for @both_approaches.
"""
NAMED_PARAMSTYLE_QUERY = "SELECT :p AS col"
PYFORMAT_PARAMSTYLE_QUERY = "SELECT %(p)s AS col"
POSITIONAL_PARAMSTYLE_QUERY = "SELECT ? AS col"
inline_type_map = {
Primitive.INT: "int_col",
Primitive.BIGINT: "bigint_col",
Primitive.SMALLINT: "small_int_col",
Primitive.FLOAT: "float_col",
Primitive.DOUBLE: "double_col",
Primitive.DECIMAL: "decimal_col",
Primitive.STRING: "string_col",
Primitive.BOOL: "boolean_col",
Primitive.DATE: "date_col",
Primitive.TIMESTAMP: "timestamp_col",
Primitive.NONE: "null_col",
}
def _get_inline_table_column(self, value):
return self.inline_type_map[Primitive(value)]
@pytest.fixture(scope="class")
def inline_table(self):
"""This table is necessary to verify that a parameter sent with INLINE
approach can actually write to its analogous data type.
For example, a Python Decimal(), when rendered inline, should be able
to read/write into a DECIMAL column in Databricks
Note that this fixture doesn't clean itself up. So the table will remain
in the schema for use by subsequent test runs.
"""
query = """
CREATE TABLE IF NOT EXISTS pysql_e2e_inline_param_test_table (
null_col INT,
int_col INT,
bigint_col BIGINT,
small_int_col SMALLINT,
float_col FLOAT,
double_col DOUBLE,
decimal_col DECIMAL(10, 2),
string_col STRING,
boolean_col BOOLEAN,
date_col DATE,
timestamp_col TIMESTAMP
) USING DELTA
"""
with self.connection() as conn:
with conn.cursor() as cursor:
cursor.execute(query)
@contextmanager
def patch_server_supports_native_params(self, supports_native_params: bool = True):
"""Applies a patch so we can test the connector's behaviour under different SPARK_CLI_SERVICE_PROTOCOL_VERSION conditions."""
with patch(
"databricks.sql.client.Connection.server_parameterized_queries_enabled",
return_value=supports_native_params,
) as mock_parameterized_queries_enabled:
try:
yield mock_parameterized_queries_enabled
finally:
pass
def _inline_roundtrip(self, params: dict, paramstyle: ParamStyle):
"""This INSERT, SELECT, DELETE dance is necessary because simply selecting
```
"SELECT %(param)s"
```
in INLINE mode would always return a str and the nature of the test is to
confirm that types are maintained.
:paramstyle:
This is a no-op but is included to make the test-code easier to read.
"""
target_column = self._get_inline_table_column(params.get("p"))
INSERT_QUERY = f"INSERT INTO pysql_e2e_inline_param_test_table (`{target_column}`) VALUES (%(p)s)"
SELECT_QUERY = f"SELECT {target_column} `col` FROM pysql_e2e_inline_param_test_table LIMIT 1"
DELETE_QUERY = "DELETE FROM pysql_e2e_inline_param_test_table"
with self.connection(extra_params={"use_inline_params": True}) as conn:
with conn.cursor() as cursor:
cursor.execute(INSERT_QUERY, parameters=params)
with conn.cursor() as cursor:
to_return = cursor.execute(SELECT_QUERY).fetchone()
with conn.cursor() as cursor:
cursor.execute(DELETE_QUERY)
return to_return
def _native_roundtrip(
self,
parameters: Union[Dict, List[Dict]],
paramstyle: ParamStyle,
parameter_structure: ParameterStructure,
):
if parameter_structure == ParameterStructure.POSITIONAL:
_query = self.POSITIONAL_PARAMSTYLE_QUERY
elif paramstyle == ParamStyle.NAMED:
_query = self.NAMED_PARAMSTYLE_QUERY
elif paramstyle == ParamStyle.PYFORMAT:
_query = self.PYFORMAT_PARAMSTYLE_QUERY
with self.connection(extra_params={"use_inline_params": False}) as conn:
with conn.cursor() as cursor:
cursor.execute(_query, parameters=parameters)
return cursor.fetchone()
def _get_one_result(
self,
params,
approach: ParameterApproach = ParameterApproach.NONE,
paramstyle: ParamStyle = ParamStyle.NONE,
parameter_structure: ParameterStructure = ParameterStructure.NONE,
):
"""When approach is INLINE then we use %(param)s paramstyle and a connection with use_inline_params=True
When approach is NATIVE then we use :param paramstyle and a connection with use_inline_params=False
"""
if approach == ParameterApproach.INLINE:
# inline mode always uses ParamStyle.PYFORMAT
# inline mode doesn't support positional parameters
return self._inline_roundtrip(params, paramstyle=ParamStyle.PYFORMAT)
elif approach == ParameterApproach.NATIVE:
# native mode can use either ParamStyle.NAMED or ParamStyle.PYFORMAT
# native mode can use either ParameterStructure.NAMED or ParameterStructure.POSITIONAL
return self._native_roundtrip(
params, paramstyle=paramstyle, parameter_structure=parameter_structure
)
def _quantize(self, input: Union[float, int], place_value=2) -> Decimal:
return Decimal(str(input)).quantize(Decimal("0." + "0" * place_value))
def _eq(self, actual, expected: Primitive):
"""This is a helper function to make the test code more readable.
If primitive is Primitive.DOUBLE than an extra quantize step is performed before
making the assertion.
"""
if expected in (Primitive.DOUBLE, Primitive.FLOAT):
return self._quantize(actual) == self._quantize(expected.value)
return actual == expected.value
@pytest.mark.parametrize("primitive", Primitive)
@pytest.mark.parametrize(
"approach,paramstyle,parameter_structure", approach_paramstyle_combinations
)
def test_primitive_single(
self,
approach,
paramstyle,
parameter_structure,
primitive: Primitive,
inline_table,
):
"""When ParameterApproach.INLINE is passed, inferrence will not be used.
When ParameterApproach.NATIVE is passed, primitive inputs will be inferred.
"""
if parameter_structure == ParameterStructure.NAMED:
params = {"p": primitive.value}
elif parameter_structure == ParameterStructure.POSITIONAL:
params = [primitive.value]
result = self._get_one_result(params, approach, paramstyle, parameter_structure)
assert self._eq(result.col, primitive)
@pytest.mark.parametrize(
"parameter_structure", (ParameterStructure.NAMED, ParameterStructure.POSITIONAL)
)
@pytest.mark.parametrize(
"primitive,dbsql_parameter_cls",
[
(Primitive.NONE, VoidParameter),
(Primitive.BOOL, BooleanParameter),
(Primitive.INT, IntegerParameter),
(Primitive.BIGINT, BigIntegerParameter),
(Primitive.STRING, StringParameter),
(Primitive.DECIMAL, DecimalParameter),
(Primitive.DATE, DateParameter),
(Primitive.TIMESTAMP, TimestampParameter),
(Primitive.DOUBLE, DoubleParameter),
(Primitive.FLOAT, FloatParameter),
(Primitive.SMALLINT, SmallIntParameter),
(PrimitiveExtra.TIMESTAMP_NTZ, TimestampNTZParameter),
(PrimitiveExtra.TINYINT, TinyIntParameter)
],
)
def test_dbsqlparameter_single(
self,
primitive: Primitive,
dbsql_parameter_cls: Type[TDbsqlParameter],
parameter_structure: ParameterStructure,
):
dbsql_param = dbsql_parameter_cls(
value=primitive.value, # type: ignore
name="p" if parameter_structure == ParameterStructure.NAMED else None,
)
params = [dbsql_param]
result = self._get_one_result(
params, ParameterApproach.NATIVE, ParamStyle.NAMED, parameter_structure
)
assert self._eq(result.col, primitive)
@pytest.mark.parametrize("use_inline_params", (True, False, "silent"))
def test_use_inline_off_by_default_with_warning(self, use_inline_params, caplog):
"""
use_inline_params should be False by default.
If a user explicitly sets use_inline_params, don't warn them about it.
"""
extra_args = (
{"use_inline_params": use_inline_params} if use_inline_params else {}
)
with self.connection(extra_params=extra_args) as conn:
with conn.cursor() as cursor:
with self.patch_server_supports_native_params(
supports_native_params=True
):
cursor.execute("SELECT %(p)s", parameters={"p": 1})
if use_inline_params is True:
assert (
"Consider using native parameters." in caplog.text
), "Log message should be suppressed"
elif use_inline_params == "silent":
assert (
"Consider using native parameters." not in caplog.text
), "Log message should not be supressed"
def test_positional_native_params_with_defaults(self):
query = "SELECT ? col"
with self.cursor() as cursor:
result = cursor.execute(query, parameters=[1]).fetchone()
assert result.col == 1
@pytest.mark.parametrize(
"params",
(
[
StringParameter(value="foo"),
StringParameter(value="bar"),
StringParameter(value="baz"),
],
["foo", "bar", "baz"],
),
)
def test_positional_native_multiple(self, params):
query = "SELECT ? `foo`, ? `bar`, ? `baz`"
with self.cursor(extra_params={"use_inline_params": False}) as cursor:
result = cursor.execute(query, params).fetchone()
expected = [i.value if isinstance(i, DbsqlParameterBase) else i for i in params]
outcome = [result.foo, result.bar, result.baz]
assert set(outcome) == set(expected)
def test_readme_example(self):
with self.cursor() as cursor:
result = cursor.execute('SELECT :param `p`, * FROM RANGE(10)', {"param": "foo"}).fetchall()
assert len(result) == 10
assert result[0].p == "foo"
class TestInlineParameterSyntax(PySQLPytestTestCase):
"""The inline parameter approach uses pyformat markers"""
def test_params_as_dict(self):
query = "SELECT %(foo)s foo, %(bar)s bar, %(baz)s baz"
params = {"foo": 1, "bar": 2, "baz": 3}
with self.connection(extra_params={"use_inline_params": True}) as conn:
with conn.cursor() as cursor:
result = cursor.execute(query, parameters=params).fetchone()
assert result.foo == 1
assert result.bar == 2
assert result.baz == 3
def test_params_as_sequence(self):
"""One side-effect of ParamEscaper using Python string interpolation to inline the values
is that it can work with "ordinal" parameters, but only if a user writes parameter markers
that are not defined with PEP-249. This test exists to prove that it works in the ideal case.
"""
# `%s` is not a valid paramstyle per PEP-249
query = "SELECT %s foo, %s bar, %s baz"
params = (1, 2, 3)
with self.connection(extra_params={"use_inline_params": True}) as conn:
with conn.cursor() as cursor:
result = cursor.execute(query, parameters=params).fetchone()
assert result.foo == 1
assert result.bar == 2
assert result.baz == 3
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
just proves that's the case.
"""
query = "SELECT 'samsonite', %s WHERE 'samsonite' LIKE '%sonite'"
params = ["luggage"]
with self.cursor(extra_params={"use_inline_params": True}) as cursor:
with pytest.raises(
TypeError, match="not enough arguments for format string"
):
cursor.execute(query, parameters=params)
def test_inline_named_dont_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
just proves that's the case.
"""
query = """
with base as (SELECT 'x(one)sonite' as `col_1`)
SELECT col_1 FROM base WHERE col_1 LIKE CONCAT(%(one)s, 'onite')
"""
params = {"one": "%(one)s"}
with self.cursor(extra_params={"use_inline_params": True}) as cursor:
result = cursor.execute(query, parameters=params).fetchone()
print("hello")
def test_native_ordinals_dont_break_sql(self):
"""This test accompanies test_inline_ordinals_can_break_sql to prove that ordinal
parameters work in native mode for the exact same query, if we use the right marker `?`
"""
query = "SELECT 'samsonite', ? WHERE 'samsonite' LIKE '%sonite'"
params = ["luggage"]
with self.cursor(extra_params={"use_inline_params": False}) as cursor:
result = cursor.execute(query, parameters=params).fetchone()
assert result.samsonite == "samsonite"
assert result.luggage == "luggage"
def test_inline_like_wildcard_breaks(self):
"""One flaw with the ParameterEscaper is that it fails if a query contains
a SQL LIKE wildcard %. This test proves that's the case.
"""
query = "SELECT 1 `col` WHERE 'foo' LIKE '%'"
params ={"param": 'bar'}
with self.cursor(extra_params={"use_inline_params": True}) as cursor:
with pytest.raises(ValueError, match="unsupported format character"):
result = cursor.execute(query, parameters=params).fetchone()
def test_native_like_wildcard_works(self):
"""This is a mirror of test_inline_like_wildcard_breaks that proves that LIKE
wildcards work under the native approach.
"""
query = "SELECT 1 `col` WHERE 'foo' LIKE '%'"
params ={"param": 'bar'}
with self.cursor(extra_params={"use_inline_params": False}) as cursor:
result = cursor.execute(query, parameters=params).fetchone()
assert result.col == 1