Types: Unlock supporting timezone-aware DateTime fields

CHANGES.md

@@ -13,6 +13,7 @@
   CrateDB dialect table options.
 - Fixed SQL rendering of special DDL table options in `CrateDDLCompiler`.
   Before, configuring `crate_"translog.durability"` was not possible.
+- Unlocked supporting timezone-aware `DateTime` fields
 
 ## 2024/06/13 0.37.0
 - Added support for CrateDB's [FLOAT_VECTOR] data type and its accompanying

src/sqlalchemy_cratedb/compiler.py

@@ -225,7 +225,7 @@ def visit_SMALLINT(self, type_, **kw):
         return 'SHORT'
 
     def visit_datetime(self, type_, **kw):
-        return 'TIMESTAMP'
+        return self.visit_TIMESTAMP(type_, **kw)
 
     def visit_date(self, type_, **kw):
         return 'TIMESTAMP'
@@ -245,6 +245,16 @@ def visit_FLOAT_VECTOR(self, type_, **kw):
             raise ValueError("FloatVector must be initialized with dimension size")
         return f"FLOAT_VECTOR({dimensions})"
 
+    def visit_TIMESTAMP(self, type_, **kw):
+        """
+        Support for `TIMESTAMP WITH|WITHOUT TIME ZONE`.
+
+        From `sqlalchemy.dialects.postgresql.base.PGTypeCompiler`.
+        """
+        return "TIMESTAMP %s" % (
+            (type_.timezone and "WITH" or "WITHOUT") + " TIME ZONE",
+        )
+
 
 class CrateCompiler(compiler.SQLCompiler):
 

src/sqlalchemy_cratedb/dialect.py

@@ -32,16 +32,15 @@
     CrateDDLCompiler,
     CrateIdentifierPreparer,
 )
-from crate.client.exceptions import TimezoneUnawareException
 from .sa_version import SA_VERSION, SA_1_4, SA_2_0
 from .type import FloatVector, ObjectArray, ObjectType
 
 TYPES_MAP = {
     "boolean": sqltypes.Boolean,
     "short": sqltypes.SmallInteger,
     "smallint": sqltypes.SmallInteger,
-    "timestamp": sqltypes.TIMESTAMP,
-    "timestamp with time zone": sqltypes.TIMESTAMP,
+    "timestamp": sqltypes.TIMESTAMP(timezone=False),
+    "timestamp with time zone": sqltypes.TIMESTAMP(timezone=True),
     "object": ObjectType,
     "integer": sqltypes.Integer,
     "long": sqltypes.NUMERIC,
@@ -62,8 +61,8 @@
     TYPES_MAP["boolean_array"] = ARRAY(sqltypes.Boolean)
     TYPES_MAP["short_array"] = ARRAY(sqltypes.SmallInteger)
     TYPES_MAP["smallint_array"] = ARRAY(sqltypes.SmallInteger)
-    TYPES_MAP["timestamp_array"] = ARRAY(sqltypes.TIMESTAMP)
-    TYPES_MAP["timestamp with time zone_array"] = ARRAY(sqltypes.TIMESTAMP)
+    TYPES_MAP["timestamp_array"] = ARRAY(sqltypes.TIMESTAMP(timezone=False))
+    TYPES_MAP["timestamp with time zone_array"] = ARRAY(sqltypes.TIMESTAMP(timezone=True))
     TYPES_MAP["long_array"] = ARRAY(sqltypes.NUMERIC)
     TYPES_MAP["bigint_array"] = ARRAY(sqltypes.NUMERIC)
     TYPES_MAP["double_array"] = ARRAY(sqltypes.DECIMAL)
@@ -114,15 +113,11 @@ def process(value):
 
 class DateTime(sqltypes.DateTime):
 
-    TZ_ERROR_MSG = "Timezone aware datetime objects are not supported"
-
     def bind_processor(self, dialect):
         def process(value):
             if value is not None:
                 assert isinstance(value, datetime)
-                if value.tzinfo is not None:
-                    raise TimezoneUnawareException(DateTime.TZ_ERROR_MSG)
-                return value.strftime('%Y-%m-%dT%H:%M:%S.%fZ')
+                return value.strftime('%Y-%m-%dT%H:%M:%S.%f%z')
             return value
         return process
 
@@ -152,8 +147,9 @@ def process(value):
 
 
 colspecs = {
+    sqltypes.Date: Date,
     sqltypes.DateTime: DateTime,
-    sqltypes.Date: Date
+    sqltypes.TIMESTAMP: DateTime,
 }
 
 

tests/create_table_test.py

@@ -67,8 +67,10 @@ class User(self.Base):
              '\n\tlong_col1 LONG, \n\tlong_col2 LONG, '
              '\n\tbool_col BOOLEAN, '
              '\n\tshort_col SHORT, '
-             '\n\tdatetime_col TIMESTAMP, \n\tdate_col TIMESTAMP, '
-             '\n\tfloat_col FLOAT, \n\tdouble_col DOUBLE, '
+             '\n\tdatetime_col TIMESTAMP WITHOUT TIME ZONE, '
+             '\n\tdate_col TIMESTAMP, '
+             '\n\tfloat_col FLOAT, '
+             '\n\tdouble_col DOUBLE, '
              '\n\tPRIMARY KEY (string_col)\n)\n\n'),
             ())
 
