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)