♻️ REFACTOR: Move to sqlachemy 1.4 (and v2 API)

aiida/backends/utils.py

@@ -38,7 +38,7 @@ def create_sqlalchemy_engine(profile, **kwargs):
         name=profile.database_name
     )
     return create_engine(
-        engine_url, json_serializer=json.dumps, json_deserializer=json.loads, encoding='utf-8', **kwargs
+        engine_url, json_serializer=json.dumps, json_deserializer=json.loads, encoding='utf-8', future=True, **kwargs
     )
 
 

aiida/orm/querybuilder.py

@@ -1952,18 +1952,12 @@ def _build(self):
 
         ################ LAST BUT NOT LEAST ############################
         # pop the entity that I added to start the query
-        self._query._entities.pop(0)  # pylint: disable=protected-access
-
-        # Dirty solution coming up:
-        # Sqlalchemy is by default de-duplicating results if possible.
-        # This can lead to strange results, as shown in:
-        # https://github.com/aiidateam/aiida-core/issues/1600
-        # essentially qb.count() != len(qb.all()) in some cases.
-        # We also addressed this with sqlachemy:
-        # https://github.com/sqlalchemy/sqlalchemy/issues/4395#event-2002418814
-        # where the following solution was sanctioned:
-        self._query._has_mapper_entities = False  # pylint: disable=protected-access
-        # We should monitor SQLAlchemy, for when a solution is officially supported by the API!
+        self._query._raw_columns.pop(0)
+
+        # Note: according to https://github.com/sqlalchemy/sqlalchemy/commit/0d1efeec475621b5c2c2aca0632b02edef54c1a6
+        # automatic uniquing of rows is turned off for the new 2.0 style of ORM querying,
+        # so this sqlalchemy < 1.4 hack is no longer necessary
+        # self._query._has_mapper_entities = False
 
         # Make a list that helps the projection postprocessing
         self._attrkeys_as_in_sql_result = {

environment.yml

@@ -35,7 +35,7 @@ dependencies:
 - pyyaml~=5.4
 - simplejson~=3.16
 - sqlalchemy-utils~=0.37.2
-- sqlalchemy~=1.3.10
+- sqlalchemy~=1.4.0
 - tabulate~=0.8.5
 - typing-extensions
 - tqdm~=4.45

requirements/requirements-py-3.7.txt

@@ -5,6 +5,8 @@ alabaster==0.7.12
 aldjemy==1.0.0
 alembic==1.6.5
 aniso8601==9.0.1
+appdirs==1.4.4
+appnope==0.1.2
 archive-path==0.2.1
 argon2-cffi==20.1.0
 ase==3.22.0
@@ -15,6 +17,7 @@ Babel==2.9.1
 backcall==0.2.0
 bcrypt==3.2.0
 beautifulsoup4==4.9.3
+black==20.8b1
 bleach==4.0.0
 certifi==2021.5.30
 cffi==1.14.6
@@ -41,6 +44,7 @@ Flask-Cors==3.0.10
 Flask-RESTful==0.3.9
 future==0.18.2
 graphviz==0.17
+greenlet==1.1.1
 idna==3.2
 imagesize==1.2.0
 importlib-metadata==4.6.3
@@ -65,16 +69,18 @@ Mako==1.1.4
 MarkupSafe==2.0.1
 matplotlib==3.4.2
 matplotlib-inline==0.1.2
+mccabe==0.6.1
 mistune==0.8.4
 monty==2021.7.8
 mpmath==1.2.1
 multidict==5.1.0
+mypy-extensions==0.4.3
 nbclient==0.5.3
 nbconvert==6.1.0
 nbformat==5.1.3
 nest-asyncio==1.4.3
 networkx==2.6.2
-notebook==6.4.2
+notebook==6.4.3
 numpy==1.21.1
 packaging==21.0
 palettable==3.3.0
@@ -83,6 +89,7 @@ pandas==1.3.1
 pandocfilters==1.4.3
 paramiko==2.7.2
 parso==0.8.2
+pathspec==0.8.0
 pexpect==4.8.0
 pg8000==1.21.0
 pgsu==0.2.0
@@ -93,7 +100,7 @@ plotly==5.1.0
 pluggy==0.13.1
 plumpy==0.20.0
 prometheus-client==0.11.0
-prompt-toolkit==3.0.14
+prompt-toolkit==3.0.19
 psutil==5.8.0
 psycopg2-binary==2.8.6
 ptyprocess==0.7.0
@@ -113,10 +120,11 @@ pytest==6.2.4
 pytest-asyncio==0.15.1
 pytest-benchmark==3.4.1
 pytest-cov==2.10.1
+pytest-datadir==1.3.1
+pytest-regressions==2.2.0
 pytest-rerunfailures==9.1.1
 pytest-timeout==1.4.2
 python-dateutil==2.8.2
-pytest-regressions==2.2.0
 python-editor==1.0.4
 python-memcached==1.59
 pytray==0.3.2
@@ -151,7 +159,7 @@ sphinxcontrib-jsmath==1.0.1
 sphinxcontrib-qthelp==1.0.3
 sphinxcontrib-serializinghtml==1.1.5
 sphinxext-rediraffe==0.2.7
-SQLAlchemy==1.3.24
+SQLAlchemy==1.4.22
 sqlalchemy-diff==0.1.5
 SQLAlchemy-Utils==0.37.8
 sqlparse==0.4.1
@@ -164,6 +172,7 @@ toml==0.10.2
 tornado==6.1
 tqdm==4.62.0
 traitlets==5.0.5
+typed-ast==1.4.1
 typing-extensions==3.10.0.0
 tzlocal==2.1
 uncertainties==3.1.6
@@ -174,5 +183,6 @@ webencodings==0.5.1
 Werkzeug==2.0.1
 widgetsnbextension==3.5.1
 wrapt==1.11.2
+yapf==0.30.0
 yarl==1.6.3
 zipp==3.5.0

requirements/requirements-py-3.8.txt

@@ -5,6 +5,7 @@ alabaster==0.7.12
 aldjemy==1.0.0
 alembic==1.6.5
 aniso8601==9.0.1
+appnope==0.1.2
 archive-path==0.2.1
 argon2-cffi==20.1.0
 ase==3.22.0
@@ -41,6 +42,7 @@ Flask-Cors==3.0.10
 Flask-RESTful==0.3.9
 future==0.18.2
 graphviz==0.17
+greenlet==1.1.1
 idna==3.2
 imagesize==1.2.0
 importlib-metadata==4.6.3
@@ -74,7 +76,7 @@ nbconvert==6.1.0
 nbformat==5.1.3
 nest-asyncio==1.4.3
 networkx==2.6.2
-notebook==6.4.2
+notebook==6.4.3
 numpy==1.21.1
 packaging==21.0
 palettable==3.3.0
@@ -113,10 +115,11 @@ pytest==6.2.4
 pytest-asyncio==0.15.1
 pytest-benchmark==3.4.1
 pytest-cov==2.10.1
+pytest-datadir==1.3.1
+pytest-regressions==2.2.0
 pytest-rerunfailures==9.1.1
 pytest-timeout==1.4.2
 python-dateutil==2.8.2
-pytest-regressions==2.2.0
 python-editor==1.0.4
 python-memcached==1.59
 pytray==0.3.2
@@ -151,7 +154,7 @@ sphinxcontrib-jsmath==1.0.1
 sphinxcontrib-qthelp==1.0.3
 sphinxcontrib-serializinghtml==1.1.5
 sphinxext-rediraffe==0.2.7
-SQLAlchemy==1.3.24
+SQLAlchemy==1.4.22
 sqlalchemy-diff==0.1.5
 SQLAlchemy-Utils==0.37.8
 sqlparse==0.4.1

requirements/requirements-py-3.9.txt

@@ -5,6 +5,7 @@ alabaster==0.7.12
 aldjemy==1.0.0
 alembic==1.6.5
 aniso8601==9.0.1
+appnope==0.1.2
 archive-path==0.2.1
 argon2-cffi==20.1.0
 ase==3.22.0
@@ -41,6 +42,7 @@ Flask-Cors==3.0.10
 Flask-RESTful==0.3.9
 future==0.18.2
 graphviz==0.17
+greenlet==1.1.1
 idna==3.2
 imagesize==1.2.0
 importlib-metadata==4.6.3
@@ -74,7 +76,7 @@ nbconvert==6.1.0
 nbformat==5.1.3
 nest-asyncio==1.4.3
 networkx==2.6.2
-notebook==6.4.2
+notebook==6.4.3
 numpy==1.21.1
 packaging==21.0
 palettable==3.3.0
@@ -113,16 +115,17 @@ pytest==6.2.4
 pytest-asyncio==0.15.1
 pytest-benchmark==3.4.1
 pytest-cov==2.10.1
+pytest-datadir==1.3.1
+pytest-regressions==2.2.0
 pytest-rerunfailures==9.1.1
 pytest-timeout==1.4.2
 python-dateutil==2.8.2
-pytest-regressions==2.2.0
 python-editor==1.0.4
 python-memcached==1.59
 pytray==0.3.2
 pytz==2021.1
 PyYAML==5.4.1
-pyzmq==22.1.1
+pyzmq==22.2.1
 qtconsole==5.1.1
 QtPy==1.9.0
 requests==2.26.0
@@ -151,7 +154,7 @@ sphinxcontrib-jsmath==1.0.1
 sphinxcontrib-qthelp==1.0.3
 sphinxcontrib-serializinghtml==1.1.5
 sphinxext-rediraffe==0.2.7
-SQLAlchemy==1.3.24
+SQLAlchemy==1.4.22
 sqlalchemy-diff==0.1.5
 SQLAlchemy-Utils==0.37.8
 sqlparse==0.4.1

setup.json

@@ -49,7 +49,7 @@
         "pyyaml~=5.4",
         "simplejson~=3.16",
         "sqlalchemy-utils~=0.37.2",
-        "sqlalchemy~=1.3.10",
+        "sqlalchemy~=1.4.0",
         "tabulate~=0.8.5",
         "typing-extensions; python_version < '3.8'",
         "tqdm~=4.45",

tests/orm/test_querybuilder.py

@@ -24,6 +24,23 @@
 @pytest.mark.usefixtures('clear_database_before_test')
 class TestQueryBuilder:
 
+    def test_single_node_count(self):
+        """Test the count method for single entity."""
+        orm.Data().store()
+        builder = orm.QueryBuilder()
+        builder.append(orm.Data)
+        self.assertEqual(builder.count(), 1)
+
+    def test_joined_node_count(self):
+        """Test the count method for joined nodes."""
+        node1 = orm.Data().store()
+        node2 = orm.Data().store()
+        node1.backend_entity.add_incoming(node2.backend_entity, link_type=LinkType.INPUT_CALC, link_label='link_1')
+        builder = orm.QueryBuilder()
+        builder.append(orm.Data)
+        builder.append(orm.Data)
+        self.assertEqual(builder.count(), 1)
+
     def test_date_filters_support(self):
         """Verify that `datetime.date` is supported in filters."""
         from aiida.common import timezone