Support extended date and time types

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Support all YDB extended date and time types, including `Interval64` and reflection
 * Fix SQLAlchemy 1.4 compatibility and test the full dialect against SQLAlchemy 1.4.54
 * Add opt-in support for YDB's native UUID type via `sa.UUID` or `types.YqlUUID` while keeping `sa.Uuid` stored as `Utf8`
 

docs/types.rst

@@ -128,6 +128,11 @@ The following table shows the complete mapping between YDB native types, YDB SQL
      -
      - ``datetime.datetime``
      - Extended timestamp range
+   * - ``Interval64``
+     - :class:`~ydb_sqlalchemy.sqlalchemy.datetime_types.YqlInterval64`
+     -
+     - ``datetime.timedelta``
+     - Extended interval range
    * - ``Json``
      - :class:`~ydb_sqlalchemy.sqlalchemy.json.YqlJSON`
      - ``JSON``
@@ -275,7 +280,7 @@ YDB provides several date and time types:
 
    from ydb_sqlalchemy.sqlalchemy.types import (
        YqlDate, YqlDateTime, YqlTimestamp,
-       YqlDate32, YqlDateTime64, YqlTimestamp64
+       YqlDate32, YqlDateTime64, YqlTimestamp64, YqlInterval64
    )
    from sqlalchemy import DateTime
    import datetime
@@ -303,6 +308,9 @@ YDB provides several date and time types:
        # Timestamp64 - extended range with microsecond precision
        extended_timestamp = Column(YqlTimestamp64)
 
+       # Interval64 - extended range with microsecond precision
+       retention = Column(YqlInterval64)
+
        # Standard SQLAlchemy DateTime also works
        updated_at = Column(DateTime)
 
@@ -318,12 +326,13 @@ YDB provides several date and time types:
        precise_datetime=now,
        precise_time=now,
        extended_timestamp=now,
+       retention=datetime.timedelta(days=50000),
        updated_at=now
    ))
 
 For detailed API reference, see:
 :class:`~ydb_sqlalchemy.sqlalchemy.datetime_types.YqlDate`, :class:`~ydb_sqlalchemy.sqlalchemy.datetime_types.YqlDateTime`, :class:`~ydb_sqlalchemy.sqlalchemy.datetime_types.YqlTimestamp`,
-:class:`~ydb_sqlalchemy.sqlalchemy.datetime_types.YqlDate32`, :class:`~ydb_sqlalchemy.sqlalchemy.datetime_types.YqlDateTime64`, :class:`~ydb_sqlalchemy.sqlalchemy.datetime_types.YqlTimestamp64`.
+:class:`~ydb_sqlalchemy.sqlalchemy.datetime_types.YqlDate32`, :class:`~ydb_sqlalchemy.sqlalchemy.datetime_types.YqlDateTime64`, :class:`~ydb_sqlalchemy.sqlalchemy.datetime_types.YqlTimestamp64`, :class:`~ydb_sqlalchemy.sqlalchemy.datetime_types.YqlInterval64`.
 
 Struct Type
 -----------

tests/integration/test_core.py

@@ -18,6 +18,7 @@
 from ydb_sqlalchemy.sqlalchemy import types
 
 _UUID_TABLE_NAME = f"test_uuid_types_{uuid.uuid4().hex[:8]}"
+_EXTENDED_DATETIME_TABLE_NAME = f"test_extended_datetime_types_{uuid.uuid4().hex[:8]}"
 
 if sa.__version__ >= "2.":
     from sqlalchemy import NullPool
@@ -266,6 +267,15 @@ def define_tables(cls, metadata: sa.MetaData):
             Column("uuid_native", types.YqlUUID),
             Column("uuid_str", sa.Uuid if not ydb_sa.OLD_SA else sa.String),
         )
+        Table(
+            _EXTENDED_DATETIME_TABLE_NAME,
+            metadata,
+            Column("id", Integer, primary_key=True),
+            Column("date32", types.YqlDate32),
+            Column("datetime64", types.YqlDateTime64),
+            Column("timestamp64", types.YqlTimestamp64),
+            Column("interval64", types.YqlInterval64),
+        )
 
     def test_primitive_types(self, connection):
         table = self.tables.test_primitive_types
@@ -391,6 +401,28 @@ def test_native_uuid_reflection(self, connection):
         reflected_type = reflected_metadata.tables[table.name].c.uuid_native.type
         assert isinstance(reflected_type, types.YqlUUID)
 
