def parse_data(content, name, date): # pylint: disable=unused-argument
if content is None:
return ''
content_type, content_string = content.split(',') # pylint: disable=unused-variable
decoded = base64.b64decode(content_string)
fd, path = tempfile.mkstemp()
try:
with os.fdopen(fd, 'wb') as tmp:
# do stuff with temp file
tmp.write(decoded)
try:
from aiida import load_dbenv, is_dbenv_loaded
if not is_dbenv_loaded():
load_dbenv()
from aiida.orm.importexport import import_data
#from aiida.common import exceptions
import_data(path)
except Exception:
msg = 'an exception occurred while importing the archive {}'.format(
name)
msg += traceback.format_exc()
else:
msg = 'Success: imported archive {}'.format(name)
finally:
os.remove(path)
print(msg)
return msg
def create_profile(self):
"""
Set AiiDA to use the test config dir and create a default profile there
Warning: the AiiDA dbenv must not be loaded when this is called!
"""
if is_dbenv_loaded():
raise FixtureError(
'AiiDA dbenv can not be loaded while creating a test profile')
if not self.__is_running_on_test_db:
self.create_aiida_db()
from aiida.cmdline.verdilib import setup
if not self.root_dir:
self.create_root_dir()
print(self.root_dir, self.config_dir)
aiida_cfg.AIIDA_CONFIG_FOLDER = self.config_dir
backend_settings.AIIDADB_PROFILE = None
aiida_cfg.create_base_dirs()
profile_name = 'test_profile'
setup(profile=profile_name,
only_config=False,
non_interactive=True,
**self.profile)
aiida_cfg.set_default_profile('verdi', profile_name)
aiida_cfg.set_default_profile('daemon', profile_name)
self.__is_running_on_test_profile = True
def aiida(self, line='', local_ns=None):
"""Load AiiDA in ipython (checking if it was already loaded), and
inserts in the namespace the main AiiDA classes (the same that are
loaded in ``verdi shell``.
Usage::
%aiida [optional parameters]
.. todo:: implement parameters, e.g. for the profile to load.
"""
import aiida
from aiida import is_dbenv_loaded, load_dbenv
self.is_warning = False
if is_dbenv_loaded():
self.current_state = "Note! AiiDA DB environment already loaded! I do not reload it again."
self.is_warning = True
else:
load_dbenv()
self.current_state = "Loaded AiiDA DB environment."
from aiida.cmdline.commands.shell import Shell
user_ns = Shell().get_start_namespace()
for k, v in user_ns.iteritems():
add_to_ns(local_ns, k, v)
return self
def get_data_aiida(cif_uuid, plot_info):
"""Query the AiiDA database"""
from aiida import load_dbenv, is_dbenv_loaded
from aiida.backends import settings
if not is_dbenv_loaded():
load_dbenv(profile=settings.AIIDADB_PROFILE)
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.data.parameter import ParameterData
from aiida.orm.data.cif import CifData
qb = QueryBuilder()
qb.append(CifData,
filters={'uuid': {
'==': cif_uuid
}},
tag='cifs',
project='*')
qb.append(
ParameterData,
descendant_of='cifs',
project='*',
)
nresults = qb.count()
if nresults == 0:
plot_info.text = "No matching COF found."
return None
return qb.one()
def __init__(self, input_plugin, text='Select code:', **kwargs):
""" Dropdown for Codes for one input plugin.
:param input_plugin: Input plugin of codes to show
:type input_plugin: str
:param text: Text to display before dropdown
:type text: str
"""
self.input_plugin = input_plugin
self.codes = {}
self.label = ipw.Label(value=text)
self.dropdown = ipw.Dropdown(options=[], disabled=True)
self.output = ipw.Output()
children = [ipw.HBox([self.label, self.dropdown, self.output])]
super(CodeDropdown, self).__init__(children=children, **kwargs)
from aiida import load_dbenv, is_dbenv_loaded
from aiida.backends import settings
if not is_dbenv_loaded():
load_dbenv(profile=settings.AIIDADB_PROFILE)
self.refresh()
def get_start_namespace(self):
"""Load all default and custom modules"""
from aiida import load_dbenv, is_dbenv_loaded
from aiida.backends import settings
if not is_dbenv_loaded():
load_dbenv(profile=settings.AIIDADB_PROFILE)
from aiida.common.setup import get_property
user_ns = {}
# load default modules
for app_mod, model_name, alias in default_modules_list:
user_ns[alias] = getattr(__import__(app_mod, {}, {},
model_name), model_name)
# load custom modules
custom_modules_list = [(str(e[0]), str(e[2])) for e in
[p.rpartition('.') for p in get_property(
'verdishell.modules', default="").split(
':')]
if e[1] == '.']
for app_mod, model_name in custom_modules_list:
try:
user_ns[model_name] = getattr(
__import__(app_mod, {}, {}, model_name), model_name)
except AttributeError:
# if the module does not exist, we ignore it
pass
return user_ns
def schema():
"""view the validation schema"""
if not is_dbenv_loaded():
load_dbenv()
from aiida.orm import DataFactory
schema = DataFactory('crystal17.structsettings').data_schema
edict.pprint(schema, depth=None, print_func=click.echo)
def get_data_aiida(projections, sliders_dict, quantities, plot_info):
"""Query the AiiDA database"""
from aiida import load_dbenv, is_dbenv_loaded
from aiida.backends import settings
if not is_dbenv_loaded():
load_dbenv(profile=settings.AIIDADB_PROFILE)
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.data.parameter import ParameterData
filters = {}
def add_range_filter(bounds, label):
# a bit of cheating until this is resolved
# https://github.com/aiidateam/aiida_core/issues/1389
# filters['attributes.'+label] = {'>=':bounds[0]}
filters["attributes." + label] = {
"and": [{
">=": bounds[0]
}, {
"<": bounds[1]
}]
}
for k, v in sliders_dict.items():
# Note: filtering is costly, avoid if possible
if not v.value == quantities[k]["range"]:
add_range_filter(v.value, k)
qb = QueryBuilder()
qb.append(
ParameterData,
filters=filters,
project=["attributes." + p
for p in projections] + ["uuid", "extras.cif_uuid"],
)
nresults = qb.count()
if nresults == 0:
plot_info.text = "No matching frameworks found."
return data_empty
plot_info.text = "{} frameworks found. Plotting...".format(nresults)
# x,y position
x, y, clrs, uuids, names, cif_uuids = zip(*qb.all())
plot_info.text = "{} frameworks queried".format(nresults)
x = map(float, x)
y = map(float, y)
cif_uuids = map(str, cif_uuids)
uuids = map(str, uuids)
if projections[2] == "group":
# clrs = map(lambda clr: bondtypes.index(clr), clrs)
clrs = map(str, clrs)
else:
clrs = map(float, clrs)
return dict(x=x, y=y, uuid=cif_uuids, color=clrs, name=names)
def test_data_loaded(self):
"""
Check that the data is indeed in the DB when calling load_node
"""
from aiida.orm import Computer
from aiida.orm import load_node
self.assertTrue(is_dbenv_loaded())
self.assertEqual(load_node(self.data_pk).uuid, self.data.uuid)
self.assertEqual(Computer.get('localhost').uuid, self.computer.uuid)
def show(pk, content):
"""show the contents of a basis set"""
if not is_dbenv_loaded():
load_dbenv()
from aiida.orm import load_node
node = load_node(pk)
if not isinstance(node, get_data_class('crystal17.basisset')):
click.echo("The node was not of type 'crystal17.basisset'", err=True)
else:
edict.pprint(node.metadata, depth=None, print_func=click.echo)
if content:
click.echo("---")
click.echo(node.content)
def create_aiida_db(self):
"""Create the necessary database on the temporary postgres instance"""
if is_dbenv_loaded():
raise FixtureError('AiiDA dbenv can not be loaded while creating a test db environment')
if not self.db_params:
self.create_db_cluster()
self.postgres = Postgres(interactive=False, quiet=True)
self.postgres.dbinfo = self.db_params
self.postgres.determine_setup()
self.db_params = self.postgres.dbinfo
if not self.postgres.pg_execute:
raise FixtureError('Could not connect to the test postgres instance')
self.postgres.create_dbuser(self.db_user, self.db_pass)
self.postgres.create_db(self.db_user, self.db_name)
self.__is_running_on_test_db = True
def show(pk, symmetries):
"""show the contents of a StructSettingsData"""
if not is_dbenv_loaded():
load_dbenv()
from aiida.orm import load_node
from aiida.orm import DataFactory
node = load_node(pk)
if not isinstance(node, DataFactory('crystal17.structsettings')):
click.echo("The node was not of type 'crystal17.structsettings'",
err=True)
elif symmetries:
edict.pprint(node.data, print_func=click.echo, round_floats=5)
else:
edict.pprint(dict(node.iterattrs()), print_func=click.echo)
def transition_json_column(profile=None):
"""
Migrate the TEXT column containing JSON into JSON columns
"""
print("\nChanging various columns to JSON format.")
if not is_dbenv_loaded():
transition_load_db_env(profile=profile)
table_col = [
('db_dbauthinfo', 'metadata'), ('db_dbauthinfo', 'auth_params'),
('db_dbcomputer', 'metadata'), ('db_dbcomputer', 'transport_params'),
('db_dblog', 'metadata')
]
inspector = reflection.Inspector.from_engine(sa.get_scoped_session().bind)
sql = ("ALTER TABLE {table} ALTER COLUMN {column} TYPE JSONB "
"USING {column}::JSONB")
session = sa.get_scoped_session()
with session.begin(subtransactions=True):
for table, col in table_col:
table_cols = inspector.get_columns(table)
col_type_list = [_["type"] for _ in table_cols if _["name"] == col]
if len(col_type_list) != 1:
raise Exception("Problem with table {} and column {}. Either"
"the column doesn't exist or multiple "
"occurrences were found.".format(table, col))
if isinstance(col_type_list[0], JSON):
print(
"Column {} of table {} is already in JSON format. "
"Proceeding with the next table & column.".format(
col, table))
continue
print("Changing column {} of table {} in JSON format.".format(
table, col))
session.execute(sql.format(table=table, column=col))
session.commit()
def __init__(self, text="Upload Structure", node_class=None, **kwargs):
""" Upload a structure and store it in AiiDA database.
:param text: Text to display before upload button
:type text: str
:param node_class: AiiDA node class for storing the structure.
Possible values: 'StructureData', 'CifData' or None (let the user decide).
Note: If your workflows require a specific node class, better fix it here.
"""
self.file_upload = FileUploadWidget(text)
self.viewer = nglview.NGLWidget()
self.btn_store = ipw.Button(description='Store in AiiDA',
disabled=True)
self.structure_description = ipw.Text(
placeholder="Description (optional)")
self.structure_ase = None
self.structure_node = None
self.data_format = ipw.RadioButtons(
options=['StructureData', 'CifData'], description='Data type:')
if node_class is None:
store = ipw.HBox(
[self.btn_store, self.data_format, self.structure_description])
else:
self.data_format.value = node_class
store = ipw.HBox([self.btn_store, self.structure_description])
children = [self.file_upload, self.viewer, store]
super(StructureUploadWidget, self).__init__(children=children,
**kwargs)
self.file_upload.observe(self._on_file_upload, names='data')
self.btn_store.on_click(self._on_click_store)
from aiida import load_dbenv, is_dbenv_loaded
from aiida.backends import settings
if not is_dbenv_loaded():
load_dbenv(profile=settings.AIIDADB_PROFILE)
def test_create_use_destroy_profile(self):
"""
Test temporary test profile creation
* The profile gets created, the dbenv loaded
* Data can be stored in the db
* reset_db deletes all data added after profile creation
* destroy_all removes all traces of the test run
"""
from aiida.common import setup as aiida_cfg
from aiida import is_dbenv_loaded
with Capturing() as output:
self.fixture_manager.create_profile()
self.assertTrue(self.fixture_manager.root_dir_ok, msg=output)
self.assertTrue(self.fixture_manager.config_dir_ok, msg=output)
self.assertTrue(self.fixture_manager.repo_ok, msg=output)
self.assertEqual(aiida_cfg.AIIDA_CONFIG_FOLDER,
self.fixture_manager.config_dir,
msg=output)
self.assertTrue(is_dbenv_loaded())
from aiida.orm import DataFactory, load_node
data = DataFactory('parameter')(dict={'key': 'value'})
data.store()
data_pk = data.pk
self.assertTrue(load_node(data_pk))
with self.assertRaises(FixtureError):
self.test_create_aiida_db()
self.fixture_manager.reset_db()
with self.assertRaises(Exception):
load_node(data_pk)
temp_dir = self.fixture_manager.root_dir
self.fixture_manager.destroy_all()
with self.assertRaises(Exception):
self.fixture_manager.postgres.db_exists(
self.fixture_manager.db_name)
self.assertFalse(os.path.exists(temp_dir))
self.assertIsNone(self.fixture_manager.root_dir)
self.assertIsNone(self.fixture_manager.pg_cluster)
def list():
"""Display all MultiplyParameters nodes"""
from aiida import is_dbenv_loaded, load_dbenv
if not is_dbenv_loaded():
load_dbenv()
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm import DataFactory
MultiplyParameters = DataFactory('template.factors')
qb = QueryBuilder()
qb.append(MultiplyParameters)
results = qb.all()
vsep = '\t'
s = ""
for result in results:
obj = result[0]
s += "{}, pk: {}\n".format(str(obj), obj.pk)
sys.stdout.write(s)
def list(): # pylint: disable=redefined-builtin
"""
Display all KkrstructureData nodes
"""
from aiida import is_dbenv_loaded, load_dbenv
if not is_dbenv_loaded():
load_dbenv()
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm import DataFactory
KKrStructure = DataFactory('kkr.kkrstructure')
qb = QueryBuilder()
qb.append(KkrStructure)
results = qb.all()
s = ""
for result in results:
obj = result[0]
s += "{}, pk: {}\n".format(str(obj), obj.pk)
sys.stdout.write(s)
def transition_settings(profile=None):
"""
Migrate the DbAttribute table into the attributes column of db_dbnode.
"""
if not is_dbenv_loaded():
transition_load_db_env(profile=profile)
from aiida.utils import timezone
class DbSetting(Base):
__tablename__ = "db_dbsetting"
id = Column(Integer, primary_key=True)
key = Column(String(255), index=True, nullable=False)
datatype = Column(String(10), index=True, nullable=False)
tval = Column(String(255), default='', nullable=True)
fval = Column(Float, default=None, nullable=True)
ival = Column(Integer, default=None, nullable=True)
bval = Column(Boolean, default=None, nullable=True)
dval = Column(DateTime(timezone=True), default=None, nullable=True)
val = Column(JSONB, default={})
description = Column(String(255), default='', nullable=True)
time = Column(DateTime(timezone=True), default=timezone.now)
print("\nStarting migration of settings.")
inspector = reflection.Inspector.from_engine(sa.session.bind)
settings_table_cols = inspector.get_columns(SETTINGS_TABLE_NAME)
col_names = [_["name"] for _ in settings_table_cols]
if SETTINGS_VAL_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(
SETTINGS_VAL_COL_NAME, SETTINGS_TABLE_NAME))
return
with sa.session.begin(subtransactions=True):
print("Creating columns..")
sa.session.execute('ALTER TABLE db_dbsetting ADD COLUMN val '
'JSONB DEFAULT \'{}\'')
sa.session.commit()
with sa.session.begin(subtransactions=True):
total_settings = sa.session.query(DbSetting).all()
for settings in total_settings:
dt = settings.datatype
val = None
if dt == "txt":
val = settings.tval
elif dt == "float":
val = settings.fval
if math.isnan(val):
val = 'NaN'
elif dt == "int":
val = settings.ival
elif dt == "bool":
val = settings.bval
elif dt == "date":
val = settings.dval
settings.val = val
flag_modified(settings, "val")
sa.session.flush()
sa.session.commit()
with sa.session.begin(subtransactions=True):
for col_name in ["datatype", "tval", "fval", "ival", "bval", "dval"]:
sql = ("ALTER TABLE {table} DROP COLUMN {column}")
sa.session.execute(
sql.format(table=SETTINGS_TABLE_NAME, column=col_name))
sa.session.commit()
print("Migration of settings finished.")
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.")
# Works run by the daemon (using submit)
from aiida import load_dbenv, is_dbenv_loaded
if not is_dbenv_loaded():
load_dbenv()
from aiida.orm import CalculationFactory, DataFactory, WorkflowFactory
from aiida.work.run import run, submit
from aiida.orm.data.structure import StructureData
from aiida.orm.data.base import Str, Float, Bool, Int
KpointsData = DataFactory("array.kpoints")
ParameterData = DataFactory('parameter')
# Define structure
import numpy as np
cell = [[3.1900000, 0.0000000, 0.0000000], [-1.5950000, 2.7626210, 0.0000000],
[0.0000000, 0.0000000, 5.1890000]]
scaled_positions = [(0.6666669, 0.3333334, 0.0000000),
(0.3333331, 0.6666663, 0.5000000),
(0.6666669, 0.3333334, 0.3750000),
(0.3333331, 0.6666663, 0.8750000)]
symbols = ['Ga', 'Ga', 'N', 'N']
positions = np.dot(scaled_positions, cell)
structure = StructureData(cell=cell)
for i, scaled_position in enumerate(scaled_positions):
def load_dbenv_if_not_loaded():
from aiida import load_dbenv, is_dbenv_loaded
if not is_dbenv_loaded():
load_dbenv()
def run(self, *args):
from aiida.backends.utils import load_dbenv, is_dbenv_loaded
if not is_dbenv_loaded():
load_dbenv()
import argparse
from aiida.cmdline.commands.shell import default_modules_list
import aiida.orm.autogroup
from aiida.orm.autogroup import Autogroup
parser = argparse.ArgumentParser(
prog=self.get_full_command_name(),
description='Execute an AiiDA script.')
parser.add_argument('-g',
'--group',
type=bool,
default=True,
help='Enables the autogrouping, default = True')
parser.add_argument('-n',
'--groupname',
type=str,
default=None,
help='Specify the name of the auto group')
# parser.add_argument('-o','--grouponly', type=str, nargs='+', default=['all'],
# help='Limit the grouping to specific classes (by default, all classes are grouped')
parser.add_argument(
'-e',
'--exclude',
type=str,
nargs='+',
default=[],
help=('Autogroup only specific calculation classes.'
" Select them by their module name."))
parser.add_argument(
'-E',
'--excludesubclasses',
type=str,
nargs='+',
default=[],
help=('Autogroup only specific calculation classes.'
" Select them by their module name."))
parser.add_argument('-i',
'--include',
type=str,
nargs='+',
default=['all'],
help=('Autogroup only specific data classes.'
" Select them by their module name."))
parser.add_argument('-I',
'--includesubclasses',
type=str,
nargs='+',
default=[],
help=('Autogroup only specific code classes.'
" Select them by their module name."))
parser.add_argument('scriptname',
metavar='ScriptName',
type=str,
help='The name of the script you want to execute')
parser.add_argument('new_args',
metavar='ARGS',
nargs=argparse.REMAINDER,
type=str,
help='Further parameters to pass to the script')
parsed_args = parser.parse_args(args)
# Prepare the environment for the script to be run
globals_dict = {
'__builtins__': globals()['__builtins__'],
'__name__': '__main__',
'__file__': parsed_args.scriptname,
'__doc__': None,
'__package__': None
}
## dynamically load modules (the same of verdi shell) - but in
## globals_dict, not in the current environment
for app_mod, model_name, alias in default_modules_list:
globals_dict["{}".format(alias)] = getattr(
__import__(app_mod, {}, {}, model_name), model_name)
if parsed_args.group:
automatic_group_name = parsed_args.groupname
if automatic_group_name is None:
import datetime
now = datetime.datetime.now()
automatic_group_name = "Verdi autogroup on " + now.strftime(
"%Y-%m-%d %H:%M:%S")
aiida_verdilib_autogroup = Autogroup()
aiida_verdilib_autogroup.set_exclude(parsed_args.exclude)
aiida_verdilib_autogroup.set_include(parsed_args.include)
aiida_verdilib_autogroup.set_exclude_with_subclasses(
parsed_args.excludesubclasses)
aiida_verdilib_autogroup.set_include_with_subclasses(
parsed_args.includesubclasses)
aiida_verdilib_autogroup.set_group_name(automatic_group_name)
## Note: this is also set in the exec environment!
## This is the intended behavior
aiida.orm.autogroup.current_autogroup = aiida_verdilib_autogroup
try:
f = open(parsed_args.scriptname)
except IOError:
print >> sys.stderr, "{}: Unable to load file '{}'".format(
self.get_full_command_name(), parsed_args.scriptname)
sys.exit(1)
else:
try:
# Must add also argv[0]
new_argv = [parsed_args.scriptname] + parsed_args.new_args
with update_environment(new_argv=new_argv):
# Add local folder to sys.path
sys.path.insert(0, os.path.abspath(os.curdir))
# Pass only globals_dict
exec(f, globals_dict)
# print sys.argv
except SystemExit as e:
## Script called sys.exit()
# print sys.argv, "(sys.exit {})".format(e.message)
## Note: remember to re-raise, the exception to have
## the error code properly returned at the end!
raise
finally:
f.close()
def setup(profile, only_config, non_interactive=False, **kwargs):
'''
setup an aiida profile and aiida user (and the aiida default user).
:param profile: Profile name
:param only_config: do not create a new user
:param non_interactive: do not prompt for configuration values, fail if not all values are given as kwargs.
:param backend: one of 'django', 'sqlalchemy'
:param email: valid email address for the user
:param db_host: hostname for the database
:param db_port: port to connect to the database
:param db_user: name of the db user
:param db_pass: password of the db user
'''
from aiida.common.setup import (create_base_dirs, create_configuration,
set_default_profile, DEFAULT_UMASK,
create_config_noninteractive)
from aiida.backends.profile import BACKEND_SQLA, BACKEND_DJANGO
from aiida.backends.utils import set_backend_type
from aiida.common.exceptions import InvalidOperation
# ~ cmdline_args = list(args)
# ~ only_user_config = False
# ~ try:
# ~ cmdline_args.remove('--only-config')
# ~ only_user_config = True
# ~ except ValueError:
# ~ # Parameter not provided
# ~ pass
only_user_config = only_config
# ~ if cmdline_args:
# ~ print >> sys.stderr, "Unknown parameters on the command line: "
# ~ print >> sys.stderr, ", ".join(cmdline_args)
# ~ sys.exit(1)
# create the directories to store the configuration files
create_base_dirs()
# gprofile = 'default' if profile is None else profile
# ~ gprofile = profile if settings_profile.AIIDADB_PROFILE is None \
# ~ else settings_profile.AIIDADB_PROFILE
if settings_profile.AIIDADB_PROFILE and profile:
sys.exit(
'the profile argument cannot be used if verdi is called with -p option: {} and {}'
.format(settings_profile.AIIDADB_PROFILE, profile))
gprofile = settings_profile.AIIDADB_PROFILE or profile
if gprofile == profile:
settings_profile.AIIDADB_PROFILE = profile
if not settings_profile.AIIDADB_PROFILE:
settings_profile.AIIDADB_PROFILE = 'default'
# used internally later
gprofile = settings_profile.AIIDADB_PROFILE
created_conf = None
# ask and store the configuration of the DB
if non_interactive:
try:
created_conf = create_config_noninteractive(
profile=gprofile,
backend=kwargs['backend'],
email=kwargs['email'],
db_host=kwargs['db_host'],
db_port=kwargs['db_port'],
db_name=kwargs['db_name'],
db_user=kwargs['db_user'],
db_pass=kwargs.get('db_pass', ''),
repo=kwargs['repo'],
force_overwrite=kwargs.get('force_overwrite', False))
except ValueError as e:
click.echo("Error during configuation: {}".format(e.message),
err=True)
sys.exit(1)
except KeyError as e:
click.echo(
"--non-interactive requires all values to be given on the commandline! Missing argument: {}"
.format(e.message),
err=True)
sys.exit(1)
else:
try:
created_conf = create_configuration(profile=gprofile)
except ValueError as e:
print >> sys.stderr, "Error during configuration: {}".format(
e.message)
sys.exit(1)
# set default DB profiles
set_default_profile('verdi', gprofile, force_rewrite=False)
set_default_profile('daemon', gprofile, force_rewrite=False)
if only_user_config:
print(
"Only user configuration requested, "
"skipping the migrate command")
else:
print "Executing now a migrate command..."
backend_choice = created_conf['AIIDADB_BACKEND']
if backend_choice == BACKEND_DJANGO:
print("...for Django backend")
# The correct profile is selected within load_dbenv.
# Setting os.umask here since sqlite database gets created in
# this step.
old_umask = os.umask(DEFAULT_UMASK)
# This check should be done more properly
# try:
# backend_type = get_backend_type()
# except KeyError:
# backend_type = None
#
# if backend_type is not None and backend_type != BACKEND_DJANGO:
# raise InvalidOperation("An already existing database found"
# "and a different than the selected"
# "backend was used for its "
# "management.")
try:
pass_to_django_manage([execname, 'migrate'], profile=gprofile)
finally:
os.umask(old_umask)
set_backend_type(BACKEND_DJANGO)
elif backend_choice == BACKEND_SQLA:
print("...for SQLAlchemy backend")
from aiida import is_dbenv_loaded
from aiida.backends import settings
from aiida.backends.sqlalchemy.utils import (
_load_dbenv_noschemacheck, check_schema_version)
from aiida.backends.profile import load_profile
# We avoid calling load_dbenv since we want to force the schema
# migration
if not is_dbenv_loaded():
settings.LOAD_DBENV_CALLED = True
# This is going to set global variables in settings, including
# settings.BACKEND
load_profile()
_load_dbenv_noschemacheck()
# Perform the needed migration quietly
check_schema_version(force_migration=True)
set_backend_type(BACKEND_SQLA)
else:
raise InvalidOperation("Not supported backend selected.")
print "Database was created successfully"
# I create here the default user
print "Loading new environment..."
if only_user_config:
from aiida.backends.utils import load_dbenv, is_dbenv_loaded
# db environment has not been loaded in this case
if not is_dbenv_loaded():
load_dbenv()
from aiida.common.setup import DEFAULT_AIIDA_USER
from aiida.orm.user import User as AiiDAUser
if not AiiDAUser.search_for_users(email=DEFAULT_AIIDA_USER):
print "Installing default AiiDA user..."
nuser = AiiDAUser(email=DEFAULT_AIIDA_USER)
nuser.first_name = "AiiDA"
nuser.last_name = "Daemon"
nuser.is_staff = True
nuser.is_active = True
nuser.is_superuser = True
nuser.force_save()
from aiida.common.utils import get_configured_user_email
email = get_configured_user_email()
print "Starting user configuration for {}...".format(email)
if email == DEFAULT_AIIDA_USER:
print "You set up AiiDA using the default Daemon email ({}),".format(
email)
print "therefore no further user configuration will be asked."
else:
# Ask to configure the new user
if not non_interactive:
user.configure.main(args=[email])
else:
# or don't ask
aiida.cmdline.commands.user.do_configure(
email=kwargs['email'],
first_name=kwargs.get('first_name'),
last_name=kwargs.get('last_name'),
institution=kwargs.get('institution'),
no_password=True,
non_interactive=non_interactive,
force=True)
print "Setup finished."
def run(self, *args):
from aiida.common.setup import (create_base_dirs, create_configuration,
set_default_profile, DEFAULT_UMASK)
from aiida.backends.profile import BACKEND_SQLA, BACKEND_DJANGO
from aiida.backends.utils import set_backend_type, get_backend_type
from aiida.common.exceptions import InvalidOperation
cmdline_args = list(args)
only_user_config = False
try:
cmdline_args.remove('--only-config')
only_user_config = True
except ValueError:
# Parameter not provided
pass
if cmdline_args:
print >> sys.stderr, "Unknown parameters on the command line: "
print >> sys.stderr, ", ".join(cmdline_args)
sys.exit(1)
# create the directories to store the configuration files
create_base_dirs()
# gprofile = 'default' if profile is None else profile
gprofile = 'default' if settings_profile.AIIDADB_PROFILE is None \
else settings_profile.AIIDADB_PROFILE
created_conf = None
# ask and store the configuration of the DB
try:
created_conf = create_configuration(profile=gprofile)
except ValueError as e:
print >> sys.stderr, "Error during configuration: {}".format(
e.message)
sys.exit(1)
# set default DB profiles
set_default_profile('verdi', gprofile, force_rewrite=False)
set_default_profile('daemon', gprofile, force_rewrite=False)
if only_user_config:
print(
"Only user configuration requested, "
"skipping the migrate command")
else:
print "Executing now a migrate command..."
backend_choice = created_conf['AIIDADB_BACKEND']
if backend_choice == BACKEND_DJANGO:
print("...for Django backend")
# The correct profile is selected within load_dbenv.
# Setting os.umask here since sqlite database gets created in
# this step.
old_umask = os.umask(DEFAULT_UMASK)
# This check should be done more properly
# try:
# backend_type = get_backend_type()
# except KeyError:
# backend_type = None
#
# if backend_type is not None and backend_type != BACKEND_DJANGO:
# raise InvalidOperation("An already existing database found"
# "and a different than the selected"
# "backend was used for its "
# "management.")
try:
pass_to_django_manage([execname, 'migrate'],
profile=gprofile)
finally:
os.umask(old_umask)
set_backend_type(BACKEND_DJANGO)
elif backend_choice == BACKEND_SQLA:
print("...for SQLAlchemy backend")
from aiida.backends.sqlalchemy.models.base import Base
from aiida.backends.sqlalchemy.utils import (get_engine,
install_tc)
from aiida.common.setup import get_profile_config
from aiida import is_dbenv_loaded, load_dbenv
if not is_dbenv_loaded():
load_dbenv()
# This check should be done more properly
# try:
# backend_type = get_backend_type()
# except KeyError:
# backend_type = None
#
# if backend_type is not None and backend_type != BACKEND_SQLA:
# raise InvalidOperation("An already existing database found"
# "and a different than the selected"
# "backend was used for its "
# "management.")
# Those import are necessary for SQLAlchemy to correctly create
# the needed database tables.
from aiida.backends.sqlalchemy.models.authinfo import (
DbAuthInfo)
from aiida.backends.sqlalchemy.models.comment import DbComment
from aiida.backends.sqlalchemy.models.computer import (
DbComputer)
from aiida.backends.sqlalchemy.models.group import (
DbGroup, table_groups_nodes)
from aiida.backends.sqlalchemy.models.lock import DbLock
from aiida.backends.sqlalchemy.models.log import DbLog
from aiida.backends.sqlalchemy.models.node import (DbLink,
DbNode,
DbPath,
DbCalcState)
from aiida.backends.sqlalchemy.models.user import DbUser
from aiida.backends.sqlalchemy.models.workflow import (
DbWorkflow, DbWorkflowData, DbWorkflowStep)
from aiida.backends.sqlalchemy.models.settings import DbSetting
connection = get_engine(get_profile_config(gprofile))
Base.metadata.create_all(connection)
install_tc(connection)
set_backend_type(BACKEND_SQLA)
else:
raise InvalidOperation("Not supported backend selected.")
print "Database was created successfully"
# I create here the default user
print "Loading new environment..."
if only_user_config:
from aiida.backends.utils import load_dbenv, is_dbenv_loaded
# db environment has not been loaded in this case
if not is_dbenv_loaded():
load_dbenv()
from aiida.common.setup import DEFAULT_AIIDA_USER
from aiida.orm.user import User as AiiDAUser
if not AiiDAUser.search_for_users(email=DEFAULT_AIIDA_USER):
print "Installing default AiiDA user..."
nuser = AiiDAUser(email=DEFAULT_AIIDA_USER)
nuser.first_name = "AiiDA"
nuser.last_name = "Daemon"
nuser.is_staff = True
nuser.is_active = True
nuser.is_superuser = True
nuser.force_save()
from aiida.common.utils import get_configured_user_email
email = get_configured_user_email()
print "Starting user configuration for {}...".format(email)
if email == DEFAULT_AIIDA_USER:
print "You set up AiiDA using the default Daemon email ({}),".format(
email)
print "therefore no further user configuration will be asked."
else:
# Ask to configure the new user
User().user_configure(email)
print "Install finished."
def setup_profile(profile,
only_config,
set_default=False,
non_interactive=False,
**kwargs):
"""
Setup an AiiDA profile and AiiDA user (and the AiiDA default user).
:param profile: Profile name
:param only_config: do not create a new user
:param set_default: set the new profile as the default
:param non_interactive: do not prompt for configuration values, fail if not all values are given as kwargs.
:param backend: one of 'django', 'sqlalchemy'
:param email: valid email address for the user
:param db_host: hostname for the database
:param db_port: port to connect to the database
:param db_user: name of the db user
:param db_pass: password of the db user
"""
from aiida.backends import settings
from aiida.backends.profile import BACKEND_SQLA, BACKEND_DJANGO
from aiida.backends.utils import set_backend_type
from aiida.cmdline import EXECNAME
from aiida.cmdline.commands import cmd_user
from aiida.common.exceptions import InvalidOperation
from aiida.common.setup import (create_base_dirs, create_configuration,
set_default_profile, DEFAULT_UMASK,
create_config_noninteractive)
only_user_config = only_config
# Create the directories to store the configuration files
create_base_dirs()
if settings.AIIDADB_PROFILE and profile:
sys.exit(
'the profile argument cannot be used if verdi is called with -p option: {} and {}'
.format(settings.AIIDADB_PROFILE, profile))
gprofile = settings.AIIDADB_PROFILE or profile
if gprofile == profile:
settings.AIIDADB_PROFILE = profile
if not settings.AIIDADB_PROFILE:
settings.AIIDADB_PROFILE = 'default'
# used internally later
gprofile = settings.AIIDADB_PROFILE
created_conf = None
# ask and store the configuration of the DB
if non_interactive:
try:
created_conf = create_config_noninteractive(
profile=gprofile,
backend=kwargs['backend'],
email=kwargs['email'],
db_host=kwargs['db_host'],
db_port=kwargs['db_port'],
db_name=kwargs['db_name'],
db_user=kwargs['db_user'],
db_pass=kwargs.get('db_pass', ''),
repo=kwargs['repo'],
force_overwrite=kwargs.get('force_overwrite', False))
except ValueError as exception:
click.echo("Error during configuation: {}".format(
exception.message),
err=True)
sys.exit(1)
except KeyError as exception:
import traceback
click.echo(traceback.format_exc())
click.echo(
"--non-interactive requires all values to be given on the commandline! Missing argument: {}"
.format(exception.message),
err=True)
sys.exit(1)
else:
try:
created_conf = create_configuration(profile=gprofile)
except ValueError as exception:
print >> sys.stderr, "Error during configuration: {}".format(
exception.message)
sys.exit(1)
# Set default DB profile
set_default_profile(gprofile, force_rewrite=False)
if only_user_config:
print("Only user configuration requested, "
"skipping the migrate command")
else:
print("Executing now a migrate command...")
backend_choice = created_conf['AIIDADB_BACKEND']
if backend_choice == BACKEND_DJANGO:
print("...for Django backend")
# The correct profile is selected within load_dbenv.
# Setting os.umask here since sqlite database gets created in
# this step.
old_umask = os.umask(DEFAULT_UMASK)
# This check should be done more properly
# try:
# backend_type = get_backend_type()
# except KeyError:
# backend_type = None
#
# if backend_type is not None and backend_type != BACKEND_DJANGO:
# raise InvalidOperation("An already existing database found"
# "and a different than the selected"
# "backend was used for its "
# "management.")
try:
from aiida.backends.djsite.utils import pass_to_django_manage
pass_to_django_manage([EXECNAME, 'migrate'], profile=gprofile)
finally:
os.umask(old_umask)
set_backend_type(BACKEND_DJANGO)
elif backend_choice == BACKEND_SQLA:
print("...for SQLAlchemy backend")
from aiida import is_dbenv_loaded
from aiida.backends.sqlalchemy.utils import _load_dbenv_noschemacheck, check_schema_version
from aiida.backends.profile import load_profile
# We avoid calling load_dbenv since we want to force the schema
# migration
if not is_dbenv_loaded():
settings.LOAD_DBENV_CALLED = True
# This is going to set global variables in settings, including settings.BACKEND
load_profile()
_load_dbenv_noschemacheck()
# Perform the needed migration quietly
check_schema_version(force_migration=True)
set_backend_type(BACKEND_SQLA)
else:
raise InvalidOperation("Not supported backend selected.")
print("Database was created successfully")
# I create here the default user
print("Loading new environment...")
if only_user_config:
from aiida.backends.utils import load_dbenv, is_dbenv_loaded
# db environment has not been loaded in this case
if not is_dbenv_loaded():
load_dbenv()
from aiida.common.setup import DEFAULT_AIIDA_USER
from aiida.orm.backend import construct_backend
backend = construct_backend()
if not backend.users.find(email=DEFAULT_AIIDA_USER):
print("Installing default AiiDA user...")
nuser = backend.users.create(email=DEFAULT_AIIDA_USER,
first_name="AiiDA",
last_name="Daemon")
nuser.is_active = True
nuser.store()
from aiida.common.utils import get_configured_user_email
email = get_configured_user_email()
print("Starting user configuration for {}...".format(email))
if email == DEFAULT_AIIDA_USER:
print("You set up AiiDA using the default Daemon email ({}),".format(
email))
print("therefore no further user configuration will be asked.")
else:
if non_interactive:
# Here we map the keyword arguments onto the command line arguments
# for verdi user configure. We have to be careful that there the
# argument names are the same as those int he kwargs dict
commands = [kwargs['email'], '--non-interactive']
for arg in ('first_name', 'last_name', 'institution'):
value = kwargs.get(arg, None)
if value is not None:
commands.extend(
('--{}'.format(arg.replace('_', '-')), str(value)))
else:
commands = [email]
# Ask to configure the user
try:
# pylint: disable=no-value-for-parameter
cmd_user.configure(commands)
except SystemExit:
# Have to catch this as the configure command will do a sys.exit()
pass
if set_default:
set_default_profile(profile, force_rewrite=True)
print("Setup finished.")
def load_dbenv_if_not_loaded():
if not is_dbenv_loaded():
load_dbenv()