Skip to content

Commit 8ee13c9

Browse files
dennnniskadennnniska
authored
Test for parametrized queries for all types (#17175)
1 parent cb504b4 commit 8ee13c9

File tree

3 files changed

+363
-0
lines changed

3 files changed

+363
-0
lines changed

‎ydb/tests/datashard/parametrized_queries/test_parametrized_queries.py

Lines changed: 340 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,340 @@
1+
import pytest
2+
import ydb
3+
4+
from datetime import datetime
5+
from ydb.tests.sql.lib.test_base import TestBase
6+
from ydb.tests.datashard.lib.create_table import create_table_sql_request, create_ttl_sql_request
7+
from ydb.tests.datashard.lib.types_of_variables import (
8+
cleanup_type_name,
9+
pk_types,
10+
non_pk_types,
11+
index_first,
12+
index_second,
13+
ttl_types,
14+
index_first_sync,
15+
index_second_sync,
16+
index_three_sync,
17+
index_three_sync_not_Bool,
18+
index_four_sync,
19+
index_zero_sync,
20+
)
21+
22+
# https://github.com/ydb-platform/ydb/issues/17178
23+
unsupported_types = [
24+
"Decimal(15,0)",
25+
"Decimal(22,9)",
26+
"Decimal(35,10)",
27+
"Date32",
28+
"Datetime64",
29+
"Timestamp64",
30+
"Interval64",
31+
]
32+
33+
34+
def filter_unsupported(type_map):
35+
return {k: v for k, v in type_map.items() if k not in unsupported_types}
36+
37+
38+
primitive_type = {
39+
"Int64": ydb.PrimitiveType.Int64,
40+
"Uint64": ydb.PrimitiveType.Uint64,
41+
"Int32": ydb.PrimitiveType.Int32,
42+
"Uint32": ydb.PrimitiveType.Uint32,
43+
"Int16": ydb.PrimitiveType.Int16,
44+
"Uint16": ydb.PrimitiveType.Uint16,
45+
"Int8": ydb.PrimitiveType.Int8,
46+
"Uint8": ydb.PrimitiveType.Uint8,
47+
"Bool": ydb.PrimitiveType.Bool,
48+
"DyNumber": ydb.PrimitiveType.DyNumber,
49+
"String": ydb.PrimitiveType.String,
50+
"Utf8": ydb.PrimitiveType.Utf8,
51+
"UUID": ydb.PrimitiveType.UUID,
52+
"Date": ydb.PrimitiveType.Date,
53+
"Datetime": ydb.PrimitiveType.Datetime,
54+
"Timestamp": ydb.PrimitiveType.Timestamp,
55+
"Interval": ydb.PrimitiveType.Interval,
56+
"Float": ydb.PrimitiveType.Float,
57+
"Double": ydb.PrimitiveType.Double,
58+
"Json": ydb.PrimitiveType.Json,
59+
"JsonDocument": ydb.PrimitiveType.JsonDocument,
60+
"Yson": ydb.PrimitiveType.Yson,
61+
}
62+
63+
64+
class TestParametrizedQueries(TestBase):
65+
@pytest.mark.parametrize(
66+
"table_name, pk_types, all_types, index, ttl, unique, sync",
67+
[
68+
(
69+
"table_index_4_UNIQUE_SYNC",
70+
filter_unsupported(pk_types),
71+
{},
72+
filter_unsupported(index_four_sync),
73+
"",
74+
"UNIQUE",
75+
"SYNC",
76+
),
77+
(
78+
"table_index_3_UNIQUE_SYNC",
79+
filter_unsupported(pk_types),
80+
{},
81+
filter_unsupported(index_three_sync_not_Bool),
82+
"",
83+
"UNIQUE",
84+
"SYNC",
85+
),
86+
(
87+
"table_index_2_UNIQUE_SYNC",
88+
filter_unsupported(pk_types),
89+
{},
90+
filter_unsupported(index_second_sync),
91+
"",
92+
"UNIQUE",
93+
"SYNC",
94+
),
95+
(
96+
"table_index_1_UNIQUE_SYNC",
97+
filter_unsupported(pk_types),
98+
{},
99+
filter_unsupported(index_first_sync),
100+
"",
101+
"UNIQUE",
102+
"SYNC",
103+
),
104+
(
105+
"table_index_0_UNIQUE_SYNC",
106+
filter_unsupported(pk_types),
107+
{},
108+
filter_unsupported(index_zero_sync),
109+
"",
110+
"UNIQUE",
111+
"SYNC",
112+
),
113+
(
114+
"table_index_4__SYNC",
115+
filter_unsupported(pk_types),
116+
{},
117+
filter_unsupported(index_four_sync),
118+
"",
119+
"",
120+
"SYNC",
121+
),
122+
(
123+
"table_index_3__SYNC",
124+
filter_unsupported(pk_types),
125+
{},
126+
filter_unsupported(index_three_sync),
127+
"",
128+
"",
129+
"SYNC",
130+
),
131+
(
132+
"table_index_2__SYNC",
133+
filter_unsupported(pk_types),
134+
{},
135+
filter_unsupported(index_second_sync),
136+
"",
137+
"",
138+
"SYNC",
139+
),
140+
(
141+
"table_index_1__SYNC",
142+
filter_unsupported(pk_types),
143+
{},
144+
filter_unsupported(index_first_sync),
145+
"",
146+
"",
147+
"SYNC",
148+
),
149+
(
150+
"table_index_0__SYNC",
151+
filter_unsupported(pk_types),
152+
{},
153+
filter_unsupported(index_zero_sync),
154+
"",
155+
"",
156+
"SYNC",
157+
),
158+
(
159+
"table_index_1__ASYNC",
160+
filter_unsupported(pk_types),
161+
{},
162+
filter_unsupported(index_second),
163+
"",
164+
"",
165+
"ASYNC",
166+
),
167+
(
168+
"table_index_0__ASYNC",
169+
filter_unsupported(pk_types),
170+
{},
171+
filter_unsupported(index_first),
172+
"",
173+
"",
174+
"ASYNC",
175+
),
176+
(
177+
"table_all_types",
178+
filter_unsupported(pk_types),
179+
{**filter_unsupported(pk_types), **filter_unsupported(non_pk_types)},
180+
{},
181+
"",
182+
"",
183+
"",
184+
),
185+
("table_ttl_DyNumber", filter_unsupported(pk_types), {}, {}, "DyNumber", "", ""),
186+
("table_ttl_Uint32", filter_unsupported(pk_types), {}, {}, "Uint32", "", ""),
187+
("table_ttl_Uint64", filter_unsupported(pk_types), {}, {}, "Uint64", "", ""),
188+
("table_ttl_Datetime", filter_unsupported(pk_types), {}, {}, "Datetime", "", ""),
189+
("table_ttl_Timestamp", filter_unsupported(pk_types), {}, {}, "Timestamp", "", ""),
190+
("table_ttl_Date", filter_unsupported(pk_types), {}, {}, "Date", "", ""),
191+
],
192+
)
193+
def test_parametrized_queries(
194+
self,
195+
table_name: str,
196+
pk_types: dict[str, str],
197+
all_types: dict[str, str],
198+
index: dict[str, str],
199+
ttl: str,
200+
unique: str,
201+
sync: str,
202+
):
203+
self.create_table(table_name, pk_types, all_types, index, ttl, unique, sync)
204+
self.insert(table_name, all_types, pk_types, index, ttl)
205+
self.select_after_insert(table_name, all_types, pk_types, index, ttl)
206+
207+
def create_table(
208+
self,
209+
table_name: str,
210+
pk_types: dict[str, str],
211+
all_types: dict[str, str],
212+
index: dict[str, str],
213+
ttl: str,
214+
unique: str,
215+
sync: str,
216+
):
217+
columns = {"pk_": pk_types.keys(), "col_": all_types.keys(), "col_index_": index.keys(), "ttl_": [ttl]}
218+
pk_columns = {"pk_": pk_types.keys()}
219+
index_columns = {"col_index_": index.keys()}
220+
sql_create_table = create_table_sql_request(table_name, columns, pk_columns, index_columns, unique, sync)
221+
self.query(sql_create_table)
222+
if ttl != "":
223+
sql_ttl = create_ttl_sql_request(
224+
f"ttl_{cleanup_type_name(ttl)}",
225+
{"P18262D": ""},
226+
"SECONDS" if ttl == "Uint32" or ttl == "Uint64" or ttl == "DyNumber" else "",
227+
table_name,
228+
)
229+
self.query(sql_ttl)
230+
231+
def insert(
232+
self, table_name: str, all_types: dict[str, str], pk_types: dict[str, str], index: dict[str, str], ttl: str
233+
):
234+
number_of_columns = len(pk_types) + len(all_types) + len(index)
235+
236+
if ttl != "":
237+
number_of_columns += 1
238+
for count in range(1, number_of_columns + 1):
239+
self.create_insert(table_name, count, all_types, pk_types, index, ttl)
240+
241+
def create_insert(
242+
self,
243+
table_name: str,
244+
value: int,
245+
all_types: dict[str, str],
246+
pk_types: dict[str, str],
247+
index: dict[str, str],
248+
ttl: str,
249+
):
250+
insert_sql = f"""
251+
{" ".join([f"DECLARE $pk_{cleanup_type_name(type_name)} AS {type_name};" for type_name in pk_types.keys()])}
252+
{" ".join([f"DECLARE $col_{cleanup_type_name(type_name)} AS {type_name};" for type_name in all_types.keys()])}
253+
{" ".join([f"DECLARE $col_index_{cleanup_type_name(type_name)} AS {type_name};" for type_name in index.keys()])}
254+
{f"DECLARE $ttl_{ttl} AS {ttl};" if ttl != "" else ""}
255+
256+
INSERT INTO {table_name}(
257+
{", ".join(["pk_" + cleanup_type_name(type_name) for type_name in pk_types.keys()])}{", " if len(all_types) != 0 else ""}
258+
{", ".join(["col_" + cleanup_type_name(type_name) for type_name in all_types.keys()])}{", " if len(index) != 0 else ""}
259+
{", ".join(["col_index_" + cleanup_type_name(type_name) for type_name in index.keys()])}{", " if len(ttl) != 0 else ""}
260+
{f"ttl_{ttl}" if ttl != "" else ""}
261+
)
262+
VALUES(
263+
{", ".join([f"$pk_{cleanup_type_name(type_name)}" for type_name in pk_types.keys()])}{", " if len(all_types) != 0 else ""}
264+
{", ".join([f"$col_{cleanup_type_name(type_name)}" for type_name in all_types.keys()])}{", " if len(index) != 0 else ""}
265+
{", ".join([f"$col_index_{cleanup_type_name(type_name)}" for type_name in index.keys()])}{", " if len(ttl) != 0 else ""}
266+
{f"$ttl_{ttl}" if ttl != "" else ""}
267+
);
268+
"""
269+
parameters = {}
270+
for type_name in pk_types.keys():
271+
parameters[f"$pk_{cleanup_type_name(type_name)}"] = self.create_type_value(pk_types, type_name, value)
272+
for type_name in all_types.keys():
273+
parameters[f"$col_{cleanup_type_name(type_name)}"] = self.create_type_value(all_types, type_name, value)
274+
for type_name in index.keys():
275+
parameters[f"$col_index_{cleanup_type_name(type_name)}"] = self.create_type_value(index, type_name, value)
276+
if ttl != "":
277+
parameters[f"$ttl_{cleanup_type_name(ttl)}"] = self.create_type_value(ttl_types, ttl, value)
278+
self.query(insert_sql, parameters=parameters)
279+
280+
def select_after_insert(
281+
self, table_name: str, all_types: dict[str, str], pk_types: dict[str, str], index: dict[str, str], ttl: str
282+
):
283+
number_of_columns = len(pk_types) + len(all_types) + len(index)
284+
285+
if ttl != "":
286+
number_of_columns += 1
287+
288+
for count in range(1, number_of_columns + 1):
289+
create_all_type = []
290+
create_all_type_declare = []
291+
for type_name in all_types.keys():
292+
if type_name != "Json" and type_name != "Yson" and type_name != "JsonDocument":
293+
create_all_type.append(f"col_{cleanup_type_name(type_name)}=$col_{cleanup_type_name(type_name)}")
294+
create_all_type_declare.append(f"DECLARE $col_{cleanup_type_name(type_name)} AS {type_name};")
295+
296+
sql_select = f"""
297+
{" ".join([f"DECLARE $pk_{cleanup_type_name(type_name)} AS {type_name};" for type_name in pk_types.keys()])}
298+
{" ".join(create_all_type_declare)}
299+
{" ".join([f"DECLARE $col_index_{cleanup_type_name(type_name)} AS {type_name};" for type_name in index.keys()])}
300+
{f"DECLARE $ttl_{ttl} AS {ttl};" if ttl != "" else ""}
301+
302+
SELECT COUNT(*) as count FROM `{table_name}` WHERE
303+
{" and ".join([f"pk_{cleanup_type_name(type_name)}=$pk_{cleanup_type_name(type_name)}" for type_name in pk_types.keys()])}
304+
{" and " if len(index) != 0 else ""}
305+
{" and ".join([f"col_index_{cleanup_type_name(type_name)}=$col_index_{cleanup_type_name(type_name)}" for type_name in index.keys()])}
306+
{" and " if len(create_all_type) != 0 else ""}
307+
{" and ".join(create_all_type)}
308+
{f" and ttl_{ttl}=$ttl_{ttl}" if ttl != "" else ""}
309+
"""
310+
parameters = {}
311+
for type_name in pk_types.keys():
312+
parameters[f"$pk_{cleanup_type_name(type_name)}"] = self.create_type_value(pk_types, type_name, count)
313+
for type_name in all_types.keys():
314+
parameters[f"$col_{cleanup_type_name(type_name)}"] = self.create_type_value(all_types, type_name, count)
315+
for type_name in index.keys():
316+
parameters[f"$col_index_{cleanup_type_name(type_name)}"] = self.create_type_value(
317+
index, type_name, count
318+
)
319+
if ttl != "":
320+
parameters[f"$ttl_{cleanup_type_name(ttl)}"] = self.create_type_value(ttl_types, ttl, count)
321+
rows = self.query(sql_select, parameters=parameters)
322+
assert (
323+
len(rows) == 1 and rows[0].count == 1
324+
), f"Expected one rows, faild in {count} value, table {table_name}"
325+
326+
rows = self.query(f"SELECT COUNT(*) as count FROM `{table_name}`")
327+
assert (
328+
len(rows) == 1 and rows[0].count == number_of_columns
329+
), f"Expected {number_of_columns} rows, after select all line"
330+
331+
def create_type_value(self, key, type_name, value):
332+
if "Decimal" in type_name:
333+
return key[type_name](value)
334+
if type_name == "String" or type_name == "Yson":
335+
return ydb.TypedValue(key[type_name](value).encode(), primitive_type[type_name])
336+
if type_name == "DyNumber":
337+
return ydb.TypedValue(str(key[type_name](value)), primitive_type[type_name])
338+
if type_name == "Datetime64" or type_name == "Datetime":
339+
return ydb.TypedValue(int(datetime.timestamp(key[type_name](value))), primitive_type[type_name])
340+
return ydb.TypedValue(key[type_name](value), primitive_type[type_name])

‎ydb/tests/datashard/parametrized_queries/ya.make

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
PY3TEST()
2+
ENV(YDB_DRIVER_BINARY="ydb/apps/ydbd/ydbd")
3+
4+
FORK_SUBTESTS()
5+
SPLIT_FACTOR(19)
6+
SIZE(MEDIUM)
7+
8+
TEST_SRCS(
9+
test_parametrized_queries.py
10+
)
11+
12+
PEERDIR(
13+
ydb/tests/sql/lib
14+
ydb/tests/datashard/lib
15+
)
16+
17+
DEPENDS(
18+
ydb/apps/ydb
19+
ydb/apps/ydbd
20+
)
21+
22+
END()

‎ydb/tests/datashard/ya.make

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ RECURSE(
55
copy_table
66
lib
77
partitioning
8+
parametrized_queries
89
split_merge
910
secondary_index
1011
s3

0 commit comments

Comments
 (0)