@@ -286,7 +288,7 @@ class DummyTable(self.Base):
         fake_cursor.execute.assert_called_with(
             ('\nCREATE TABLE t (\n\t'
              'pk STRING NOT NULL, \n\t'
-             'a TIMESTAMP DEFAULT now(), \n\t'
+             'a TIMESTAMP WITHOUT TIME ZONE DEFAULT now(), \n\t'
              'PRIMARY KEY (pk)\n)\n\n'), ())
 
     def test_column_server_default_string(self):
@@ -312,7 +314,7 @@ class DummyTable(self.Base):
         fake_cursor.execute.assert_called_with(
             ('\nCREATE TABLE t (\n\t'
              'pk STRING NOT NULL, \n\t'
-             'a TIMESTAMP DEFAULT now(), \n\t'
+             'a TIMESTAMP WITHOUT TIME ZONE DEFAULT now(), \n\t'
              'PRIMARY KEY (pk)\n)\n\n'), ())
 
     def test_column_server_default_text_constant(self):

tests/datetime_test.py

@@ -21,21 +21,27 @@
 
 from __future__ import absolute_import
 
-from datetime import datetime, tzinfo, timedelta
+import datetime as dt
 from unittest import TestCase, skipIf
 from unittest.mock import patch, MagicMock
 
+import pytest
 import sqlalchemy as sa
-from sqlalchemy.exc import DBAPIError
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import Session, sessionmaker
 
 from sqlalchemy_cratedb import SA_VERSION, SA_1_4
+from sqlalchemy_cratedb.dialect import DateTime
 
 try:
     from sqlalchemy.orm import declarative_base
 except ImportError:
     from sqlalchemy.ext.declarative import declarative_base
 
+try:
+    import zoneinfo
+except ImportError:
+    from backports import zoneinfo
+
 from crate.client.cursor import Cursor
 
 
@@ -44,16 +50,14 @@
 FakeCursor.return_value = fake_cursor
 
 
-class CST(tzinfo):
-    """
-    Timezone object for CST
-    """
-
-    def utcoffset(self, date_time):
-        return timedelta(seconds=-3600)
-
-    def dst(self, date_time):
-        return timedelta(seconds=-7200)
+INPUT_DATE = dt.date(2009, 5, 13)
+INPUT_DATETIME_NOTZ = dt.datetime(2009, 5, 13, 19, 00, 30, 123456)
+INPUT_DATETIME_TZ = dt.datetime(2009, 5, 13, 19, 00, 30, 123456, tzinfo=zoneinfo.ZoneInfo("Europe/Kyiv"))
+OUTPUT_DATE = INPUT_DATE
+OUTPUT_TIMETZ_NOTZ = dt.time(19, 00, 30, 123000)
+OUTPUT_TIMETZ_TZ = dt.time(16, 00, 30, 123000)
+OUTPUT_DATETIME_NOTZ = dt.datetime(2009, 5, 13, 19, 00, 30, 123000)
+OUTPUT_DATETIME_TZ = dt.datetime(2009, 5, 13, 16, 00, 30, 123000)
 
 
 @skipIf(SA_VERSION < SA_1_4, "SQLAlchemy 1.3 suddenly has problems with these test cases")
@@ -68,7 +72,7 @@ class Character(Base):
             __tablename__ = 'characters'
             name = sa.Column(sa.String, primary_key=True)
             date = sa.Column(sa.Date)
-            timestamp = sa.Column(sa.DateTime)
+            datetime = sa.Column(sa.DateTime)
 
         fake_cursor.description = (
             ('characters_name', None, None, None, None, None, None),
@@ -87,9 +91,96 @@ def test_date_can_handle_datetime(self):
         ]
         self.session.query(self.Character).first()
 
-    def test_date_cannot_handle_tz_aware_datetime(self):
+    def test_date_can_handle_tz_aware_datetime(self):
         character = self.Character()
         character.name = "Athur"
-        character.timestamp = datetime(2009, 5, 13, 19, 19, 30, tzinfo=CST())
+        character.datetime = INPUT_DATETIME_NOTZ
         self.session.add(character)
