Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Speed up creation of Nodes in the AiiDA ORM #2214

Merged
merged 2 commits into from
Nov 30, 2018
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 3 additions & 5 deletions .ci/test_fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,15 +32,12 @@ def setUp(self):

def test_create_db_cluster(self):
self.fixture_manager.create_db_cluster()
self.assertTrue(
pgtest.is_server_running(self.fixture_manager.pg_cluster.cluster))
self.assertTrue(pgtest.is_server_running(self.fixture_manager.pg_cluster.cluster))

def test_create_aiida_db(self):
self.fixture_manager.create_db_cluster()
self.fixture_manager.create_aiida_db()
self.assertTrue(
self.fixture_manager.postgres.db_exists(
self.fixture_manager.db_name))
self.assertTrue(self.fixture_manager.postgres.db_exists(self.fixture_manager.db_name))

def test_create_use_destroy_profile(self):
"""
Expand All @@ -55,6 +52,7 @@ def test_create_use_destroy_profile(self):
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)
Expand Down
2 changes: 1 addition & 1 deletion .ci/test_plugin_testcase.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ def test_data_loaded(self):

self.assertTrue(is_dbenv_loaded())
self.assertEqual(orm.load_node(self.data_pk).uuid, self.data.uuid)
self.assertEqual(orm.Computer.objects(self.backend).get(name='localhost').uuid, self.computer.uuid)
self.assertEqual(orm.Computer.objects.get(name='localhost').uuid, self.computer.uuid)

def test_tear_down(self):
"""
Expand Down
1 change: 0 additions & 1 deletion .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -189,7 +189,6 @@
aiida/common/links.py|
aiida/common/orbital/__init__.py|
aiida/common/orbital/realhydrogen.py|
aiida/common/profile.py|
aiida/common/test_extendeddicts.py|
aiida/common/test_logging.py|
aiida/control/tests/test_postgres.py|
Expand Down
20 changes: 10 additions & 10 deletions aiida/backends/djsite/db/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
# load_dbenv() function).
SCHEMA_VERSION = migrations.current_schema_version()


class AiidaQuerySet(QuerySet):
def iterator(self):
for obj in super(AiidaQuerySet, self).iterator():
Expand All @@ -54,7 +55,6 @@ def __iter__(self):

return (x.get_aiida_class() for x in super(AiidaQuerySet, self).__iter__())


def __getitem__(self, key):
"""Get item for [] operator

Expand Down Expand Up @@ -120,8 +120,8 @@ class DbUser(AbstractBaseUser, PermissionsMixin):

def get_aiida_class(self):
from aiida.orm.implementation.django.users import DjangoUser
from aiida.orm.backends import construct_backend
return DjangoUser.from_dbmodel(self, construct_backend())
from aiida.manage import get_manager
return DjangoUser.from_dbmodel(self, get_manager().get_backend())


@python_2_unicode_compatible
Expand Down Expand Up @@ -745,7 +745,7 @@ def create_value(cls, key, value, subspecifier_value=None,
bulk_create() call).
"""
import datetime

import aiida.utils.json as json
from aiida.utils.timezone import is_naive, make_aware, get_current_timezone

Expand Down Expand Up @@ -1315,8 +1315,8 @@ def __str__(self):

def get_aiida_class(self):
from aiida.orm.implementation.django.groups import DjangoGroup
from aiida.orm.backends import construct_backend
return DjangoGroup.from_dbmodel(self, construct_backend())
from aiida.manage import get_manager
return DjangoGroup.from_dbmodel(self, get_manager().get_backend())


@python_2_unicode_compatible
Expand Down Expand Up @@ -1399,8 +1399,8 @@ def get_dbcomputer(cls, computer):

def get_aiida_class(self):
from aiida.orm.implementation.django.computer import DjangoComputer
from aiida.orm.backends import construct_backend
return DjangoComputer.from_dbmodel(self, construct_backend())
from aiida.manage import get_manager
return DjangoComputer.from_dbmodel(self, get_manager().get_backend())

def _get_val_from_metadata(self, key):
import aiida.utils.json as json
Expand Down Expand Up @@ -1452,8 +1452,8 @@ def __str__(self):

def get_aiida_class(self):
from aiida.orm.implementation.django.authinfo import DjangoAuthInfo
from aiida.orm.backends import construct_backend
return DjangoAuthInfo.from_dbmodel(self, construct_backend())
from aiida.manage import get_manager
return DjangoAuthInfo.from_dbmodel(self, get_manager().get_backend())


@python_2_unicode_compatible
Expand Down
4 changes: 2 additions & 2 deletions aiida/backends/djsite/db/subtests/generic.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ def test_deletion(self):
workdir='/tmp/aiida').store()

# # This should be possible, because nothing is using this computer
self.backend.computers.delete(newcomputer.id)
orm.Computer.objects.delete(newcomputer.id)

