Date32, Datetime64 and Timestamp64 support

Author: vgvoleg

requirements.txt

@@ -1,3 +1,3 @@
 sqlalchemy >= 1.4.0, < 3.0.0
-ydb >= 3.18.8
+ydb >= 3.21.6
 ydb-dbapi >= 0.1.10

test/test_suite.py

@@ -1,4 +1,5 @@
 import ctypes
+import datetime
 import decimal
 
 import pytest
@@ -424,6 +425,16 @@ class DateTest(_DateTest):
     run_dispose_bind = "once"
 
 
+class Date32Test(_DateTest):
+    run_dispose_bind = "once"
+    datatype = ydb_sa_types.YqlDate32
+    data = datetime.date(1969, 1, 1)
+
+    @pytest.mark.skip("Default binding for DATE is not compatible with Date32")
+    def test_select_direct(self, connection):
+        pass
+
+
 class DateTimeMicrosecondsTest(_DateTimeMicrosecondsTest):
     run_dispose_bind = "once"
 
@@ -432,10 +443,30 @@ class DateTimeTest(_DateTimeTest):
     run_dispose_bind = "once"
 
 
+class DateTime64Test(_DateTimeTest):
+    datatype = ydb_sa_types.YqlDateTime64
+    data = datetime.datetime(1969, 10, 15, 12, 57, 18)
+    run_dispose_bind = "once"
+
+    @pytest.mark.skip("Default binding for DATETIME is not compatible with DateTime64")
+    def test_select_direct(self, connection):
+        pass
+
+
 class TimestampMicrosecondsTest(_TimestampMicrosecondsTest):
     run_dispose_bind = "once"
 
 
+class Timestamp64MicrosecondsTest(_TimestampMicrosecondsTest):
+    run_dispose_bind = "once"
+    datatype = ydb_sa_types.YqlTimestamp64
+    data = datetime.datetime(1969, 10, 15, 12, 57, 18, 396)
+
+    @pytest.mark.skip("Default binding for TIMESTAMP is not compatible with Timestamp64")
+    def test_select_direct(self, connection):
+        pass
+
+
 @pytest.mark.skip("unsupported Time data type")
 class TimeTest(_TimeTest):
     pass

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -61,6 +61,9 @@ def upsert(table):
     ydb.DecimalType: sa.DECIMAL,
     ydb.PrimitiveType.Yson: sa.TEXT,
     ydb.PrimitiveType.Date: sa.DATE,
+    ydb.PrimitiveType.Date32: sa.DATE,
+    ydb.PrimitiveType.Timestamp64: sa.TIMESTAMP,
+    ydb.PrimitiveType.Datetime64: sa.DATETIME,
     ydb.PrimitiveType.Datetime: sa.DATETIME,
     ydb.PrimitiveType.Timestamp: sa.TIMESTAMP,
     ydb.PrimitiveType.Interval: sa.INTEGER,
@@ -437,6 +440,8 @@ def do_execute(
     ) -> None:
         operation, parameters = self._prepare_ydb_query(statement, context, parameters, execute_many=False)
         is_ddl = context.isddl if context is not None else False
+        print(operation)
+        print(parameters)
         if is_ddl:
             cursor.execute_scheme(operation, parameters)
         else:

ydb_sqlalchemy/sqlalchemy/compiler/base.py

@@ -135,6 +135,15 @@ def visit_DATETIME(self, type_: sa.DATETIME, **kw):
     def visit_TIMESTAMP(self, type_: sa.TIMESTAMP, **kw):
         return "Timestamp"
 
+    def visit_date32(self, type_: types.YqlDate32, **kw):
+        return "Date32"
+
+    def visit_timestamp64(self, type_: types.YqlTimestamp64, **kw):
+        return "Timestamp64"
+
+    def visit_datetime64(self, type_: types.YqlDateTime64, **kw):
+        return "DateTime64"
+
     def visit_list_type(self, type_: types.ListType, **kw):
         inner = self.process(type_.item_type, **kw)
         return f"List<{inner}>"
@@ -193,6 +202,12 @@ def get_ydb_type(
         elif isinstance(type_, types.YqlJSON.YqlJSONPathType):
             ydb_type = ydb.PrimitiveType.Utf8
         # Json
+        elif isinstance(type_, types.YqlDate32):
+            ydb_type = ydb.PrimitiveType.Date32
+        elif isinstance(type_, types.YqlTimestamp64):
+            ydb_type = ydb.PrimitiveType.Timestamp64
+        elif isinstance(type_, types.YqlDateTime64):
+            ydb_type = ydb.PrimitiveType.Datetime64
         elif isinstance(type_, sa.DATETIME):
             ydb_type = ydb.PrimitiveType.Datetime
         elif isinstance(type_, sa.TIMESTAMP):

ydb_sqlalchemy/sqlalchemy/datetime_types.py

@@ -36,3 +36,23 @@ def process(value: Optional[datetime.datetime]) -> Optional[int]:
             return int(value.timestamp())
 
         return process
+
+
+class YqlDate32(YqlDate):
+    __visit_name__ = "date32"
+
+    def literal_processor(self, dialect):
+        parent = super().literal_processor(dialect)
+
+        def process(value):
+            return f"Date32({parent(value)})"
+
+        return process
+
+
+class YqlTimestamp64(YqlTimestamp):
+    __visit_name__ = "timestamp64"
+
+
+class YqlDateTime64(YqlDateTime):
+    __visit_name__ = "datetime64"

ydb_sqlalchemy/sqlalchemy/types.py

@@ -11,7 +11,7 @@
 from sqlalchemy import ARRAY, exc, types
 from sqlalchemy.sql import type_api
 
-from .datetime_types import YqlDate, YqlDateTime, YqlTimestamp  # noqa: F401
+from .datetime_types import YqlDate, YqlDateTime, YqlTimestamp, YqlDate32, YqlTimestamp64, YqlDateTime64  # noqa: F401
 from .json import YqlJSON  # noqa: F401