Use Datetime for general DateTime

Author: LuckySting

test/test_core.py

@@ -213,8 +213,8 @@ def define_tables(cls, metadata: sa.MetaData):
         Table(
             "test_datetime_types",
             metadata,
-            Column("datetime", sa.DateTime, primary_key=True),
-            Column("datetime_tz", sa.DateTime(timezone=True)),
+            Column("datetime", sa.DATETIME, primary_key=True),
+            Column("datetime_tz", sa.DATETIME(timezone=True)),
             Column("timestamp", sa.TIMESTAMP),
             Column("timestamp_tz", sa.TIMESTAMP(timezone=True)),
             Column("date", sa.Date),
@@ -260,33 +260,35 @@ def test_datetime_types(self, connection: sa.Connection):
         )
 
         result = connection.execute(stmt).fetchone()
-        assert result == (b"Timestamp", b"Datetime", b"Timestamp")
+        assert result == (b"Datetime", b"Datetime", b"Timestamp")
 
     def test_datetime_types_timezone(self, connection: sa.Connection):
         table = self.tables.test_datetime_types
+        tzinfo = datetime.timezone(datetime.timedelta(hours=3, minutes=42))
 
-        now_dt = datetime.datetime.now()
-        now_dt_tz = now_dt.replace(tzinfo=datetime.timezone(datetime.timedelta(hours=3, minutes=42)))
-        today = now_dt.date()
+        timestamp_value = datetime.datetime.now()
+        timestamp_value_tz = timestamp_value.replace(tzinfo=tzinfo)
+        datetime_value = timestamp_value.replace(microsecond=0)
+        datetime_value_tz = timestamp_value_tz.replace(microsecond=0)
+        today = timestamp_value.date()
 
         statement = sa.insert(table).values(
-            datetime=now_dt,
-            datetime_tz=now_dt_tz,
-            timestamp=now_dt,
-            timestamp_tz=now_dt_tz,
+            datetime=datetime_value,
+            datetime_tz=datetime_value_tz,
+            timestamp=timestamp_value,
+            timestamp_tz=timestamp_value_tz,
             date=today,
             # interval=datetime.timedelta(minutes=45),
         )
         connection.execute(statement)
 
         row = connection.execute(sa.select(table)).fetchone()
 
-        now_dt_tz_utc = now_dt.replace(tzinfo=datetime.timezone.utc) - datetime.timedelta(hours=3, minutes=42)
         assert row == (
-            now_dt,
-            now_dt_tz_utc,  # YDB doesn't store timezone, so it is always utc
-            now_dt,
-            now_dt_tz_utc,  # YDB doesn't store timezone, so it is always utc
+            datetime_value,
+            datetime_value_tz.astimezone(datetime.timezone.utc),  # YDB doesn't store timezone, so it is always utc
+            timestamp_value,
+            timestamp_value_tz.astimezone(datetime.timezone.utc),  # YDB doesn't store timezone, so it is always utc
             today,
         )
 

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -132,8 +132,8 @@ 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):
-        return self.visit_TIMESTAMP(type_, **kw)
+    def visit_datetime(self, type_: sa.DATETIME, **kw):
+        return self.visit_DATETIME(type_, **kw)
 
     def visit_DATETIME(self, type_: sa.DATETIME, **kw):
         return "DateTime"
@@ -204,7 +204,7 @@ def get_ydb_type(
         elif isinstance(type_, sa.TIMESTAMP):
             ydb_type = ydb.PrimitiveType.Timestamp
         elif isinstance(type_, sa.DateTime):
-            ydb_type = ydb.PrimitiveType.Timestamp
+            ydb_type = ydb.PrimitiveType.Datetime
         elif isinstance(type_, sa.Date):
             ydb_type = ydb.PrimitiveType.Date
         elif isinstance(type_, sa.BINARY):
@@ -549,7 +549,7 @@ def upsert(table):
     ydb.PrimitiveType.Yson: sa.TEXT,
     ydb.PrimitiveType.Date: sa.DATE,
     ydb.PrimitiveType.Datetime: sa.DATETIME,
-    ydb.PrimitiveType.Timestamp: sa.DATETIME,
+    ydb.PrimitiveType.Timestamp: sa.TIMESTAMP,
     ydb.PrimitiveType.Interval: sa.INTEGER,
     ydb.PrimitiveType.Bool: sa.BOOLEAN,
     ydb.PrimitiveType.DyNumber: sa.TEXT,
@@ -619,7 +619,7 @@ class YqlDialect(StrCompileDialect):
     colspecs = {
         sa.types.JSON: types.YqlJSON,
         sa.types.JSON.JSONPathType: types.YqlJSON.YqlJSONPathType,
-        sa.types.DateTime: types.YqlTimestamp,
+        sa.types.DateTime: types.YqlDateTime,
         sa.types.DATETIME: types.YqlDateTime,
         sa.types.TIMESTAMP: types.YqlTimestamp,
     }

ydb_sqlalchemy/sqlalchemy/datetime_types.py

@@ -11,18 +11,20 @@ def result_processor(self, dialect: Dialect, coltype: str) -> Optional[_ResultPr
         def process(value: Optional[datetime.datetime]) -> Optional[datetime.datetime]:
             if value is None:
                 return None
+            if not self.timezone:
+                return value
             return value.replace(tzinfo=datetime.timezone.utc)
 
-        if self.timezone:
-            return process
-        return None
+        return process
 
 
 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
+            if not self.timezone:  # if timezone is disabled, consider it as utc
+                value = value.replace(tzinfo=datetime.timezone.utc)
             return int(value.timestamp())
 
         return process