Skip to content

Commit 8e11f9f

Browse files
authored
Impl compatibility tests for parametrized writes and reads (#26529)
1 parent 94fc437 commit 8e11f9f

File tree

3 files changed

+148
-31
lines changed

3 files changed

+148
-31
lines changed

ydb/tests/compatibility/test_data_type.py

Lines changed: 111 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,20 @@
11
import pytest
22
import math
33

4+
from decimal import Decimal
45
from datetime import datetime, timedelta
56
from ydb.tests.library.compatibility.fixtures import RestartToAnotherVersionFixture
67
from ydb.tests.oss.ydb_sdk_import import ydb
78
from ydb.tests.datashard.lib.create_table import create_table_sql_request
8-
from ydb.tests.datashard.lib.types_of_variables import pk_types, non_pk_types, cleanup_type_name, format_sql_value, types_not_supported_yet_in_columnshard
9+
from ydb.tests.datashard.lib.types_of_variables import (
10+
pk_types,
11+
non_pk_types,
12+
cleanup_type_name,
13+
format_sql_value,
14+
types_not_supported_yet_in_columnshard,
15+
non_comparable_types,
16+
primitive_type,
17+
)
918

1019

1120
class TestDataType(RestartToAnotherVersionFixture):
@@ -159,6 +168,90 @@ def create_table(self):
159168
for query in querys:
160169
session_pool.execute_with_retries(query)
161170

171+
def create_typed_value(self, type_name, value):
172+
if "Decimal" in type_name:
173+
prec, scale = type_name.replace("Decimal(", "").replace(")", "").split(",")
174+
return ydb.TypedValue(Decimal(value), ydb.DecimalType(int(prec), int(scale)))
175+
if type_name == "String" or type_name == "Yson":
176+
return ydb.TypedValue(value.encode(), primitive_type[type_name])
177+
if type_name == "DyNumber":
178+
return ydb.TypedValue(str(value), primitive_type[type_name])
179+
if type_name == "Datetime64" or type_name == "Datetime":
180+
return ydb.TypedValue(int(value.timestamp()), primitive_type[type_name])
181+
return ydb.TypedValue(value, primitive_type[type_name])
182+
183+
def parametrized_write_data(self):
184+
queries_with_parameters = []
185+
186+
for i in range(self.count_table):
187+
query = f"""
188+
{";".join([f"DECLARE $pk_{cleanup_type_name(name)} AS {name}" for name in self.pk_types[i].keys()])};
189+
{";".join([f"DECLARE $col_{cleanup_type_name(name)} AS {name}" for name in self.all_types.keys()])};
190+
191+
UPSERT INTO {self.table_names[i]} (
192+
{", ".join([f"pk_{cleanup_type_name(name)}" for name in self.pk_types[i].keys()])},
193+
{", ".join([f"col_{cleanup_type_name(name)}" for name in self.all_types.keys()])}
194+
)
195+
VALUES (
196+
{", ".join([f"$pk_{cleanup_type_name(name)}" for name in self.pk_types[i].keys()])},
197+
{", ".join([f"$col_{cleanup_type_name(name)}" for name in self.all_types.keys()])}
198+
);
199+
"""
200+
201+
for row in range(1, self.count_rows + 1):
202+
parameters = {}
203+
for type_name, lamb in self.pk_types[i].items():
204+
parameters[f"$pk_{cleanup_type_name(type_name)}"] = self.create_typed_value(type_name, lamb(row))
205+
for type_name, lamb in self.all_types.items():
206+
parameters[f"$col_{cleanup_type_name(type_name)}"] = self.create_typed_value(type_name, lamb(row))
207+
queries_with_parameters.append((query, parameters))
208+
209+
with ydb.QuerySessionPool(self.driver) as session_pool:
210+
for query, parameters in queries_with_parameters:
211+
session_pool.execute_with_retries(query, parameters)
212+
213+
def parametrized_check_table(self):
214+
queries_with_parameters = []
215+
comparable_types = [(type_name, lamb) for type_name, lamb in self.all_types.items() if type_name not in non_comparable_types]
216+
217+
for i in range(self.count_table):
218+
query = f"""
219+
{";".join([f"DECLARE $pk_{cleanup_type_name(name)} AS {name}" for name in self.pk_types[i].keys()])};
220+
{";".join([f"DECLARE $col_{cleanup_type_name(name)} AS {name}" for name, _ in comparable_types])};
221+
222+
SELECT * FROM {self.table_names[i]} WHERE
223+
{" AND ".join([f"pk_{cleanup_type_name(name)} = $pk_{cleanup_type_name(name)}" for name in self.pk_types[i].keys()])}
224+
{" AND " if len(comparable_types) > 0 else ""}
225+
{" AND ".join([f"col_{cleanup_type_name(name)} = $col_{cleanup_type_name(name)}" for name, _ in comparable_types])}
226+
"""
227+
228+
for row in range(1, self.count_rows + 1):
229+
parameters = {}
230+
for type_name, lamb in self.pk_types[i].items():
231+
parameters[f"$pk_{cleanup_type_name(type_name)}"] = self.create_typed_value(type_name, lamb(row))
232+
for type_name, lamb in comparable_types:
233+
parameters[f"$col_{cleanup_type_name(type_name)}"] = self.create_typed_value(type_name, lamb(row))
234+
queries_with_parameters.append((query, parameters))
235+
236+
with ydb.QuerySessionPool(self.driver) as session_pool:
237+
for query_index, (query, parameters) in enumerate(queries_with_parameters):
238+
table_index = query_index // self.count_rows
239+
row_num = (query_index % self.count_rows) + 1
240+
241+
result_rows = []
242+
243+
result_sets = session_pool.execute_with_retries(query, parameters)
244+
for result_set in result_sets:
245+
for row in result_set.rows:
246+
result_rows.append(row)
247+
248+
assert len(result_rows) == 1
249+
250+
for row in result_rows:
251+
for prefix in self.columns[table_index].keys():
252+
for type_name in self.columns[table_index][prefix]:
253+
self.assert_type(type_name, row_num, row[f"{prefix}{cleanup_type_name(type_name)}"])
254+
162255
@pytest.mark.parametrize("store_type", ["ROW", "COLUMN"], indirect=True)
163256
def test_data_type(self):
164257
if any("Decimal" in type_name for type_name in self.all_types.keys()) and self.store_type == "COLUMN":
@@ -175,3 +268,20 @@ def test_data_type(self):
175268
self.check_table()
176269
self.write_data()
177270
self.check_table()
271+
272+
@pytest.mark.parametrize("store_type", ["ROW", "COLUMN"], indirect=True)
273+
def test_parametrized_data_type(self):
274+
if any("Decimal" in type_name for type_name in self.all_types.keys()) and self.store_type == "COLUMN":
275+
if (min(self.versions) < (25, 1)):
276+
pytest.skip("Decimal types are not supported in columnshard in this version")
277+
278+
self.create_table()
279+
280+
self.parametrized_write_data()
281+
self.parametrized_check_table()
282+
283+
self.change_cluster_version()
284+
285+
self.parametrized_check_table()
286+
self.parametrized_write_data()
287+
self.parametrized_check_table()

ydb/tests/datashard/lib/types_of_variables.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from uuid import UUID
22
from datetime import datetime
3+
import ydb
34

45

56
def format_sql_value(value, type_name, unwrap_after_cast: bool = False):
@@ -254,3 +255,38 @@ def cleanup_type_name(type_name):
254255
"UUID",
255256
"Interval"
256257
}
258+
259+
non_comparable_types = {
260+
"Yson",
261+
"Json",
262+
"JsonDocument",
263+
}
264+
265+
primitive_type = {
266+
"Int64": ydb.PrimitiveType.Int64,
267+
"Uint64": ydb.PrimitiveType.Uint64,
268+
"Int32": ydb.PrimitiveType.Int32,
269+
"Uint32": ydb.PrimitiveType.Uint32,
270+
"Int16": ydb.PrimitiveType.Int16,
271+
"Uint16": ydb.PrimitiveType.Uint16,
272+
"Int8": ydb.PrimitiveType.Int8,
273+
"Uint8": ydb.PrimitiveType.Uint8,
274+
"Bool": ydb.PrimitiveType.Bool,
275+
"DyNumber": ydb.PrimitiveType.DyNumber,
276+
"String": ydb.PrimitiveType.String,
277+
"Utf8": ydb.PrimitiveType.Utf8,
278+
"UUID": ydb.PrimitiveType.UUID,
279+
"Date": ydb.PrimitiveType.Date,
280+
"Datetime": ydb.PrimitiveType.Datetime,
281+
"Timestamp": ydb.PrimitiveType.Timestamp,
282+
"Interval": ydb.PrimitiveType.Interval,
283+
"Float": ydb.PrimitiveType.Float,
284+
"Double": ydb.PrimitiveType.Double,
285+
"Json": ydb.PrimitiveType.Json,
286+
"JsonDocument": ydb.PrimitiveType.JsonDocument,
287+
"Yson": ydb.PrimitiveType.Yson,
288+
"Date32": ydb.PrimitiveType.Date32,
289+
"Datetime64": ydb.PrimitiveType.Datetime64,
290+
"Timestamp64": ydb.PrimitiveType.Timestamp64,
291+
"Interval64": ydb.PrimitiveType.Interval64,
292+
}

