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.

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

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:
@@ -106,6 +107,7 @@ def __init__(self,
         """
 
         self._converter = converter
+        self.time_zone = time_zone
 
         if client:
             self.client = client
@@ -135,10 +137,12 @@ def cursor(self, cursor=None, **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/cursor.py

@@ -18,10 +18,12 @@
 # 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 Converter, DefaultTypeConverter
+from .converter import Converter, DefaultTypeConverter, CrateDatatypeIdentifier
 from .exceptions import ProgrammingError
 import warnings
+import typing as t
 
 
 class Cursor(object):
@@ -31,13 +33,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):
         """
@@ -238,3 +242,64 @@ def get_default_converter() -> Converter:
         Return the standard converter instance.
         """
         return DefaultTypeConverter()
+
+    @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.
+
+        It supports different ways to populate. Some examples::
+
+        - ``datetime.timezone.utc``
+        - ``datetime.timezone(datetime.timedelta(hours=7), name="MST")``
+        - ``pytz.timezone("Australia/Sydney")``
+        - ``+0530`` (UTC offset in string format)
+        """
+
+        # 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(CrateDatatypeIdentifier.TIMESTAMP, _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

@@ -363,6 +363,60 @@ Proof that the converter works correctly, ``B\'0110\'`` should be converted to
     [6]
 
 
+``TIMESTAMP`` conversion with time zone
+=======================================
+
+Based on the data type converter machinery, the driver offers a convenient
+interface to make it return timezone-aware ``datetime`` objects, using the
+desired time zone.
+
+For your reference, 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'))]
+
+There are different ways to populate ``time_zone``. Some examples::
+
+- ``datetime.timezone.utc``
+- ``datetime.timezone(datetime.timedelta(hours=7), name="MST")``
+- ``pytz.timezone("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>)]
+
+    >>> 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))

src/crate/client/test_cursor.py

@@ -19,11 +19,13 @@
 # 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
+import datetime
 from ipaddress import IPv4Address
 from unittest import TestCase
 from unittest.mock import MagicMock
 
+import pytz
+
 from crate.client import connect
 from crate.client.converter import CrateDatatypeIdentifier
 from crate.client.cursor import Cursor
@@ -33,6 +35,84 @@
 
 class CursorTest(TestCase):
 
+    @staticmethod
+    def get_mocked_connection():
+        client = MagicMock(spec=Client)
+        return connect(client=client)
+
+    def test_create_with_timezone_as_datetime_object(self):
+        """
+        Verify the cursor returns timezone-aware `datetime` objects when requested to.
+        Switching the time zone at runtime on the cursor object is possible.
+        Here: Use a `datetime.timezone` instance.
+        """
+
+        connection = self.get_mocked_connection()
+
+        tz_mst = datetime.timezone(datetime.timedelta(hours=7), name="MST")
+        cursor = connection.cursor(time_zone=tz_mst)
+
+        self.assertEqual(cursor.time_zone.tzname(None), "MST")
+        self.assertEqual(cursor.time_zone.utcoffset(None), datetime.timedelta(seconds=25200))
+
+        cursor.time_zone = datetime.timezone.utc
+        self.assertEqual(cursor.time_zone.tzname(None), "UTC")
+        self.assertEqual(cursor.time_zone.utcoffset(None), datetime.timedelta(0))
+
+    def test_create_with_timezone_as_pytz_object(self):
+        """
+        Verify the cursor returns timezone-aware `datetime` objects when requested to.
+        Here: Use a `pytz.timezone` instance.
+        """
+        connection = self.get_mocked_connection()
+        cursor = connection.cursor(time_zone=pytz.timezone('Australia/Sydney'))
+        self.assertEqual(cursor.time_zone.tzname(None), "Australia/Sydney")
+
+        # Apparently, when using `pytz`, the timezone object does not return an offset.
+        # Nevertheless, it works, as demonstrated per doctest in `cursor.txt`.
+        self.assertEqual(cursor.time_zone.utcoffset(None), None)
+
+    def test_create_with_timezone_as_utc_offset_success(self):
+        """
+        Verify the cursor returns timezone-aware `datetime` objects when requested to.
+        Here: Use a UTC offset in string format.
+        """
+        connection = self.get_mocked_connection()
+        cursor = connection.cursor(time_zone="+0530")
+        self.assertEqual(cursor.time_zone.tzname(None), "+0530")
+        self.assertEqual(cursor.time_zone.utcoffset(None), datetime.timedelta(seconds=19800))
+
+        connection = self.get_mocked_connection()
+        cursor = connection.cursor(time_zone="-1145")
+        self.assertEqual(cursor.time_zone.tzname(None), "-1145")
+        self.assertEqual(cursor.time_zone.utcoffset(None), datetime.timedelta(days=-1, seconds=44100))
+
+    def test_create_with_timezone_as_utc_offset_failure(self):
+        """
+        Verify the cursor croaks when trying to create it with invalid UTC offset strings.
+        """
+        connection = self.get_mocked_connection()
+        with self.assertRaises(AssertionError) as ex:
+            connection.cursor(time_zone="foobar")
+        self.assertEqual(str(ex.exception), "Time zone 'foobar' is given in invalid UTC offset format")
+
+        connection = self.get_mocked_connection()
+        with self.assertRaises(ValueError) as ex:
+            connection.cursor(time_zone="+abcd")
+        self.assertEqual(str(ex.exception), "Time zone '+abcd' is given in invalid UTC offset format: "
+                                            "invalid literal for int() with base 10: '+ab'")
+
+    def test_create_with_timezone_connection_cursor_precedence(self):
+        """
+        Verify that the time zone specified on the cursor object instance
+        takes precedence over the one specified on the connection instance.
+        """
+        client = MagicMock(spec=Client)
+        connection = connect(client=client, time_zone=pytz.timezone('Australia/Sydney'))
+        cursor = connection.cursor(time_zone="+0530")
+        self.assertEqual(cursor.time_zone.tzname(None), "+0530")
+        self.assertEqual(cursor.time_zone.utcoffset(None), datetime.timedelta(seconds=19800))
+
     def test_execute_with_args(self):
         client = MagicMock(spec=Client)
         conn = connect(client=client)
