Use "Datetime" ydb type for sa.DATETIME

Author: LuckySting

test/test_core.py

@@ -253,6 +253,16 @@ def test_integer_types(self, connection):
         assert result == (b"Uint8", b"Uint16", b"Uint32", b"Uint64", b"Int8", b"Int16", b"Int32", b"Int64")
 
     def test_datetime_types(self, connection: sa.Connection):
+        stmt = sa.Select(
+            sa.func.FormatType(sa.func.TypeOf(sa.bindparam("p_datetime", datetime.datetime.now(), sa.DateTime))),
+            sa.func.FormatType(sa.func.TypeOf(sa.bindparam("p_DATETIME", datetime.datetime.now(), sa.DATETIME))),
+            sa.func.FormatType(sa.func.TypeOf(sa.bindparam("p_TIMESTAMP", datetime.datetime.now(), sa.TIMESTAMP))),
+        )
+
+        result = connection.execute(stmt).fetchone()
+        assert result == (b"Timestamp", b"Datetime", b"Timestamp")
+
+    def test_datetime_types_timezone(self, connection: sa.Connection):
         table = self.tables.test_datetime_types
 
         now_dt = datetime.datetime.now()

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -132,7 +132,13 @@ def visit_BINARY(self, type_: sa.BINARY, **kw):
     def visit_BLOB(self, type_: sa.BLOB, **kw):
         return "String"
 
-    def visit_DATETIME(self, type_: sa.TIMESTAMP, **kw):
+    def visit_datetime(self, type_: sa.TIMESTAMP, **kw):
+        return self.visit_TIMESTAMP(type_, **kw)
+
+    def visit_DATETIME(self, type_: sa.DATETIME, **kw):
+        return "DateTime"
+
+    def visit_TIMESTAMP(self, type_: sa.TIMESTAMP, **kw):
         return "Timestamp"
 
     def visit_list_type(self, type_: types.ListType, **kw):
@@ -193,7 +199,10 @@ def get_ydb_type(
         elif isinstance(type_, types.YqlJSON.YqlJSONPathType):
             ydb_type = ydb.PrimitiveType.Utf8
         # Json
-
+        elif isinstance(type_, sa.DATETIME):
+            ydb_type = ydb.PrimitiveType.Datetime
+        elif isinstance(type_, sa.TIMESTAMP):
+            ydb_type = ydb.PrimitiveType.Timestamp
         elif isinstance(type_, sa.DateTime):
             ydb_type = ydb.PrimitiveType.Timestamp
         elif isinstance(type_, sa.Date):
@@ -610,7 +619,9 @@ class YqlDialect(StrCompileDialect):
     colspecs = {
         sa.types.JSON: types.YqlJSON,
         sa.types.JSON.JSONPathType: types.YqlJSON.YqlJSONPathType,
-        sa.types.DateTime: types.YqlDateTime,
+        sa.types.DateTime: types.YqlTimestamp,
+        sa.types.DATETIME: types.YqlDateTime,
+        sa.types.TIMESTAMP: types.YqlTimestamp,
     }
 
     connection_characteristics = util.immutabledict(

ydb_sqlalchemy/sqlalchemy/datetime_types.py

@@ -3,10 +3,10 @@
 
 from sqlalchemy import Dialect
 from sqlalchemy import types as sqltypes
-from sqlalchemy.sql.type_api import _ResultProcessorType
+from sqlalchemy.sql.type_api import _BindProcessorType, _ResultProcessorType
 
 
-class YqlDateTime(sqltypes.DateTime):
+class YqlTimestamp(sqltypes.DateTime):
     def result_processor(self, dialect: Dialect, coltype: str) -> Optional[_ResultProcessorType[datetime.datetime]]:
         def process(value: Optional[datetime.datetime]) -> Optional[datetime.datetime]:
             if value is None:
@@ -16,3 +16,13 @@ def process(value: Optional[datetime.datetime]) -> Optional[datetime.datetime]:
         if self.timezone:
             return process
         return None
+
+
+class YqlDateTime(YqlTimestamp):
+    def bind_processor(self, dialect: Dialect) -> Optional[_BindProcessorType[datetime.datetime]]:
+        def process(value: Optional[datetime.datetime]) -> Optional[int]:
+            if value is None:
+                return None
+            return int(value.timestamp())
+
+        return process

ydb_sqlalchemy/sqlalchemy/types.py

@@ -3,7 +3,7 @@
 from sqlalchemy import ARRAY, ColumnElement, exc, types
 from sqlalchemy.sql import type_api
 
-from .datetime_types import YqlDateTime  # noqa: F401
+from .datetime_types import YqlDateTime, YqlTimestamp  # noqa: F401
 from .json import YqlJSON  # noqa: F401