diff --git a/src/crate/client/connection.py b/src/crate/client/connection.py index 271a9a734..2d9d6f8ad 100644 --- a/src/crate/client/connection.py +++ b/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,14 @@ 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. """ self._converter = converter + self.time_zone = time_zone if client: self.client = client @@ -135,10 +141,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") diff --git a/src/crate/client/cursor.py b/src/crate/client/cursor.py index 69f68a7ec..fdb1ad767 100644 --- a/src/crate/client/cursor.py +++ b/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}") diff --git a/src/crate/client/doctests/cursor.txt b/src/crate/client/doctests/cursor.txt index 07b9e3d6b..03b27b191 100644 --- a/src/crate/client/doctests/cursor.txt +++ b/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=)] + + >>> 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() diff --git a/src/crate/client/test_connection.py b/src/crate/client/test_connection.py index 5faa46a89..078f88c85 100644 --- a/src/crate/client/test_connection.py +++ b/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)) diff --git a/src/crate/client/test_cursor.py b/src/crate/client/test_cursor.py index 2ede6ebd0..bf55749c8 100644 --- a/src/crate/client/test_cursor.py +++ b/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()