@@ -82,7 +162,7 @@ def test_execute_with_converter(self):
             [
                 'foo',
                 IPv4Address('10.10.10.1'),
-                datetime(2022, 7, 18, 18, 10, 36, 758000),
+                datetime.datetime(2022, 7, 18, 18, 10, 36, 758000),
                 6,
             ],
             [
@@ -145,3 +225,58 @@ def test_execute_nested_array_with_converter(self):
             'foo',
             [[IPv4Address('10.10.10.1'), IPv4Address('10.10.10.2')], [IPv4Address('10.10.10.3')], [], None],
         ])
+
+    def test_execute_with_timezone(self):
+        client = ClientMocked()
+        conn = connect(client=client)
+
+        # Create a `Cursor` object with `time_zone`.
+        tz_mst = datetime.timezone(datetime.timedelta(hours=7), name="MST")
+        c = conn.cursor(time_zone=tz_mst)
+
+        # Make up a response using CrateDB data type `TIMESTAMP`.
+        conn.client.set_next_response({
+            "col_types": [4, 11],
+            "cols": ["name", "timestamp"],
+            "rows": [
+                ["foo", 1658167836758],
+                [None, None],
+            ],
+            "rowcount": 1,
+            "duration": 123
+        })
+
+        # Run execution and verify the returned `datetime` object is timezone-aware,
+        # using the designated timezone object.
+        c.execute("")
+        result = c.fetchall()
+        self.assertEqual(result, [
+            [
+                'foo',
+                datetime.datetime(2022, 7, 19, 1, 10, 36, 758000,
+                                  tzinfo=datetime.timezone(datetime.timedelta(seconds=25200), 'MST')),
+            ],
+            [
+                None,
+                None,
+            ],
+        ])
+        self.assertEqual(result[0][1].tzname(), "MST")
+
+        # Change timezone and verify the returned `datetime` object is using it.
+        c.time_zone = datetime.timezone.utc
+        c.execute("")
+        result = c.fetchall()
+        self.assertEqual(result, [
+            [
+                'foo',
+                datetime.datetime(2022, 7, 18, 18, 10, 36, 758000, tzinfo=datetime.timezone.utc),
+            ],
+            [
+                None,
+                None,
+            ],
+        ])
+        self.assertEqual(result[0][1].tzname(), "UTC")
+
+        conn.close()