def transition_db(profile=None, group_size=1000, delete_table=False):
"""
Migrate the attributes, extra, and some other columns to use JSONMigrate
the attributes, extra, and some other columns to use JSONB column type.
"""
cont = query_yes_no(
"Starting complete database transition. Be sure to "
"backup your database before continuing, and that no "
"one else is using it. Do you want to continue?", "no")
if not cont:
print("Answered no. Exiting")
sys.exit(0)
transition_attributes(profile=profile,
group_size=group_size,
delete_table=delete_table)
transition_extras(profile=profile,
group_size=group_size,
delete_table=delete_table)
create_gin_index()
modify_link_table()
transition_settings(profile=profile)
transition_json_column(profile=profile)
set_correct_schema_version_and_backend()
print("\nDatabase transition finished.")
def check_schema_version(force_migration=False, alembic_cfg=None):
"""
Check if the version stored in the database is the same of the version
of the code.
:note: if the DbSetting table does not exist, this function does not
fail. The reason is to avoid to have problems before running the first
migrate call.
:note: if no version is found, the version is set to the version of the
code. This is useful to have the code automatically set the DB version
at the first code execution.
:raise ConfigurationError: if the two schema versions do not match.
Otherwise, just return.
"""
import sys
from aiida.common.utils import query_yes_no
from aiida.backends import sqlalchemy as sa
from aiida.backends.settings import IN_DOC_MODE
# Early exit if we compile the documentation since the schema
# check is not needed and it creates problems with the sqlalchemy
# migrations
if IN_DOC_MODE:
return
# If an alembic configuration file is given then use that one.
if alembic_cfg is None:
alembic_cfg = get_alembic_conf()
# Getting the version of the code and the database
# Reusing the existing engine (initialized by AiiDA)
with sa.engine.begin() as connection:
alembic_cfg.attributes['connection'] = connection
code_schema_version = get_migration_head(alembic_cfg)
db_schema_version = get_db_schema_version(alembic_cfg)
if code_schema_version != db_schema_version:
if db_schema_version is None:
print("It is time to perform your first SQLAlchemy migration.")
else:
print("The code schema version is {}, but the version stored in "
"the database is {}.".format(code_schema_version,
db_schema_version))
if force_migration or query_yes_no(
"Would you like to migrate to the "
"latest version?", "yes"):
print("Migrating to the last version")
# Reusing the existing engine (initialized by AiiDA)
with sa.engine.begin() as connection:
alembic_cfg.attributes['connection'] = connection
command.upgrade(alembic_cfg, "head")
else:
print("No migration is performed. Exiting since database is out "
"of sync with the code.")
sys.exit(1)
def create_dir(self, question, dir_path):
final_path = utils.query_string(question, dir_path)
if not os.path.exists(final_path):
if utils.query_yes_no("The path {} doesn't exist. Should it be "
"created?".format(final_path), "yes"):
try:
os.makedirs(final_path)
except OSError:
self._logger.error("Error creating the path "
"{}.".format(final_path))
raise
return final_path
def test_query_yes_no(self):
"""
This method tests the query_yes_no method behaves as expected. To
perform this, a lambda function is used to simulate the user input.
"""
from aiida.utils.capturing import Capturing
# Capture the sysout for the following code
with Capturing():
# Check the yes
utils.raw_input = lambda _: "y"
self.assertTrue(utils.query_yes_no("", "yes"))
utils.raw_input = lambda _: "yes"
self.assertTrue(utils.query_yes_no("", "yes"))
# Check the no
utils.raw_input = lambda _: "no"
self.assertFalse(utils.query_yes_no("", "yes"))
utils.raw_input = lambda _: "n"
self.assertFalse(utils.query_yes_no("", "yes"))
# Check the empty default value that should
# lead to an error
with self.assertRaises(ValueError):
utils.query_yes_no("", "")
# Check that a None default value and no answer from
# the user should lead to the repetition of the query until
# it is answered properly
self.seq = -1
answers = ["", "", "", "yes"]
utils.raw_input = lambda _: answers[self.array_counter()]
self.assertTrue(utils.query_yes_no("", None))
self.assertEqual(self.seq, len(answers) - 1)
# Check that the default answer is returned
# when the user doesn't give an answer
utils.raw_input = lambda _: ""
self.assertTrue(utils.query_yes_no("", "yes"))
utils.raw_input = lambda _: ""
self.assertFalse(utils.query_yes_no("", "no"))
def transition_attributes(profile=None,
group_size=1000,
debug=False,
delete_table=False):
"""
Migrate the DbAttribute table into the attributes column of db_dbnode.
"""
if not is_dbenv_loaded():
transition_load_db_env(profile=profile)
class DbAttribute(Base):
"""
DbAttribute table, use only for migration purposes.
"""
__tablename__ = ATTR_TABLE_NAME
id = Column(Integer, primary_key=True)
key = Column(String(1024), nullable=False)
datatype = Column(String(10), nullable=False)
tval = Column(Text, nullable=False)
fval = Column(Float)
ival = Column(Integer)
bval = Column(Boolean)
dval = Column(DateTime(timezone=True))
dbnode_id = Column(Integer, ForeignKey('db_dbnode.id'), nullable=False)
dbnode = relationship('DbNode', backref='old_attrs')
print("\nStarting migration of attributes")
inspector = reflection.Inspector.from_engine(sa.session.bind)
table_names = inspector.get_table_names()
if NODE_TABLE_NAME not in table_names:
raise Exception(
"There is no {} table in the database. Transition"
"to SQLAlchemy can not be done. Exiting".format(NODE_TABLE_NAME))
node_table_cols = inspector.get_columns(NODE_TABLE_NAME)
col_names = [_["name"] for _ in node_table_cols]
if ATTR_COL_NAME in col_names:
print(
"Column named {} found at the {} table of the database. I assume "
"that the migration of the attributes has already been done and "
"therefore I proceed with the next migration step.".format(
ATTR_COL_NAME, NODE_TABLE_NAME))
return
with sa.session.begin(subtransactions=True):
print("Creating columns..")
sa.session.execute('ALTER TABLE db_dbnode ADD COLUMN attributes '
'JSONB DEFAULT \'{}\'')
from aiida.backends.sqlalchemy.models.node import DbNode
total_nodes = sa.session.query(func.count(DbNode.id)).scalar()
total_groups = int(math.ceil(total_nodes / float(group_size)))
error = False
for i in xrange(total_groups):
print("Migrating group {} of {}".format(i, total_groups))
nodes = DbNode.query.options(
subqueryload('old_attrs'),
load_only('id', 'attributes')).order_by(
DbNode.id)[i * group_size:(i + 1) * group_size]
for node in nodes:
attrs, err_ = attributes_to_dict(
sorted(node.old_attrs, key=lambda a: a.key))
error |= err_
node.attributes = attrs
sa.session.add(node)
# Remove the db_dbnode from sqlalchemy, to allow the GC to do its
# job.
sa.session.flush()
sa.session.expunge_all()
del nodes
gc.collect()
if error:
cont = query_yes_no(
"There has been some errors during the "
"migration. Do you want to continue?", "no")
if not cont:
sa.session.rollback()
sys.exit(-1)
if delete_table:
sa.session.execute('DROP TABLE db_dbattribute')
sa.session.commit()
print("Migration of attributes finished.")
def run(self):
conf_backup_folder_abs = self.create_dir(
"Please provide the backup folder by providing the full path.",
os.path.join(expanduser(AIIDA_CONFIG_FOLDER),
self._conf_backup_folder_rel))
file_backup_folder_abs = self.create_dir(
"Please provide the destination folder of the backup (normally in "
"the previously provided backup folder).",
os.path.join(conf_backup_folder_abs, self._file_backup_folder_rel))
# The template backup configuration file
template_conf_path = os.path.join(
os.path.dirname(os.path.abspath(__file__)),
self._backup_info_tmpl_filename)
# Copy the sample configuration file to the backup folder
try:
shutil.copy(template_conf_path, conf_backup_folder_abs)
except Exception:
self._logger.error(
"Error copying the file {} ".format(template_conf_path) +
"to the directory {}.".format(conf_backup_folder_abs))
raise
if utils.query_yes_no("A sample configuration file was copied to {}. "
"Would you like to ".format(
conf_backup_folder_abs) +
"see the configuration parameters explanation?",
default="yes"):
self.print_info()
# Construct the path to the backup configuration file
final_conf_filepath = os.path.join(conf_backup_folder_abs,
self._backup_info_filename)
# If the backup parameters are configured now
if utils.query_yes_no("Would you like to configure the backup " +
"configuration file now?", default="yes"):
# Ask questions to properly setup the backup variables
backup_variables = self.construct_backup_variables(
file_backup_folder_abs)
with open(final_conf_filepath, 'w') as backup_info_file:
json.dump(backup_variables, backup_info_file)
# If the backup parameters are configured manually
else:
sys.stdout.write(
"Please rename the file {} ".format(
self._backup_info_tmpl_filename) +
"found in {} to ".format(conf_backup_folder_abs) +
"{} and ".format(self._backup_info_filename) +
"change the backup parameters accordingly.\n")
sys.stdout.write(
"Please adapt the startup script accordingly to point to the " +
"correct backup configuration file. For the moment, it points " +
"to {}\n".format(os.path.join(conf_backup_folder_abs,
self._backup_info_filename)))
# The contents of the startup script
script_content = \
"""#!/usr/bin/env runaiida
import logging
from aiida.common.additions.backup_script.backup import Backup
# Create the backup instance
backup_inst = Backup(backup_info_filepath="{}", additional_back_time_mins = 2)
# Define the backup logging level
backup_inst._logger.setLevel(logging.INFO)
# Start the backup
backup_inst.run()
""".format(final_conf_filepath)
# Script full path
script_path = os.path.join(conf_backup_folder_abs,
self._script_filename)
# Write the contents to the script
with open(script_path, 'w') as script_file:
script_file.write(script_content)
# Set the right permissions
try:
st = os.stat(script_path)
os.chmod(script_path, st.st_mode | stat.S_IEXEC)
except OSError:
self._logger.error("Problem setting the right permissions to the " +
"script {}.".format(script_path))
raise
def calculation_cleanworkdir(self, *args):
"""
Clean the working directory of calculations by removing all the content of the
associated RemoteFolder node. Calculations can be identified by pk with the -k flag
or by specifying limits on the modification times with -p/-o flags
"""
import argparse
parser = argparse.ArgumentParser(
prog=self.get_full_command_name(),
description="""
Clean all content of all output remote folders of calculations,
passed as a list of pks, or identified by modification time.
If a list of calculation PKs is not passed with the -k option, one or both
of the -p and -o options has to be specified. If both are specified, a logical
AND is done between the two, i.e. the calculations that will be cleaned have been
modified AFTER [-p option] days from now but BEFORE [-o option] days from now.
Passing the -f option will prevent the confirmation dialog from being prompted.
"""
)
parser.add_argument(
'-k', '--pk', metavar='PK', type=int, nargs='+', dest='pk',
help='The principal key (PK) of the calculations of which to clean the work directory'
)
parser.add_argument(
'-f', '--force', action='store_true',
help='Force the cleaning (no prompt)'
)
parser.add_argument(
'-p', '--past-days', metavar='N', type=int, action='store', dest='past_days',
help='Include calculations that have been modified within the last N days',
)
parser.add_argument(
'-o', '--older-than', metavar='N', type=int, action='store', dest='older_than',
help='Include calculations that have been modified more than N days ago',
)
parser.add_argument(
'-c', '--computers', metavar='label', nargs='+', type=str, action='store', dest='computer',
help='Include only calculations that were ran on these computers'
)
if not is_dbenv_loaded():
load_dbenv()
from aiida.backends.utils import get_automatic_user
from aiida.backends.utils import get_authinfo
from aiida.common.utils import query_yes_no
from aiida.orm.computer import Computer as OrmComputer
from aiida.orm.user import User as OrmUser
from aiida.orm.calculation import Calculation as OrmCalculation
from aiida.orm.querybuilder import QueryBuilder
from aiida.utils import timezone
import datetime
parsed_args = parser.parse_args(args)
# If a pk is given then the -o & -p options should not be specified
if parsed_args.pk is not None:
if (parsed_args.past_days is not None or parsed_args.older_than is not None):
print("You cannot specify both a list of calculation pks and the -p or -o options")
return
# If no pk is given then at least one of the -o & -p options should be specified
else:
if (parsed_args.past_days is None and parsed_args.older_than is None):
print("You should specify at least a list of calculations or the -p, -o options")
return
qb_user_filters = dict()
user = OrmUser(dbuser=get_automatic_user())
qb_user_filters["email"] = user.email
qb_computer_filters = dict()
if parsed_args.computer is not None:
qb_computer_filters["name"] = {"in": parsed_args.computer}
qb_calc_filters = dict()
if parsed_args.past_days is not None:
pd_ts = timezone.now() - datetime.timedelta(days=parsed_args.past_days)
qb_calc_filters["mtime"] = {">": pd_ts}
if parsed_args.older_than is not None:
ot_ts = timezone.now() - datetime.timedelta(days=parsed_args.older_than)
qb_calc_filters["mtime"] = {"<": ot_ts}
if parsed_args.pk is not None:
print("parsed_args.pk: ", parsed_args.pk)
qb_calc_filters["id"] = {"in": parsed_args.pk}
qb = QueryBuilder()
qb.append(OrmCalculation, tag="calc",
filters=qb_calc_filters,
project=["id", "uuid", "attributes.remote_workdir"])
qb.append(OrmComputer, computer_of="calc", tag="computer",
project=["*"],
filters=qb_computer_filters)
qb.append(OrmUser, creator_of="calc", tag="user",
project=["*"],
filters=qb_user_filters)
no_of_calcs = qb.count()
if no_of_calcs == 0:
print("No calculations found with the given criteria.")
return
print("Found {} calculations with the given criteria.".format(
no_of_calcs))
if not parsed_args.force:
if not query_yes_no("Are you sure you want to clean the work "
"directory?", "no"):
return
# get the uuids of all calculations matching the filters
calc_list_data = qb.dict()
# get all computers associated to the calc uuids above, and load them
# we group them by uuid to avoid computer duplicates
comp_uuid_to_computers = {_["computer"]["*"].uuid: _["computer"]["*"] for _ in calc_list_data}
# now build a dictionary with the info of folders to delete
remotes = {}
for computer in comp_uuid_to_computers.values():
# initialize a key of info for a given computer
remotes[computer.name] = {'transport': get_authinfo(
computer=computer, aiidauser=user._dbuser).get_transport(),
'computer': computer,
}
# select the calc pks done on this computer
this_calc_pks = [_["calc"]["id"] for _ in calc_list_data
if _["computer"]["*"].id == computer.id]
this_calc_uuids = [unicode(_["calc"]["uuid"])
for _ in calc_list_data
if _["computer"]["*"].id == computer.id]
remote_workdirs = [_["calc"]["attributes.remote_workdir"]
for _ in calc_list_data
if _["calc"]["id"] in this_calc_pks
if _["calc"]["attributes.remote_workdir"]
is not None]
remotes[computer.name]['remotes'] = remote_workdirs
remotes[computer.name]['uuids'] = this_calc_uuids
# now proceed to cleaning
for computer, dic in remotes.iteritems():
print("Cleaning the work directory on computer {}.".format(computer))
counter = 0
t = dic['transport']
with t:
remote_user = remote_user = t.whoami()
aiida_workdir = dic['computer'].get_workdir().format(
username=remote_user)
t.chdir(aiida_workdir)
# Hardcoding the sharding equal to 3 parts!
existing_folders = t.glob('*/*/*')
folders_to_delete = [i for i in existing_folders if
i.replace("/", "") in dic['uuids']]
for folder in folders_to_delete:
t.rmtree(folder)
counter += 1
if counter % 20 == 0 and counter > 0:
print("Deleted work directories: {}".format(counter))
print("{} remote folder(s) cleaned.".format(counter))
def create_configuration(profile='default'):
"""
:param profile: The profile to be configured
:return: The populated profile that was also stored.
"""
import readline
from aiida.common.exceptions import ConfigurationError
from validate_email import validate_email
from aiida.common.utils import query_yes_no
aiida_dir = os.path.expanduser(AIIDA_CONFIG_FOLDER)
print("Setting up profile {}.".format(profile))
is_test_profile = False
if profile.startswith(TEST_KEYWORD):
print(
"This is a test profile. All the data that will be stored under "
"this profile are subjected to possible deletion or "
"modification (repository and database data).")
is_test_profile = True
try:
confs = get_config()
except ConfigurationError:
# No configuration file found
confs = {}
# first time creation check
try:
confs['profiles']
except KeyError:
confs['profiles'] = {}
# load the old configuration for the given profile
try:
this_existing_confs = confs['profiles'][profile]
except KeyError:
this_existing_confs = {}
# if there is an existing configuration, print it and ask if the user wants
# to modify it.
updating_existing_prof = False
if this_existing_confs:
print(
"The following configuration found corresponding to "
"profile {}.".format(profile))
for k, v in this_existing_confs.iteritems():
if key_explanation.has_key(k):
print("{}: {}".format(key_explanation.get(k), v))
else:
print("{}: {}".format(k, v))
answ = query_yes_no("Would you like to change it?", "no")
# If the user doesn't want to change it, we abandon
if answ is False:
return this_existing_confs
# Otherwise, we continue.
else:
updating_existing_prof = True
this_new_confs = {}
try:
# Defining the backend to be used
aiida_backend = this_existing_confs.get('AIIDADB_BACKEND')
if updating_existing_prof:
print(
"The backend of already stored profiles can not be "
"changed. The current backend is {}.".format(aiida_backend))
this_new_confs['AIIDADB_BACKEND'] = aiida_backend
else:
backend_possibilities = ['django', 'sqlalchemy']
if len(backend_possibilities) > 0:
valid_aiida_backend = False
while not valid_aiida_backend:
backend_ans = raw_input(
'AiiDA backend (available: {} - sqlalchemy is in beta mode): '
.format(', '.join(backend_possibilities)))
if backend_ans in backend_possibilities:
valid_aiida_backend = True
else:
print "* ERROR! Invalid backend inserted."
print(
"* The available middlewares are {}".format(
', '.join(backend_possibilities)))
this_new_confs['AIIDADB_BACKEND'] = backend_ans
aiida_backend = backend_ans
# Setting the email
valid_email = False
readline.set_startup_hook(lambda: readline.insert_text(
this_existing_confs.get(DEFAULT_USER_CONFIG_FIELD,
DEFAULT_AIIDA_USER)))
while not valid_email:
this_new_confs[DEFAULT_USER_CONFIG_FIELD] = raw_input(
'Default user email: ')
valid_email = validate_email(
this_new_confs[DEFAULT_USER_CONFIG_FIELD])
if not valid_email:
print "** Invalid email provided!"
# Setting the database engine
db_possibilities = []
if aiida_backend == 'django':
db_possibilities.extend(['postgresql_psycopg2', 'mysql'])
elif aiida_backend == 'sqlalchemy':
db_possibilities.extend(['postgresql_psycopg2'])
if len(db_possibilities) > 0:
db_engine = this_existing_confs.get('AIIDADB_ENGINE',
db_possibilities[0])
readline.set_startup_hook(lambda: readline.insert_text(db_engine))
valid_db_engine = False
while not valid_db_engine:
db_engine_ans = raw_input(
'Database engine (available: {} - mysql is deprecated): '.
format(', '.join(db_possibilities)))
if db_engine_ans in db_possibilities:
valid_db_engine = True
else:
print "* ERROR! Invalid database engine inserted."
print("* The available engines are {}".format(
', '.join(db_possibilities)))
this_new_confs['AIIDADB_ENGINE'] = db_engine_ans
if 'postgresql_psycopg2' in this_new_confs['AIIDADB_ENGINE']:
this_new_confs['AIIDADB_ENGINE'] = 'postgresql_psycopg2'
old_host = this_existing_confs.get('AIIDADB_HOST', 'localhost')
if not old_host:
old_host = 'localhost'
readline.set_startup_hook(lambda: readline.insert_text(old_host))
this_new_confs['AIIDADB_HOST'] = raw_input('PostgreSQL host: ')
old_port = this_existing_confs.get('AIIDADB_PORT', '5432')
if not old_port:
old_port = '5432'
readline.set_startup_hook(lambda: readline.insert_text(old_port))
this_new_confs['AIIDADB_PORT'] = raw_input('PostgreSQL port: ')
readline.set_startup_hook(lambda: readline.insert_text(
this_existing_confs.get('AIIDADB_NAME')))
db_name = ''
while True:
db_name = raw_input('AiiDA Database name: ')
if is_test_profile and db_name.startswith(TEST_KEYWORD):
break
if (not is_test_profile
and not db_name.startswith(TEST_KEYWORD)):
break
print(
"The test databases should start with the prefix {} and "
"the non-test databases should not have this prefix.".
format(TEST_KEYWORD))
this_new_confs['AIIDADB_NAME'] = db_name
old_user = this_existing_confs.get('AIIDADB_USER', 'aiida')
if not old_user:
old_user = '******'
readline.set_startup_hook(lambda: readline.insert_text(old_user))
this_new_confs['AIIDADB_USER'] = raw_input('AiiDA Database user: '******'AIIDADB_PASS')))
this_new_confs['AIIDADB_PASS'] = raw_input(
'AiiDA Database password: '******'mysql' in this_new_confs['AIIDADB_ENGINE']:
this_new_confs['AIIDADB_ENGINE'] = 'mysql'
old_host = this_existing_confs.get('AIIDADB_HOST', 'localhost')
if not old_host:
old_host = 'localhost'
readline.set_startup_hook(lambda: readline.insert_text(old_host))
this_new_confs['AIIDADB_HOST'] = raw_input('mySQL host: ')
old_port = this_existing_confs.get('AIIDADB_PORT', '3306')
if not old_port:
old_port = '3306'
readline.set_startup_hook(lambda: readline.insert_text(old_port))
this_new_confs['AIIDADB_PORT'] = raw_input('mySQL port: ')
readline.set_startup_hook(lambda: readline.insert_text(
this_existing_confs.get('AIIDADB_NAME')))
db_name = ''
while True:
db_name = raw_input('AiiDA Database name: ')
if is_test_profile and db_name.startswith(TEST_KEYWORD):
break
if (not is_test_profile
and not db_name.startswith(TEST_KEYWORD)):
break
print(
"The test databases should start with the prefix {} and "
"the non-test databases should not have this prefix.".
format(TEST_KEYWORD))
this_new_confs['AIIDADB_NAME'] = db_name
old_user = this_existing_confs.get('AIIDADB_USER', 'aiida')
if not old_user:
old_user = '******'
readline.set_startup_hook(lambda: readline.insert_text(old_user))
this_new_confs['AIIDADB_USER'] = raw_input('AiiDA Database user: '******'AIIDADB_PASS')))
this_new_confs['AIIDADB_PASS'] = raw_input(
'AiiDA Database password: '******'mysql', 'postgres')")
# This part for the time being is a bit oddly written
# it should change in the future to add the possibility of having a
# remote repository. Atm, I act as only a local repo is possible
existing_repo = this_existing_confs.get(
'AIIDADB_REPOSITORY_URI',
os.path.join(aiida_dir, "repository-{}/".format(profile)))
default_protocol = 'file://'
if existing_repo.startswith(default_protocol):
existing_repo = existing_repo[len(default_protocol):]
readline.set_startup_hook(lambda: readline.insert_text(existing_repo))
new_repo_path = raw_input('AiiDA repository directory: ')
# Constructing the repo path
new_repo_path = os.path.expanduser(new_repo_path)
if not os.path.isabs(new_repo_path):
raise ValueError("You must specify an absolute path")
# Check if the new repository is a test repository and if it already
# exists.
if is_test_profile:
if TEST_KEYWORD not in new_repo_path:
raise ValueError("The test prefix {} should be contained only"
"in repository names related to test "
"profiles.".format(TEST_KEYWORD))
if os.path.isdir(new_repo_path):
print(
"The repository {} already exists. It will be used for "
"tests. Any content may be deleted.".format(new_repo_path))
else:
if TEST_KEYWORD in new_repo_path:
raise ValueError("Only test profiles can have a test "
"repository. Your repository contains the "
"test prefix {}.".format(TEST_KEYWORD))
if not os.path.isdir(new_repo_path):
print("The repository {} will be created.".format(new_repo_path))
old_umask = os.umask(DEFAULT_UMASK)
try:
os.makedirs(new_repo_path)
finally:
os.umask(old_umask)
this_new_confs['AIIDADB_REPOSITORY_URI'] = 'file://' + new_repo_path
confs['profiles'][profile] = this_new_confs
backup_config()
store_config(confs)
return this_new_confs
finally:
readline.set_startup_hook(lambda: readline.insert_text(""))
def calculation_cleanworkdir(self, *args):
"""
Clean all the content of all the output remote folders of calculations,
passed as a list of pks, or identified by modification time.
If a list of calculation PKs is not passed through -c option, one of
the option -p or -u has to be specified (if both are given, a logical
AND is done between the 2 - you clean out calculations modified AFTER
[-p option] days from now but BEFORE [-o option] days from now).
If you also pass the -f option, no confirmation will be asked.
"""
import argparse
parser = argparse.ArgumentParser(
prog=self.get_full_command_name(),
description="Clean work directory (i.e. remote folder) of AiiDA "
"calculations.")
parser.add_argument("-k",
"--pk",
metavar="PK",
type=int,
nargs="+",
help="The principal key (PK) of the calculations "
"to clean the workdir of",
dest="pk")
parser.add_argument("-f",
"--force",
action="store_true",
help="Force the cleaning (no prompt)")
parser.add_argument("-p",
"--past-days",
metavar="N",
help="Add a filter to clean workdir of "
"calculations modified during the past N "
"days",
type=int,
action="store",
dest="past_days")
parser.add_argument("-o",
"--older-than",
metavar="N",
help="Add a filter to clean workdir of "
"calculations that have been modified on a "
"date before N days ago",
type=int,
action="store",
dest="older_than")
parser.add_argument("-c",
"--computers",
metavar="label",
nargs="+",
help="Add a filter to clean workdir of "
"calculations on this computer(s) only",
type=str,
action="store",
dest="computer")
if not is_dbenv_loaded():
load_dbenv()
from aiida.backends.utils import get_automatic_user
from aiida.backends.utils import get_authinfo
from aiida.common.utils import query_yes_no
from aiida.orm.computer import Computer as OrmComputer
from aiida.orm.user import User as OrmUser
from aiida.orm.calculation import Calculation as OrmCalculation
from aiida.orm.querybuilder import QueryBuilder
from aiida.utils import timezone
import datetime
parsed_args = parser.parse_args(args)
# If a pk is given then the -o & -p options should not be specified
if parsed_args.pk is not None:
if ((parsed_args.past_days is not None)
or (parsed_args.older_than is not None)):
print(
"You cannot specify both a list of calculation pks and "
"the -p or -o options")
return
# If no pk is given then at least one of the -o & -p options should be
# specified
else:
if ((parsed_args.past_days is None)
and (parsed_args.older_than is None)):
print(
"You should specify at least a list of calculations or "
"the -p, -o options")
return
# At this point we know that either the pk or the -p -o options are
# specified
# We also check that not both -o & -p options are specified
if ((parsed_args.past_days is not None)
and (parsed_args.older_than is not None)):
print(
"Not both of the -p, -o options can be specified in the "
"same time")
return
qb_user_filters = dict()
user = OrmUser(dbuser=get_automatic_user())
qb_user_filters["email"] = user.email
qb_computer_filters = dict()
if parsed_args.computer is not None:
qb_computer_filters["name"] = {"in": parsed_args.computer}
qb_calc_filters = dict()
if parsed_args.past_days is not None:
pd_ts = timezone.now() - datetime.timedelta(
days=parsed_args.past_days)
qb_calc_filters["mtime"] = {">": pd_ts}
if parsed_args.older_than is not None:
ot_ts = timezone.now() - datetime.timedelta(
days=parsed_args.older_than)
qb_calc_filters["mtime"] = {"<": ot_ts}
if parsed_args.pk is not None:
print("parsed_args.pk: ", parsed_args.pk)
qb_calc_filters["id"] = {"in": parsed_args.pk}
qb = QueryBuilder()
qb.append(OrmCalculation,
tag="calc",
filters=qb_calc_filters,
project=["id", "uuid", "attributes.remote_workdir"])
qb.append(OrmComputer,
computer_of="calc",
project=["*"],
filters=qb_computer_filters)
qb.append(OrmUser,
creator_of="calc",
project=["*"],
filters=qb_user_filters)
no_of_calcs = qb.count()
if no_of_calcs == 0:
print("No calculations found with the given criteria.")
return
print("Found {} calculations with the given criteria.".format(
no_of_calcs))
if not parsed_args.force:
if not query_yes_no(
"Are you sure you want to clean the work "
"directory?", "no"):
return
# get the uuids of all calculations matching the filters
calc_list_data = qb.dict()
# get all computers associated to the calc uuids above, and load them
# we group them by uuid to avoid computer duplicates
comp_uuid_to_computers = {
_["computer"]["*"].uuid: _["computer"]["*"]
for _ in calc_list_data
}
# now build a dictionary with the info of folders to delete
remotes = {}
for computer in comp_uuid_to_computers.values():
# initialize a key of info for a given computer
remotes[computer.name] = {
'transport':
get_authinfo(computer=computer,
aiidauser=user._dbuser).get_transport(),
'computer':
computer,
}
# select the calc pks done on this computer
this_calc_pks = [
_["calc"]["id"] for _ in calc_list_data
if _["computer"]["*"].id == computer.id
]
this_calc_uuids = [
unicode(_["calc"]["uuid"]) for _ in calc_list_data
if _["computer"]["*"].id == computer.id
]
remote_workdirs = [
_["calc"]["attributes.remote_workdir"] for _ in calc_list_data
if _["calc"]["id"] in this_calc_pks
if _["calc"]["attributes.remote_workdir"] is not None
]
remotes[computer.name]['remotes'] = remote_workdirs
remotes[computer.name]['uuids'] = this_calc_uuids
# now proceed to cleaning
for computer, dic in remotes.iteritems():
print(
"Cleaning the work directory on computer {}.".format(computer))
counter = 0
t = dic['transport']
with t:
remote_user = remote_user = t.whoami()
aiida_workdir = dic['computer'].get_workdir().format(
username=remote_user)
t.chdir(aiida_workdir)
# Hardcoding the sharding equal to 3 parts!
existing_folders = t.glob('*/*/*')
folders_to_delete = [
i for i in existing_folders
if i.replace("/", "") in dic['uuids']
]
for folder in folders_to_delete:
t.rmtree(folder)
counter += 1
if counter % 20 == 0 and counter > 0:
print("Deleted work directories: {}".format(counter))
print("{} remote folder(s) cleaned.".format(counter))
def configure_user(self, *args):
"""
Configure the user that can run the daemon.
"""
if not is_dbenv_loaded():
from aiida.backends.utils import load_dbenv
load_dbenv(process='daemon')
if args:
print >> sys.stderr, (
"No arguments allowed for the '{}' command.".format(
self.get_full_command_name()))
sys.exit(1)
from aiida.utils import timezone
from aiida.backends.utils import get_daemon_user, set_daemon_user
from aiida.common.utils import (get_configured_user_email,
query_yes_no, query_string)
from aiida.daemon.timestamps import get_most_recent_daemon_timestamp
from aiida.common.utils import str_timedelta
from aiida.orm.user import User
old_daemon_user = get_daemon_user()
this_user = get_configured_user_email()
print("> Current default user: {}".format(this_user))
print("> Currently configured user who can run the daemon: {}".format(
old_daemon_user))
if old_daemon_user == this_user:
print(
" (therefore, at the moment you are the user who can run "
"the daemon)")
pid = self.get_daemon_pid()
if pid is not None:
print("The daemon is running! I will not proceed.")
sys.exit(1)
else:
print(" (therefore, you cannot run the daemon, at the moment)")
most_recent_timestamp = get_most_recent_daemon_timestamp()
print "*" * 76
print "* {:72s} *".format("WARNING! Change this setting only if you "
"are sure of what you are doing.")
print "* {:72s} *".format("Moreover, make sure that the "
"daemon is stopped.")
if most_recent_timestamp is not None:
timestamp_delta = timezone.now() - most_recent_timestamp
last_check_string = (
"[The most recent timestamp from the daemon was {}]".format(
str_timedelta(timestamp_delta)))
print "* {:72s} *".format(last_check_string)
print "*" * 76
answer = query_yes_no(
"Are you really sure that you want to change "
"the daemon user?",
default="no")
if not answer:
sys.exit(0)
print ""
print "Enter below the email of the new user who can run the daemon."
new_daemon_user_email = query_string("New daemon user: "******"ERROR! The user you specified ({}) does "
"not exist in the database!!".format(new_daemon_user_email))
print("The available users are {}".format(
[_.email for _ in User.search_for_users()]))
sys.exit(1)
set_daemon_user(new_daemon_user_email)
print "The new user that can run the daemon is now {} {}.".format(
found_users[0].first_name, found_users[0].last_name)