+    def test_extended_datetime_types_and_reflection(self, connection):
+        table = self.tables[_EXTENDED_DATETIME_TABLE_NAME]
+        values = {
+            "id": 1,
+            "date32": datetime.date(1969, 1, 1),
+            "datetime64": datetime.datetime(1969, 10, 15, 12, 57, 18),
+            "timestamp64": datetime.datetime(1969, 10, 15, 12, 57, 18, 396),
+            "interval64": datetime.timedelta(days=-50_000, microseconds=123),
+        }
+
+        connection.execute(sa.insert(table).values(**values))
+        assert connection.execute(sa.select(table)).one()._mapping == values
+
+        reflected_metadata = sa.MetaData()
+        reflected_metadata.reflect(connection, only=[table.name])
+        reflected_columns = reflected_metadata.tables[table.name].c
+
+        assert isinstance(reflected_columns.date32.type, types.YqlDate32)
+        assert isinstance(reflected_columns.datetime64.type, types.YqlDateTime64)
+        assert isinstance(reflected_columns.timestamp64.type, types.YqlTimestamp64)
+        assert isinstance(reflected_columns.interval64.type, types.YqlInterval64)
+
 
 class TestWithClause(TablesTest):
     __backend__ = True

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -69,12 +69,13 @@ 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.Date32: types.YqlDate32,
+    ydb.PrimitiveType.Timestamp64: types.YqlTimestamp64,
+    ydb.PrimitiveType.Datetime64: types.YqlDateTime64,
     ydb.PrimitiveType.Datetime: sa.DATETIME,
     ydb.PrimitiveType.Timestamp: sa.TIMESTAMP,
     ydb.PrimitiveType.Interval: sa.INTEGER,
+    ydb.PrimitiveType.Interval64: types.YqlInterval64,
     ydb.PrimitiveType.Bool: sa.BOOLEAN,
     ydb.PrimitiveType.DyNumber: sa.TEXT,
     ydb.PrimitiveType.UUID: types.YqlUUID,

ydb_sqlalchemy/sqlalchemy/compiler/base.py

@@ -154,6 +154,9 @@ def visit_timestamp64(self, type_: types.YqlTimestamp64, **kw):
     def visit_datetime64(self, type_: types.YqlDateTime64, **kw):
         return "DateTime64"
 
+    def visit_interval64(self, type_: types.YqlInterval64, **kw):
+        return "Interval64"
+
     def visit_list_type(self, type_: types.ListType, **kw):
         inner = self.process(type_.item_type, **kw)
         return f"List<{inner}>"
@@ -177,6 +180,12 @@ def visit_struct_type(self, type_: types.StructType, **kw):
     def get_ydb_type(
         self, type_: sa.types.TypeEngine, is_optional: bool
     ) -> Union[ydb.PrimitiveType, ydb.AbstractTypeBuilder]:
+        if isinstance(type_, types.YqlInterval64):
+            ydb_type = ydb.PrimitiveType.Interval64
+            if is_optional:
+                return ydb.OptionalType(ydb_type)
+            return ydb_type
+
         if isinstance(type_, sa.TypeDecorator):
             type_ = type_.impl
 

ydb_sqlalchemy/sqlalchemy/compiler/sa14.py

@@ -6,6 +6,7 @@
 from sqlalchemy.sql import literal_column
 from sqlalchemy.util.compat import inspect_getfullargspec
 
