From ad1a431f33a6e57d8b6867447ecdfd8ff41bc8f5 Mon Sep 17 00:00:00 2001
From: Sebastiaan Huber <mail@sphuber.net>
Date: Mon, 5 Aug 2024 11:09:28 +0200
Subject: [PATCH] CLI: Validate storage in `verdi storage version` (#6551)

The `verdi storage version`, in addition to printing the version of the
code's and storage's schema, now also validates the storage. If the
storage is corrupt or cannot be reached, the command returns the exit
code 3. If the storage and code schema versions are incompatible, exit
code 4 is returned. This way this command serves as an alternative to
running `verdi storage migrate` as a way to check whether a profile
needs to be migrated. The `verdi storage migrate` command needs to
perform checks such as whether the daemon is running and so is always
going to be slower.
---
 src/aiida/cmdline/commands/cmd_storage.py | 37 +++++++++++++++++++----
 tests/cmdline/commands/test_storage.py    | 20 ++++++++++++
 2 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/src/aiida/cmdline/commands/cmd_storage.py b/src/aiida/cmdline/commands/cmd_storage.py
index c4fd17601..f6f64b755 100644
--- a/src/aiida/cmdline/commands/cmd_storage.py
+++ b/src/aiida/cmdline/commands/cmd_storage.py
@@ -8,6 +8,8 @@
 ###########################################################################
 """`verdi storage` commands."""
 
+import sys
+
 import click
 from click_spinner import spinner
 
@@ -24,14 +26,37 @@ def verdi_storage():
 
 @verdi_storage.command('version')
 def storage_version():
-    """Print the current version of the storage schema."""
+    """Print the current version of the storage schema.
+
+    The command returns the following exit codes:
+
+    * 0: If the storage schema is equal and compatible to the schema version of the code
+    * 3: If the storage cannot be reached or is corrupt
+    * 4: If the storage schema is compatible with the code schema version and probably needs to be migrated.
+    """
     from aiida import get_profile
+    from aiida.common.exceptions import CorruptStorage, IncompatibleStorageSchema, UnreachableStorage
 
-    profile = get_profile()
-    head_version = profile.storage_cls.version_head()
-    profile_version = profile.storage_cls.version_profile(profile)
-    echo.echo(f'Latest storage schema version: {head_version!r}')
-    echo.echo(f'Storage schema version of {profile.name!r}: {profile_version!r}')
+    try:
+        profile = get_profile()
+        head_version = profile.storage_cls.version_head()
+        profile_version = profile.storage_cls.version_profile(profile)
+        echo.echo(f'Latest storage schema version: {head_version!r}')
+        echo.echo(f'Storage schema version of {profile.name!r}: {profile_version!r}')
+    except Exception as exception:
+        echo.echo_critical(f'Failed to determine the storage version: {exception}')
+
+    try:
+        profile.storage_cls(profile)
+    except (CorruptStorage, UnreachableStorage) as exception:
+        echo.echo_error(f'The storage cannot be reached or is corrupt: {exception}')
+        sys.exit(3)
+    except IncompatibleStorageSchema:
+        echo.echo_error(
+            f'The storage schema version {profile_version} is incompatible with the code version {head_version}.'
+            'Run `verdi storage migrate` to migrate the storage.'
+        )
+        sys.exit(4)
 
 
 @verdi_storage.command('migrate')
diff --git a/tests/cmdline/commands/test_storage.py b/tests/cmdline/commands/test_storage.py
index eb629d013..59698343e 100644
--- a/tests/cmdline/commands/test_storage.py
+++ b/tests/cmdline/commands/test_storage.py
@@ -23,6 +23,26 @@ def tests_storage_version(run_cli_command):
     assert version in result.output
 
 
+@pytest.mark.parametrize(
+    'exception_cls, exit_code',
+    (
+        (exceptions.CorruptStorage, 3),
+        (exceptions.UnreachableStorage, 3),
+        (exceptions.IncompatibleStorageSchema, 4),
+    ),
+)
+def tests_storage_version_non_zero_exit_code(aiida_profile, run_cli_command, monkeypatch, exception_cls, exit_code):
+    """Test the ``verdi storage version`` command when it returns a non-zero exit code."""
+
+    def validate_storage(self):
+        raise exception_cls()
+
+    with monkeypatch.context() as context:
+        context.setattr(aiida_profile.storage_cls.migrator, 'validate_storage', validate_storage)
+        result = run_cli_command(cmd_storage.storage_version, raises=True)
+        assert result.exit_code == exit_code
+
+
 def tests_storage_info(aiida_localhost, run_cli_command):
     """Test the ``verdi storage info`` command with the ``--detailed`` option."""
     from aiida import orm