-        self.assertRaises(DBAPIError, self.session.commit)
+
+
+Base = declarative_base()
+
+
+class FooBar(Base):
+    __tablename__ = "foobar"
+    name = sa.Column(sa.String, primary_key=True)
+    date = sa.Column(sa.Date)
+    datetime_notz = sa.Column(DateTime(timezone=False))
+    datetime_tz = sa.Column(DateTime(timezone=True))
+
+
+@pytest.fixture
+def session(cratedb_service):
+    engine = cratedb_service.database.engine
+    session = sessionmaker(bind=engine)()
+
+    Base.metadata.drop_all(engine, checkfirst=True)
+    Base.metadata.create_all(engine, checkfirst=True)
+    return session
+
+
+@pytest.mark.skipif(SA_VERSION < SA_1_4, reason="Test case not supported on SQLAlchemy 1.3")
+def test_datetime_notz(session):
+    """
+    An integration test for `sa.Date` and `sa.DateTime`, not using timezones.
+    """
+
+    # Insert record.
+    foo_item = FooBar(
+        name="foo",
+        date=INPUT_DATE,
+        datetime_notz=INPUT_DATETIME_NOTZ,
+        datetime_tz=INPUT_DATETIME_NOTZ,
+    )
+    session.add(foo_item)
+    session.commit()
+    session.execute(sa.text("REFRESH TABLE foobar"))
+
+    # Query record.
+    result = session.execute(sa.select(
+        FooBar.name, FooBar.date, FooBar.datetime_notz, FooBar.datetime_tz)).mappings().first()
+
+    # Compare outcome.
+    assert result["date"] == OUTPUT_DATE
+    assert result["datetime_notz"] == OUTPUT_DATETIME_NOTZ
+    assert result["datetime_notz"].tzname() is None
+    assert result["datetime_notz"].timetz() == OUTPUT_TIMETZ_NOTZ
+    assert result["datetime_notz"].tzinfo is None
+    assert result["datetime_tz"] == OUTPUT_DATETIME_NOTZ
+    assert result["datetime_tz"].tzname() is None
+    assert result["datetime_tz"].timetz() == OUTPUT_TIMETZ_NOTZ
+    assert result["datetime_tz"].tzinfo is None
+
+
+@pytest.mark.skipif(SA_VERSION < SA_1_4, reason="Test case not supported on SQLAlchemy 1.3")
+def test_datetime_tz(session):
+    """
+    An integration test for `sa.Date` and `sa.DateTime`, now using timezones.
+    """
+
+    # Insert record.
+    foo_item = FooBar(
+        name="foo",
+        date=INPUT_DATE,
+        datetime_notz=INPUT_DATETIME_TZ,
+        datetime_tz=INPUT_DATETIME_TZ,
+    )
+    session.add(foo_item)
+    session.commit()
+    session.execute(sa.text("REFRESH TABLE foobar"))
+
+    # Query record.
+    session.expunge(foo_item)
+    result = session.execute(sa.select(
+        FooBar.name, FooBar.date, FooBar.datetime_notz, FooBar.datetime_tz)).mappings().first()
+
+    # Compare outcome.
+    assert result["date"] == OUTPUT_DATE
+    assert result["datetime_notz"] == OUTPUT_DATETIME_NOTZ
+    assert result["datetime_notz"].tzname() is None
+    assert result["datetime_notz"].timetz() == OUTPUT_TIMETZ_NOTZ
+    assert result["datetime_notz"].tzinfo is None
+    assert result["datetime_tz"] == OUTPUT_DATETIME_TZ
+    assert result["datetime_tz"].tzname() is None
+    assert result["datetime_tz"].timetz() == OUTPUT_TIMETZ_TZ
+    assert result["datetime_tz"].tzinfo is None

tests/update_test.py

@@ -81,7 +81,7 @@ def test_onupdate_is_triggered(self):
         args = args[1]
         self.assertEqual(expected_stmt, stmt)
         self.assertEqual(40, args[0])
-        dt = datetime.strptime(args[1], '%Y-%m-%dT%H:%M:%S.%fZ')
+        dt = datetime.strptime(args[1], '%Y-%m-%dT%H:%M:%S.%f')
         self.assertIsInstance(dt, datetime)
         self.assertGreater(dt, now)
         self.assertEqual('Arthur', args[2])
@@ -110,6 +110,6 @@ def test_bulk_update(self):
         self.assertEqual(expected_stmt, stmt)
         self.assertEqual('Julia', args[0])
         self.assertEqual({'favorite_book': 'Romeo & Juliet'}, args[1])
-        dt = datetime.strptime(args[2], '%Y-%m-%dT%H:%M:%S.%fZ')
+        dt = datetime.strptime(args[2], '%Y-%m-%dT%H:%M:%S.%f')
         self.assertIsInstance(dt, datetime)
         self.assertGreater(dt, before_update_time)