Skip to content

Commit 2715867

Browse files
damncabbagedamncabbage
authored
Merge pull request #277 from cipherstash/rh/test/python/jsonb-operators
Adds JSONB tests to the Python integration suite
2 parents 7215bdd + 3916860 commit 2715867

8 files changed

+1074
-0
lines changed

‎tests/python/tests/test_jsonb_array_elements.py‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
import json
2+
import os
3+
import psycopg
4+
import random
5+
from itertools import product
6+
7+
import pytest
8+
9+
conn_params = {
10+
"user": os.environ.get("CS_DATABASE__USERNAME"),
11+
"password": os.environ.get("CS_DATABASE__PASSWORD"),
12+
"dbname": os.environ.get("CS_DATABASE__NAME"),
13+
"host": os.environ.get("CS_DATABASE__HOST"),
14+
"port": 6432,
15+
}
16+
17+
connection_str = psycopg.conninfo.make_conninfo(**conn_params)
18+
19+
print("Connection to Proxy with {}".format(connection_str))
20+
21+
# Common test data
22+
val = {
23+
"key": "value",
24+
"number": 42,
25+
"array_number": [3, 2, 1],
26+
"array_string": ["hello", "world"],
27+
"nested": {"foo": "bar", "number": 1312},
28+
}
29+
30+
31+
def test_numbers():
32+
select_jsonb("encrypted_jsonb", val, "$.array_number[@]",
33+
[3, 2, 1])
34+
35+
36+
def test_strings():
37+
select_jsonb("encrypted_jsonb", val, "$.array_string[@]",
38+
["hello", "world"])
39+
40+
41+
def test_with_unknown():
42+
select_jsonb("encrypted_jsonb", val, "$.nonexistent",
43+
[])
44+
45+
46+
def select_jsonb(column, value, selector, expected, alias=None):
47+
alias = "AS {}".format(alias) if alias is not None else ""
48+
tests = [
49+
("jsonb_array_elements(jsonb_path_query({}, '{}')) {}".format(column, selector, alias), []),
50+
("jsonb_array_elements(jsonb_path_query({}, %s)) {}".format(column, alias), [selector]),
51+
]
52+
53+
for (select_fragment, params) in tests:
54+
print("Testing fragment: {}, params: {}, expecting: {}".format(
55+
select_fragment, params, expected))
56+
57+
for (binary, prepare) in product([None, True], repeat=2):
58+
execute(json.dumps(value), column,
59+
select_fragment=select_fragment,
60+
select_params=params,
61+
expected=expected,
62+
binary=binary,
63+
prepare=prepare)
64+
65+
66+
def make_id():
67+
return random.randrange(1, 1000000000)
68+
69+
70+
def execute(val, column, binary=None, prepare=None, expected=None,
71+
select_fragment=None, select_params=[]):
72+
with psycopg.connect(connection_str, autocommit=True) as conn:
73+
74+
with conn.cursor() as cursor:
75+
76+
with conn.transaction():
77+
id = make_id()
78+
79+
print("... for column {}, with binary: {}, prepare: {}".format(
80+
column, binary, prepare))
81+
82+
sql = "INSERT INTO encrypted (id, {}) VALUES (%s, %s)".format(
83+
column)
84+
85+
cursor.execute(sql, [id, val], binary=binary, prepare=prepare)
86+
87+
sql = "SELECT {} FROM encrypted WHERE id = %s".format(
88+
select_fragment)
89+
cursor.execute(
90+
sql, (select_params + [id]),
91+
binary=binary, prepare=prepare)
92+
93+
rows = list(map(
94+
lambda row_tuple: row_tuple[0], cursor.fetchall()))
95+
96+
rows.sort()
97+
expected.sort()
98+
99+
assert rows == expected

