From 14b83d9850c4e1a31c5ef034ddf22ddd53848566 Mon Sep 17 00:00:00 2001 From: "leo@pegasus" Date: Tue, 6 Dec 2016 12:12:59 +0100 Subject: [PATCH 1/2] [Working on #178] Wrote test for the path queries of the QueryBuilder --- aiida/backends/sqlalchemy/tests/query.py | 142 +++++++++++++++++++++++ 1 file changed, 142 insertions(+) diff --git a/aiida/backends/sqlalchemy/tests/query.py b/aiida/backends/sqlalchemy/tests/query.py index 5d6aebdd2e..67f75620e5 100644 --- a/aiida/backends/sqlalchemy/tests/query.py +++ b/aiida/backends/sqlalchemy/tests/query.py @@ -28,3 +28,145 @@ def test_clsf_sqla(self): class QueryBuilderJoinsTestsSQLA(SqlAlchemyTests, QueryBuilderJoinsTests): pass + + + +class QueryBuilderPathSQLA(SqlAlchemyTests): + def test_query_path(self): + + from aiida.orm.querybuilder import QueryBuilder + from aiida.orm import Node + + n1 = Node() + n1.store() + n2 = Node() + n2.store() + n3 = Node() + n3.store() + n4 = Node() + n4.store() + n5 = Node() + n5.store() + n6 = Node() + n6.store() + n7 = Node() + n7.store() + n8 = Node() + n8.store() + n9 = Node() + n9.store() + + # I create a strange graph, inserting links in a order + # such that I often have to create the transitive closure + # between two graphs + n3.add_link_from(n2) + n2.add_link_from(n1) + n5.add_link_from(n3) + n5.add_link_from(n4) + n4.add_link_from(n2) + + n7.add_link_from(n6) + n8.add_link_from(n7) + + # Yet, no links from 1 to 8 + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n1.pk}, tag='anc' + ).append(Node, descendant_of='anc', filters={'id':n8.pk} + ).count(), 0) + + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n1.pk}, tag='anc' + ).append(Node, descendant_of_beta='anc', filters={'id':n8.pk} + ).count(), 0) + + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n8.pk}, tag='desc' + ).append(Node, ancestor_of='desc', filters={'id':n1.pk} + ).count(), 0) + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n8.pk}, tag='desc' + ).append(Node, ancestor_of_beta='desc', filters={'id':n1.pk} + ).count(), 0) + + + + n6.add_link_from(n5) + # Yet, now 2 links from 1 to 8 + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n1.pk}, tag='anc' + ).append(Node, descendant_of='anc', filters={'id':n8.pk} + ).count(), 2 + ) + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n1.pk}, tag='anc' + ).append(Node, descendant_of_beta='anc', filters={'id':n8.pk} + ).count(), 2 + ) + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n8.pk}, tag='desc' + ).append(Node, ancestor_of='desc', filters={'id':n1.pk} + ).count(), 2) + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n8.pk}, tag='desc' + ).append(Node, ancestor_of_beta='desc', filters={'id':n1.pk} + ).count(), 2) + + n7.add_link_from(n9) + # Still two links... + + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n1.pk}, tag='anc' + ).append(Node, descendant_of='anc', filters={'id':n8.pk} + ).count(), 2 + ) + + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n1.pk}, tag='anc' + ).append(Node, descendant_of_beta='anc', filters={'id':n8.pk} + ).count(), 2 + ) + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n8.pk}, tag='desc' + ).append(Node, ancestor_of='desc', filters={'id':n1.pk} + ).count(), 2) + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n8.pk}, tag='desc' + ).append(Node, ancestor_of_beta='desc', filters={'id':n1.pk} + ).count(), 2) + + n9.add_link_from(n6) + # And now there should be 4 nodes + + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n1.pk}, tag='anc' + ).append(Node, descendant_of='anc', filters={'id':n8.pk} + ).count(), 4) + + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n1.pk}, tag='anc' + ).append(Node, descendant_of_beta='anc', filters={'id':n8.pk} + ).count(), 4) + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n8.pk}, tag='desc' + ).append(Node, ancestor_of='desc', filters={'id':n1.pk} + ).count(), 4) + self.assertEquals( + QueryBuilder().append( + Node, filters={'id':n8.pk}, tag='desc' + ).append(Node, ancestor_of_beta='desc', filters={'id':n1.pk} + ).count(), 4) From d7c9df0efce3548f09954f6923ee09a30ef3f795 Mon Sep 17 00:00:00 2001 From: "leo@pegasus" Date: Wed, 7 Dec 2016 18:39:17 +0100 Subject: [PATCH 2/2] Making QueryBuilder for Django more stable by using atomic transactions. Also removed the return of the path in DbPathBeta --- .../backends/querybuild/querybuilder_base.py | 25 ++++--------- .../querybuild/querybuilder_django.py | 29 +++++++++++++++ aiida/backends/querybuild/querybuilder_sa.py | 37 +++++++++++++++++++ aiida/backends/sqlalchemy/models/__init__.py | 4 +- 4 files changed, 75 insertions(+), 20 deletions(-) diff --git a/aiida/backends/querybuild/querybuilder_base.py b/aiida/backends/querybuild/querybuilder_base.py index b0217c33b2..955c332064 100644 --- a/aiida/backends/querybuild/querybuilder_base.py +++ b/aiida/backends/querybuild/querybuilder_base.py @@ -1975,6 +1975,7 @@ def distinct(self): return self + @abstractmethod def _yield_per(self, batch_size): """ :param count: Number of rows to yield per step @@ -1983,34 +1984,22 @@ def _yield_per(self, batch_size): :returns: a generator """ - try: - return self.get_query().yield_per(batch_size) - except Exception as e: - # exception was raised. Rollback the session - self._get_session().rollback() - raise e + pass + + @abstractmethod def _all(self): - try: - return self.get_query().all() - except Exception as e: - # exception was raised. Rollback the session - self._get_session().rollback() - raise e + pass + @abstractmethod def _first(self): """ Executes query in the backend asking for one instance. :returns: One row of aiida results """ - try: - return self.get_query().first() - except Exception as e: - # exception was raised. Rollback the session - self._get_session().rollback() - raise e + pass def first(self): diff --git a/aiida/backends/querybuild/querybuilder_django.py b/aiida/backends/querybuild/querybuilder_django.py index 42bcf6d492..73180c6eeb 100644 --- a/aiida/backends/querybuild/querybuilder_django.py +++ b/aiida/backends/querybuild/querybuilder_django.py @@ -306,3 +306,32 @@ def _get_projectable_attribute( entity = case([(exists_stmt, select_stmt), ], else_=None) return entity + + + def _yield_per(self, batch_size): + """ + :param count: Number of rows to yield per step + + Yields *count* rows at a time + + :returns: a generator + """ + from django.db import transaction + with transaction.atomic(): + return self.get_query().yield_per(batch_size) + + + def _all(self): + from django.db import transaction + with transaction.atomic(): + return self.get_query().all() + + def _first(self): + """ + Executes query in the backend asking for one instance. + + :returns: One row of aiida results + """ + from django.db import transaction + with transaction.atomic(): + return self.get_query().first() diff --git a/aiida/backends/querybuild/querybuilder_sa.py b/aiida/backends/querybuild/querybuilder_sa.py index 7cf87202e9..58a40047b6 100644 --- a/aiida/backends/querybuild/querybuilder_sa.py +++ b/aiida/backends/querybuild/querybuilder_sa.py @@ -239,3 +239,40 @@ def _get_aiida_res(self, key, res): return returnval + def _yield_per(self, batch_size): + """ + :param count: Number of rows to yield per step + + Yields *count* rows at a time + + :returns: a generator + """ + try: + return self.get_query().yield_per(batch_size) + except Exception as e: + # exception was raised. Rollback the session + self._get_session().rollback() + raise e + + + def _all(self): + try: + return self.get_query().all() + except Exception as e: + # exception was raised. Rollback the session + self._get_session().rollback() + raise e + + def _first(self): + """ + Executes query in the backend asking for one instance. + + :returns: One row of aiida results + """ + try: + return self.get_query().first() + except Exception as e: + # exception was raised. Rollback the session + self._get_session().rollback() + raise e + diff --git a/aiida/backends/sqlalchemy/models/__init__.py b/aiida/backends/sqlalchemy/models/__init__.py index dd77becf93..c9ed36c261 100644 --- a/aiida/backends/sqlalchemy/models/__init__.py +++ b/aiida/backends/sqlalchemy/models/__init__.py @@ -43,7 +43,7 @@ DbNode.id.label('ancestor_id'), DbNode.id.label('descendant_id'), cast(-1, Integer).label('depth'), - array([DbNode.id]).label('path') #Arrays can only be used with postgres + #~ array([DbNode.id]).label('path') #Arrays can only be used with postgres ]).select_from(DbNode).cte(recursive=True) #, name="incl_aliased3") @@ -52,7 +52,7 @@ walk.c.ancestor_id, node_aliased.id, walk.c.depth + cast(1, Integer), - (walk.c.path+array([node_aliased.id])).label('path') + #~ (walk.c.path+array([node_aliased.id])).label('path') #, As above, but if arrays are supported # This is the way to reconstruct the path (the sequence of nodes traversed) ]).select_from(