from __future__ import annotations
from collections import defaultdict
from collections.abc import Iterator, Sequence
from contextlib import ExitStack
from enum import Enum
from functools import reduce
import io
import os.path
from pathlib import Path
import re
import time
from time import sleep
from typing import Any, TypedDict, cast
from unittest.mock import patch
import click
from packaging.version import Version
from . import __version__, lgr
from .consts import (
DRAFT,
DandiInstance,
dandiset_identifier_regex,
dandiset_metadata_file,
)
from .dandiapi import DandiAPIClient, RemoteAsset
from .dandiset import Dandiset
from .exceptions import NotFoundError, UploadError
from .files import (
DandiFile,
DandisetMetadataFile,
LocalAsset,
LocalDirectoryAsset,
ZarrAsset,
)
from .misctypes import Digest
from .support import pyout as pyouts
from .support.pyout import naturalsize
from .utils import ensure_datetime, path_is_subpath, pluralize
from .validate_types import Severity
[docs]class Uploaded(TypedDict):
size: int
errors: list[str]
[docs]class UploadExisting(str, Enum):
ERROR = "error"
SKIP = "skip"
FORCE = "force"
OVERWRITE = "overwrite"
REFRESH = "refresh"
def __str__(self) -> str:
return self.value
[docs]class UploadValidation(str, Enum):
REQUIRE = "require"
SKIP = "skip"
IGNORE = "ignore"
def __str__(self) -> str:
return self.value
[docs]def upload(
paths: Sequence[str | Path] | None = None,
existing: UploadExisting = UploadExisting.REFRESH,
validation: UploadValidation = UploadValidation.REQUIRE,
dandi_instance: str | DandiInstance = "dandi",
allow_any_path: bool = False,
upload_dandiset_metadata: bool = False,
devel_debug: bool = False,
jobs: int | None = None,
jobs_per_file: int | None = None,
sync: bool = False,
) -> None:
if paths:
paths = [Path(p).absolute() for p in paths]
dandiset = Dandiset.find(os.path.commonpath(paths))
else:
dandiset = Dandiset.find(None)
if dandiset is None:
raise RuntimeError(
f"Found no {dandiset_metadata_file} anywhere in common ancestor of"
" paths. Use 'dandi download' or 'organize' first."
)
with ExitStack() as stack:
# We need to use the client as a context manager in order to ensure the
# session gets properly closed. Otherwise, pytest sometimes complains
# under obscure conditions.
client = stack.enter_context(DandiAPIClient.for_dandi_instance(dandi_instance))
client.check_schema_version()
client.dandi_authenticate()
if os.environ.get("DANDI_DEVEL_INSTRUMENT_REQUESTS_SUPERLEN"):
from requests.utils import super_len
def check_len(obj: io.IOBase, name: Any) -> None:
first = True
for i in range(10):
if first:
first = False
else:
lgr.debug("- Will sleep and then stat %r again", name)
sleep(1)
try:
st = os.stat(name)
except OSError as e:
lgr.debug(
"- Attempt to stat %r failed with %s: %s",
name,
type(e).__name__,
e,
)
stat_size = None
else:
lgr.debug("- stat(%r) = %r", name, st)
stat_size = st.st_size
try:
fileno = obj.fileno()
except Exception:
lgr.debug(
"- I/O object for %r has no fileno; cannot fstat", name
)
fstat_size = None
else:
try:
st = os.fstat(fileno)
except OSError as e:
lgr.debug(
"- Attempt to fstat %r failed with %s: %s",
name,
type(e).__name__,
e,
)
fstat_size = None
else:
lgr.debug("- fstat(%r) = %r", name, st)
fstat_size = st.st_size
if stat_size not in (None, 0):
raise RuntimeError(
f"requests.utils.super_len() reported size of 0 for"
f" {name!r}, but os.stat() reported size"
f" {stat_size} bytes {i + 1} tries later"
)
if fstat_size not in (None, 0):
raise RuntimeError(
f"requests.utils.super_len() reported size of 0 for"
f" {name!r}, but os.fstat() reported size"
f" {fstat_size} bytes {i + 1} tries later"
)
lgr.debug(
"- Size of %r still appears to be 0 after 10 rounds of"
" stat'ing; returning size 0",
name,
)
def new_super_len(o: Any) -> int:
try:
n = super_len(o)
except Exception:
lgr.debug(
"requests.utils.super_len() failed on %r:", o, exc_info=True
)
raise
else:
lgr.debug("requests.utils.super_len() reported %d for %r", n, o)
if (
n == 0
and isinstance(o, io.IOBase)
and (name := getattr(o, "name", None)) not in (None, "")
):
lgr.debug(
"- Size of 0 is suspicious; double-checking that NFS isn't lying"
)
check_len(o, name)
return cast(int, n)
stack.enter_context(patch("requests.models.super_len", new_super_len))
ds_identifier = dandiset.identifier
remote_dandiset = client.get_dandiset(ds_identifier, DRAFT)
if not re.match(dandiset_identifier_regex, str(ds_identifier)):
raise ValueError(
f"Dandiset identifier {ds_identifier} does not follow expected "
f"convention {dandiset_identifier_regex!r}."
)
# Avoid heavy import by importing within function:
from .pynwb_utils import ignore_benign_pynwb_warnings
ignore_benign_pynwb_warnings() # so validate doesn't whine
if not paths:
paths = [dandiset.path]
# DO NOT FACTOR OUT THIS VARIABLE! It stores any
# BIDSDatasetDescriptionAsset instances for the Dandiset, which need to
# remain alive until we're done working with all BIDS assets.
assets = dandiset.assets(allow_all=allow_any_path)
dandi_files: list[DandiFile] = []
# Build the list step by step so as not to confuse mypy
dandi_files.append(dandiset.metadata_file())
dandi_files.extend(
assets.under_paths(Path(p).relative_to(dandiset.path) for p in paths)
)
lgr.info(f"Found {len(dandi_files)} files to consider")
# We will keep a shared set of "being processed" paths so
# we could limit the number of them until
# https://github.com/pyout/pyout/issues/87
# properly addressed
process_paths: set[str] = set()
uploaded_paths: dict[str, Uploaded] = defaultdict(
lambda: {"size": 0, "errors": []}
)
upload_err: Exception | None = None
validate_ok = True
# TODO: we might want to always yield a full record so no field is not
# provided to pyout to cause it to halt
def process_path(dfile: DandiFile) -> Iterator[dict]:
"""
Parameters
----------
dfile: DandiFile
Yields
------
dict
Records for pyout
"""
nonlocal upload_err, validate_ok
strpath = str(dfile.filepath)
try:
if not isinstance(dfile, LocalDirectoryAsset):
try:
yield {"size": dfile.size}
except FileNotFoundError:
raise UploadError("File not found")
except Exception as exc:
# without limiting [:50] it might cause some pyout indigestion
raise UploadError(str(exc)[:50])
#
# Validate first, so we do not bother server at all if not kosher
#
# TODO: enable back validation of dandiset.yaml
if (
isinstance(dfile, LocalAsset)
and validation != UploadValidation.SKIP
):
yield {"status": "pre-validating"}
validation_statuses = dfile.get_validation_errors()
validation_errors = [
s for s in validation_statuses if s.severity == Severity.ERROR
]
yield {"errors": len(validation_errors)}
# TODO: split for dandi, pynwb errors
if validation_errors:
if validation is UploadValidation.REQUIRE:
lgr.warning(
"%r had %d validation errors preventing its upload:",
strpath,
len(validation_errors),
)
for i, e in enumerate(validation_errors, start=1):
lgr.warning(" Error %d: %s", i, e)
validate_ok = False
raise UploadError("failed validation")
else:
yield {"status": "validated"}
else:
# yielding empty causes pyout to get stuck or crash
# https://github.com/pyout/pyout/issues/91
# yield {"errors": '',}
pass
#
# Special handling for dandiset.yaml
# Yarik hates it but that is life for now. TODO
#
if isinstance(dfile, DandisetMetadataFile):
# TODO This is a temporary measure to avoid breaking web UI
# dandiset metadata schema assumptions. All edits should happen
# online.
if upload_dandiset_metadata:
yield {"status": "updating metadata"}
assert dandiset is not None
assert dandiset.metadata is not None
remote_dandiset.set_raw_metadata(dandiset.metadata)
yield {"status": "updated metadata"}
else:
yield skip_file("should be edited online")
return
assert isinstance(dfile, LocalAsset)
#
# Compute checksums
#
file_etag: Digest | None
if isinstance(dfile, ZarrAsset):
file_etag = None
else:
yield {"status": "digesting"}
try:
file_etag = dfile.get_digest()
except Exception as exc:
raise UploadError("failed to compute digest: %s" % str(exc))
try:
extant = remote_dandiset.get_asset_by_path(dfile.path)
except NotFoundError:
extant = None
else:
assert extant is not None
replace, out = check_replace_asset(
local_asset=dfile,
remote_asset=extant,
existing=existing,
local_etag=file_etag,
)
yield out
if not replace:
return
#
# Extract metadata - delayed since takes time, but is done before
# actual upload, so we could skip if this fails
#
# Extract metadata before actual upload and skip if fails
# TODO: allow for for non-nwb files to skip this step
# ad-hoc for dandiset.yaml for now
yield {"status": "extracting metadata"}
try:
metadata = dfile.get_metadata(
digest=file_etag, ignore_errors=allow_any_path
).model_dump(mode="json", exclude_none=True)
except Exception as e:
raise UploadError("failed to extract metadata: %s" % str(e))
#
# Upload file
#
yield {"status": "uploading"}
validating = False
for r in dfile.iter_upload(
remote_dandiset, metadata, jobs=jobs_per_file, replacing=extant
):
r.pop("asset", None) # to keep pyout from choking
if r["status"] == "uploading":
uploaded_paths[strpath]["size"] = r.pop("current")
yield r
elif r["status"] == "post-validating":
# Only yield the first "post-validating" status
if not validating:
yield r
validating = True
else:
yield r
yield {"status": "done"}
except Exception as exc:
if upload_err is None:
upload_err = exc
if devel_debug:
raise
lgr.exception("Error uploading %s:", strpath)
# Custom formatting for some exceptions we know to extract
# user-meaningful message
message = str(exc)
uploaded_paths[strpath]["errors"].append(message)
yield error_file(message)
finally:
process_paths.remove(strpath)
# We will again use pyout to provide a neat table summarizing our progress
# with upload etc
# for the upload speeds we need to provide a custom aggregate
t0 = time.time()
def upload_agg(*ignored: Any) -> str:
dt = time.time() - t0
# to help avoiding dict length changes during upload
# might be not a proper solution
# see https://github.com/dandi/dandi-cli/issues/502 for more info
uploaded_recs = list(uploaded_paths.values())
total = sum(v["size"] for v in uploaded_recs)
if not total:
return ""
speed = total / dt if dt else 0
return "%s/s" % naturalsize(speed)
pyout_style = pyouts.get_style(hide_if_missing=False)
pyout_style["progress"]["aggregate"] = upload_agg
rec_fields = ["path", "size", "errors", "progress", "status", "message"]
out = pyouts.LogSafeTabular(
style=pyout_style, columns=rec_fields, max_workers=jobs or 5
)
with out:
for dfile in dandi_files:
while len(process_paths) >= 10:
lgr.log(2, "Sleep waiting for some paths to finish processing")
time.sleep(0.5)
process_paths.add(str(dfile.filepath))
rec: dict[Any, Any]
if isinstance(dfile, DandisetMetadataFile):
rec = {"path": dandiset_metadata_file}
else:
assert isinstance(dfile, LocalAsset)
rec = {"path": dfile.path}
try:
if devel_debug:
# DEBUG: do serially
for v in process_path(dfile):
print(str(v), flush=True)
else:
rec[tuple(rec_fields[1:])] = process_path(dfile)
except ValueError as exc:
rec.update(error_file(exc))
out(rec)
if not validate_ok:
lgr.warning(
"One or more assets failed validation. Consult the logfile for"
" details."
)
if upload_err is not None:
try:
import etelemetry
latest_version = etelemetry.get_project("dandi/dandi-cli")["version"]
except Exception:
pass
else:
if Version(latest_version) > Version(__version__):
lgr.warning(
"Upload failed, and you are not using the latest"
" version of dandi. We suggest upgrading dandi to v%s"
" and trying again.",
latest_version,
)
raise upload_err
if sync:
relpaths: list[str] = []
for p in paths:
rp = os.path.relpath(p, dandiset.path)
relpaths.append("" if rp == "." else rp)
path_prefix = reduce(os.path.commonprefix, relpaths) # type: ignore[arg-type]
to_delete = []
for asset in remote_dandiset.get_assets_with_path_prefix(path_prefix):
if any(
p == "" or path_is_subpath(asset.path, p) for p in relpaths
) and not os.path.lexists(Path(dandiset.path, asset.path)):
to_delete.append(asset)
if to_delete and click.confirm(
f"Delete {pluralize(len(to_delete), 'asset')} on server?"
):
for asset in to_delete:
asset.delete()
[docs]def check_replace_asset(
local_asset: LocalAsset,
remote_asset: RemoteAsset,
existing: UploadExisting,
local_etag: Digest | None,
) -> tuple[bool, dict[str, str]]:
# Returns a (replace asset, message to yield) tuple
if isinstance(local_asset, ZarrAsset):
return (True, {"message": "exists - reuploading"})
assert local_etag is not None
metadata = remote_asset.get_raw_metadata()
local_mtime = local_asset.modified
remote_mtime_str = metadata.get("blobDateModified")
# TODO: Should this error if the digest is missing?
remote_etag = metadata.get("digest", {}).get(local_etag.algorithm.value)
if remote_mtime_str is not None:
remote_mtime = ensure_datetime(remote_mtime_str)
remote_file_status = (
"same"
if remote_etag == local_etag.value and remote_mtime == local_mtime
else (
"newer"
if remote_mtime > local_mtime
else ("older" if remote_mtime < local_mtime else "diff")
)
)
else:
remote_mtime = None
remote_file_status = "no mtime"
exists_msg = f"exists ({remote_file_status})"
if existing is UploadExisting.ERROR:
# as promised -- not gentle at all!
raise FileExistsError(exists_msg)
if existing is UploadExisting.SKIP:
return (False, skip_file(exists_msg))
# Logic below only for overwrite and reupload
if existing is UploadExisting.OVERWRITE:
if remote_etag == local_etag.value:
return (False, skip_file(exists_msg))
elif existing is UploadExisting.REFRESH:
if remote_etag == local_etag.value:
return (False, skip_file("file exists"))
elif remote_mtime is not None and remote_mtime >= local_mtime:
return (False, skip_file(exists_msg))
elif existing is UploadExisting.FORCE:
pass
else:
raise AssertionError(f"Unhandled UploadExisting member: {existing!r}")
return (True, {"message": f"{exists_msg} - reuploading"})
[docs]def skip_file(msg: Any) -> dict[str, str]:
return {"status": "skipped", "message": str(msg)}
[docs]def error_file(msg: Any) -> dict[str, str]:
return {"status": "ERROR", "message": str(msg)}