WIP

Author: jjnesbitt

dandi/files/bases.py

@@ -1,28 +1,27 @@
 from __future__ import annotations
 
+import os
+import re
 from abc import ABC, abstractmethod
 from collections import deque
 from collections.abc import Iterator
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass
 from datetime import datetime
-import os
 from pathlib import Path
-import re
 from threading import Lock
 from typing import IO, Any, Generic
 from xml.etree.ElementTree import fromstring
 
 import dandischema
+import requests
 from dandischema.consts import DANDI_SCHEMA_VERSION
 from dandischema.digests.dandietag import DandiETag
-from dandischema.models import BareAsset, CommonModel
+from dandischema.models import BareAsset, CommonModel, get_schema_version
 from dandischema.models import Dandiset as DandisetMeta
-from dandischema.models import get_schema_version
 from packaging.version import Version
 from pydantic import ValidationError
 from pydantic_core import ErrorDetails
-import requests
 
 import dandi
 from dandi.dandiapi import RemoteAsset, RemoteDandiset, RESTFullAPIClient
@@ -320,6 +319,7 @@ def iter_upload(
         metadata: dict[str, Any],
         jobs: int | None = None,
         replacing: RemoteAsset | None = None,
+        zarr_id: str = "",
     ) -> Iterator[dict]:
         """
         Upload the file as an asset with the given metadata to the given
@@ -335,6 +335,9 @@ def iter_upload(
         :param RemoteAsset replacing:
             If set, replace the given asset, which must have the same path as
             the new asset
+        :param str zarr_id:
+            The zarr id being uploaded to, if applicable. If set, this will
+            modify the request params to the archive.
         :returns:
             A generator of `dict`\\s containing at least a ``"status"`` key.
             Upon successful upload, the last `dict` will have a status of
@@ -361,17 +364,23 @@ def iter_upload(
         lgr.debug("%s: Beginning upload", asset_path)
         total_size = pre_upload_size_check(self.filepath)
         try:
-            resp = client.post(
-                "/uploads/initialize/",
-                json={
-                    "contentSize": total_size,
-                    "digest": {
-                        "algorithm": "dandi:dandi-etag",
-                        "value": filetag,
-                    },
-                    "dandiset": dandiset.identifier,
+            body = {
+                "contentSize": total_size,
+                "digest": {
+                    "algorithm": "dandi:dandi-etag",
+                    "value": filetag,
                 },
-            )
+            }
+
+            if zarr_id:
+                body["zarr"] = {
+                    "chunk_key": asset_path,
+                    "zarr_id": zarr_id,
+                }
+            else:
+                body["dandiset"] = dandiset.identifier
+
+            resp = client.post("/uploads/initialize/", json=body)
         except requests.HTTPError as e:
             if e.response is not None and e.response.status_code == 409:
                 lgr.debug("%s: Blob already exists on server", asset_path)
@@ -583,8 +592,8 @@ def get_validation_errors(
                 )
 
         # Avoid circular imports by importing within function:
-        from .bids import NWBBIDSAsset
         from ..organize import validate_organized_path
+        from .bids import NWBBIDSAsset
 
         if not isinstance(self, NWBBIDSAsset) and self.dandiset_path is not None:
             errors.extend(

dandi/files/zarr.py

@@ -1,23 +1,23 @@
 from __future__ import annotations
 
+import json
+import os
+import os.path
+import urllib.parse
 from base64 import b64encode
 from collections.abc import Generator, Iterator
 from concurrent.futures import Future, ThreadPoolExecutor, as_completed
 from contextlib import closing
 from dataclasses import dataclass, field, replace
 from datetime import datetime
 from enum import Enum
-import json
-import os
-import os.path
 from pathlib import Path
 from time import sleep
 from typing import Any, Optional
-import urllib.parse
 
+import requests
 from dandischema.models import BareAsset, DigestType
 from pydantic import BaseModel, ConfigDict, ValidationError
-import requests
 from zarr_checksum.tree import ZarrChecksumTree
 
 from dandi import __version__ as dandi_version
@@ -47,7 +47,6 @@
     pre_upload_size_check,
 )
 
-from .bases import LocalDirectoryAsset, LocalFileAsset
 from ..validate._types import (
     ORIGIN_VALIDATION_DANDI_ZARR,
     Origin,
@@ -58,6 +57,7 @@
     ValidationResult,
     Validator,
 )
+from .bases import LocalDirectoryAsset, LocalFileAsset
 
 lgr = get_logger()
 
@@ -764,7 +764,10 @@ def mkzarr() -> str:
                             path=it.entry_path,
                         )
                         for status in large_asset.iter_upload(
-                            dandiset, {"path": it.entry_path}, jobs=jobs
+                            dandiset,
+                            {"path": it.entry_path},
+                            jobs=jobs,
+                            zarr_id=zarr_id,
                         ):
                             if status.get("status") == "done":
                                 changed = True

dandi/tests/test_files.py

@@ -4,7 +4,7 @@
 import os
 from pathlib import Path
 import subprocess
-from unittest.mock import ANY
+from unittest.mock import ANY, patch
 
 from dandischema.models import get_schema_version
 import numpy as np
@@ -536,6 +536,36 @@ def test_upload_zarr_entry_content_type(new_dandiset, tmp_path):
     assert r.headers["Content-Type"] == "application/json"
 
 
+@pytest.mark.ai_generated
+def test_upload_zarr_large_chunks(new_dandiset, tmp_path):
+    """Chunks above ZARR_LARGE_CHUNK_THRESHOLD are uploaded via LocalFileAsset.iter_upload."""
+    filepath = tmp_path / "example.zarr"
+    zarr.save(filepath, np.arange(1000), np.arange(1000, 0, -1))
+    zf = dandi_file(filepath)
+    assert isinstance(zf, ZarrAsset)
+
+    from ..files.bases import LocalFileAsset
+
+    real_iter_upload = LocalFileAsset.iter_upload
+    called_paths: list[str] = []
+
+    def spy_iter_upload(self, *args, **kwargs):
+        called_paths.append(self.path)
+        yield from real_iter_upload(self, *args, **kwargs)
+
+    # Set threshold to 0 so every chunk is treated as "large"
+    with (
+        patch("dandi.files.zarr.ZARR_LARGE_CHUNK_THRESHOLD", 0),
+        patch.object(LocalFileAsset, "iter_upload", spy_iter_upload),
+    ):
+        asset = zf.upload(new_dandiset.dandiset, {})
+
+    assert isinstance(asset, RemoteZarrAsset)
+    # Every chunk file in the zarr should have been routed through LocalFileAsset.iter_upload
+    remote_entries = {str(e) for e in asset.iterfiles()}
+    assert remote_entries == set(called_paths)
+
+
 def test_validate_deep_zarr(tmp_path: Path) -> None:
     zarr_path = tmp_path / "foo.zarr"
     zarr.save(zarr_path, np.arange(1000), np.arange(1000, 0, -1))