From 9b4bf720ada6790646576fb839430be02d7bb750 Mon Sep 17 00:00:00 2001
From: Jacob Nesbitt <jjnesbitt2@gmail.com>
Date: Mon, 20 Mar 2023 13:37:33 -0400
Subject: [PATCH] Add management command for asset metadata extract

---
 .../management/commands/extract_metadata.py   | 49 +++++++++++++++++++
 setup.py                                      |  2 +
 2 files changed, 51 insertions(+)
 create mode 100644 dandiapi/api/management/commands/extract_metadata.py

diff --git a/dandiapi/api/management/commands/extract_metadata.py b/dandiapi/api/management/commands/extract_metadata.py
new file mode 100644
index 000000000..d8d52d2a8
--- /dev/null
+++ b/dandiapi/api/management/commands/extract_metadata.py
@@ -0,0 +1,49 @@
+from pathlib import Path
+
+from dandi.dandiapi import RemoteReadableAsset
+from dandi.metadata import nwb2asset
+import djclick as click
+from tqdm import tqdm
+
+from dandiapi.api.models import Asset, Dandiset
+
+
+@click.group()
+def group():
+    pass
+
+
+def extract_asset_metadata(asset: Asset):
+    readable_asset = RemoteReadableAsset(
+        asset.s3_url, size=asset.size, mtime=asset.modified, name=Path(asset.path).name
+    )
+    asset.metadata = nwb2asset(readable_asset).json_dict()
+    asset.save()
+
+
+def extract_dandiset_assets(dandiset: Dandiset):
+    for asset in tqdm(dandiset.draft_version.assets.all()):
+        extract_asset_metadata(asset)
+
+
+@group.command(help='Re-extracts the metadata of this asset')
+@click.argument('asset_id')
+def asset(asset_id: str):
+    asset = Asset.objects.get(asset_id=asset_id)
+    extract_asset_metadata(asset)
+
+
+@group.command(
+    help='Re-extracts the metadata of all assets in the draft version of the provided dandiset'
+)
+@click.argument('dandiset_id')
+def dandiset(dandiset_id: str):
+    dandiset = Dandiset.objects.get(id=int(dandiset_id))
+    extract_dandiset_assets(dandiset)
+
+
+@group.command(help='Re-extracts the metadata of all assets in all draft versions')
+def all():
+    for dandiset in Dandiset.objects.all():
+        print(f'DANDISET: {dandiset.identifier}')
+        extract_dandiset_assets(dandiset)
diff --git a/setup.py b/setup.py
index 1e59ca07c..57050520e 100644
--- a/setup.py
+++ b/setup.py
@@ -77,6 +77,8 @@
             'ipython',
             'tox',
             'memray',
+            'dandi',
+            'fsspec[http]',
         ],
         'test': [
             'factory-boy',