Skip to content

Commit

Permalink
Fix for the link import problem under SQLA (#1769)
Browse files Browse the repository at this point in the history
  • Loading branch information
@szoupanos @sphuber
szoupanos authored and sphuber committed Jul 20, 2018
1 parent ee29447 commit e863140
Show file tree
Hide file tree
Showing 2 changed files with 101 additions and 25 deletions.
67 changes: 57 additions & 10 deletions aiida/backends/tests/export_and_import.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,11 @@
Tests for the export and import routines.
"""

import unittest

from aiida.backends.testbase import AiidaTestCase
from aiida.orm.importexport import import_data


class TestSpecificImport(AiidaTestCase):

def setUp(self):
Expand Down Expand Up @@ -1496,7 +1497,6 @@ def test_input_and_create_links(self):
from aiida.orm.importexport import export
from aiida.orm.calculation.work import WorkCalculation
from aiida.common.links import LinkType
from aiida.common.exceptions import NotExistent

tmp_folder = tempfile.mkdtemp()

Expand Down Expand Up @@ -1546,13 +1546,11 @@ def test_input_and_create_links_proper(self):
import os, shutil, tempfile

from aiida.orm.data.base import Int
from aiida.orm import Node, Data
from aiida.orm import Data
from aiida.orm.importexport import export
from aiida.orm.calculation import Calculation
from aiida.orm.calculation.inline import InlineCalculation
from aiida.orm.calculation.work import WorkCalculation
from aiida.common.links import LinkType
from aiida.common.exceptions import NotExistent
from aiida.orm.querybuilder import QueryBuilder
tmp_folder = tempfile.mkdtemp()

Expand Down Expand Up @@ -1614,14 +1612,9 @@ def test_links_for_workflows(self):
import os, shutil, tempfile

from aiida.orm.data.base import Int
from aiida.orm import Node, Data
from aiida.orm.importexport import export
from aiida.orm.calculation import Calculation
from aiida.orm.calculation.inline import InlineCalculation
from aiida.orm.calculation.work import WorkCalculation
from aiida.common.links import LinkType
from aiida.common.exceptions import NotExistent
from aiida.orm.querybuilder import QueryBuilder
tmp_folder = tempfile.mkdtemp()

try:
Expand Down Expand Up @@ -1660,5 +1653,59 @@ def test_links_for_workflows(self):



finally:
shutil.rmtree(tmp_folder, ignore_errors=True)


@unittest.skip("This test should be activated after PR #1764 is merged")
def test_double_return_links_for_workflows(self):
"""
This test checks that double return links to a node can be exported
and imported without problems,
"""
import os, shutil, tempfile

from aiida.orm.data.base import Int
from aiida.orm.importexport import export
from aiida.orm.calculation.work import WorkCalculation
from aiida.common.links import LinkType
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.node import Node

tmp_folder = tempfile.mkdtemp()

try:
w1 = WorkCalculation().store()
w2 = WorkCalculation().store()
i1 = Int(1).store()
o1 = Int(2).store()

w1.add_link_from(i1, 'input-i1', link_type=LinkType.INPUT)
w1.add_link_from(w2, 'call', link_type=LinkType.CALL)
o1.add_link_from(w1, 'output', link_type=LinkType.CREATE)
o1.add_link_from(w1, 'return', link_type=LinkType.RETURN)
o1.add_link_from(w2, 'return', link_type=LinkType.RETURN)

uuids_wanted = set(_.uuid for _ in (w1,o1,i1,w2))
links_wanted = [l for l in self.get_all_node_links() if l[3] in
('createlink', 'inputlink', 'returnlink')]

export_file = os.path.join(tmp_folder, 'export.tar.gz')
export([o1.dbnode, w1.dbnode, w2.dbnode, i1.dbnode],
outfile=export_file, silent=True)

self.clean_db()
self.insert_data()

import_data(export_file, silent=True)

uuids_in_db = [str(uuid) for [uuid] in
QueryBuilder().append(Node, project='uuid').all()]
self.assertEquals(sorted(uuids_wanted), sorted(uuids_in_db))

links_in_db = self.get_all_node_links()
self.assertEquals(sorted(links_wanted), sorted(links_in_db))


finally:
shutil.rmtree(tmp_folder, ignore_errors=True)
59 changes: 44 additions & 15 deletions aiida/orm/importexport.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -694,6 +694,7 @@ def import_data_dj(in_path,ignore_unknown_nodes=False,
#~ print foreign_ids_reverse_mappings
dbnode_reverse_mappings = foreign_ids_reverse_mappings[NODE_ENTITY_NAME]
#~ get_class_string(models.DbNode)]

for link in import_links:
try:
in_id = dbnode_reverse_mappings[link['input']]
Expand Down Expand Up @@ -721,13 +722,26 @@ def import_data_dj(in_path,ignore_unknown_nodes=False,
# the correct name
except KeyError:
try:
existing_input = existing_input_links[out_id, link['label']]
# If existing_input were the correct one, I would have found
# it already in the previous step!
raise ValueError("There exists already an input link "
"to node {} with label {} but it "
"does not come the expected input {}"
.format(out_id, link['label'], in_id))
# We try to get the existing input of the link that
# points to "out" and has label link['label'].
# If there is no existing_input, it means that the
# link doesn't exist and it has to be created. If
# it exists, then the only case that we can have more
# than one links with the same name entering a node
# is the case of the RETURN links of workflows/
# workchains. If it is not this case, then it is
# an error.
existing_input = existing_input_links[out_id,
link['label']]

if link['type'] != LinkType.RETURN:
raise ValueError(
"There exists already an input link to node "
"with UUID {} with label {} but it does not "
"come from the expected input with UUID {} "
"but from a node with UUID {}."
.format(link['output'], link['label'],
link['input'], existing_input))
except KeyError:
# New link
links_to_store.append(models.DbLink(
Expand Down Expand Up @@ -851,6 +865,7 @@ def import_data_sqla(in_path, ignore_unknown_nodes=False, silent=False):
from aiida.common.utils import get_object_from_string
from aiida.common.datastructures import calc_states
from aiida.orm.querybuilder import QueryBuilder
from aiida.common.links import LinkType

# Backend specific imports
from aiida.backends.sqlalchemy.models.node import DbCalcState
Expand Down Expand Up @@ -1280,8 +1295,8 @@ def import_data_sqla(in_path, ignore_unknown_nodes=False, silent=False):

# Needed for fast checks of existing links
from aiida.backends.sqlalchemy.models.node import DbLink
existing_links_raw = session.query(DbLink.input, DbLink.output,
DbLink.label).all()
existing_links_raw = session.query(
DbLink.input_id, DbLink.output_id,DbLink.label).all()
existing_links_labels = {(l[0], l[1]): l[2]
for l in existing_links_raw}
existing_input_links = {(l[1], l[2]): l[0]
Expand Down Expand Up @@ -1315,14 +1330,26 @@ def import_data_sqla(in_path, ignore_unknown_nodes=False, silent=False):
# the correct name
except KeyError:
try:
# We try to get the existing input of the link that
# points to "out" and has label link['label'].
# If there is no existing_input, it means that the
# link doesn't exist and it has to be created. If
# it exists, then the only case that we can have more
# than one links with the same name entering a node
# is the case of the RETURN links of workflows/
# workchains. If it is not this case, then it is
# an error.
existing_input = existing_input_links[out_id,
link['label']]
# If existing_input were the correct one, I would have
# it already in the previous step!
raise ValueError("There exists already an input link "
"to node {} with label {} but it "
"does not come the expected input {}"
.format(out_id, link['label'], in_id))

if link['type'] != LinkType.RETURN:
raise ValueError(
"There exists already an input link to node "
"with UUID {} with label {} but it does not "
"come from the expected input with UUID {} "
"but from a node with UUID {}."
.format(link['output'], link['label'],
link['input'], existing_input))
except KeyError:
# New link
links_to_store.append(DbLink(
Expand Down Expand Up @@ -1406,6 +1433,8 @@ def import_data_sqla(in_path, ignore_unknown_nodes=False, silent=False):
if not silent:
print "NO DBNODES TO IMPORT, SO NO GROUP CREATED"

if not silent:
print "COMMITTING EVERYTHING..."
session.commit()
except:
print "Rolling back"
Expand Down

0 comments on commit e863140

Please sign in to comment.