feat(codecs): explicit context argument, additive to key (#1556)

`encode()` and `decode()` take a `context` argument carrying `schema`, `table`, `field` and `config`, so `key` means what it means everywhere else in DataJoint: the primary key.

Passing the calling connection's `config` is what matters. Without it, `_build_path` and `_get_backend` fall back to the global `dj.config`, which in a process holding connections for several users belongs to none of them and resolves a different store silently rather than raising.

Nothing breaks. DataJoint passes `context` only to codecs whose signature declares it, the pre-2.3.4 underscore keys are still populated in `key`, and `_extract_context(key)` still accepts one argument. Codecs written before 2.3.4 keep working unchanged; `key` reverting to primary-key-only, and `config` becoming required, wait for 2.4.

Closes #1550.

Author: dimitri-yatsenko

src/datajoint/builtin_codecs/__init__.py

@@ -29,14 +29,14 @@ class GraphCodec(dj.Codec):
             def get_dtype(self, is_store: bool) -> str:
                 return "<blob>"  # Compose with blob for serialization
 
-            def encode(self, graph, *, key=None, store_name=None):
+            def encode(self, graph, *, key=None, context=None, store_name=None):
                 # Convert graph to a serializable format
                 return {
                     'nodes': list(graph.nodes(data=True)),
                     'edges': list(graph.edges(data=True)),
                 }
 
-            def decode(self, stored, *, key=None):
+            def decode(self, stored, *, key=None, context=None):
                 # Reconstruct graph from stored format
                 G = nx.Graph()
                 G.add_nodes_from(stored['nodes'])

src/datajoint/builtin_codecs/attach.py

@@ -50,7 +50,9 @@ def get_dtype(self, is_store: bool) -> str:
         """Return bytes for in-table, <hash> for in-store storage."""
         return "<hash>" if is_store else "bytes"
 
-    def encode(self, value: Any, *, key: dict | None = None, store_name: str | None = None) -> bytes:
+    def encode(
+        self, value: Any, *, key: dict | None = None, context: dict | None = None, store_name: str | None = None
+    ) -> bytes:
         """
         Read file and encode as filename + contents.
 
@@ -80,7 +82,7 @@ def encode(self, value: Any, *, key: dict | None = None, store_name: str | None
         contents = path.read_bytes()
         return filename.encode("utf-8") + b"\x00" + contents
 
-    def decode(self, stored: bytes, *, key: dict | None = None) -> str:
+    def decode(self, stored: bytes, *, key: dict | None = None, context: dict | None = None) -> str:
         """
         Extract file to download path and return local path.
 
@@ -104,7 +106,7 @@ def decode(self, stored: bytes, *, key: dict | None = None) -> str:
         contents = stored[null_pos + 1 :]
 
         # Write to download path
-        config = (key or {}).get("_config")
+        config = self._codec_config(key, context)
         if config is None:
             from ..settings import config  # type: ignore[assignment]
         assert config is not None

src/datajoint/builtin_codecs/filepath.py

@@ -74,7 +74,9 @@ def get_dtype(self, is_store: bool) -> str:
             )
         return "json"
 
-    def encode(self, value: Any, *, key: dict | None = None, store_name: str | None = None) -> dict:
+    def encode(
+        self, value: Any, *, key: dict | None = None, context: dict | None = None, store_name: str | None = None
+    ) -> dict:
         """
         Store path reference as JSON metadata.
 
@@ -104,7 +106,7 @@ def encode(self, value: Any, *, key: dict | None = None, store_name: str | None
 
         from ..hash_registry import get_store_backend
 
-        config = (key or {}).get("_config")
+        config = self._codec_config(key, context)
         if config is None:
             from ..settings import config  # type: ignore[assignment]
         assert config is not None
@@ -168,7 +170,7 @@ def encode(self, value: Any, *, key: dict | None = None, store_name: str | None
             "timestamp": datetime.now(timezone.utc).isoformat(),
         }
 
-    def decode(self, stored: dict, *, key: dict | None = None) -> Any:
+    def decode(self, stored: dict, *, key: dict | None = None, context: dict | None = None) -> Any:
         """
         Create ObjectRef handle for lazy access.
 
@@ -187,7 +189,7 @@ def decode(self, stored: dict, *, key: dict | None = None) -> Any:
         from ..objectref import ObjectRef
         from ..hash_registry import get_store_backend
 
-        config = (key or {}).get("_config")
+        config = self._codec_config(key, context)
         store_name = stored.get("store")
         backend = get_store_backend(store_name, config=config)
         return ObjectRef.from_json(stored, backend=backend)

src/datajoint/builtin_codecs/hash.py

@@ -57,7 +57,9 @@ def get_dtype(self, is_store: bool) -> str:
             raise DataJointError("<hash> requires @ (in-store storage only)")
         return "json"
 
