Skip to content

Commit 31ae103

Browse files
vgvolegvgvoleg
committed
Support new YDB types
1 parent 0fff143 commit 31ae103

File tree

2 files changed

+204
-4
lines changed

2 files changed

+204
-4
lines changed

‎tests/query/test_types.py‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
import pytest
2+
import ydb
3+
4+
from datetime import date, datetime, timedelta, timezone
5+
from decimal import Decimal
6+
from uuid import uuid4
7+
8+
9+
@pytest.mark.parametrize(
10+
"value,ydb_type_str,ydb_type",
11+
[
12+
(True, "Bool", ydb.PrimitiveType.Bool),
13+
(-125, "Int8", ydb.PrimitiveType.Int8),
14+
(None, "Int8?", ydb.OptionalType(ydb.PrimitiveType.Int8)),
15+
(-32766, "Int16", ydb.PrimitiveType.Int16),
16+
(-1123, "Int32", ydb.PrimitiveType.Int32),
17+
(-2157583648, "Int64", ydb.PrimitiveType.Int64),
18+
(255, "UInt8", ydb.PrimitiveType.Uint8),
19+
(65534, "UInt16", ydb.PrimitiveType.Uint16),
20+
(5555, "UInt32", ydb.PrimitiveType.Uint32),
21+
(2157583649, "UInt64", ydb.PrimitiveType.Uint64),
22+
(3.1415, "Double", ydb.PrimitiveType.Double),
23+
(".31415926535e1", "DyNumber", ydb.PrimitiveType.DyNumber),
24+
(Decimal("3.1415926535"), "Decimal(28, 10)", ydb.DecimalType(28, 10)),
25+
(b"Hello, YDB!", "String", ydb.PrimitiveType.String),
26+
("Hello, 🐍!", "Utf8", ydb.PrimitiveType.Utf8),
27+
('{"foo": "bar"}', "Json", ydb.PrimitiveType.Json),
28+
(b'{"foo"="bar"}', "Yson", ydb.PrimitiveType.Yson),
29+
('{"foo":"bar"}', "JsonDocument", ydb.PrimitiveType.JsonDocument),
30+
(uuid4(), "Uuid", ydb.PrimitiveType.UUID),
31+
([1, 2, 3], "List<Int8>", ydb.ListType(ydb.PrimitiveType.Int8)),
32+
({1: None, 2: None, 3: None}, "Set<Int8>", ydb.SetType(ydb.PrimitiveType.Int8)),
33+
([b"a", b"b", b"c"], "List<String>", ydb.ListType(ydb.PrimitiveType.String)),
34+
({"a": 1001, "b": 1002}, "Dict<Utf8, Int32>", ydb.DictType(ydb.PrimitiveType.Utf8, ydb.PrimitiveType.Int32)),
35+
(
36+
("a", 1001),
37+
"Tuple<Utf8, Int32>",
38+
ydb.TupleType().add_element(ydb.PrimitiveType.Utf8).add_element(ydb.PrimitiveType.Int32),
39+
),
40+
(
41+
{"foo": True, "bar": None},
42+
"Struct<foo:Bool?, bar:Int32?>",
43+
ydb.StructType()
44+
.add_member("foo", ydb.OptionalType(ydb.PrimitiveType.Bool))
45+
.add_member("bar", ydb.OptionalType(ydb.PrimitiveType.Int32)),
46+
),
47+
(100, "Date", ydb.PrimitiveType.Date),
48+
(100, "Date32", ydb.PrimitiveType.Date32),
49+
(-100, "Date32", ydb.PrimitiveType.Date32),
50+
(100, "Datetime", ydb.PrimitiveType.Datetime),
51+
(100, "Datetime64", ydb.PrimitiveType.Datetime64),
52+
(-100, "Datetime64", ydb.PrimitiveType.Datetime64),
53+
(-100, "Interval", ydb.PrimitiveType.Interval),
54+
(-100, "Interval64", ydb.PrimitiveType.Interval64),
55+
(100, "Timestamp", ydb.PrimitiveType.Timestamp),
56+
(100, "Timestamp64", ydb.PrimitiveType.Timestamp64),
57+
(-100, "Timestamp64", ydb.PrimitiveType.Timestamp64),
58+
(1511789040123456, "Timestamp", ydb.PrimitiveType.Timestamp),
59+
(1511789040123456, "Timestamp64", ydb.PrimitiveType.Timestamp64),
60+
(-1511789040123456, "Timestamp64", ydb.PrimitiveType.Timestamp64),
61+
],
62+
)
63+
def test_types(driver_sync: ydb.Driver, value, ydb_type_str, ydb_type):
64+
settings = (
65+
ydb.QueryClientSettings()
66+
.with_native_date_in_result_sets(False)
67+
.with_native_datetime_in_result_sets(False)
68+
.with_native_timestamp_in_result_sets(False)
69+
.with_native_interval_in_result_sets(False)
70+
.with_native_json_in_result_sets(False)
71+
)
72+
with ydb.QuerySessionPool(driver_sync, query_client_settings=settings) as pool:
73+
result = pool.execute_with_retries(
74+
f"DECLARE $param as {ydb_type_str}; SELECT $param as value",
75+
{"$param": (value, ydb_type)},
76+
)
77+
assert result[0].rows[0].value == value
78+
79+
80+
test_td = timedelta(microseconds=-100)
81+
test_now = datetime.utcnow()
82+
test_old_date = datetime(1221, 1, 1, 0, 0)
83+
test_today = test_now.date()
84+
test_dt_today = datetime.today()
85+
tz4h = timezone(timedelta(hours=4))
86+
87+
88+
@pytest.mark.parametrize(
89+
"value,ydb_type_str,ydb_type,result_value",
90+
[
91+
# FIXME: TypeError: 'datetime.datetime' object cannot be interpreted as an integer
92+
# (test_dt_today, "Datetime", test_dt_today),
93+
(test_today, "Date", ydb.PrimitiveType.Date, test_today),
94+
(365, "Date", ydb.PrimitiveType.Date, date(1971, 1, 1)),
95+
(-365, "Date32", ydb.PrimitiveType.Date32, date(1969, 1, 1)),
96+
(3600 * 24 * 365, "Datetime", ydb.PrimitiveType.Datetime, datetime(1971, 1, 1, 0, 0)),
97+
(3600 * 24 * 365 * (-1), "Datetime64", ydb.PrimitiveType.Datetime64, datetime(1969, 1, 1, 0, 0)),
98+
(datetime(1970, 1, 1, 4, 0, tzinfo=tz4h), "Timestamp", ydb.PrimitiveType.Timestamp, datetime(1970, 1, 1, 0, 0)),
99+
(test_td, "Interval", ydb.PrimitiveType.Interval, test_td),
100+
(test_td, "Interval64", ydb.PrimitiveType.Interval64, test_td),
101+
(test_now, "Timestamp", ydb.PrimitiveType.Timestamp, test_now),
102+
(test_old_date, "Timestamp64", ydb.PrimitiveType.Timestamp64, test_old_date),
103+
(
104+
1511789040123456,
105+
"Timestamp",
106+
ydb.PrimitiveType.Timestamp,
107+
datetime.fromisoformat("2017-11-27 13:24:00.123456"),
108+
),
109+
('{"foo": "bar"}', "Json", ydb.PrimitiveType.Json, {"foo": "bar"}),
110+
('{"foo": "bar"}', "JsonDocument", ydb.PrimitiveType.JsonDocument, {"foo": "bar"}),
111+
],
112+
)
113+
def test_types_native(driver_sync, value, ydb_type_str, ydb_type, result_value):
114+
with ydb.QuerySessionPool(driver_sync) as pool:
115+
result = pool.execute_with_retries(
116+
f"DECLARE $param as {ydb_type_str}; SELECT $param as value",
117+
{"$param": (value, ydb_type)},
118+
)
119+
assert result[0].rows[0].value == result_value

