🐛 Close broker and storage backend when verdi commands exit

src/aiida/cmdline/commands/cmd_archive.py

@@ -23,6 +23,7 @@
 from aiida.cmdline.params import arguments, options
 from aiida.cmdline.params.types import GroupParamType, PathOrUrl
 from aiida.cmdline.utils import decorators, echo
+from aiida.cmdline.utils.loaders import load_codes, load_computers, load_group, load_groups, load_nodes
 from aiida.common.exceptions import CorruptStorage, IncompatibleStorageSchema, UnreachableStorage
 from aiida.common.links import GraphTraversalRules
 from aiida.common.log import AIIDA_LOGGER
@@ -192,16 +193,16 @@ def create(
         entities = []
 
         if codes:
-            entities.extend(codes)
+            entities.extend(load_codes(codes))
 
         if computers:
-            entities.extend(computers)
+            entities.extend(load_computers(computers))
 
         if groups:
-            entities.extend(groups)
+            entities.extend(load_groups(groups))
 
         if nodes:
-            entities.extend(nodes)
+            entities.extend(load_nodes(nodes))
 
     kwargs = {
         'input_calc_forward': input_calc_forward,
@@ -396,6 +397,7 @@ def import_archive(
     else:
         set_progress_reporter(None)
 
+    group = load_group(group) if group is not None else None
     all_archives = _gather_imports(archives, webpages)
 
     # Preliminary sanity check

src/aiida/cmdline/commands/cmd_calcjob.py

@@ -19,6 +19,7 @@
 from aiida.cmdline.params import arguments, options
 from aiida.cmdline.params.types import CalculationParamType
 from aiida.cmdline.utils import decorators, echo
+from aiida.cmdline.utils.loaders import load_calculation, load_calculations, load_computers
 
 if t.TYPE_CHECKING:
     from aiida import orm
@@ -31,6 +32,7 @@ def verdi_calcjob():
 
 @verdi_calcjob.command('gotocomputer')
 @arguments.CALCULATION('calcjob', type=CalculationParamType(sub_classes=('aiida.node:process.calculation.calcjob',)))
+@decorators.with_dbenv()
 def calcjob_gotocomputer(calcjob):
     """Open a shell in the remote folder on the calcjob.
 
@@ -39,6 +41,8 @@ def calcjob_gotocomputer(calcjob):
     """
     from aiida.common.exceptions import NotExistent
 
+    calcjob = load_calculation(calcjob, param_name='calcjob')
+
     try:
         transport = calcjob.get_transport()
     except NotExistent as exception:
@@ -67,6 +71,8 @@ def calcjob_res(calcjob, fmt, keys):
     """Print data from the result output Dict node of a calcjob."""
     from aiida.cmdline.utils.echo import echo_dictionary
 
+    calcjob = load_calculation(calcjob, param_name='calcjob')
+
     try:
         results = calcjob.res.get_results()
     except ValueError as exception:
@@ -98,6 +104,8 @@ def calcjob_inputcat(calcjob, path):
     import sys
     from shutil import copyfileobj
 
+    calcjob = load_calculation(calcjob, param_name='calcjob')
+
     # Get path from the given CalcJobNode if not defined by user
     if path is None:
         path = calcjob.get_option('input_filename')
@@ -143,6 +151,8 @@ def calcjob_remotecat(calcjob: orm.CalcJobNode, path: str | None):
     import sys
     import tempfile
 
+    calcjob = load_calculation(calcjob, param_name='calcjob')
+
     remote_folder, path = get_remote_and_path(calcjob, path)
 
     with tempfile.NamedTemporaryFile() as tmp_path:
@@ -168,6 +178,8 @@ def calcjob_outputcat(calcjob, path):
     """
     import errno
     import sys
+
+    calcjob = load_calculation(calcjob, param_name='calcjob')
     from shutil import copyfileobj
 
     try:
@@ -220,6 +232,8 @@ def calcjob_inputls(calcjob, path, color):
     """
     from aiida.cmdline.utils.repository import list_repository_contents
 
+    calcjob = load_calculation(calcjob, param_name='calcjob')
+
     try:
         list_repository_contents(calcjob, path, color)
     except FileNotFoundError:
@@ -241,6 +255,8 @@ def calcjob_outputls(calcjob, path, color):
     """
     from aiida.cmdline.utils.repository import list_repository_contents
 
+    calcjob = load_calculation(calcjob, param_name='calcjob')
+
     try:
         retrieved = calcjob.outputs.retrieved
     except AttributeError:
@@ -267,6 +283,8 @@ def calcjob_cleanworkdir(calcjobs, past_days, older_than, computers, force, exit
     If both are specified, a logical AND is done between the two, i.e. the calcjobs that will be cleaned have been
     modified AFTER [-p option] days from now, but BEFORE [-o option] days from now.
     """
+    calcjobs = load_calculations(calcjobs, param_name='calcjobs')
+    computers = load_computers(computers) if computers else computers
     from aiida.orm.utils.remote import clean_mapping_remote_paths, get_calcjob_remote_paths
 
     if calcjobs:

src/aiida/cmdline/commands/cmd_code.py

@@ -26,6 +26,7 @@
 from aiida.cmdline.utils import echo, echo_tabulate
 from aiida.cmdline.utils.common import validate_output_filename
 from aiida.cmdline.utils.decorators import with_dbenv
+from aiida.cmdline.utils.loaders import load_code, load_codes, load_computer, load_entity
 from aiida.common import exceptions
 
 if TYPE_CHECKING:
@@ -90,13 +91,21 @@ def get_on_computer(ctx: click.Context) -> bool:
 
 
 def set_code_builder(ctx: click.Context, _param: Any, value: Any) -> Any:
-    """Set the code spec for defaults of following options."""
+    """Set the code spec for defaults of following options.
+
+    This callback seeds the defaults that the following options prompt with, so unlike a command body it has to
+    resolve the identifier while the command line is still being parsed, and loads the backend itself.
+    """
+    from aiida.cmdline.utils.decorators import load_backend_if_not_loaded
     from aiida.orm.utils.builders.code import CodeBuilder
 
+    load_backend_if_not_loaded()
+    code = load_entity(value, param_name=_param.name)
+
     # TODO(danielhollas): CodeBuilder is deprecated, rewrite this somehow?
     with warnings.catch_warnings(record=True):
-        ctx.code_builder = CodeBuilder.from_code(value)  # type: ignore[attr-defined]
-    return value
+        ctx.code_builder = CodeBuilder.from_code(code)  # type: ignore[attr-defined]
+    return code
 
 
 # Defining the ``COMPUTER`` option first guarantees that the user is prompted for the computer first. This is necessary
@@ -233,6 +242,7 @@ def show(code: Code):
     """Display detailed information for a code."""
     from aiida.cmdline import is_verbose
 
+    code = load_code(code)
     table = []
 
     # These are excluded from the CLI, so we add them manually
@@ -265,6 +275,7 @@ def show(code: Code):
 @with_dbenv()
 def export(code, output_file, overwrite, sort):
     """Export code to a yaml file. If no output file is given, default name is created based on the code label."""
+    code = load_code(code)
     other_args = {'sort': sort}
     fileformat = 'yaml'
 
@@ -309,6 +320,7 @@ def delete(codes, dry_run, force):
     """
     from aiida.tools import delete_nodes
 
+    codes = load_codes(codes)
     node_pks_to_delete = [code.pk for code in codes]
 
     def _dry_run_callback(pks):
@@ -328,7 +340,7 @@ def _dry_run_callback(pks):
 @with_dbenv()
 def hide(codes):
     """Hide one or more codes from `verdi code list`."""
-    for code in codes:
+    for code in load_codes(codes):
         code.is_hidden = True
         echo.echo_success(f'Code<{code.pk}> {code.full_label} hidden')
 
@@ -338,7 +350,7 @@ def hide(codes):
 @with_dbenv()
 def reveal(codes):
     """Reveal one or more hidden codes in `verdi code list`."""
-    for code in codes:
+    for code in load_codes(codes):
         code.is_hidden = False
         echo.echo_success(f'Code<{code.pk}> {code.full_label} revealed')
 
@@ -349,6 +361,7 @@ def reveal(codes):
 @with_dbenv()
 def relabel(code, label):
     """Relabel a code."""
+    code = load_code(code)
     old_label = code.full_label
 
     try:
@@ -390,6 +403,8 @@ def code_list(computer, default_calc_job_plugin, all_entries, all_users, raw, sh
     from aiida import orm
     from aiida.orm.utils.node import load_node_class
 
+    computer = load_computer(computer) if computer is not None else None
+
     if show_owner:
         echo.echo_deprecated(
             'the `-o/--show-owner` option is deprecated. To show the user use the `-P/--project` option instead, e.g., '

src/aiida/cmdline/commands/cmd_computer.py

@@ -24,6 +24,7 @@
 from aiida.cmdline.utils import echo, echo_tabulate
 from aiida.cmdline.utils.common import validate_output_filename
 from aiida.cmdline.utils.decorators import with_dbenv
+from aiida.cmdline.utils.loaders import load_computer, load_entity, load_user
 from aiida.common.exceptions import EntryPointError, ValidationError
 from aiida.plugins.entry_point import get_entry_point_names
 
@@ -265,11 +266,18 @@ def get_parameter_default(parameter, ctx):
 
 
 def set_computer_builder(ctx, param, value):
-    """Set the computer spec for defaults of following options."""
+    """Set the computer spec for defaults of following options.
+
+    This callback seeds the defaults that the following options prompt with, so unlike a command body it has to
+    resolve the identifier while the command line is still being parsed. It therefore loads the backend itself.
+    """
+    from aiida.cmdline.utils.decorators import load_backend_if_not_loaded
     from aiida.orm.utils.builders.computer import ComputerBuilder
 
-    ctx.computer_builder = ComputerBuilder.from_computer(value)
-    return value
+    load_backend_if_not_loaded()
+    computer = load_entity(value, param_name=param.name)
+    ctx.computer_builder = ComputerBuilder.from_computer(computer)
+    return computer
 
 
 @verdi_computer.command('setup')
@@ -386,6 +394,9 @@ def computer_enable(computer, user):
     """Enable the computer for the given user."""
     from aiida.common.exceptions import NotExistent
 
+    computer = load_computer(computer)
+    user = load_user(user)
+
     try:
         authinfo = computer.get_authinfo(user)
     except NotExistent:
@@ -411,6 +422,9 @@ def computer_disable(computer, user):
     """
     from aiida.common.exceptions import NotExistent
 
+    computer = load_computer(computer)
+    user = load_user(user)
+
     try:
         authinfo = computer.get_authinfo(user)
     except NotExistent:
@@ -451,6 +465,7 @@ def computer_list(all_entries, raw):
 
 @verdi_computer.command('goto')
 @arguments.COMPUTER()
+@with_dbenv()
 def computer_goto(computer):
     """Open a shell connecting to the remote computer.
 
@@ -459,6 +474,8 @@ def computer_goto(computer):
     """
     from aiida.common.exceptions import NotExistent
 
+    computer = load_computer(computer)
+
     try:
         transport = computer.get_transport()
     except NotExistent as exception:
@@ -477,6 +494,7 @@ def computer_goto(computer):
 @with_dbenv()
 def computer_show(computer):
     """Show detailed information for a computer."""
+    computer = load_computer(computer)
     table = [
         ['Label', computer.label],
         ['PK', computer.pk],
@@ -504,6 +522,7 @@ def computer_relabel(computer, label):
     """Relabel a computer."""
     from aiida.common.exceptions import UniquenessError
 
+    computer = load_computer(computer)
     old_label = computer.label
 
     if old_label == label:
@@ -542,9 +561,13 @@ def computer_test(user, print_traceback, computer):
     from aiida import orm
     from aiida.common.exceptions import NotExistent
 
+    computer = load_computer(computer)
+
     # Set a user automatically if one is not specified in the command line
     if user is None:
         user = orm.User.collection.get_default()
+    else:
+        user = load_user(user)
 
     echo.echo_report(f'Testing computer<{computer.label}> for user<{user.email}>...')
 
@@ -647,6 +670,7 @@ def computer_delete(computer, dry_run):
     from aiida.orm.querybuilder import QueryBuilder
     from aiida.tools import delete_nodes
 
+    computer = load_computer(computer)
     label = computer.label
 
     # Sofar, we can only get this info with QueryBuilder
@@ -716,11 +740,15 @@ def computer_configure():
     help='Email address of the AiiDA user for whom to configure this computer (if different from default user).'
 )
 @arguments.COMPUTER()
+@with_dbenv()
 def computer_config_show(computer, user, defaults, as_option_string):
     """Show the current configuration for a computer."""
     from aiida.common.escaping import escape_for_bash
     from aiida.transports import cli as transport_cli
 
+    computer = load_computer(computer)
+    user = load_user(user) if user is not None else None
+
     transport_cls = computer.get_transport_class()
     option_list = [
         param
@@ -778,6 +806,7 @@ def computer_export_setup(computer, output_file, overwrite, sort):
     """Export computer setup to a YAML file."""
     import yaml
 
+    computer = load_computer(computer)
     computer_setup = {
         'label': computer.label,
         'hostname': computer.hostname,
@@ -826,6 +855,9 @@ def computer_export_config(computer, output_file, user, overwrite, sort):
     """Export computer transport configuration for a user to a YAML file."""
     import yaml
 
+    computer = load_computer(computer)
+    user = load_user(user) if user is not None else None
+
     if not computer.is_configured:
         echo.echo_critical(
             f'Computer<{computer.pk}> {computer.label} configuration cannot be exported,'

src/aiida/cmdline/commands/cmd_data/cmd_array.py

@@ -10,6 +10,8 @@
 
 from aiida.cmdline.commands.cmd_data import verdi_data
 from aiida.cmdline.params import arguments, options, types
+from aiida.cmdline.utils import decorators
+from aiida.cmdline.utils.loaders import load_data
 
 
 @verdi_data.group('core.array')
@@ -20,11 +22,12 @@ def array():
 @array.command('show')
 @arguments.DATA(type=types.DataParamType(sub_classes=('aiida.data:core.array',)))
 @options.DICT_FORMAT()
+@decorators.with_dbenv()
 def array_show(data, fmt):
     """Visualize ArrayData objects."""
     from aiida.cmdline.utils.echo import echo_dictionary
 
-    for node in data:
+    for node in load_data(data):
         the_dict = {}
         for arrayname in node.get_arraynames():
             the_dict[arrayname] = node.get_array(arrayname).tolist()