‎tests/python/tests/test_jsonb_array_length.py‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
import json
2+
import os
3+
import psycopg
4+
import random
5+
from itertools import product
6+
7+
conn_params = {
8+
"user": os.environ.get("CS_DATABASE__USERNAME"),
9+
"password": os.environ.get("CS_DATABASE__PASSWORD"),
10+
"dbname": os.environ.get("CS_DATABASE__NAME"),
11+
"host": os.environ.get("CS_DATABASE__HOST"),
12+
"port": 6432,
13+
}
14+
15+
connection_str = psycopg.conninfo.make_conninfo(**conn_params)
16+
17+
print("Connection to Proxy with {}".format(connection_str))
18+
19+
# Common test data
20+
val = {
21+
"key": "value",
22+
"number": 42,
23+
"array_number": [3, 2, 1],
24+
"array_string": ["hello", "world"],
25+
"nested": {"foo": "bar", "number": 1312},
26+
}
27+
28+
29+
def test_length_with_number():
30+
select_jsonb("encrypted_jsonb", val, "$.array_number[@]", 3)
31+
32+
33+
def test_length_with_string():
34+
select_jsonb("encrypted_jsonb", val, "$.array_string[@]", 2)
35+
36+
37+
def test_with_unknown():
38+
select_jsonb("encrypted_jsonb", val, "$.nonexistent", None)
39+
40+
41+
def select_jsonb(column, value, selector, expected, alias=None):
42+
alias = "AS {}".format(alias) if alias is not None else ""
43+
tests = [
44+
("jsonb_array_length(jsonb_path_query({}, '{}')) {}".format(column, selector, alias), []),
45+
("jsonb_array_length(jsonb_path_query({}, %s)) {}".format(column, alias), [selector]),
46+
]
47+
48+
for (select_fragment, params) in tests:
49+
print("Testing fragment: {}, params: {}, expecting: {}".format(
50+
select_fragment, params, expected))
51+
52+
for (binary, prepare) in product([True, None], repeat=2):
53+
execute(json.dumps(value), column,
54+
select_fragment=select_fragment,
55+
select_params=params,
56+
expected=expected,
57+
binary=binary,
58+
prepare=prepare)
59+
60+
61+
def make_id():
62+
return random.randrange(1, 1000000000)
63+
64+
65+
def execute(val, column, binary=None, prepare=None, expected=None,
66+
select_fragment=None, select_params=[]):
67+
with psycopg.connect(connection_str, autocommit=True) as conn:
68+
69+
with conn.cursor() as cursor:
70+
71+
with conn.transaction():
72+
id = make_id()
73+
74+
print("... for column {}, with binary: {}, prepare: {}".format(
75+
column, binary, prepare))
76+
77+
sql = "INSERT INTO encrypted (id, {}) VALUES (%s, %s)".format(
78+
column)
79+
80+
cursor.execute(sql, [id, val], binary=binary, prepare=prepare)
81+
82+
sql = "SELECT id, {} FROM encrypted WHERE id = %s".format(
83+
select_fragment)
84+
cursor.execute(
85+
sql, (select_params + [id]),
86+
binary=binary, prepare=prepare)
87+
88+
row = cursor.fetchone()
89+
90+
# If expected is None, we mean that there is no result.
91+
# If expected is not None, we mean that we expect a row.
92+
if expected is None:
93+
assert row == expected
94+
else:
95+
(result_id, result) = row
96+
assert result_id == id
97+
assert result == expected

