logging moved to cmd_storage

aiidateam · Sep 4, 2024 · cc62561 · cc62561
1 parent d22a9b6
commit cc62561
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/src/aiida/cmdline/commands/cmd_storage.py b/src/aiida/cmdline/commands/cmd_storage.py
@@ -8,6 +8,7 @@
 ###########################################################################
 """`verdi storage` commands."""
 
+import logging
 import sys
 
 import click
@@ -153,7 +154,9 @@ def storage_info(detailed):
 def storage_maintain(ctx, full, no_repack, force, dry_run, compress):
     """Performs maintenance tasks on the repository."""
     from aiida.common.exceptions import LockingProfileError
+    from aiida.common.progress_reporter import set_progress_bar_tqdm, set_progress_reporter
     from aiida.manage.manager import get_manager
+    from aiida.storage.log import STORAGE_LOGGER
 
     manager = get_manager()
     profile = ctx.obj.profile
@@ -184,6 +187,11 @@ def storage_maintain(ctx, full, no_repack, force, dry_run, compress):
     if not dry_run and not force and not click.confirm('Are you sure you want continue in this mode?'):
         return
 
+    if STORAGE_LOGGER.level <= logging.REPORT:
+        set_progress_bar_tqdm(leave=STORAGE_LOGGER.level <= logging.INFO)
+    else:
+        set_progress_reporter(None)
+
     try:
         if full and no_repack:
             storage.maintain(full=full, dry_run=dry_run, do_repack=False, compress=compress)

diff --git a/src/aiida/repository/backend/disk_object_store.py b/src/aiida/repository/backend/disk_object_store.py
@@ -2,7 +2,6 @@
 
 import contextlib
 import dataclasses
-import logging
 import shutil
 import typing as t
 
@@ -168,12 +167,7 @@ def maintain(  # type: ignore[override]
         """
         from disk_objectstore import CompressMode
 
-        import aiida.common.progress_reporter as pr
-
-        if STORAGE_LOGGER.level <= logging.REPORT:
-            pr.set_progress_bar_tqdm(leave=STORAGE_LOGGER.level <= logging.INFO)
-        else:
-            pr.set_progress_reporter(None)
+        from aiida.common.progress_reporter import create_callback, get_progress_reporter
 
         if live and (do_repack or clean_storage or do_vacuum):
             overrides = {'do_repack': do_repack, 'clean_storage': clean_storage, 'do_vacuum': do_vacuum}
@@ -200,17 +194,17 @@ def maintain(  # type: ignore[override]
                 files_size = container.get_total_size().total_size_loose * BYTES_TO_MB
                 logger.report(f'Packing all loose files ({files_numb} files occupying {files_size} MB) ...')
                 if not dry_run:
-                    with pr.get_progress_reporter()(total=1) as progress:
-                        callback = pr.create_callback(progress)
+                    with get_progress_reporter()(total=1) as progress:
+                        callback = create_callback(progress)
                         container.pack_all_loose(compress=compress, callback=callback)
 
             if do_repack:
                 files_numb = container.count_objects().packed
                 files_size = container.get_total_size().total_size_packfiles_on_disk * BYTES_TO_MB
                 logger.report(f'Re-packing all pack files ({files_numb} files in packs, occupying {files_size} MB) ...')
                 if not dry_run:
-                    with pr.get_progress_reporter()(total=1) as progress:
-                        callback = pr.create_callback(progress)
+                    with get_progress_reporter()(total=1) as progress:
+                        callback = create_callback(progress)
                         container.repack(callback=callback)
 
             if clean_storage: