Skip to content

Commit

Permalink
verdi node rehash: Add aiida.node as group for --entry-point
Browse files Browse the repository at this point in the history
This will allow users to target specific node types, such as all
`CalcJobNodes` or all `WorkChainNodes`. This is especially useful if a
migration needs to be added for one of these node types where the hash
is reset and so needs to be regenerated.
  • Loading branch information
sphuber committed May 15, 2023
1 parent a2a05a6 commit 2fd0751
Showing 1 changed file with 11 additions and 8 deletions.
19 changes: 11 additions & 8 deletions aiida/cmdline/commands/cmd_node.py
Original file line number Diff line number Diff line change
Expand Up @@ -329,7 +329,7 @@ def _dry_run_callback(pks):
@click.option(
'-e',
'--entry-point',
type=PluginParamType(group=('aiida.calculations', 'aiida.data', 'aiida.workflows'), load=True),
type=PluginParamType(group=('aiida.calculations', 'aiida.data', 'aiida.node', 'aiida.workflows'), load=True),
default=None,
help='Only include nodes that are class or sub class of the class identified by this entry point.'
)
Expand All @@ -342,8 +342,15 @@ def rehash(nodes, entry_point, force):
"""
from aiida.orm import Data, ProcessNode, QueryBuilder

# If no explicit entry point is defined, rehash all nodes, which are either Data nodes or ProcessNodes
if entry_point is None:
classes = (Data, ProcessNode)
else:
classes = (entry_point,)

if not force:
echo.echo_warning('This command will recompute and overwrite the hashes of all nodes.')
class_names = ', '.join([cls.__name__ for cls in classes])
echo.echo_warning(f'This command will recompute and overwrite the hashes of all {class_names} nodes.')
echo.echo_warning('Note that this can take a lot of time for big databases.')
echo.echo_warning('')
echo.echo_warning('', nl=False)
Expand All @@ -356,16 +363,12 @@ def rehash(nodes, entry_point, force):
echo.echo('\n')
echo.echo_critical('Migration aborted, the data has not been affected.')

# If no explicit entry point is defined, rehash all nodes, which are either Data nodes or ProcessNodes
if entry_point is None:
entry_point = (Data, ProcessNode)

if nodes:
to_hash = [(node,) for node in nodes if isinstance(node, entry_point)]
to_hash = [(node,) for node in nodes if isinstance(node, classes)]
num_nodes = len(to_hash)
else:
builder = QueryBuilder()
builder.append(entry_point, tag='node')
builder.append(classes, tag='node')
to_hash = builder.all()
num_nodes = builder.count()

Expand Down

0 comments on commit 2fd0751

Please sign in to comment.