Marshalling: Restore CrateDB standard encoder

Author: amotl

src/crate/client/http.py

@@ -20,6 +20,8 @@
 # software solely pursuant to the terms of the relevant commercial agreement.
 
 
+import calendar
+import datetime as dt
 import heapq
 import io
 import logging
@@ -84,19 +86,35 @@ def super_len(o):
     return None
 
 
-def cratedb_json_encoder(obj: t.Any) -> str:
+epoch_aware = dt.datetime(1970, 1, 1, tzinfo=dt.timezone.utc)
+epoch_naive = dt.datetime(1970, 1, 1)
+
+
+def json_encoder(obj: t.Any) -> t.Union[int, str]:
     """
     Encoder function for orjson, with additional type support.
 
-    - Python's `Decimal` type.
-    - freezegun's `FakeDatetime` type.
+    - Python's `Decimal` type will be serialized to `str`.
+    - Python's `dt.datetime` and `dt.date` types will be
+      serialized to `int` after converting to milliseconds
+      since epoch.
 
     https://github.com/ijl/orjson#default
+    https://cratedb.com/docs/crate/reference/en/latest/general/ddl/data-types.html#type-timestamp
     """
     if isinstance(obj, Decimal):
         return str(obj)
-    elif hasattr(obj, "isoformat"):
-        return obj.isoformat()
+    if isinstance(obj, dt.datetime):
+        if obj.tzinfo is not None:
+            delta = obj - epoch_aware
+        else:
+            delta = obj - epoch_naive
+        return int(
+            delta.microseconds / 1000.0
+            + (delta.seconds + delta.days * 24 * 3600) * 1000.0
+        )
+    if isinstance(obj, dt.date):
+        return calendar.timegm(obj.timetuple()) * 1000
     raise TypeError
 
 
@@ -108,8 +126,12 @@ def json_dumps(obj: t.Any) -> bytes:
     """
     return orjson.dumps(
         obj,
-        default=cratedb_json_encoder,
-        option=(orjson.OPT_NON_STR_KEYS | orjson.OPT_SERIALIZE_NUMPY),
+        default=json_encoder,
+        option=(
+            orjson.OPT_PASSTHROUGH_DATETIME
+            | orjson.OPT_NON_STR_KEYS
+            | orjson.OPT_SERIALIZE_NUMPY
+        ),
     )
 
 

tests/client/test_http.py

@@ -318,7 +318,7 @@ def test_datetime_is_converted_to_ts(self, request):
         # convert string to dict
         # because the order of the keys isn't deterministic
         data = json.loads(request.call_args[1]["data"])
-        self.assertEqual(data["args"], ["2015-02-28T07:31:40"])
+        self.assertEqual(data["args"], [1425108700000])
         client.close()
 
     @patch(REQUEST, autospec=True)
@@ -329,7 +329,7 @@ def test_date_is_converted_to_ts(self, request):
         day = dt.date(2016, 4, 21)
         client.sql("insert into users (dt) values (?)", (day,))
         data = json.loads(request.call_args[1]["data"])
-        self.assertEqual(data["args"], ["2016-04-21"])
+        self.assertEqual(data["args"], [1461196800000])
         client.close()
 
     def test_socket_options_contain_keepalive(self):
@@ -725,9 +725,9 @@ class TestCrateJsonEncoder(TestCase):
     def test_naive_datetime(self):
         data = dt.datetime.fromisoformat("2023-06-26T09:24:00.123")
         result = json_dumps(data)
-        self.assertEqual(result, b'"2023-06-26T09:24:00.123000"')
+        self.assertEqual(result, b"1687771440123")
 
     def test_aware_datetime(self):
         data = dt.datetime.fromisoformat("2023-06-26T09:24:00.123+02:00")
         result = json_dumps(data)
-        self.assertEqual(result, b'"2023-06-26T09:24:00.123000+02:00"')
+        self.assertEqual(result, b"1687764240123")