forked from databricks/databricks-sql-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_param_escaper.py
More file actions
155 lines (114 loc) Β· 5.36 KB
/
test_param_escaper.py
File metadata and controls
155 lines (114 loc) Β· 5.36 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
from datetime import date, datetime
import unittest, pytest, decimal
from databricks.sql.utils import ParamEscaper, inject_parameters
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):
"""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]) == "(1,2,3,4)"
def test_escape_sequence_float(self):
assert pe.escape_sequence([1.1,2.2,3.3,4.4]) == "(1.1,2.2,3.3,4.4)"
def test_escape_sequence_string(self):
assert pe.escape_sequence(
["his", "name", "was", "robert", "palmer"]) == \
"('his','name','was','robert','palmer')"
def test_escape_sequence_sequence_of_strings(self):
# This is not valid SQL.
INPUT = [["his", "name"], ["was", "robert"], ["palmer"]]
OUTPUT = "(('his','name'),('was','robert'),('palmer'))"
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))