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()