‎tests/python/tests/test_jsonb_ops.py‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
import json
2+
import os
3+
import psycopg
4+
import random
5+
from itertools import product
6+
7+
conn_params = {
8+
"user": os.environ.get("CS_DATABASE__USERNAME"),
9+
"password": os.environ.get("CS_DATABASE__PASSWORD"),
10+
"dbname": os.environ.get("CS_DATABASE__NAME"),
11+
"host": os.environ.get("CS_DATABASE__HOST"),
12+
"port": 6432,
13+
}
14+
15+
connection_str = psycopg.conninfo.make_conninfo(**conn_params)
16+
17+
print("Connection to Proxy with {}".format(connection_str))
18+
19+
20+
def test_jsonb_contained_by():
21+
val = {"key": "value"}
22+
column = "encrypted_jsonb"
23+
select_fragment = "%s <@ encrypted_jsonb"
24+
tests = [
25+
(val, True),
26+
({"key": "different value"}, False)
27+
]
28+
29+
for (param, expected) in tests:
30+
params = [json.dumps(param)]
31+
print("Testing params: {}, expecting: {}".format(
32+
params, expected))
33+
34+
for (binary, prepare) in product([True, None], repeat=2):
35+
execute(json.dumps(val), column,
36+
select_fragment=select_fragment,
37+
select_params=params,
38+
expected=expected,
39+
binary=binary,
40+
prepare=prepare)
41+
42+
43+
def test_jsonb_contains():
44+
val = {"key": "value"}
45+
column = "encrypted_jsonb"
46+
select_fragment = "encrypted_jsonb @> %s"
47+
tests = [
48+
(val, True),
49+
({"key": "different value"}, False)
50+
]
51+
52+
for (param, expected) in tests:
53+
params = [json.dumps(param)]
54+
print("Testing params: {}, expecting: {}".format(
55+
params, expected))
56+
57+
for (binary, prepare) in product([True, None], repeat=2):
58+
execute(json.dumps(val), column,
59+
select_fragment=select_fragment,
60+
select_params=params,
61+
expected=expected,
62+
binary=binary,
63+
prepare=prepare)
64+
65+
66+
def test_jsonb_extract_simple():
67+
expected = "value"
68+
val = {"key": expected}
69+
column = "encrypted_jsonb"
70+
select_fragment_template = "{}->'{}'"
71+
accessors = [
72+
'key',
73+
'$.key', # Undocumented JSONPath selector
74+
]
75+
76+
for accessor in accessors:
77+
select_fragment = select_fragment_template.format(column, accessor)
78+
79+
print("Testing field: {}, expecting: {}".format(
80+
accessor, expected))
81+
82+
for (binary, prepare) in product([True, None], repeat=2):
83+
execute(json.dumps(val), column,
84+
select_fragment=select_fragment,
85+
select_params=[],
86+
expected=expected,
87+
binary=binary,
88+
prepare=prepare)
89+
90+
91+
def test_jsonb_extract_parameterised():
92+
val = {
93+
"string": "hello",
94+
"number": 42,
95+
"nested": {
96+
"number": 1815,
97+
"string": "world",
98+
},
99+
"array_string": ["hello", "world"],
100+
"array_number": [42, 84],
101+
}
102+
column = "encrypted_jsonb"
103+
select_fragment = "encrypted_jsonb->%s"
104+
tests = [
105+
("string", "hello"),
106+
("number", 42),
107+
("array_string", ["hello", "world"]),
108+
("array_number", [42, 84]),
109+
("nested", {"number": 1815, "string": "world"}),
110+
# TODO: Test ("nonexistentkey", None)
111+
]
112+
113+
for (param, expected) in tests:
114+
# JSONPath selectors *also* work with EQL extract, but are undocumented
115+
for accessor in [param, "$." + param]:
116+
print("Testing accessor: {}, expecting: {}".format(accessor, expected))
117+
118+
for (binary, prepare) in product([True, None], repeat=2):
119+
execute(json.dumps(val), column,
120+
select_fragment=select_fragment,
121+
select_params=[accessor],
122+
expected=expected,
123+
binary=binary,
124+
prepare=prepare)
125+
126+
127+
def make_id():
128+
return random.randrange(1, 1000000000)
129+
130+
131+
def execute(val, column, binary=None, prepare=None, expected=None,
132+
select_fragment=None, select_params=[]):
133+
with psycopg.connect(connection_str, autocommit=True) as conn:
134+
135+
with conn.cursor() as cursor:
136+
137+
with conn.transaction():
138+
id = make_id()
139+
140+
print("... for column {}, with binary: {}, prepare: {}".format(
141+
column, binary, prepare))
142+
143+
sql = "INSERT INTO encrypted (id, {}) VALUES (%s, %s)".format(
144+
column)
145+
146+
cursor.execute(sql, [id, val], binary=binary, prepare=prepare)
147+
148+
sql = "SELECT id, {} FROM encrypted WHERE id = %s".format(
149+
select_fragment)
150+
cursor.execute(
151+
sql, (select_params + [id]),
152+
binary=binary, prepare=prepare)
153+
154+
row = cursor.fetchone()
155+
156+
(result_id, result) = row
157+
expected_result = expected if expected is not None else val
158+
159+
assert result_id == id
160+
assert result == expected_result

0 commit comments

Comments
 (0)