-
Notifications
You must be signed in to change notification settings - Fork 140
Expand file tree
/
Copy pathtest_param_escaper.py
More file actions
249 lines (205 loc) ยท 8.29 KB
/
test_param_escaper.py
File metadata and controls
249 lines (205 loc) ยท 8.29 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
from datetime import date, datetime
import unittest, pytest, decimal
from typing import Any, Dict
from databricks.sql.parameters.native import dbsql_parameter_from_primitive
from databricks.sql.utils import (
ParamEscaper,
inject_parameters,
transform_paramstyle,
ParameterStructure,
)
pe = ParamEscaper()
class TestIndividualFormatters(object):
# Test individual type escapers
def test_escape_number_integer(self):
"""This behaviour falls back to Python's default string formatting of numbers"""
assert pe.escape_number(100) == 100
def test_escape_number_float(self):
"""This behaviour falls back to Python's default string formatting of numbers"""
assert pe.escape_number(100.1234) == 100.1234
def test_escape_number_decimal(self):
"""This behaviour uses the string representation of a decimal"""
assert pe.escape_decimal(decimal.Decimal("124.32")) == "124.32"
def test_escape_string_normal(self):
""" """
assert pe.escape_string("golly bob howdy") == "'golly bob howdy'"
def test_escape_string_that_includes_special_characters(self):
r"""Tests for how special characters are treated.
When passed a string, the `escape_string` method wraps it in single quotes
and escapes any special characters with a back stroke (\)
Example:
IN : his name was 'robert palmer'
OUT: 'his name was \'robert palmer\''
"""
# Testing for the presence of these characters: '"/\๐
assert (
pe.escape_string("his name was 'robert palmer'")
== r"'his name was \'robert palmer\''"
)
# These tests represent the same user input in the several ways it can be written in Python
# Each argument to `escape_string` evaluates to the same bytes. But Python lets us write it differently.
assert (
pe.escape_string('his name was "robert palmer"')
== "'his name was \"robert palmer\"'"
)
assert (
pe.escape_string('his name was "robert palmer"')
== "'his name was \"robert palmer\"'"
)
assert (
pe.escape_string("his name was {}".format('"robert palmer"'))
== "'his name was \"robert palmer\"'"
)
assert (
pe.escape_string("his name was robert / palmer")
== r"'his name was robert / palmer'"
)
# If you need to include a single backslash, use an r-string to prevent Python from raising a
# DeprecationWarning for an invalid escape sequence
assert (
pe.escape_string("his name was robert \\/ palmer")
== r"'his name was robert \\/ palmer'"
)
assert (
pe.escape_string("his name was robert \\ palmer")
== r"'his name was robert \\ palmer'"
)
assert (
pe.escape_string("his name was robert \\\\ palmer")
== r"'his name was robert \\\\ palmer'"
)
assert (
pe.escape_string("his name was robert palmer ๐")
== r"'his name was robert palmer ๐'"
)
# Adding the test from PR #56 to prove escape behaviour
assert pe.escape_string("you're") == r"'you\'re'"
# Adding this test from #51 to prove escape behaviour when the target string involves repeated SQL escape chars
assert pe.escape_string("cat\\'s meow") == r"'cat\\\'s meow'"
# Tests from the docs: https://docs.databricks.com/sql/language-manual/data-types/string-type.html
assert pe.escape_string("Spark") == "'Spark'"
assert pe.escape_string("O'Connell") == r"'O\'Connell'"
assert pe.escape_string("Some\\nText") == r"'Some\\nText'"
assert pe.escape_string("Some\\\\nText") == r"'Some\\\\nText'"
assert pe.escape_string("์์ธ์") == "'์์ธ์'"
assert pe.escape_string("\\\\") == r"'\\\\'"
def test_escape_date_time(self):
INPUT = datetime(1991, 8, 3, 21, 55)
FORMAT = "%Y-%m-%d %H:%M:%S"
OUTPUT = "'1991-08-03 21:55:00'"
assert pe.escape_datetime(INPUT, FORMAT) == OUTPUT
def test_escape_date(self):
INPUT = date(1991, 8, 3)
FORMAT = "%Y-%m-%d"
OUTPUT = "'1991-08-03'"
assert pe.escape_datetime(INPUT, FORMAT) == OUTPUT
def test_escape_sequence_integer(self):
assert pe.escape_sequence([1, 2, 3, 4]) == "ARRAY(1,2,3,4)"
def test_escape_sequence_float(self):
assert pe.escape_sequence([1.1, 2.2, 3.3, 4.4]) == "ARRAY(1.1,2.2,3.3,4.4)"
def test_escape_sequence_string(self):
assert (
pe.escape_sequence(["his", "name", "was", "robert", "palmer"])
== "ARRAY('his','name','was','robert','palmer')"
)
def test_escape_sequence_sequence_of_strings(self):
INPUT = [["his", "name"], ["was", "robert"], ["palmer"]]
OUTPUT = "ARRAY(ARRAY('his','name'),ARRAY('was','robert'),ARRAY('palmer'))"
assert pe.escape_sequence(INPUT) == OUTPUT
def test_escape_map_string_int(self):
INPUT = {"a": 1, "b": 2}
OUTPUT = "MAP('a',1,'b',2)"
assert pe.escape_mapping(INPUT) == OUTPUT
def test_escape_map_string_sequence_of_floats(self):
INPUT = {"a": [1.1, 2.2, 3.3], "b": [4.4, 5.5, 6.6]}
OUTPUT = "MAP('a',ARRAY(1.1,2.2,3.3),'b',ARRAY(4.4,5.5,6.6))"
assert pe.escape_mapping(INPUT) == OUTPUT
def test_escape_sequence_of_map_int_string(self):
INPUT = [{1: "a", 2: "foo"}, {3: "b", 4: "bar"}]
OUTPUT = "ARRAY(MAP(1,'a',2,'foo'),MAP(3,'b',4,'bar'))"
assert pe.escape_sequence(INPUT) == OUTPUT
class TestFullQueryEscaping(object):
def test_simple(self):
INPUT = """
SELECT
field1,
field2,
field3
FROM
table
WHERE
field1 = %(param1)s
"""
OUTPUT = """
SELECT
field1,
field2,
field3
FROM
table
WHERE
field1 = ';DROP ALL TABLES'
"""
args = {"param1": ";DROP ALL TABLES"}
assert inject_parameters(INPUT, pe.escape_args(args)) == OUTPUT
@unittest.skipUnless(False, "Thrift server supports native parameter binding.")
def test_only_bind_in_where_clause(self):
INPUT = """
SELECT
%(field)s,
field2,
field3
FROM table
"""
args = {"field": "Some Value"}
with pytest.raises(Exception):
inject_parameters(INPUT, pe.escape_args(args))
class TestInlineToNativeTransformer(object):
@pytest.mark.parametrize(
("label", "query", "params", "expected"),
(
("no effect", "SELECT 1", {}, "SELECT 1"),
("one marker", "%(param)s", {"param": ""}, ":param"),
(
"multiple markers",
"%(foo)s %(bar)s %(baz)s",
{"foo": None, "bar": None, "baz": None},
":foo :bar :baz",
),
(
"sql query",
"SELECT * FROM table WHERE field = %(param)s AND other_field IN (%(list)s)",
{"param": None, "list": None},
"SELECT * FROM table WHERE field = :param AND other_field IN (:list)",
),
(
"query with like wildcard",
'select * from table where field like "%"',
{},
'select * from table where field like "%"',
),
(
"query with named param and like wildcard",
'select :param from table where field like "%"',
{"param": None},
'select :param from table where field like "%"',
),
(
"query with doubled wildcards",
"select 1 where " ' like "%%"',
{"param": None},
"select 1 where " ' like "%%"',
),
),
)
def test_transformer(
self, label: str, query: str, params: Dict[str, Any], expected: str
):
_params = [
dbsql_parameter_from_primitive(value=value, name=name)
for name, value in params.items()
]
output = transform_paramstyle(
query, _params, param_structure=ParameterStructure.NAMED
)
assert output == expected