Use QuerySet.iterator to reduce memory usage

dandiapi/api/management/commands/extract_metadata.py

@@ -9,6 +9,7 @@
 from dandischema.models import get_schema_version
 from django.contrib.auth.models import User
 from django.db import transaction
+from django.db.models import QuerySet
 import djclick as click
 from tqdm import tqdm
 
@@ -77,7 +78,7 @@ def extract_asset_metadata(asset: Asset, draft_version: Version):
 
 def extract_dandiset_assets(dandiset: Dandiset):
     # Only update NWB assets which are out of date and do not belong to a published version
-    assets = dandiset.draft_version.assets.filter(
+    assets: QuerySet[Asset] = dandiset.draft_version.assets.filter(
         published=False,
         path__iendswith='.nwb',
         metadata__schemaVersion__lt=get_schema_version(),
@@ -86,7 +87,7 @@ def extract_dandiset_assets(dandiset: Dandiset):
         logger.info('No old draft NWB assets found in dandiset %s. Skipping...', dandiset)
         return
 
-    for asset in tqdm(assets):
+    for asset in tqdm(assets.iterator(), total=assets.count()):
         extract_asset_metadata(asset=asset, draft_version=dandiset.draft_version)
 
 
@@ -101,7 +102,7 @@ def asset(asset_id: str):
         )
 
     # Re-extract for every draft version
-    for version in draft_versions:
+    for version in draft_versions.iterator():
         extract_asset_metadata(asset=asset, draft_version=version)
 
 
@@ -116,6 +117,6 @@ def dandiset(dandiset_id: str):
 
 @group.command(name='all', help='Re-extracts the metadata of all assets in all draft versions')
 def all_dandisets():
-    for dandiset in Dandiset.objects.all():
+    for dandiset in Dandiset.objects.iterator():
         logger.info('DANDISET: %s', dandiset.identifier)
         extract_dandiset_assets(dandiset)