Convert TIMESTAMP columns to timezone-aware datetime objects

This is additional API convenience based on the data type converter
machinery. A `time_zone` keyword argument can be passed to both the
`connect()` method, or when creating new `Cursor` objects.

The `time_zone` attribute can also be changed at runtime on both the
`connection` and `cursor` object instances.

Examples:

- connect('localhost:4200', time_zone=pytz.timezone("Australia/Sydney"))
- connection.cursor(time_zone="+0530")

CHANGES.txt

@@ -8,6 +8,10 @@ Unreleased
 - SQLAlchemy: Added support for ``crate_index`` and ``nullable`` attributes in
   ORM column definitions.
 
+- Added support for converting ``TIMESTAMP`` columns to timezone-aware
+  ``datetime`` objects, using the new ``time_zone`` keyword argument.
+
+
 2022/12/02 0.28.0
 =================
 

docs/query.rst

@@ -246,6 +246,67 @@ converter function defined as ``lambda``, which assigns ``yes`` for boolean
     ['no']
 
 
+``TIMESTAMP`` conversion with time zone
+=======================================
+
+Based on the data type converter functionality, the driver offers a convenient
+interface to make it return timezone-aware ``datetime`` objects, using the
+desired time zone.
+
+For your reference, in the following examples, epoch 1658167836758 is
+``Mon, 18 Jul 2022 18:10:36 GMT``.
+
+::
+
+    >>> import datetime
+    >>> tz_mst = datetime.timezone(datetime.timedelta(hours=7), name="MST")
+    >>> cursor = connection.cursor(time_zone=tz_mst)
+
+    >>> cursor.execute("SELECT datetime_tz FROM locations ORDER BY name")
+
+    >>> cursor.fetchone()
+    [datetime.datetime(2022, 7, 19, 1, 10, 36, 758000, tzinfo=datetime.timezone(datetime.timedelta(seconds=25200), 'MST'))]
+
+For the ``time_zone`` keyword argument, different data types are supported.
+The available options are:
+
+- ``datetime.timezone.utc``
+- ``datetime.timezone(datetime.timedelta(hours=7), name="MST")``
+- ``pytz.timezone("Australia/Sydney")``
+- ``zoneinfo.ZoneInfo("Australia/Sydney")``
+- ``+0530`` (UTC offset in string format)
+
+Let's exercise all of them.
+
+::
+
+    >>> cursor.time_zone = datetime.timezone.utc
+    >>> cursor.execute("SELECT datetime_tz FROM locations ORDER BY name")
+    >>> cursor.fetchone()
+    [datetime.datetime(2022, 7, 18, 18, 10, 36, 758000, tzinfo=datetime.timezone.utc)]
+
+    >>> import pytz
+    >>> cursor.time_zone = pytz.timezone("Australia/Sydney")
+    >>> cursor.execute("SELECT datetime_tz FROM locations ORDER BY name")
+    >>> cursor.fetchone()
+    ['foo', datetime.datetime(2022, 7, 19, 4, 10, 36, 758000, tzinfo=<DstTzInfo 'Australia/Sydney' AEST+10:00:00 STD>)]
+
+    >>> try:
+    ...     import zoneinfo
+    ... except ImportError:
+    ...     from backports import zoneinfo
+
+    >>> cursor.time_zone = zoneinfo.ZoneInfo("Australia/Sydney")
+    >>> cursor.execute("SELECT datetime_tz FROM locations ORDER BY name")
+    >>> cursor.fetchone()
+    [datetime.datetime(2022, 7, 19, 4, 10, 36, 758000, tzinfo=zoneinfo.ZoneInfo(key='Australia/Sydney'))]
+
+    >>> cursor.time_zone = "+0530"
+    >>> cursor.execute("SELECT datetime_tz FROM locations ORDER BY name")
+    >>> cursor.fetchone()
+    [datetime.datetime(2022, 7, 18, 23, 40, 36, 758000, tzinfo=datetime.timezone(datetime.timedelta(seconds=19800), '+0530'))]
+
+
 .. _Bulk inserts: https://crate.io/docs/crate/reference/en/latest/interfaces/http.html#bulk-operations
 .. _CrateDB data type identifiers for the HTTP interface: https://crate.io/docs/crate/reference/en/latest/interfaces/http.html#column-types
 .. _Database API: http://www.python.org/dev/peps/pep-0249/

