diff --git a/CHANGES.txt b/CHANGES.txt index a37f5352..7dfdef33 100644 --- a/CHANGES.txt +++ b/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 ================= diff --git a/docs/query.rst b/docs/query.rst index b80bcb9c..36f11983 100644 --- a/docs/query.rst +++ b/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=)] + + >>> 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/ diff --git a/setup.py b/setup.py index 66035caa..7070977a 100644 --- a/setup.py +++ b/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', diff --git a/src/crate/client/connection.py b/src/crate/client/connection.py index 84e4ec8e..ef1aefef 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,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") diff --git a/src/crate/client/converter.py b/src/crate/client/converter.py index d0dad3b0..c4dbf598 100644 --- a/src/crate/client/converter.py +++ b/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: diff --git a/src/crate/client/cursor.py b/src/crate/client/cursor.py index 25b88667..c458ae1b 100644 --- a/src/crate/client/cursor.py +++ b/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}") diff --git a/src/crate/client/doctests/cursor.txt b/src/crate/client/doctests/cursor.txt index f1f2ee6b..e3a92132 100644 --- a/src/crate/client/doctests/cursor.txt +++ b/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=)] + + >>> 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() diff --git a/src/crate/client/test_connection.py b/src/crate/client/test_connection.py index 5faa46a8..078f88c8 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 020ba7f5..79e7ddd6 100644 --- a/src/crate/client/test_cursor.py +++ b/src/crate/client/test_cursor.py @@ -19,10 +19,16 @@ # 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 +try: + import zoneinfo +except ImportError: + from backports import zoneinfo + +import pytz from crate.client import connect from crate.client.converter import DataType, DefaultTypeConverter @@ -32,6 +38,93 @@ 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_zoneinfo_object(self): + """ + Verify the cursor returns timezone-aware `datetime` objects when requested to. + Here: Use a `zoneinfo.ZoneInfo` instance. + """ + connection = self.get_mocked_connection() + cursor = connection.cursor(time_zone=zoneinfo.ZoneInfo('Australia/Sydney')) + self.assertEqual(cursor.time_zone.key, 'Australia/Sydney') + + 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) @@ -78,7 +171,7 @@ def test_execute_with_converter(self): c.execute("") result = c.fetchall() self.assertEqual(result, [ - ['foo', IPv4Address('10.10.10.1'), datetime(2022, 7, 18, 18, 10, 36, 758000), 6], + ['foo', IPv4Address('10.10.10.1'), datetime.datetime(2022, 7, 18, 18, 10, 36, 758000), 6], [None, None, None, None], ]) @@ -193,3 +286,56 @@ def test_executemany_with_converter(self): # ``executemany()`` is not intended to be used with statements returning result # sets. The result will always be empty. self.assertEqual(result, []) + + 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], + ], + }) + + # 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()