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.

Sign up for GitHub

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

Fix 552 export dbpath #554

Merged
merged 3 commits into from
Jun 9, 2017
Merged

Fix 552 export dbpath #554

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
02a4c73
For #552, I removed all the explicit sqlalchemy queries in the export…
lekah Jun 2, 2017
fbdc2da
Resolves #552 by adding function for recursive query with Django (get…
lekah Jun 2, 2017
03a0f2a
Merge branch 'develop' into fix_552_export_dbpath
szoupanos Jun 8, 2017
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
96 changes: 51 additions & 45 deletions aiida/orm/importexport.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2002,9 +2002,8 @@ def fill_in_query(partial_query, originating_entity_str, current_entity_str):
fill_in_query(partial_query, current_entity_str, ref_model_name)


def export_tree_sqla(what, folder, also_parents = True, also_calc_outputs=True,
allowed_licenses=None, forbidden_licenses=None,
silent=False):
def export_tree_sqla(what, folder, also_parents=True, also_calc_outputs=True,
allowed_licenses=None, forbidden_licenses=None, silent=False):
"""
Export the DB entries passed in the 'what' list to a file tree.

Expand All @@ -2013,8 +2012,7 @@ def export_tree_sqla(what, folder, also_parents = True, also_calc_outputs=True,
:param what: a list of Django database entries; they can belong to different
models.
:param folder: a :py:class:`Folder <aiida.common.folders.Folder>` object
:param also_parents: if True, also all the parents are stored (from th
DbPath transitive closure table)
:param also_parents: if True, also all the parents are stored (from the transitive closure)
:param also_calc_outputs: if True, any output of a calculation is also exported
:param allowed_licenses: a list or a function. If a list, then checks
whether all licenses of Data nodes are in the list. If a function,
Expand Down Expand Up @@ -2064,10 +2062,11 @@ def export_tree_sqla(what, folder, also_parents = True, also_calc_outputs=True,
"aiida.backends.sqlalchemy.models.node.DbNode"]
if given_nodes:
# Also add the parents (to any level) to the query
# This is done via the ancestor relationship.
qb = QueryBuilder()
qb.append(Node, tag='low_node', filters={'id': {'in': given_nodes}})
qb.append(Node, ancestor_of='low_node', project=['id'])
additional_ids = [_ for [_] in qb.all()]
additional_ids = [_ for [_] in qb.all()] # Good way, since that works also when qb.all() returns []
given_nodes = list(set(given_nodes + additional_ids))
entries_ids_to_add[
"aiida.backends.sqlalchemy.models.node.DbNode"] = given_nodes
Expand All @@ -2080,37 +2079,27 @@ def export_tree_sqla(what, folder, also_parents = True, also_calc_outputs=True,
# selected
qb = QueryBuilder()
qb.append(Calculation, tag='high_node',
filters={'id': {'in': given_nodes}})
qb.append(Node, output_of='high_node', project=['id'])
filters={'id': {'in': given_nodes}}) # Only looking at calculations and subclasses
qb.append(Node, output_of='high_node', project=['id']) # and the outputs
additional_ids = [_ for [_] in qb.all()]
given_nodes = list(set(given_nodes + additional_ids))
given_nodes = list(set(given_nodes + additional_ids)) # What if given_nodes is a set from the start?
entries_ids_to_add[
"aiida.backends.sqlalchemy.models.node.DbNode"] = given_nodes

# Initial query to fire the generation of the export data

# print entries_ids_to_add
# for k, v in entries_ids_to_add.iteritems():
# print k, v

from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.node import Node


# The following is done for the nodes but we should also do it for the
# other type of input given data


# Here we get all the columns that we plan to project
model_name = get_class_string(models.node.DbNode)
project_cols = ["id"]
# The following gets a list of fields that we need, e.g. user, mtime, uuid, computer
entity_prop = all_fields_info[sqla_to_django_schema[model_name]].keys()

# Here we do the necessary renaming of
for prop in entity_prop:
nprop = (django_fields_to_sqla[model_name][prop]
if django_fields_to_sqla[model_name].has_key(prop)
else prop)
project_cols.append(nprop)
# project_cols contains the strings we can use to project, i.e. user_id, mtime, uuid, dbcomputer_id

entries_to_add = dict()
for k, v in entries_ids_to_add.iteritems():
Expand Down Expand Up @@ -2244,24 +2233,21 @@ def export_tree_sqla(what, folder, also_parents = True, also_calc_outputs=True,
## All 'parent' links (in this way, I can automatically export a node
## that will get automatically attached to a parent node in the end DB,
## if the parent node is already present in the DB)
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
from aiida.backends.sqlalchemy.models.node import DbLink
# authinfo = session.query(DbLink).filter(DbLink.output_id.in_(all_nodes_pk)).all()
linksquery = session.query(DbLink).filter(
DbLink.output_id.in_(all_nodes_pk)).distinct()

# print "================="
#
# print linksquery.all()

# I have to find a way to project only specific columns -
# To be seen with Leo

links_qb = QueryBuilder()
links_qb.append(Node, project=['uuid'], tag='input')
links_qb.append(Node,
project=['uuid'], tag='output', filters={'id':{'in':all_nodes_pk}},
edge_project=['label'], output_of='input')
links_uuid = list()
for link in linksquery.all():
links_uuid.append({"input": str(link.input.uuid),
"output": str(link.output.uuid),
"label": str(link.label)})

for input_uuid, output_uuid, link_label in links_qb.iterall():
links_uuid.append({
'input':str(input_uuid),
'output':str(output_uuid),
'label':str(link_label)
})


# The following has to be written more properly
if not silent:
Expand Down Expand Up @@ -2400,6 +2386,28 @@ def export_tree(what, folder, also_parents = True, also_calc_outputs=True,
BACKEND))


def get_all_parents_dj(node_pks):
"""
Get all the parents of given nodes
:param node_pks: one node pk or an iterable of node pks
:return: a list of aiida objects with all the parents of the nodes
"""
from aiida.backends.djsite.db import models

try:
the_node_pks = list(node_pks)
except TypeError:
the_node_pks = [node_pks]

parents = models.DbNode.objects.none()
q_inputs = models.DbNode.aiidaobjects.filter(outputs__pk__in=the_node_pks).distinct()
while q_inputs.count() > 0:
inputs = list(q_inputs)
parents = q_inputs | parents.all()
q_inputs = models.DbNode.aiidaobjects.filter(outputs__in=inputs).distinct()
return parents


def export_tree_dj(what, folder, also_parents = True, also_calc_outputs=True,
allowed_licenses=None, forbidden_licenses=None,
silent=False):
Expand All @@ -2411,8 +2419,7 @@ def export_tree_dj(what, folder, also_parents = True, also_calc_outputs=True,
:param what: a list of Django database entries; they can belong to different
models.
:param folder: a :py:class:`Folder <aiida.common.folders.Folder>` object
:param also_parents: if True, also all the parents are stored (from th
DbPath transitive closure table)
:param also_parents: if True, also all the parents are stored
:param also_calc_outputs: if True, any output of a calculation is also exported
:param allowed_licenses: a list or a function. If a list, then checks
whether all licenses of Data nodes are in the list. If a function,
Expand Down Expand Up @@ -2462,10 +2469,10 @@ def export_tree_dj(what, folder, also_parents = True, also_calc_outputs=True,
if given_nodes:
# Also add the parents (to any level) to the query
given_nodes = list(set(given_nodes +
list(models.DbNode.objects.filter(
children__in=given_nodes).values_list('pk', flat=True))))
list(get_all_parents_dj(given_nodes).values_list('pk', flat=True))))
entries_ids_to_add[get_class_string(models.DbNode)] = given_nodes


if also_calc_outputs:
given_nodes = entries_ids_to_add[get_class_string(models.DbNode)]

Expand Down Expand Up @@ -2864,8 +2871,7 @@ def export(what, outfile = 'export_data.aiida.tar.gz', overwrite = False,

:param what: a list of Django database entries; they can belong to different
models.
:param also_parents: if True, also all the parents are stored (from th
DbPath transitive closure table)
:param also_parents: if True, also all the parents are stored (from the transitive closure)
:param also_calc_outputs: if True, any output of a calculation is also exported
:param outfile: the filename of the file on which to export
:param overwrite: if True, overwrite the output file without asking.
Expand Down