+from .. import types
 from .base import (
     BaseYqlCompiler,
     BaseYqlDDLCompiler,
@@ -22,6 +23,9 @@ def visit_FLOAT(self, type_: sa.FLOAT, **kw):
     def get_ydb_type(
         self, type_: sa.types.TypeEngine, is_optional: bool
     ) -> Union[ydb.PrimitiveType, ydb.AbstractTypeBuilder]:
+        if isinstance(type_, types.YqlInterval64):
+            return super().get_ydb_type(type_, is_optional)
+
         if isinstance(type_, sa.TypeDecorator):
             type_ = type_.impl
 

ydb_sqlalchemy/sqlalchemy/compiler/sa20.py

@@ -5,6 +5,7 @@
 from sqlalchemy.sql import literal_column
 from sqlalchemy.util.compat import inspect_getfullargspec
 
+from .. import types
 from .base import (
     BaseYqlCompiler,
     BaseYqlDDLCompiler,
@@ -21,6 +22,9 @@ def visit_uuid(self, type_: sa.Uuid, **kw):
     def get_ydb_type(
         self, type_: sa.types.TypeEngine, is_optional: bool
     ) -> Union[ydb.PrimitiveType, ydb.AbstractTypeBuilder]:
+        if isinstance(type_, types.YqlInterval64):
+            return super().get_ydb_type(type_, is_optional)
+
         if isinstance(type_, sa.TypeDecorator):
             type_ = type_.impl
 

ydb_sqlalchemy/sqlalchemy/datetime_types.py

@@ -20,6 +20,33 @@ def process(value):
     return process
 
 
+def _interval_literal(value: datetime.timedelta) -> str:
+    total_microseconds = (value.days * 24 * 60 * 60 + value.seconds) * 1_000_000 + value.microseconds
+    sign = "-" if total_microseconds < 0 else ""
+    total_microseconds = abs(total_microseconds)
+
+    days, remainder = divmod(total_microseconds, 24 * 60 * 60 * 1_000_000)
+    hours, remainder = divmod(remainder, 60 * 60 * 1_000_000)
+    minutes, remainder = divmod(remainder, 60 * 1_000_000)
+    seconds, microseconds = divmod(remainder, 1_000_000)
+
+    result = f"{sign}P"
+    if days:
+        result += f"{days}D"
+    if hours or minutes or seconds or microseconds or not days:
+        result += "T"
+        if hours:
+            result += f"{hours}H"
+        if minutes:
+            result += f"{minutes}M"
+        if microseconds:
+            result += f"{seconds}.{microseconds:06d}S"
+        elif seconds or not (hours or minutes):
+            result += f"{seconds}S"
+
+    return f"'{result}'"
+
+
 class YqlDate(sqltypes.Date):
     def literal_processor(self, dialect):
         parent = super().literal_processor(dialect)
@@ -72,3 +99,27 @@ class YqlDateTime64(YqlDateTime):
     def literal_processor(self, dialect):
         parent = super().literal_processor(dialect)
         return _literal_processor(parent, "DateTime64")
+
+
+class YqlInterval64(sqltypes.Interval):
+    """Store ``datetime.timedelta`` values using YDB's ``Interval64`` type."""
+
+    __visit_name__ = "interval64"
+    cache_ok = True
+
+    def bind_processor(self, dialect):
+        def process(value: Optional[datetime.timedelta]) -> Optional[datetime.timedelta]:
+            return value
+
+        return process
+
+    def result_processor(self, dialect, coltype):
+        def process(value) -> Optional[datetime.timedelta]:
+            if value is None or isinstance(value, datetime.timedelta):
+                return value
+            return datetime.timedelta(microseconds=value)
+
+        return process
+
+    def literal_processor(self, dialect):
+        return _literal_processor(_interval_literal, "Interval64")

ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py

@@ -1,4 +1,4 @@
-from datetime import date, datetime
+from datetime import date, datetime, timedelta
 import uuid
 
 import pytest
@@ -57,6 +57,11 @@ def test_lambda_compilation():
             datetime(1996, 11, 19, 12, 34, 56, 789),
             "DateTime64('1996-11-19 12:34:56.000789')",
         ),
+        (
+            types.YqlInterval64(),
+            timedelta(days=2, hours=3, minutes=4, seconds=5, microseconds=6),
+            "Interval64('P2DT3H4M5.000006S')",
+        ),
     ],
 )
 def test_datetime_literal_compilation(type_, value, expected):
@@ -68,6 +73,20 @@ def test_datetime_literal_compilation(type_, value, expected):
     assert str(compiled) == expected
 
 
+def test_extended_datetime_type_mapping_and_processors():
+    dialect = YqlDialect()
+    type_compiler = dialect.type_compiler
+    interval_type = types.YqlInterval64()
+    value = timedelta(days=-50_000, microseconds=123)
+
+    assert type_compiler.process(interval_type) == "Interval64"
+    assert type_compiler.get_ydb_type(interval_type, is_optional=False) == ydb.PrimitiveType.Interval64
+    assert type_compiler.get_ydb_type(interval_type, is_optional=True).item == ydb.PrimitiveType.Interval64
+    assert interval_type.bind_processor(dialect)(value) == value
+    assert interval_type.result_processor(dialect, None)(value) == value
+    assert interval_type.result_processor(dialect, None)(123) == timedelta(microseconds=123)
+
+
 def test_binary_type():
     dialect = YqlDialect()
     expr = sa.literal(b"some bytes")

ydb_sqlalchemy/sqlalchemy/types.py

@@ -12,7 +12,15 @@
 from sqlalchemy import ARRAY, exc, Table, types
 from sqlalchemy.sql import type_api
 
-from .datetime_types import YqlDate, YqlDateTime, YqlTimestamp, YqlDate32, YqlTimestamp64, YqlDateTime64  # noqa: F401
+from .datetime_types import (  # noqa: F401
+    YqlDate,
+    YqlDate32,
+    YqlDateTime,
+    YqlDateTime64,
+    YqlInterval64,
+    YqlTimestamp,
+    YqlTimestamp64,
+)
 from .json import YqlJSON  # noqa: F401