Add timezone support

Author: LuckySting

requirements.txt

@@ -1,2 +1,2 @@
 sqlalchemy >= 2.0.7, < 3.0.0
-ydb >= 3.2.0
+ydb >= 3.11.3

test-requirements.txt

@@ -1,5 +1,5 @@
 sqlalchemy==2.0.7
-ydb==3.2.0
+ydb==3.11.3
 
 requests<2.29
 pytest==7.2.2

test/test_core.py

@@ -1,5 +1,5 @@
 import asyncio
-from datetime import date, datetime
+import datetime
 from decimal import Decimal
 from typing import NamedTuple
 
@@ -187,39 +187,42 @@ class TestTypes(TablesTest):
     __backend__ = True
 
     @classmethod
-    def define_tables(cls, metadata):
+    def define_tables(cls, metadata: sa.MetaData):
         Table(
-            "test_types",
+            "test_primitive_types",
             metadata,
-            Column("id", Integer, primary_key=True),
+            Column("int", sa.Integer, primary_key=True),
             # Column("bin", sa.BINARY),
             Column("str", sa.String),
-            Column("num", sa.Float),
-            Column("bl", sa.Boolean),
-            Column("ts", sa.TIMESTAMP),
+            Column("float", sa.Float),
+            Column("bool", sa.Boolean),
+        )
+        Table(
+            "test_datetime_types",
+            metadata,
+            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),
             # Column("interval", sa.Interval),
         )
 
-    def test_select_types(self, connection):
-        tb = self.tables.test_types
-
-        now, today = datetime.now(), date.today()
+    def test_primitive_types(self, connection):
+        table = self.tables.test_primitive_types
 
-        stm = tb.insert().values(
-            id=1,
+        statement = sa.insert(table).values(
+            int=42,
             # bin=b"abc",
             str="Hello World!",
-            num=3.5,
-            bl=True,
-            ts=now,
-            date=today,
+            float=3.5,
+            bool=True,
             # interval=timedelta(minutes=45),
         )
-        connection.execute(stm)
+        connection.execute(statement)
 
-        row = connection.execute(sa.select(tb)).fetchone()
-        assert row == (1, "Hello World!", 3.5, True, now, today)
+        row = connection.execute(sa.select(table)).fetchone()
+        assert row == (42, "Hello World!", 3.5, True)
 
     def test_integer_types(self, connection):
         stmt = sa.Select(
@@ -236,6 +239,34 @@ def test_integer_types(self, connection):
         result = connection.execute(stmt).fetchone()
         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):
+        table = self.tables.test_datetime_types
+
+        now_dt = datetime.datetime.now()
+        now_dt_tz = now_dt.replace(tzinfo=datetime.timezone(datetime.timedelta(hours=3, minutes=42)))
+        today = now_dt.date()
+
+        statement = sa.insert(table).values(
+            datetime=now_dt,
+            datetime_tz=now_dt_tz,
+            timestamp=now_dt,
+            timestamp_tz=now_dt_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
+            today,
+        )
+
 
 class TestWithClause(TablesTest):
     __backend__ = True

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -597,6 +597,7 @@ class YqlDialect(StrCompileDialect):
     colspecs = {
         sa.types.JSON: types.YqlJSON,
         sa.types.JSON.JSONPathType: types.YqlJSON.YqlJSONPathType,
+        sa.types.DateTime: types.YqlDateTime,
     }
 
     construct_arguments = [

ydb_sqlalchemy/sqlalchemy/datetime_types.py

@@ -0,0 +1,18 @@
+import datetime
+from typing import Optional
+
+from sqlalchemy import Dialect
+from sqlalchemy import types as sqltypes
+from sqlalchemy.sql.type_api import _ResultProcessorType
+
+
+class YqlDateTime(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:
+                return None
+            return value.replace(tzinfo=datetime.timezone.utc)
+
+        if self.timezone:
+            return process
+        return None

ydb_sqlalchemy/sqlalchemy/types.py

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