Support new YDB types

Author: vgvoleg

ydb/types.py

@@ -40,6 +40,19 @@ def _to_date(pb: ydb_value_pb2.Value, value: typing.Union[date, int]) -> None:
         pb.uint32_value = value
 
 
+def _from_date32(x: ydb_value_pb2.Value, table_client_settings: table.TableClientSettings) -> typing.Union[date, int]:
+    if table_client_settings is not None and table_client_settings._native_date_in_result_sets:
+        return _EPOCH.date() + timedelta(days=x.int32_value)
+    return x.int32_value
+
+
+def _to_date32(pb: ydb_value_pb2.Value, value: typing.Union[date, int]) -> None:
+    if isinstance(value, date):
+        pb.int32_value = (value - _EPOCH.date()).days
+    else:
+        pb.int32_value = value
+
+
 def _from_datetime_number(
     x: typing.Union[float, datetime], table_client_settings: table.TableClientSettings
 ) -> datetime:
@@ -63,6 +76,10 @@ def _from_uuid(pb: ydb_value_pb2.Value, value: uuid.UUID):
     pb.high_128 = struct.unpack("Q", value.bytes_le[8:16])[0]
 
 
+def _timedelta_to_microseconds(value: timedelta) -> int:
+    return (value.days * _SECONDS_IN_DAY + value.seconds) * 1000000 + value.microseconds
+
+
 def _from_interval(
     value_pb: ydb_value_pb2.Value, table_client_settings: table.TableClientSettings
 ) -> typing.Union[timedelta, int]:
@@ -71,10 +88,6 @@ def _from_interval(
     return value_pb.int64_value
 
 
-def _timedelta_to_microseconds(value: timedelta) -> int:
-    return (value.days * _SECONDS_IN_DAY + value.seconds) * 1000000 + value.microseconds
-
-
 def _to_interval(pb: ydb_value_pb2.Value, value: typing.Union[timedelta, int]):
     if isinstance(value, timedelta):
         pb.int64_value = _timedelta_to_microseconds(value)
@@ -101,6 +114,25 @@ def _to_timestamp(pb: ydb_value_pb2.Value, value: typing.Union[datetime, int]):
         pb.uint64_value = value
 
 
+def _from_timestamp64(
+    value_pb: ydb_value_pb2.Value, table_client_settings: table.TableClientSettings
+) -> typing.Union[datetime, int]:
+    if table_client_settings is not None and table_client_settings._native_timestamp_in_result_sets:
+        return _EPOCH + timedelta(microseconds=value_pb.int64_value)
+    return value_pb.int64_value
+
+
+def _to_timestamp64(pb: ydb_value_pb2.Value, value: typing.Union[datetime, int]):
+    if isinstance(value, datetime):
+        if value.tzinfo:
+            epoch = _EPOCH_UTC
+        else:
+            epoch = _EPOCH
+        pb.int64_value = _timedelta_to_microseconds(value - epoch)
+    else:
+        pb.int64_value = value
+
+
 @enum.unique
 class PrimitiveType(enum.Enum):
     """
@@ -133,23 +165,46 @@ class PrimitiveType(enum.Enum):
         _from_date,
         _to_date,
     )
+    Date32 = (
+        _apis.primitive_types.DATE32,
+        None,
+        _from_date32,
+        _to_date32,
+    )
     Datetime = (
         _apis.primitive_types.DATETIME,
         "uint32_value",
         _from_datetime_number,
     )
+    Datetime64 = (
+        _apis.primitive_types.DATETIME64,
+        "int32_value",
+        _from_datetime_number,
+    )
     Timestamp = (
         _apis.primitive_types.TIMESTAMP,
         None,
         _from_timestamp,
         _to_timestamp,
     )
+    Timestamp64 = (
+        _apis.primitive_types.TIMESTAMP64,
+        None,
+        _from_timestamp,
+        _to_timestamp,
+    )
     Interval = (
         _apis.primitive_types.INTERVAL,
         None,
         _from_interval,
         _to_interval,
     )
+    Interval64 = (
+        _apis.primitive_types.INTERVAL64,
+        None,
+        _from_interval,
+        _to_interval,
+    )
 
     DyNumber = _apis.primitive_types.DYNUMBER, "text_value"