calc_params = {
'computer': self.computer,
Expand All @@ -47,7 +47,7 @@ def test_deletion(self):
# This should fail, because there is at least a calculation
# using this computer (the one created just above)
with self.assertRaises(InvalidOperation):
self.backend.computers.delete(self.computer.id)
orm.Computer.objects.delete(self.computer.id)


class TestGroupsDjango(AiidaTestCase):
Expand Down
35 changes: 1 addition & 34 deletions aiida/backends/djsite/db/testbase.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,38 +40,9 @@ class DjangoTests(AiidaTestImplementation):
def setUpClass_method(self):
self.clean_db()
self.backend = DjangoBackend()
self.insert_data()

def setUp_method(self):
pass

def tearDown_method(self):
pass

def insert_data(self):
"""
Insert default data into the DB.
"""
from django.core.exceptions import ObjectDoesNotExist # pylint: disable=import-error, no-name-in-module

from aiida.backends.djsite.db.models import DbUser
from aiida.common.utils import get_configured_user_email
# We create the user only once:
# Otherwise, get_automatic_user() will fail when the
# user is recreated because it caches the user!
# In any case, store it in self.user though
try:
self.user = DbUser.objects.get(email=get_configured_user_email())
except ObjectDoesNotExist:
self.user = DbUser.objects.create_user(get_configured_user_email(), 'fakepwd')
# Reqired by the calling class
self.user_email = self.user.email

super(DjangoTests, self).insert_data()

def clean_db(self):
from aiida.backends.djsite.db.models import (DbComputer, DbUser, DbWorkflow, DbWorkflowStep, DbWorkflowData)
from aiida.common.utils import get_configured_user_email

# Complicated way to make sure we 'unwind' all the relationships
# between workflows and their children.
Expand Down Expand Up @@ -99,9 +70,7 @@ def clean_db(self):

DbNode.objects.all().delete() # pylint: disable=no-member

# I delete all the users except the default user.
# See discussion in setUpClass
DbUser.objects.exclude(email=get_configured_user_email()).delete()
DbUser.objects.all().delete() # pylint: disable=no-member

DbComputer.objects.all().delete()

Expand All @@ -123,8 +92,6 @@ def tearDownClass_method(self):
"the repository. Repository path: "
"{}".format(REPOSITORY_PATH))

self.clean_db()

# I clean the test repository
shutil.rmtree(REPOSITORY_PATH, ignore_errors=True)
os.makedirs(REPOSITORY_PATH)
4 changes: 2 additions & 2 deletions aiida/backends/sqlalchemy/models/authinfo.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,5 +55,5 @@ def __str__(self):

def get_aiida_class(self):
from aiida.orm.implementation.sqlalchemy.authinfo import SqlaAuthInfo
from aiida.orm.backends import construct_backend
return SqlaAuthInfo.from_dbmodel(dbmodel=self, backend=construct_backend())
from aiida.manage import get_manager
return SqlaAuthInfo.from_dbmodel(dbmodel=self, backend=get_manager().get_backend())
4 changes: 2 additions & 2 deletions aiida/backends/sqlalchemy/models/user.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,5 +50,5 @@ def __str__(self):

def get_aiida_class(self):
from aiida.orm.implementation.sqlalchemy.users import SqlaUser
from aiida.orm.backends import construct_backend
return SqlaUser.from_dbmodel(self, construct_backend())
from aiida.manage import get_manager
return SqlaUser.from_dbmodel(self, get_manager().get_backend())
3 changes: 1 addition & 2 deletions aiida/backends/sqlalchemy/tests/nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -307,7 +307,6 @@ def test_multiple_node_creation(self):
"""
from aiida.backends.sqlalchemy.models.node import DbNode
from aiida.common.utils import get_new_uuid
from aiida.orm.implementation.sqlalchemy import users as users
import aiida.backends.sqlalchemy

backend = self.backend
Expand Down Expand Up @@ -336,7 +335,7 @@ def test_multiple_node_creation(self):
"UUID in the session/DB.")

# Get the automatic user
dbuser = orm.User.objects(self.backend).get_default().backend_entity.dbmodel
dbuser = orm.User.objects.get_default().backend_entity.dbmodel
# Create a new node but now add it to the session
node_uuid = get_new_uuid()
node = DbNode(user=dbuser, uuid=node_uuid, type=None)
Expand Down
16 changes: 9 additions & 7 deletions aiida/backends/sqlalchemy/tests/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,11 @@
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
import os
from sqlalchemy.orm import sessionmaker

import aiida.backends
from aiida.backends.testbase import AiidaTestCase
from aiida.common.utils import get_configured_user_email
from aiida.manage import get_manager


class TestSessionSqla(AiidaTestCase):
Expand Down Expand Up @@ -65,7 +64,8 @@ def test_session_update_and_expiration_1(self):

session = aiida.backends.sqlalchemy.get_scoped_session()

user = self.backend.users.create(email=get_configured_user_email())
email = get_manager().get_profile().default_user_email
user = self.backend.users.create(email=email)
session.add(user.dbmodel)
session.commit()

Expand All @@ -89,14 +89,14 @@ def test_session_update_and_expiration_2(self):
expire_on_commit=True & committing computer and code objects with
their built-in store function.
"""
from aiida.backends.sqlalchemy.models.user import DbUser
from aiida.orm.data.code import Code