‎ydb/types.py‎

Lines changed: 85 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,19 @@ def _to_date(pb: ydb_value_pb2.Value, value: typing.Union[date, int]) -> None:
4040
pb.uint32_value = value
4141

4242

43+
def _from_date32(x: ydb_value_pb2.Value, table_client_settings: table.TableClientSettings) -> typing.Union[date, int]:
44+
if table_client_settings is not None and table_client_settings._native_date_in_result_sets:
45+
return _EPOCH.date() + timedelta(days=x.int32_value)
46+
return x.int32_value
47+
48+
49+
def _to_date32(pb: ydb_value_pb2.Value, value: typing.Union[date, int]) -> None:
50+
if isinstance(value, date):
51+
pb.int32_value = (value - _EPOCH.date()).days
52+
else:
53+
pb.int32_value = value
54+
55+
4356
def _from_datetime_number(
4457
x: typing.Union[float, datetime], table_client_settings: table.TableClientSettings
4558
) -> datetime:
@@ -63,6 +76,10 @@ def _from_uuid(pb: ydb_value_pb2.Value, value: uuid.UUID):
6376
pb.high_128 = struct.unpack("Q", value.bytes_le[8:16])[0]
6477

6578

79+
def _timedelta_to_microseconds(value: timedelta) -> int:
80+
return (value.days * _SECONDS_IN_DAY + value.seconds) * 1000000 + value.microseconds
81+
82+
6683
def _from_interval(
6784
value_pb: ydb_value_pb2.Value, table_client_settings: table.TableClientSettings
6885
) -> typing.Union[timedelta, int]:
@@ -71,10 +88,6 @@ def _from_interval(
7188
return value_pb.int64_value
7289

7390

74-
def _timedelta_to_microseconds(value: timedelta) -> int:
75-
return (value.days * _SECONDS_IN_DAY + value.seconds) * 1000000 + value.microseconds
76-
77-
7891
def _to_interval(pb: ydb_value_pb2.Value, value: typing.Union[timedelta, int]):
7992
if isinstance(value, timedelta):
8093
pb.int64_value = _timedelta_to_microseconds(value)
@@ -101,6 +114,25 @@ def _to_timestamp(pb: ydb_value_pb2.Value, value: typing.Union[datetime, int]):
101114
pb.uint64_value = value
102115

103116

117+
def _from_timestamp64(
118+
value_pb: ydb_value_pb2.Value, table_client_settings: table.TableClientSettings
119+
) -> typing.Union[datetime, int]:
120+
if table_client_settings is not None and table_client_settings._native_timestamp_in_result_sets:
121+
return _EPOCH + timedelta(microseconds=value_pb.int64_value)
122+
return value_pb.int64_value
123+
124+
125+
def _to_timestamp64(pb: ydb_value_pb2.Value, value: typing.Union[datetime, int]):
126+
if isinstance(value, datetime):
127+
if value.tzinfo:
128+
epoch = _EPOCH_UTC
129+
else:
130+
epoch = _EPOCH
131+
pb.int64_value = _timedelta_to_microseconds(value - epoch)
132+
else:
133+
pb.int64_value = value
134+
135+
104136
@enum.unique
105137
class PrimitiveType(enum.Enum):
106138
"""
@@ -133,23 +165,46 @@ class PrimitiveType(enum.Enum):
133165
_from_date,
134166
_to_date,
135167
)
168+
Date32 = (
169+
_apis.primitive_types.DATE32,
170+
None,
171+
_from_date32,
172+
_to_date32,
173+
)
136174
Datetime = (
137175
_apis.primitive_types.DATETIME,
138176
"uint32_value",
139177
_from_datetime_number,
140178
)
179+
Datetime64 = (
180+
_apis.primitive_types.DATETIME64,
181+
"int64_value",
182+
_from_datetime_number,
183+
)
141184
Timestamp = (
142185
_apis.primitive_types.TIMESTAMP,
143186
None,
144187
_from_timestamp,
145188
_to_timestamp,
146189
)
190+
Timestamp64 = (
191+
_apis.primitive_types.TIMESTAMP64,
192+
None,
193+
_from_timestamp64,
194+
_to_timestamp64,
195+
)
147196
Interval = (
148197
_apis.primitive_types.INTERVAL,
149198
None,
150199
_from_interval,
151200
_to_interval,
152201
)
202+
Interval64 = (
203+
_apis.primitive_types.INTERVAL64,
204+
None,
205+
_from_interval,
206+
_to_interval,
207+
)
153208

154209
DyNumber = _apis.primitive_types.DYNUMBER, "text_value"
155210

@@ -366,6 +421,32 @@ def __str__(self):
366421
return self._repr
367422

368423

424+
class SetType(AbstractTypeBuilder):
425+
__slots__ = ("__repr", "__proto")
426+
427+
def __init__(
428+
self,
429+
key_type: typing.Union[AbstractTypeBuilder, PrimitiveType],
430+
):
431+
"""
432+
:param key_type: Key type builder
433+
"""
434+
self._repr = "Set<%s>" % (str(key_type))
435+
self._proto = _apis.ydb_value.Type(
436+
dict_type=_apis.ydb_value.DictType(
437+
key=key_type.proto,
438+
payload=_apis.ydb_value.Type(void_type=struct_pb2.NULL_VALUE),
439+
)
440+
)
441+
442+
@property
443+
def proto(self):
444+
return self._proto
445+
446+
def __str__(self):
447+
return self._repr
448+
449+
369450
class TupleType(AbstractTypeBuilder):
370451
__slots__ = ("__elements_repr", "__proto")
371452

0 commit comments

Comments
 (0)