unify json usage, make sure datetime is handled (#1444)

* unify json usage, make sure datetime is handled

* cleanup and simplify serializer

Author: shcheklein

.github/workflows/tests-studio.yml

@@ -69,17 +69,6 @@ jobs:
           ref: ${{ env.STUDIO_BRANCH }}
           token: ${{ secrets.ITERATIVE_STUDIO_READ_ACCESS_TOKEN }}
 
-      - name: Check out repository
-        uses: actions/checkout@v5
-        with:
-          path: './backend/datachain'
-          fetch-depth: 0
-
-      - name: Install FFmpeg
-        run: |
-          sudo apt update
-          sudo apt install -y ffmpeg
-
       - name: Set up Python ${{ matrix.pyv }}
         uses: actions/setup-python@v6
         with:
@@ -94,6 +83,23 @@ jobs:
             backend/datachain_server/pyproject.toml
             backend/datachain/pyproject.toml
 
+      - name: Update DataChain requirement in Studio
+        env:
+          DATACHAIN_BRANCH: ${{ env.BRANCH }}
+        working-directory: backend
+        run: make update_datachain_deps "$DATACHAIN_BRANCH"
+
+      - name: Check out repository
+        uses: actions/checkout@v5
+        with:
+          path: './backend/datachain'
+          fetch-depth: 0
+
+      - name: Install FFmpeg
+        run: |
+          sudo apt update
+          sudo apt install -y ffmpeg
+
       - name: Install dependencies
         run: uv pip install --system ./backend/datachain_server[tests] ./backend/datachain[tests]
 

src/datachain/catalog/catalog.py

@@ -1,5 +1,4 @@
 import io
-import json
 import logging
 import os
 import os.path
@@ -21,6 +20,7 @@
 from sqlalchemy import Column
 from tqdm.auto import tqdm
 
+from datachain import json
 from datachain.cache import Cache
 from datachain.client import Client
 from datachain.dataset import (

src/datachain/data_storage/metastore.py

@@ -1,5 +1,4 @@
 import copy
-import json
 import logging
 import os
 from abc import ABC, abstractmethod
@@ -27,6 +26,7 @@
 )
 from sqlalchemy.sql import func as f
 
+from datachain import json
 from datachain.catalog.dependency import DatasetDependencyNode
 from datachain.checkpoint import Checkpoint
 from datachain.data_storage import JobQueryType, JobStatus
@@ -53,7 +53,6 @@
 from datachain.job import Job
 from datachain.namespace import Namespace
 from datachain.project import Project
-from datachain.utils import JSONSerialize
 
 if TYPE_CHECKING:
     from sqlalchemy import CTE, Delete, Insert, Select, Subquery, Update
@@ -1194,7 +1193,7 @@ def update_dataset_version(
                         f"Field '{field}' must be a list, got {type(value).__name__}"
                     )
                 else:
-                    values[field] = json.dumps(value, cls=JSONSerialize)
+                    values[field] = json.dumps(value, serialize_bytes=True)
                 version_values["_preview_data"] = value
             else:
                 values[field] = value

src/datachain/data_storage/serializer.py

@@ -1,9 +1,9 @@
 import base64
-import json
 from abc import abstractmethod
 from collections.abc import Callable
 from typing import Any, ClassVar
 
+from datachain import json
 from datachain.plugins import ensure_plugins_loaded
 
 

src/datachain/data_storage/warehouse.py

@@ -10,9 +10,9 @@
 
 import attrs
 import sqlalchemy as sa
-import ujson as json
 from sqlalchemy.sql.expression import true
 
+from datachain import json
 from datachain.client import Client
 from datachain.data_storage.schema import convert_rows_custom_column_types
 from datachain.data_storage.serializer import Serializable

src/datachain/dataset.py

@@ -1,5 +1,4 @@
 import builtins
-import json
 from dataclasses import dataclass, fields
 from datetime import datetime
 from functools import cached_property
@@ -9,7 +8,7 @@
 from packaging.specifiers import SpecifierSet
 from packaging.version import Version
 
-from datachain import semver
+from datachain import json, semver
 from datachain.error import DatasetVersionNotFoundError, InvalidDatasetNameError
 from datachain.namespace import Namespace
 from datachain.project import Project

src/datachain/hash_utils.py

@@ -1,12 +1,13 @@
 import hashlib
 import inspect
-import json
 import textwrap
 from collections.abc import Sequence
 from typing import TypeAlias, TypeVar
 
 from sqlalchemy.sql.elements import ClauseElement, ColumnElement
 
+from datachain import json
+
 T = TypeVar("T", bound=ColumnElement)
 ColumnLike: TypeAlias = str | T
 
@@ -72,7 +73,9 @@ def hash_column_elements(columns: ColumnLike | Sequence[ColumnLike]) -> str:
         columns = (columns,)
 
     serialized = [serialize_column_element(c) for c in columns]
-    json_str = json.dumps(serialized, sort_keys=True)  # stable JSON
+    json_str = json.dumps(
+        serialized, sort_keys=True, separators=(", ", ": ")
+    )  # stable JSON
     return hashlib.sha256(json_str.encode("utf-8")).hexdigest()
 
 

src/datachain/job.py

@@ -1,9 +1,10 @@
-import json
 import uuid
 from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, TypeVar
 
+from datachain import json
+
 J = TypeVar("J", bound="Job")
 
 

src/datachain/json.py

@@ -0,0 +1,138 @@
+"""DataChain JSON utilities.
+
+This module wraps :mod:`ujson` so we can guarantee consistent handling
+of values that the encoder does not support out of the box (for example
+``datetime`` objects or ``bytes``).
+All code inside DataChain should import this module instead of using
+:mod:`ujson` directly.
+"""
+
+import datetime as _dt
+import json as _json
+import uuid as _uuid
+from collections.abc import Callable
+from typing import Any
+
+import ujson as _ujson
+
+__all__ = [
+    "JSONDecodeError",
+    "dump",
+    "dumps",
+    "load",
+    "loads",
+]
+
+JSONDecodeError = (_ujson.JSONDecodeError, _json.JSONDecodeError)
+
+_SENTINEL = object()
+_Default = Callable[[Any], Any]
+DEFAULT_PREVIEW_BYTES = 1024
+
+
+# To make it looks like Pydantic's ISO format with 'Z' for UTC
+# It is minor but nice to have consistency
+def _format_datetime(value: _dt.datetime) -> str:
+    iso = value.isoformat()
+
+    offset = value.utcoffset()
+    if value.tzinfo is None or offset is None:
+        return iso
+
+    if offset == _dt.timedelta(0) and iso.endswith(("+00:00", "-00:00")):
+        return iso[:-6] + "Z"
+
+    return iso
+
+
+def _format_time(value: _dt.time) -> str:
+    iso = value.isoformat()
+
+    offset = value.utcoffset()
+    if value.tzinfo is None or offset is None:
+        return iso
+
+    if offset == _dt.timedelta(0) and iso.endswith(("+00:00", "-00:00")):
+        return iso[:-6] + "Z"
+
+    return iso
+
+
+def _coerce(value: Any, serialize_bytes: bool) -> Any:
+    """Return a JSON-serializable representation for supported extra types."""
+
+    if isinstance(value, _dt.datetime):
+        return _format_datetime(value)
+    if isinstance(value, _dt.date):
+        return value.isoformat()
+    if isinstance(value, _dt.time):
+        return _format_time(value)
+    if isinstance(value, _uuid.UUID):
+        return str(value)
+    if serialize_bytes and isinstance(value, (bytes, bytearray)):
+        return list(bytes(value)[:DEFAULT_PREVIEW_BYTES])
+    return _SENTINEL
+
+
+def _base_default(value: Any, serialize_bytes: bool) -> Any:
+    converted = _coerce(value, serialize_bytes)
+    if converted is not _SENTINEL:
+        return converted
+    raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")
+
+
+def _build_default(user_default: _Default | None, serialize_bytes: bool) -> _Default:
+    if user_default is None:
+        return lambda value: _base_default(value, serialize_bytes)
+
+    def combined(value: Any) -> Any:
+        converted = _coerce(value, serialize_bytes)
+        if converted is not _SENTINEL:
+            return converted
+        return user_default(value)
+
+    return combined
+
+
+def dumps(
+    obj: Any,
+    *,
+    default: _Default | None = None,
+    serialize_bytes: bool = False,
+    **kwargs: Any,
+) -> str:
+    """Serialize *obj* to a JSON-formatted ``str``."""
+
+    if serialize_bytes:
+        return _json.dumps(obj, default=_build_default(default, True), **kwargs)
+
+    return _ujson.dumps(obj, default=_build_default(default, False), **kwargs)
+
+
+def dump(
+    obj: Any,
+    fp,
+    *,
+    default: _Default | None = None,
+    serialize_bytes: bool = False,
+    **kwargs: Any,
+) -> None:
+    """Serialize *obj* as a JSON formatted stream to *fp*."""
+
+    if serialize_bytes:
+        _json.dump(obj, fp, default=_build_default(default, True), **kwargs)
+        return
+
+    _ujson.dump(obj, fp, default=_build_default(default, False), **kwargs)
+
+
+def loads(s: str | bytes | bytearray, **kwargs: Any) -> Any:
+    """Deserialize *s* to a Python object."""
+
+    return _ujson.loads(s, **kwargs)
+
+
+def load(fp, **kwargs: Any) -> Any:
+    """Deserialize JSON content from *fp* to a Python object."""
+
+    return loads(fp.read(), **kwargs)

src/datachain/lib/arrow.py

@@ -3,11 +3,11 @@
 from typing import TYPE_CHECKING, Any
 
 import pyarrow as pa
-import ujson as json
 from pyarrow._csv import ParseOptions
 from pyarrow.dataset import CsvFileFormat, dataset
 from tqdm.auto import tqdm
 
+from datachain import json
 from datachain.fs.reference import ReferenceFileSystem
 from datachain.lib.data_model import dict_to_data_model
 from datachain.lib.file import ArrowRow, File