session = aiida.backends.sqlalchemy.get_scoped_session()

self.set_connection(expire_on_commit=True)

user = self.backend.users.create(email=get_configured_user_email())
email = get_manager().get_profile().default_user_email
user = self.backend.users.create(email=email)
session.add(user.dbmodel)
session.commit()

Expand Down Expand Up @@ -124,7 +124,8 @@ def test_session_update_and_expiration_3(self):

session = aiida.backends.sqlalchemy.get_scoped_session()

user = self.backend.users.create(email=get_configured_user_email())
email = get_manager().get_profile().default_user_email
user = self.backend.users.create(email=email)
session.add(user.dbmodel)
session.commit()

Expand Down Expand Up @@ -154,7 +155,8 @@ def test_session_update_and_expiration_4(self):

session = aiida.backends.sqlalchemy.get_scoped_session()

user = self.backend.users.create(email=get_configured_user_email())
email = get_manager().get_profile().default_user_email
user = self.backend.users.create(email=email)
session.add(user.dbmodel)
session.commit()

Expand Down
27 changes: 2 additions & 25 deletions aiida/backends/sqlalchemy/tests/testbase.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
from aiida.backends.sqlalchemy.models.user import DbUser
from aiida.backends.sqlalchemy.utils import install_tc
from aiida.backends.testimplbase import AiidaTestImplementation
from aiida.common.utils import get_configured_user_email
from aiida.manage import get_manager
from aiida.orm.implementation.sqlalchemy.backend import SqlaBackend

# Querying for expired objects automatically doesn't seem to work.
Expand Down Expand Up @@ -65,34 +65,13 @@ def setUpClass_method(self):
else:
self.clean_db()
self.backend = SqlaBackend()
self.insert_data()

def setUp_method(self):
pass

def tearDown_method(self):
pass

def insert_data(self):
"""
Insert default data into the DB.
"""
email = get_configured_user_email()

has_user = DbUser.query.filter(DbUser.email == email).first()
if not has_user:
self.user = DbUser(get_configured_user_email(), "foo", "bar",
"tests")
self.test_session.add(self.user)
self.test_session.commit()
else:
self.user = has_user

# Required by the calling class
self.user_email = self.user.email

super(SqlAlchemyTests, self).insert_data()

@staticmethod
def inject_computer(f):
@functools.wraps(f)
Expand Down Expand Up @@ -140,7 +119,7 @@ def clean_db(self):
# delete computers and users
self.test_session.query(DbNode).delete()

# # Delete the users
# Delete the users
self.test_session.query(DbUser).delete()

# Delete the computers
Expand All @@ -162,8 +141,6 @@ def tearDownClass_method(self):
"the repository. Repository path: "
"{}".format(REPOSITORY_PATH))

self.clean_db()

self.test_session.close()
self.test_session = None

Expand Down
12 changes: 7 additions & 5 deletions aiida/backends/testbase.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
from aiida.backends.tests import get_db_test_list
from aiida.common.exceptions import ConfigurationError, TestsNotAllowedError, InternalError
from aiida.common.utils import classproperty
from aiida import work
from aiida.manage import get_manager, reset_manager


def check_if_tests_can_run():
Expand Down Expand Up @@ -89,16 +89,15 @@ def get_backend_class(cls):

@classmethod
def setUpClass(cls, *args, **kwargs):
from aiida.orm.backends import construct_backend

# Note: this will raise an exception, that will be seen as a test
# failure. To be safe, you should do the same check also in the tearDownClass
# to avoid that it is run
check_if_tests_can_run()

cls.__backend_instance = cls.get_backend_class()()
cls.__backend_instance.setUpClass_method(*args, **kwargs)
cls.backend = construct_backend()
cls.backend = cls.__backend_instance.backend
cls.insert_data()

cls._class_was_setup = True

Expand All @@ -113,7 +112,7 @@ def tearDown(self):
self.__backend_instance.tearDown_method()
# Clean up the loop we created in set up.
# Call this after the instance tear down just in case it uses the loop
work.AiiDAManager.reset()
reset_manager()
loop = ioloop.IOLoop.current()
if not loop._closing:
loop.close()
Expand All @@ -135,6 +134,8 @@ def clean_db(cls):
raise InvalidOperation("You cannot call clean_db before running the setUpClass")

cls.__backend_instance.clean_db()
# Make sure to reset the backend when cleaning the database
reset_manager()

@classproperty
def computer(cls):
Expand All @@ -155,6 +156,7 @@ def tearDownClass(cls, *args, **kwargs):
# Double check for double security to avoid to run the tearDown
# if this is not a test profile
check_if_tests_can_run()
cls.clean_db()
cls.__backend_instance.tearDownClass_method(*args, **kwargs)

def assertClickResultNoException(self, cli_result):
Expand Down
Loading