datachain 0.30.5__py3-none-any.whl → 0.39.0__py3-none-any.whl

datachain/data_storage/schema.py

@@ -1,12 +1,6 @@
 import inspect
 from collections.abc import Iterable, Iterator, Sequence
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Generic,
-    Optional,
-    TypeVar,
-)
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 import sqlalchemy as sa
 from sqlalchemy.sql import func as f
@@ -17,7 +11,6 @@ from datachain.sql.types import (
     JSON,
     Boolean,
     DateTime,
-    Int,
     Int64,
     SQLType,
     String,
@@ -51,7 +44,7 @@ def dedup_columns(columns: Iterable[sa.Column]) -> list[sa.Column]:
     """
     c_set: dict[str, sa.Column] = {}
     for c in columns:
-        if (ec := c_set.get(c.name, None)) is not None:
+        if (ec := c_set.get(c.name)) is not None:
             if str(ec.type) != str(c.type):
                 raise ValueError(
                     f"conflicting types for column {c.name}:{c.type!s} and {ec.type!s}"
@@ -96,11 +89,11 @@ class DirExpansion:
     def __init__(self, column: str):
         self.column = column
 
-    def col_name(self, name: str, column: Optional[str] = None) -> str:
+    def col_name(self, name: str, column: str | None = None) -> str:
         column = column or self.column
         return col_name(name, column)
 
-    def c(self, query, name: str, column: Optional[str] = None) -> str:
+    def c(self, query, name: str, column: str | None = None) -> str:
         return getattr(query.c, self.col_name(name, column=column))
 
     def base_select(self, q):
@@ -161,7 +154,7 @@ class DataTable:
         self,
         name: str,
         engine: "DatabaseEngine",
-        column_types: Optional[dict[str, SQLType]] = None,
+        column_types: dict[str, SQLType] | None = None,
         column: str = "file",
     ):
         self.name: str = name
@@ -172,12 +165,12 @@ class DataTable:
     @staticmethod
     def copy_column(
         column: sa.Column,
-        primary_key: Optional[bool] = None,
-        index: Optional[bool] = None,
-        nullable: Optional[bool] = None,
-        default: Optional[Any] = None,
-        server_default: Optional[Any] = None,
-        unique: Optional[bool] = None,
+        primary_key: bool | None = None,
+        index: bool | None = None,
+        nullable: bool | None = None,
+        default: Any | None = None,
+        server_default: Any | None = None,
+        unique: bool | None = None,
     ) -> sa.Column:
         """
         Copy a sqlalchemy Column object intended for use as a signal column.
@@ -206,8 +199,8 @@ class DataTable:
     def new_table(
         cls,
         name: str,
-        columns: Sequence["sa.Column"] = (),
-        metadata: Optional["sa.MetaData"] = None,
+        columns: Sequence[sa.Column] = (),
+        metadata: sa.MetaData | None = None,
     ):
         # copy columns, since reusing the same objects from another table
         # may raise an error
@@ -218,7 +211,7 @@ class DataTable:
             metadata = sa.MetaData()
         return sa.Table(name, metadata, *columns)
 
-    def get_table(self) -> "sa.Table":
+    def get_table(self) -> sa.Table:
         table = self.engine.get_table(self.name)
 
         column_types = self.column_types | {c.name: c.type for c in self.sys_columns()}
@@ -233,19 +226,19 @@ class DataTable:
     def columns(self) -> "ReadOnlyColumnCollection[str, sa.Column[Any]]":
         return self.table.columns
 
-    def col_name(self, name: str, column: Optional[str] = None) -> str:
+    def col_name(self, name: str, column: str | None = None) -> str:
         column = column or self.column
         return col_name(name, column)
 
-    def without_object(self, column_name: str, column: Optional[str] = None) -> str:
+    def without_object(self, column_name: str, column: str | None = None) -> str:
         column = column or self.column
         return column_name.removeprefix(f"{column}{DEFAULT_DELIMITER}")
 
-    def c(self, name: str, column: Optional[str] = None):
+    def c(self, name: str, column: str | None = None):
         return getattr(self.columns, self.col_name(name, column=column))
 
     @property
-    def table(self) -> "sa.Table":
+    def table(self) -> sa.Table:
         return self.get_table()
 
     def apply_conditions(self, query: "Executable") -> "Executable":
@@ -275,7 +268,7 @@ class DataTable:
     @classmethod
     def sys_columns(cls):
         return [
-            sa.Column("sys__id", Int, primary_key=True),
+            sa.Column("sys__id", UInt64, primary_key=True),
             sa.Column(
                 "sys__rand", UInt64, nullable=False, server_default=f.abs(f.random())
             ),
@@ -303,7 +296,7 @@ PARTITION_COLUMN_ID = "partition_id"
 partition_col_names = [PARTITION_COLUMN_ID]
 
 
-def partition_columns() -> Sequence["sa.Column"]:
+def partition_columns() -> Sequence[sa.Column]:
     return [
         sa.Column(PARTITION_COLUMN_ID, sa.Integer),
     ]

datachain/data_storage/serializer.py

@@ -1,29 +1,119 @@
 import base64
-import pickle
 from abc import abstractmethod
 from collections.abc import Callable
-from typing import Any
+from typing import Any, ClassVar
+
+from datachain import json
+from datachain.plugins import ensure_plugins_loaded
+
+
+class CallableRegistry:
+    _registry: ClassVar[dict[str, Callable]] = {}
+
+    @classmethod
+    def register(cls, callable_obj: Callable, name: str) -> str:
+        cls._registry[name] = callable_obj
+        return name
+
+    @classmethod
+    def get(cls, name: str) -> Callable:
+        return cls._registry[name]
 
 
 class Serializable:
+    @classmethod
+    @abstractmethod
+    def serialize_callable_name(cls) -> str:
+        """Return the registered name used for this class' factory callable."""
+
     @abstractmethod
     def clone_params(self) -> tuple[Callable[..., Any], list[Any], dict[str, Any]]:
-        """
-        Returns the class, args, and kwargs needed to instantiate a cloned copy
-        of this instance for use in separate processes or machines.
-        """
+        """Return (callable, args, kwargs) necessary to recreate this object."""
+
+    def _prepare(self, params: tuple) -> dict:
+        callable, args, kwargs = params
+        callable_name = callable.__self__.serialize_callable_name()
+        return {
+            "callable": callable_name,
+            "args": args,
+            "kwargs": {
+                k: self._prepare(v) if isinstance(v, tuple) else v
+                for k, v in kwargs.items()
+            },
+        }
 
     def serialize(self) -> str:
-        """
-        Returns a string representation of clone params.
-        This is useful for storing the state of an object in environment variable.
-        """
-        return base64.b64encode(pickle.dumps(self.clone_params())).decode()
+        """Return a base64-encoded JSON string with registered callable + params."""
+        _ensure_default_callables_registered()
+        data = self.clone_params()
+        return base64.b64encode(json.dumps(self._prepare(data)).encode()).decode()
 
 
 def deserialize(s: str) -> Serializable:
+    """Deserialize from base64-encoded JSON using only registered callables.
+
+    Nested serialized objects are instantiated automatically except for those
+    passed via clone parameter tuples (keys ending with ``_clone_params``),
+    which must remain as (callable, args, kwargs) for later factory usage.
     """
-    Returns a new instance of the class represented by the string.
-    """
-    (f, args, kwargs) = pickle.loads(base64.b64decode(s.encode()))  # noqa: S301
-    return f(*args, **kwargs)
+    ensure_plugins_loaded()
+    _ensure_default_callables_registered()
+    decoded = base64.b64decode(s.encode())
+    data = json.loads(decoded.decode())
+
+    def _is_serialized(obj: Any) -> bool:
+        return isinstance(obj, dict) and {"callable", "args", "kwargs"}.issubset(
+            obj.keys()
+        )
+
+    def _reconstruct(obj: Any, nested: bool = False) -> Any:
+        if not _is_serialized(obj):
+            return obj
+        callable_name: str = obj["callable"]
+        args: list[Any] = obj["args"]
+        kwargs: dict[str, Any] = obj["kwargs"]
+        # Recurse only inside kwargs because serialize() only nests through kwargs
+        for k, v in list(kwargs.items()):
+            if _is_serialized(v):
+                kwargs[k] = _reconstruct(v, True)
+        callable_obj = CallableRegistry.get(callable_name)
+        if nested:
+            return (callable_obj, args, kwargs)
+        # Otherwise instantiate
+        return callable_obj(*args, **kwargs)
+
+    if not _is_serialized(data):
+        raise ValueError("Invalid serialized data format")
+    return _reconstruct(data, False)
+
+
+class _DefaultsState:
+    registered = False
+
+
+def _ensure_default_callables_registered() -> None:
+    if _DefaultsState.registered:
+        return
+
+    from datachain.data_storage.sqlite import (
+        SQLiteDatabaseEngine,
+        SQLiteMetastore,
+        SQLiteWarehouse,
+    )
+
+    # Register (idempotent by name overwrite is fine) using class-level
+    # serialization names to avoid hard-coded literals here.
+    CallableRegistry.register(
+        SQLiteDatabaseEngine.from_db_file,
+        SQLiteDatabaseEngine.serialize_callable_name(),
+    )
+    CallableRegistry.register(
+        SQLiteMetastore.init_after_clone,
+        SQLiteMetastore.serialize_callable_name(),
+    )
+    CallableRegistry.register(
+        SQLiteWarehouse.init_after_clone,
+        SQLiteWarehouse.serialize_callable_name(),
+    )
+
+    _DefaultsState.registered = True