setup.py

@@ -61,7 +61,8 @@ def read(path):
     install_requires=['urllib3>=1.9,<2'],
     extras_require=dict(
         sqlalchemy=['sqlalchemy>=1.0,<1.5',
-                    'geojson>=2.5.0,<3'],
+                    'geojson>=2.5.0,<3',
+                    'backports.zoneinfo<1; python_version<"3.9"'],
         test=['tox>=3,<4',
               'zope.testing>=4,<5',
               'zope.testrunner>=5,<6',

src/crate/client/connection.py

@@ -47,6 +47,7 @@ def __init__(self,
                  socket_tcp_keepintvl=None,
                  socket_tcp_keepcnt=None,
                  converter=None,
+                 time_zone=None,
                  ):
         """
         :param servers:
@@ -103,9 +104,28 @@ def __init__(self,
         :param converter:
             (optional, defaults to ``None``)
             A `Converter` object to propagate to newly created `Cursor` objects.
+        :param time_zone:
+            (optional, defaults to ``None``)
+            A time zone specifier used for returning `TIMESTAMP` types as
+            timezone-aware native Python `datetime` objects.
+
+            Different data types are supported. Available options are:
+
+            - ``datetime.timezone.utc``
+            - ``datetime.timezone(datetime.timedelta(hours=7), name="MST")``
+            - ``pytz.timezone("Australia/Sydney")``
+            - ``zoneinfo.ZoneInfo("Australia/Sydney")``
+            - ``+0530`` (UTC offset in string format)
+
+            When `time_zone` is `None`, the returned `datetime` objects are
+            "naive", without any `tzinfo`, converted using ``datetime.utcfromtimestamp(...)``.
+
+            When `time_zone` is given, the returned `datetime` objects are "aware",
+            with `tzinfo` set, converted using ``datetime.fromtimestamp(..., tz=...)``.
         """
 
         self._converter = converter
+        self.time_zone = time_zone
 
         if client:
             self.client = client
@@ -135,10 +155,12 @@ def cursor(self, **kwargs) -> Cursor:
         Return a new Cursor Object using the connection.
         """
         converter = kwargs.pop("converter", self._converter)
+        time_zone = kwargs.pop("time_zone", self.time_zone)
         if not self._closed:
             return Cursor(
                 connection=self,
                 converter=converter,
+                time_zone=time_zone,
             )
         else:
             raise ProgrammingError("Connection closed")

src/crate/client/converter.py

@@ -123,6 +123,9 @@ def convert(value: Any) -> Optional[List[Any]]:
 
         return convert
 
+    def set(self, type_: DataType, converter: ConverterFunction):
+        self._mappings[type_] = converter
+
 
 class DefaultTypeConverter(Converter):
     def __init__(self, more_mappings: Optional[ConverterMapping] = None) -> None:

src/crate/client/cursor.py

@@ -18,8 +18,11 @@
 # However, if you have executed another commercial license agreement
 # with Crate these terms will supersede the license and you may use the
 # software solely pursuant to the terms of the relevant commercial agreement.
+from datetime import datetime, timedelta, timezone
 
+from .converter import DataType
 import warnings
+import typing as t
 
 from .converter import Converter
 from .exceptions import ProgrammingError
@@ -32,13 +35,15 @@ class Cursor(object):
     """
     lastrowid = None  # currently not supported
 
-    def __init__(self, connection, converter: Converter):
+    def __init__(self, connection, converter: Converter, **kwargs):
         self.arraysize = 1
         self.connection = connection
         self._converter = converter
         self._closed = False
         self._result = None
         self.rows = None
+        self._time_zone = None
+        self.time_zone = kwargs.get("time_zone")
 
     def execute(self, sql, parameters=None, bulk_parameters=None):
         """
@@ -241,3 +246,72 @@ def _convert_rows(self):
                 convert(value)
                 for convert, value in zip(converters, row)
             ]
+
+    @property
+    def time_zone(self):
+        """
+        Get the current time zone.
+        """
+        return self._time_zone
+
+    @time_zone.setter
+    def time_zone(self, tz):
+        """
+        Set the time zone.
+
+        Different data types are supported. Available options are:
+
+        - ``datetime.timezone.utc``
+        - ``datetime.timezone(datetime.timedelta(hours=7), name="MST")``
+        - ``pytz.timezone("Australia/Sydney")``
+        - ``zoneinfo.ZoneInfo("Australia/Sydney")``
+        - ``+0530`` (UTC offset in string format)
+
+        When `time_zone` is `None`, the returned `datetime` objects are
+        "naive", without any `tzinfo`, converted using ``datetime.utcfromtimestamp(...)``.
+
+        When `time_zone` is given, the returned `datetime` objects are "aware",
+        with `tzinfo` set, converted using ``datetime.fromtimestamp(..., tz=...)``.
+        """
+
+        # Do nothing when time zone is reset.
+        if tz is None:
+            self._time_zone = None
+            return
+
+        # Requesting datetime-aware `datetime` objects needs the data type converter.
+        # Implicitly create one, when needed.
+        if self._converter is None:
+            self._converter = Converter()
+
+        # When the time zone is given as a string, assume UTC offset format, e.g. `+0530`.
+        if isinstance(tz, str):
+            tz = self._timezone_from_utc_offset(tz)
+
+        self._time_zone = tz
+
+        def _to_datetime_with_tz(value: t.Optional[float]) -> t.Optional[datetime]:
+            """
+            Convert CrateDB's `TIMESTAMP` value to a native Python `datetime`
+            object, with timezone-awareness.
+            """
+            if value is None:
+                return None
+            return datetime.fromtimestamp(value / 1e3, tz=self._time_zone)
+
+        # Register converter function for `TIMESTAMP` type.
+        self._converter.set(DataType.TIMESTAMP_WITH_TZ, _to_datetime_with_tz)
+        self._converter.set(DataType.TIMESTAMP_WITHOUT_TZ, _to_datetime_with_tz)
+
+    @staticmethod
+    def _timezone_from_utc_offset(tz) -> timezone:
+        """
+        Convert UTC offset in string format (e.g. `+0530`) into `datetime.timezone` object.
+        """
+        assert len(tz) == 5, f"Time zone '{tz}' is given in invalid UTC offset format"
+        try:
+            hours = int(tz[:3])
+            minutes = int(tz[0] + tz[3:])
+            return timezone(timedelta(hours=hours, minutes=minutes), name=tz)
+        except Exception as ex:
+            raise ValueError(f"Time zone '{tz}' is given in invalid UTC offset format: {ex}")

src/crate/client/doctests/cursor.txt

@@ -364,6 +364,76 @@ Proof that the converter works correctly, ``B\'0110\'`` should be converted to
     [6]
 
 
+``TIMESTAMP`` conversion with time zone
+=======================================
+
+Based on the data type converter functionality, the driver offers a convenient
+interface to make it return timezone-aware ``datetime`` objects, using the
+desired time zone.
+
+For your reference, in the following examples, epoch 1658167836758 is
+``Mon, 18 Jul 2022 18:10:36 GMT``.
+
+::
+
+    >>> import datetime
+    >>> tz_mst = datetime.timezone(datetime.timedelta(hours=7), name="MST")
+    >>> cursor = connection.cursor(time_zone=tz_mst)
+
+    >>> connection.client.set_next_response({
+    ...     "col_types": [4, 11],
+    ...     "rows":[ [ "foo", 1658167836758 ] ],
+    ...     "cols":[ "name", "timestamp" ],
+    ...     "rowcount":1,
+    ...     "duration":123
+    ... })
+
+    >>> cursor.execute('')
+
+    >>> cursor.fetchone()
+    ['foo', datetime.datetime(2022, 7, 19, 1, 10, 36, 758000, tzinfo=datetime.timezone(datetime.timedelta(seconds=25200), 'MST'))]
+
+For the ``time_zone`` keyword argument, different data types are supported.
+The available options are:
+
+- ``datetime.timezone.utc``
+- ``datetime.timezone(datetime.timedelta(hours=7), name="MST")``
+- ``pytz.timezone("Australia/Sydney")``
+- ``zoneinfo.ZoneInfo("Australia/Sydney")``
+- ``+0530`` (UTC offset in string format)
+
+Let's exercise all of them::
+
+    >>> cursor.time_zone = datetime.timezone.utc
+    >>> cursor.execute('')
+    >>> cursor.fetchone()
+    ['foo', datetime.datetime(2022, 7, 18, 18, 10, 36, 758000, tzinfo=datetime.timezone.utc)]
+
+    >>> import pytz
+    >>> cursor.time_zone = pytz.timezone("Australia/Sydney")
+    >>> cursor.execute('')
+    >>> cursor.fetchone()
+    ['foo', datetime.datetime(2022, 7, 19, 4, 10, 36, 758000, tzinfo=<DstTzInfo 'Australia/Sydney' AEST+10:00:00 STD>)]
+
+    >>> try:
+    ...     import zoneinfo
+    ... except ImportError:
+    ...     from backports import zoneinfo
+    >>> cursor.time_zone = zoneinfo.ZoneInfo("Australia/Sydney")
+    >>> cursor.execute('')
+    >>> record = cursor.fetchone()
+    >>> record
+    ['foo', datetime.datetime(2022, 7, 19, 4, 10, 36, 758000, ...zoneinfo.ZoneInfo(key='Australia/Sydney'))]
+
+    >>> record[1].tzname()
+    'AEST'
+
+    >>> cursor.time_zone = "+0530"
+    >>> cursor.execute('')
+    >>> cursor.fetchone()
+    ['foo', datetime.datetime(2022, 7, 18, 23, 40, 36, 758000, tzinfo=datetime.timezone(datetime.timedelta(seconds=19800), '+0530'))]
+
+
 .. Hidden: close connection
 
     >>> connection.close()

src/crate/client/test_connection.py

@@ -1,3 +1,5 @@
+import datetime
+
 from .http import Client
 from crate.client import connect
 from unittest import TestCase
@@ -23,7 +25,25 @@ def test_invalid_server_version(self):
         self.assertEqual((0, 0, 0), connection.lowest_server_version.version)
         connection.close()
 
-    def test_with_is_supported(self):
+    def test_context_manager(self):
         with connect('localhost:4200') as conn:
             pass
         self.assertEqual(conn._closed, True)
+
+    def test_with_timezone(self):
+        """
+        Verify the cursor objects will return timezone-aware `datetime` objects when requested to.
+        When switching the time zone at runtime on the connection object, only new cursor objects
+        will inherit the new time zone.
+        """
+
+        tz_mst = datetime.timezone(datetime.timedelta(hours=7), name="MST")
+        connection = connect('localhost:4200', time_zone=tz_mst)
+        cursor = connection.cursor()
+        self.assertEqual(cursor.time_zone.tzname(None), "MST")
+        self.assertEqual(cursor.time_zone.utcoffset(None), datetime.timedelta(seconds=25200))
+
+        connection.time_zone = datetime.timezone.utc
+        cursor = connection.cursor()
+        self.assertEqual(cursor.time_zone.tzname(None), "UTC")
+        self.assertEqual(cursor.time_zone.utcoffset(None), datetime.timedelta(0))