-    def encode(self, value: bytes, *, key: dict | None = None, store_name: str | None = None) -> dict:
+    def encode(
+        self, value: bytes, *, key: dict | None = None, context: dict | None = None, store_name: str | None = None
+    ) -> dict:
         """
         Store content and return metadata.
 
@@ -66,7 +68,11 @@ def encode(self, value: bytes, *, key: dict | None = None, store_name: str | Non
         value : bytes
             Raw bytes to store.
         key : dict, optional
-            Context dict with ``_schema`` for path isolation.
+            Primary key values. Before 2.3.4 this also carried the schema name
+            under ``_schema``; that is still read when ``context`` omits it.
+        context : dict, optional
+            Connection context. ``schema`` isolates the stored path and
+            ``config`` resolves the store.
         store_name : str, optional
             Store to use. If None, uses default store.
 
@@ -77,11 +83,14 @@ def encode(self, value: bytes, *, key: dict | None = None, store_name: str | Non
         """
         from ..hash_registry import put_hash
 
-        schema_name = (key or {}).get("_schema", "unknown")
-        config = (key or {}).get("_config")
+        # Same precedence as `_extract_context`. Reading `key` alone would store
+        # under "unknown" for any caller that passes `context` with a key holding
+        # only primary-key values -- which is every caller once 2.4 lands.
+        schema_name = (context or {}).get("schema", (key or {}).get("_schema", "unknown"))
+        config = self._codec_config(key, context)
         return put_hash(value, schema_name=schema_name, store_name=store_name, config=config)
 
-    def decode(self, stored: dict, *, key: dict | None = None) -> bytes:
+    def decode(self, stored: dict, *, key: dict | None = None, context: dict | None = None) -> bytes:
         """
         Retrieve content using stored metadata.
 
@@ -99,7 +108,7 @@ def decode(self, stored: dict, *, key: dict | None = None) -> bytes:
         """
         from ..hash_registry import get_hash
 
-        config = (key or {}).get("_config")
+        config = self._codec_config(key, context)
         return get_hash(stored, config=config)
 
     def validate(self, value: Any) -> None:

src/datajoint/builtin_codecs/npy.py

@@ -312,6 +312,7 @@ def encode(
         value: Any,
         *,
         key: dict | None = None,
+        context: dict | None = None,
         store_name: str | None = None,
     ) -> dict:
         """
@@ -337,8 +338,8 @@ def encode(
         import numpy as np
 
         # Extract context using inherited helper
-        schema, table, field, primary_key = self._extract_context(key)
-        config = (key or {}).get("_config")
+        schema, table, field, primary_key = self._extract_context(key, context)
+        config = self._codec_config(key, context)
 
         # Build schema-addressed storage path
         path, _ = self._build_path(schema, table, field, primary_key, ext=".npy", store_name=store_name, config=config)
@@ -360,7 +361,7 @@ def encode(
             "shape": list(value.shape),
         }
 
-    def decode(self, stored: dict, *, key: dict | None = None) -> NpyRef:
+    def decode(self, stored: dict, *, key: dict | None = None, context: dict | None = None) -> NpyRef:
         """
         Create lazy NpyRef from stored metadata.
 
@@ -376,6 +377,6 @@ def decode(self, stored: dict, *, key: dict | None = None) -> NpyRef:
         NpyRef
             Lazy array reference with metadata access and numpy integration.
         """
-        config = (key or {}).get("_config")
+        config = self._codec_config(key, context)
         backend = self._get_backend(stored.get("store"), config=config)
         return NpyRef(stored, backend)

src/datajoint/builtin_codecs/object.py

@@ -81,6 +81,7 @@ def encode(
         value: Any,
         *,
         key: dict | None = None,
+        context: dict | None = None,
         store_name: str | None = None,
     ) -> dict:
         """