ydb/tests/datashard/parametrized_queries/test_parametrized_queries.py

Lines changed: 1 addition & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
index_three_sync_not_Bool,
1818
index_four_sync,
1919
index_zero_sync,
20+
primitive_type,
2021
)
2122

2223
# https://github.com/ydb-platform/ydb/issues/17178
@@ -31,36 +32,6 @@ def filter_unsupported(type_map):
3132
return {k: v for k, v in type_map.items() if k not in unsupported_types}
3233

3334

34-
primitive_type = {
35-
"Int64": ydb.PrimitiveType.Int64,
36-
"Uint64": ydb.PrimitiveType.Uint64,
37-
"Int32": ydb.PrimitiveType.Int32,
38-
"Uint32": ydb.PrimitiveType.Uint32,
39-
"Int16": ydb.PrimitiveType.Int16,
40-
"Uint16": ydb.PrimitiveType.Uint16,
41-
"Int8": ydb.PrimitiveType.Int8,
42-
"Uint8": ydb.PrimitiveType.Uint8,
43-
"Bool": ydb.PrimitiveType.Bool,
44-
"DyNumber": ydb.PrimitiveType.DyNumber,
45-
"String": ydb.PrimitiveType.String,
46-
"Utf8": ydb.PrimitiveType.Utf8,
47-
"UUID": ydb.PrimitiveType.UUID,
48-
"Date": ydb.PrimitiveType.Date,
49-
"Datetime": ydb.PrimitiveType.Datetime,
50-
"Timestamp": ydb.PrimitiveType.Timestamp,
51-
"Interval": ydb.PrimitiveType.Interval,
52-
"Float": ydb.PrimitiveType.Float,
53-
"Double": ydb.PrimitiveType.Double,
54-
"Json": ydb.PrimitiveType.Json,
55-
"JsonDocument": ydb.PrimitiveType.JsonDocument,
56-
"Yson": ydb.PrimitiveType.Yson,
57-
"Date32": ydb.PrimitiveType.Date32,
58-
"Datetime64": ydb.PrimitiveType.Datetime64,
59-
"Timestamp64": ydb.PrimitiveType.Timestamp64,
60-
"Interval64": ydb.PrimitiveType.Interval64,
61-
}
62-
63-
6435
class TestParametrizedQueries(TestBase):
6536
@pytest.mark.parametrize(
6637
"table_name, pk_types, all_types, index, ttl, unique, sync",

0 commit comments

Comments
 (0)