Update to crate-2.0.0, which uses orjson for JSON marshalling

Author: amotl

CHANGES.md

@@ -3,6 +3,7 @@
 ## Unreleased
 - Fixed connectivity for `jobstats collect`
 - Refactored code and improved CLI interface of `ctk info` vs. `ctk cfr`
+- Dependencies: Updated to `crate-2.0.0`, which uses `orjson` for JSON marshalling
 
 ## 2025/01/13 v0.0.30
 - Dependencies: Minimize dependencies of core installation,

cratedb_toolkit/cfr/systable.py

@@ -32,7 +32,6 @@
 from cratedb_toolkit.info.core import InfoContainer
 from cratedb_toolkit.util import DatabaseAdapter
 from cratedb_toolkit.util.cli import error_logger
-from cratedb_toolkit.util.sqlalchemy import patch_encoder
 
 logger = logging.getLogger(__name__)
 
@@ -256,6 +255,3 @@ def load_table(self, path: Path) -> "pl.DataFrame":
             return pl.read_parquet(path)
         else:
             raise NotImplementedError(f"Input format not implemented: {path.suffix}")
-
-
-patch_encoder()

cratedb_toolkit/util/sqlalchemy.py

@@ -4,57 +4,8 @@
 TODO: Refactor to `crate` or `sqlalchemy-cratedb` packages.
 """
 
-import calendar
-import datetime as dt
-import json
-from decimal import Decimal
-from uuid import UUID
-
-from sqlalchemy_cratedb.dialect import TYPES_MAP
-
-try:
-    import numpy as np
-
-    has_numpy = True
-except ImportError:
-    has_numpy = False
-
 from sqlalchemy import types as sqltypes
-
-
-def patch_encoder():
-    import crate.client.http
-
-    crate.client.http.CrateJsonEncoder = CrateJsonEncoderWithNumPy
-
-
-class CrateJsonEncoderWithNumPy(json.JSONEncoder):
-    epoch_aware = dt.datetime(1970, 1, 1, tzinfo=dt.timezone.utc)
-    epoch_naive = dt.datetime(1970, 1, 1)
-
-    def default(self, o):
-        # Vanilla CrateDB Python.
-        if isinstance(o, (Decimal, UUID)):
-            return str(o)
-        if isinstance(o, dt.datetime):
-            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, dt.date):
-            return calendar.timegm(o.timetuple()) * 1000
-
-        # NumPy ndarray and friends.
-        # https://stackoverflow.com/a/49677241
-        if has_numpy:
-            if isinstance(o, np.integer):
-                return int(o)
-            elif isinstance(o, np.floating):
-                return float(o)
-            elif isinstance(o, np.ndarray):
-                return o.tolist()
-        return json.JSONEncoder.default(self, o)
+from sqlalchemy_cratedb.dialect import TYPES_MAP
 
 
 def patch_types_map():

pyproject.toml

@@ -91,6 +91,7 @@ dependencies = [
   "colorama<1",
   "colorlog",
   "crash",
+  "crate>=2.0.0.dev2",
   "cratedb-sqlparse==0.0.10",
   'importlib-metadata; python_version < "3.8"',
   'importlib-resources; python_version < "3.9"',
@@ -100,7 +101,7 @@ dependencies = [
   "python-dotenv<2",
   "python-slugify<9",
   "pyyaml<7",
-  "sqlalchemy-cratedb>=0.40,<1",
+  "sqlalchemy-cratedb @ git+https://github.com/crate/sqlalchemy-cratedb.git@crate-orjson",
   "sqlparse<0.6",
   "tqdm<5",
   'typing-extensions<5; python_version <= "3.7"',

tests/sqlalchemy/test_patch.py

@@ -1,10 +1,10 @@
 import datetime
-import json
 
 import pytest
 import sqlalchemy as sa
+from crate.client.http import cratedb_json_encoder
+from orjson import orjson
 
-from cratedb_toolkit.util.sqlalchemy import CrateJsonEncoderWithNumPy
 from tests.conftest import TESTDRIVE_DATA_SCHEMA
 
 
@@ -48,18 +48,36 @@ def test_inspector_patched(database, needs_sqlalchemy2):
 def test_json_encoder_date():
     """
     Verify the extended JSON encoder also accepts Python's `date` types.
+
+    TODO: Move to different test file, as this no longer requires
+          monkeypatching after using orjson for marshalling.
     """
     data = {"date": datetime.date(2024, 6, 4)}
-    encoded = json.dumps(data, cls=CrateJsonEncoderWithNumPy)
-    assert encoded == '{"date": 1717459200000}'
+    # TODO: Refactor core marshalling method out of crate-python's
+    #       `_create_sql_payload`, to make it reusable here.
+    encoded = orjson.dumps(
+        data,
+        default=cratedb_json_encoder,
+        option=(orjson.OPT_PASSTHROUGH_DATETIME | orjson.OPT_SERIALIZE_NUMPY),
+    )
+    assert encoded == b'{"date":1717459200000}'
 
 
 def test_json_encoder_numpy():
     """
     Verify the extended JSON encoder also accepts NumPy types.
+
+    TODO: Move to different test file, as this no longer requires
+          monkeypatching after using orjson for marshalling.
     """
     np = pytest.importorskip("numpy")
 
     data = {"scalar-int": np.float32(42.42).astype(int), "scalar-float": np.float32(42.42), "ndarray": np.ndarray([1])}
-    encoded = json.dumps(data, cls=CrateJsonEncoderWithNumPy)
-    assert encoded == """{"scalar-int": 42, "scalar-float": 42.41999816894531, "ndarray": [2.08e-322]}"""
+    # TODO: Refactor core marshalling method out of crate-python's
+    #       `_create_sql_payload`, to make it reusable here.
+    encoded = orjson.dumps(
+        data,
+        default=cratedb_json_encoder,
+        option=(orjson.OPT_PASSTHROUGH_DATETIME | orjson.OPT_SERIALIZE_NUMPY),
+    )
+    assert encoded == b"""{"scalar-int":42,"scalar-float":42.42,"ndarray":[2.08e-322]}"""