# -*- coding: utf-8 -*-
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved. #
# This file is part of the AiiDA code. #
# #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida_core #
# For further information on the license, see the LICENSE.txt file #
# For further information please visit http://www.aiida.net #
###########################################################################
from aiida.backends.testbase import AiidaTestCase
[docs]class TestQueryBuilderSQLA(AiidaTestCase):
[docs] def test_clsf_sqla(self):
from aiida.orm import Group, User, Computer, Node, Data, Calculation
from aiida.backends.sqlalchemy.models.node import DbNode
from aiida.backends.sqlalchemy.models.group import DbGroup
from aiida.backends.sqlalchemy.models.user import DbUser
from aiida.backends.sqlalchemy.models.computer import DbComputer
from aiida.orm.querybuilder import QueryBuilder
qb = QueryBuilder()
for AiidaCls, ORMCls, typestr in zip(
(Group, User, Computer, Node, Data, Calculation),
(DbGroup, DbUser, DbComputer, DbNode, DbNode, DbNode),
(None, None, None, Node._query_type_string, Data._query_type_string, Calculation._query_type_string)):
cls, clstype, query_type_string = qb._get_ormclass(AiidaCls, None)
self.assertEqual(cls, ORMCls)
self.assertEqual(query_type_string, typestr)
[docs]class QueryBuilderLimitOffsetsTestSQLA(AiidaTestCase):
[docs] def test_ordering_limits_offsets_of_results_for_SQLA(self):
from aiida.orm import Node
from aiida.orm.querybuilder import QueryBuilder
# Creating 10 nodes with an attribute that can be ordered
for i in range(10):
n = Node()
n._set_attr('foo', i)
n.store()
qb = QueryBuilder().append(
Node, project='attributes.foo'
).order_by(
{Node:{'attributes.foo':{'cast':'i'}}}
)
res = list(zip(*qb.all())[0])
self.assertEqual(res, range(10))
# Now applying an offset:
qb.offset(5)
res = list(zip(*qb.all())[0])
self.assertEqual(res, range(5,10))
# Now also applying a limit:
qb.limit(3)
res = list(zip(*qb.all())[0])
self.assertEqual(res, range(5,8))