@@ -105,8 +106,8 @@ def encode(
         from pathlib import Path
 
         # Extract context using inherited helper
-        schema, table, field, primary_key = self._extract_context(key)
-        config = (key or {}).get("_config")
+        schema, table, field, primary_key = self._extract_context(key, context)
+        config = self._codec_config(key, context)
 
         # Check for pre-computed metadata (from staged insert)
         if isinstance(value, dict) and "path" in value:
@@ -177,7 +178,7 @@ def encode(
 
         return metadata
 
-    def decode(self, stored: dict, *, key: dict | None = None) -> Any:
+    def decode(self, stored: dict, *, key: dict | None = None, context: dict | None = None) -> Any:
         """
         Create ObjectRef handle for lazy access.
 
@@ -195,7 +196,7 @@ def decode(self, stored: dict, *, key: dict | None = None) -> Any:
         """
         from ..objectref import ObjectRef
 
-        config = (key or {}).get("_config")
+        config = self._codec_config(key, context)
         backend = self._get_backend(stored.get("store"), config=config)
         return ObjectRef.from_json(stored, backend=backend)
 

src/datajoint/builtin_codecs/schema.py

@@ -4,6 +4,8 @@
 
 from __future__ import annotations
 
+import warnings
+
 from ..codecs import Codec
 from ..errors import DataJointError
 
@@ -24,14 +26,27 @@ class SchemaCodec(Codec, register=False):
         - ``validate()``: Validate input values
 
     Helper Methods:
-        - ``_extract_context()``: Parse key dict into schema/table/field/pk
+        - ``_extract_context()``: Parse key/context into schema/table/field/pk
+        - ``_codec_config()``: Read the calling connection's config
         - ``_build_path()``: Construct storage path from context
         - ``_get_backend()``: Get storage backend by name
 
-    Both helpers take a ``config`` and fall back to the global ``dj.config``
-    without one. Read it off ``key["_config"]`` and pass it through, as below: it
-    is the calling connection's config, and in a process holding connections for
-    several users the global one belongs to none of them.
+    ``_build_path`` and ``_get_backend`` take a ``config`` and fall back to the
+    global ``dj.config`` without one. Always pass the calling connection's
+    config: in a process holding connections for several users, the global one
+    belongs to none of them, and the fallback resolves a different store
+    silently rather than raising.
+
+    Since 2.3.4 that config arrives in an explicit ``context`` argument rather
+    than hidden among the primary key values. **Declare ``context=None`` in
+    ``encode`` and ``decode`` and pass it to the helpers**, as the example below
+    does.
+
+    The old underscore keys in ``key`` still work and are still populated, so a
+    codec written before 2.3.4 keeps running. They are removed in 2.4.
+    ``_extract_context`` warns when it falls back to them; ``_codec_config`` and
+    a direct ``key["_config"]`` do not, so the warning is a prompt rather than a
+    complete inventory of what still reads them.
 
     Comparison with Hash-addressed:
         - **Schema-addressed** (this): Path from schema structure, no dedup
@@ -42,9 +57,9 @@ class SchemaCodec(Codec, register=False):
         class MyCodec(SchemaCodec):
             name = "my"
 
-            def encode(self, value, *, key=None, store_name=None):
-                schema, table, field, pk = self._extract_context(key)
-                config = (key or {}).get("_config")
+            def encode(self, value, *, key=None, context=None, store_name=None):
+                schema, table, field, pk = self._extract_context(key, context)
+                config = self._codec_config(key, context)
                 path, _ = self._build_path(
                     schema, table, field, pk, ext=".dat",
                     store_name=store_name, config=config,
@@ -53,8 +68,8 @@ def encode(self, value, *, key=None, store_name=None):
                 backend.put_buffer(serialize(value), path)
                 return {"path": path, "store": store_name, ...}
 
-            def decode(self, stored, *, key=None):
-                config = (key or {}).get("_config")
+            def decode(self, stored, *, key=None, context=None):
+                config = self._codec_config(key, context)
                 backend = self._get_backend(stored.get("store"), config=config)
                 return MyRef(stored, backend)
 
@@ -88,25 +103,41 @@ def get_dtype(self, is_store: bool) -> str:
             raise DataJointError(f"<{self.name}> requires @ (store only)")
         return "json"
 
-    def _extract_context(self, key: dict | None) -> tuple[str, str, str, dict]:
+    def _extract_context(self, key: dict | None, context: dict | None = None) -> tuple[str, str, str, dict]:
         """
-        Extract schema, table, field, and primary key from context dict.
+        Extract schema, table, field, and primary key.
 
         Parameters
         ----------
         key : dict or None
-            Context dict with ``_schema``, ``_table``, ``_field``,
-            and primary key values.
+            Primary key values. Before 2.3.4 this also carried connection
+            context under ``_schema``, ``_table``, ``_field`` and ``_config``;
+            those keys are still populated and still read, with a
+            ``DeprecationWarning``, when ``context`` is not supplied.
+        context : dict or None
+            Connection context with ``schema``, ``table``, ``field`` and
+            ``config``. Pass the ``context`` argument your ``encode``/``decode``
+            received.
 
         Returns
         -------
         tuple[str, str, str, dict]
             ``(schema, table, field, primary_key)``
         """
         key = dict(key) if key else {}
-        schema = key.pop("_schema", "unknown")
-        table = key.pop("_table", "unknown")
-        field = key.pop("_field", "data")
+        if context is None and any(k.startswith("_") for k in key):
+            warnings.warn(
+                "Reading connection context from the `key` dict is deprecated and will "
+                "be removed in DataJoint 2.4. Accept a `context` argument in encode()/decode() "
+                "and pass it to _extract_context(key, context). See "
+                "https://github.com/datajoint/datajoint-python/issues/1550",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        context = context or {}
+        schema = context.get("schema", key.pop("_schema", "unknown"))
+        table = context.get("table", key.pop("_table", "unknown"))
+        field = context.get("field", key.pop("_field", "data"))
         primary_key = {k: v for k, v in key.items() if not k.startswith("_")}
         return schema, table, field, primary_key