Allow handling "aware" datetime values when inserting or updating

CHANGES.txt

@@ -11,6 +11,8 @@ Unreleased
 - SQLAlchemy DDL: Allow setting ``server_default`` on columns to enable
   server-generated defaults.
 
+- Allow handling "aware" datetime values with time zone info when inserting or updating.
+
 
 2023/04/18 0.31.1
 =================

docs/by-example/client.rst

@@ -139,6 +139,24 @@ Refresh locations:
 
     >>> cursor.execute("REFRESH TABLE locations")
 
+Updating Data
+=============
+
+Both when inserting or updating data, values for ``TIMESTAMP`` columns can be obtained
+in different formats. Both scalar strings and datetime objects are supported.
+
+    >>> import datetime as dt
+    >>> timestamp_full = "2023-06-26T09:24:00.123+02:00"
+    >>> timestamp_date = "2023-06-26"
+    >>> datetime_aware = dt.datetime.fromisoformat("2023-06-26T09:24:00.123+02:00")
+    >>> datetime_naive = dt.datetime.fromisoformat("2023-06-26T09:24:00.123")
+    >>> datetime_date = dt.date.fromisoformat("2023-06-26")
+    >>> cursor.execute("UPDATE locations SET date=? WHERE name='Cloverleaf'", (timestamp_full, ))
+    >>> cursor.execute("UPDATE locations SET date=? WHERE name='Cloverleaf'", (timestamp_date, ))
+    >>> cursor.execute("UPDATE locations SET date=? WHERE name='Cloverleaf'", (datetime_aware, ))
+    >>> cursor.execute("UPDATE locations SET date=? WHERE name='Cloverleaf'", (datetime_naive, ))
+    >>> cursor.execute("UPDATE locations SET date=? WHERE name='Cloverleaf'", (datetime_date, ))
+
 Selecting Data
 ==============
 

src/crate/client/http.py

@@ -33,7 +33,7 @@
 from urllib.parse import urlparse
 from base64 import b64encode
 from time import time
-from datetime import datetime, date
+from datetime import datetime, date, timezone
 from decimal import Decimal
 from urllib3 import connection_from_url
 from urllib3.connection import HTTPConnection
@@ -82,13 +82,17 @@ def super_len(o):
 
 class CrateJsonEncoder(json.JSONEncoder):
 
-    epoch = datetime(1970, 1, 1)
+    epoch_aware = datetime(1970, 1, 1, tzinfo=timezone.utc)
+    epoch_naive = datetime(1970, 1, 1)
 
     def default(self, o):
         if isinstance(o, Decimal):
             return str(o)
         if isinstance(o, datetime):
-            delta = o - self.epoch
+            if o.tzinfo is not None:
+                delta = o - self.epoch_aware
+            else:
+                delta = o - self.epoch_naive
             return int(delta.microseconds / 1000.0 +
                        (delta.seconds + delta.days * 24 * 3600) * 1000.0)
         if isinstance(o, date):

src/crate/client/test_http.py

@@ -40,7 +40,7 @@
 from urllib.parse import urlparse, parse_qs
 from setuptools.ssl_support import find_ca_bundle
 
-from .http import Client, _get_socket_opts, _remove_certs_for_non_https
+from .http import Client, CrateJsonEncoder, _get_socket_opts, _remove_certs_for_non_https
 from .exceptions import ConnectionError, ProgrammingError
 
 
@@ -626,3 +626,16 @@ def test_username(self):
         self.assertEqual(TestingHTTPServer.SHARED['usernameFromXUser'], 'testDBUser')
         self.assertEqual(TestingHTTPServer.SHARED['username'], 'testDBUser')
         self.assertEqual(TestingHTTPServer.SHARED['password'], 'test:password')
+
+
+class TestCrateJsonEncoder(TestCase):
+
+    def test_naive_datetime(self):
+        data = dt.datetime.fromisoformat("2023-06-26T09:24:00.123")
+        result = json.dumps(data, cls=CrateJsonEncoder)
+        self.assertEqual(result, "1687771440123")
+
+    def test_aware_datetime(self):
+        data = dt.datetime.fromisoformat("2023-06-26T09:24:00.123+02:00")
+        result = json.dumps(data, cls=CrateJsonEncoder)
+        self.assertEqual(result, "1687764240123")

src/crate/client/tests.py

@@ -51,6 +51,7 @@
     RetryOnTimeoutServerTest,
     RequestsCaBundleTest,
     TestUsernameSentAsHeader,
+    TestCrateJsonEncoder,
     TestDefaultSchemaHeader,
 )
 from .sqlalchemy.tests import test_suite as sqlalchemy_test_suite
@@ -341,6 +342,7 @@ def test_suite():
     suite.addTest(unittest.makeSuite(RetryOnTimeoutServerTest))
     suite.addTest(unittest.makeSuite(RequestsCaBundleTest))
     suite.addTest(unittest.makeSuite(TestUsernameSentAsHeader))
+    suite.addTest(unittest.makeSuite(TestCrateJsonEncoder))
     suite.addTest(unittest.makeSuite(TestDefaultSchemaHeader))
     suite.addTest(sqlalchemy_test_suite())
     suite.addTest(doctest.DocTestSuite('crate.client.connection'))