Merged in leonidkahle/aiida_core/query-n-test (pull request aiidateam…

…#113)

Query n test

aiida/backends/querybuild/querybuilder_base.py

@@ -1977,6 +1977,7 @@ def distinct(self):
  return self
 
 
+ @abstractmethod
  def _yield_per(self, batch_size):
  """
  :param count: Number of rows to yield per step
@@ -1985,34 +1986,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):

aiida/backends/querybuild/querybuilder_django.py

@@ -307,3 +307,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()

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
+

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(

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)