Migrate to compatible with aiida-core 2.0 (#331)

* Update to compatible with aiida-core 2.0 Following changes are made: - deprecate 3.7 support - Use aiida API to search and update extras field - drop django related code - Add support to python 3.10 - Add badge for dep matrix Co-authored-by: Casper Welzel Andersen <casper.w.andersen@sintef.no>
aiidateam · Sep 6, 2022 · b3caba2 · b3caba2
1 parent cb1b59f
commit b3caba2
Show file tree

Hide file tree

Showing 33 changed files with 159 additions and 153 deletions.
diff --git a/.ci/aiida-version.json b/.ci/aiida-version.json
@@ -1,6 +1,6 @@
 {
   "schemaVersion": 1,
   "label": "AiiDA",
-  "message": "v1.6.8",
+  "message": "v2.0.1",
   "color": "brightgreen"
 }
diff --git a/.github/aiida/profile.yaml b/.github/aiida/profile.yaml
@@ -18,3 +18,4 @@ broker_password: guest
 broker_host: 127.0.0.1
 broker_port: 5672
 broker_virtual_host: ''
+test_profile: true
diff --git a/.github/mongo/test_structures_mongo.json b/.github/mongo/test_structures_mongo.json
diff --git a/.github/workflows/ci_dependabot.yml b/.github/workflows/ci_dependabot.yml
@@ -98,4 +98,3 @@ jobs:
 
     - name: Information
       run: 'echo "${{ steps.cpr.outputs.pull-request-operation }} PR #${{ steps.cpr.outputs.pull-request-number }}: ${{ steps.cpr.outputs.pull-request-url }}"'
-
diff --git a/.github/workflows/ci_tests.yml b/.github/workflows/ci_tests.yml
@@ -44,16 +44,16 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [3.7, 3.8, 3.9]
-        backend: ['django', 'sqlalchemy']
+        python-version: ['3.8', '3.9', '3.10']
+        backend: ['psql_dos']
 
     services:
       mongo:
         image: mongo:4
         ports:
           - 27017:27017
       postgres:
-        image: postgres:10
+        image: postgres:12
         env:
           POSTGRES_DB: test_${{ matrix.backend }}
           POSTGRES_PASSWORD: test
@@ -80,7 +80,6 @@ jobs:
         python -m pip install -U pip
         pip install -U setuptools
         pip install -e .[testing]
-        reentry scan
 
     - name: Setup up environment for AiiDA
       env:
@@ -133,9 +132,9 @@ jobs:
         ports:
         - 27017:27017
       postgres:
-        image: postgres:10
+        image: postgres:12
         env:
-          POSTGRES_DB: test_django
+          POSTGRES_DB: test_psql_dos
           POSTGRES_PASSWORD: test
         options: >-
           --health-cmd pg_isready
@@ -165,19 +164,18 @@ jobs:
         done < requirements.txt
 
         pip install ${aiida_core}
-        reentry scan
 
         pip install pymongo
 
     - name: Setup up environment for AiiDA
       env:
-        AIIDA_TEST_BACKEND: django
+        AIIDA_TEST_BACKEND: psql_dos
       run: .github/aiida/setup_aiida.sh
 
     - name: Load test data
       run: |
         if [ "${{ matrix.database }}" == "aiida" ]; then
-          verdi import --migration --non-interactive .github/aiida/optimade.aiida
+          verdi archive import --migration .github/aiida/optimade.aiida
         else
             .github/mongo/load_data.py
         fi
@@ -209,9 +207,9 @@ jobs:
 
     services:
       postgres:
-        image: postgres:10
+        image: postgres:12
         env:
-          POSTGRES_DB: test_django
+          POSTGRES_DB: test_psql_dos
           POSTGRES_PASSWORD: test
         options: >-
           --health-cmd pg_isready
@@ -241,15 +239,14 @@ jobs:
         done < requirements.txt
 
         pip install ${aiida_core}
-        reentry scan
 
     - name: Setup up environment for AiiDA
       env:
-        AIIDA_TEST_BACKEND: django
+        AIIDA_TEST_BACKEND: psql_dos
       run: .github/aiida/setup_aiida.sh
 
     - name: Load test data (AiiDA)
-      run: verdi import --migration --non-interactive .github/aiida/optimade.aiida
+      run: verdi archive import --migration .github/aiida/optimade.aiida
 
     - name: Build docker image
       run: docker-compose -f profiles/docker-compose-mongo.yml build --build-arg CONFIG_FILE=".github/mongo/ci_config.json"
@@ -279,9 +276,9 @@ jobs:
 
     services:
       postgres:
-        image: postgres:10
+        image: postgres:12
         env:
-          POSTGRES_DB: test_django
+          POSTGRES_DB: test_psql_dos
           POSTGRES_PASSWORD: test
         options: >-
           --health-cmd pg_isready
@@ -311,11 +308,10 @@ jobs:
         done < requirements.txt
 
         pip install ${aiida_core}
-        reentry scan
 
     - name: Setup up environment for AiiDA
       env:
-        AIIDA_TEST_BACKEND: django
+        AIIDA_TEST_BACKEND: psql_dos
       run: .github/aiida/setup_aiida.sh
 
     - name: Build docker image

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -6,6 +6,7 @@ repos:
     exclude: README.md
   - id: check-yaml
   - id: check-json
+  - id: end-of-file-fixer
 
 - repo: https://github.com/ambv/black
   rev: 22.6.0

diff --git a/Dockerfile b/Dockerfile
@@ -1,10 +1,10 @@
-FROM python:3.8
+FROM python:3.9
 
 WORKDIR /app
 
 # Install specific optimade and aiida-core versions
 ARG OPTIMADE_TOOLS_VERSION=0.18.0
-ARG AIIDA_VERSION=1.6.8
+ARG AIIDA_VERSION=2.0.1
 
 # Copy repo contents
 COPY setup.py setup.json README.md requirements*.txt ./
@@ -13,7 +13,6 @@ COPY aiida_optimade ./aiida_optimade
 RUN pip install -U pip setuptools wheel \
     && pip install optimade==${OPTIMADE_TOOLS_VERSION} \
     && pip install aiida-core==${AIIDA_VERSION} \
-    && reentry scan \
     && pip install -e .
 
 COPY .docker/run.sh ./

diff --git a/README.md b/README.md
@@ -1,5 +1,12 @@
 # OPTIMADE API implementation for AiiDA
 
+The compatibility matrix below assumes the user always install the latest patch release of the specified minor version, which is recommended.
+
+| Plugin | AiiDA | Python | Specification |
+|-|-|-|-|
+| `v1.0 < v2.0` | ![Compatibility for v1.0][AiiDA v2 range] |  [![PyPI pyversions](https://img.shields.io/pypi/pyversions/aiida-optimade)](https://pypi.org/project/aiida-optimade) | ![OPTIMADE API compatibility][OPTIMADE from OPT] |
+| `v0.18 <= v0.20` | ![Compatibility for v0][AiiDA v1 range] |  [![PyPI pyversions][Python v3.7-v3.9]](https://pypi.org/project/aiida-optimade/0.20.0/) | ![OPTIMADE API compatibility][OPTIMADE from OPT] |
+
 | Latest release | Build status | Activity |
 |:--------------:|:------------:|:--------:|
 | [![AiiDA](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/aiidateam/aiida-optimade/develop/.ci/aiida-version.json)](https://github.com/aiidateam/aiida-core/)<br>[![PyPI](https://img.shields.io/pypi/v/aiida-optimade)](https://pypi.org/project/aiida-optimade/)<br>[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/aiida-optimade)](https://pypi.org/project/aiida-optimade/)<br>[![OPTIMADE](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/Materials-Consortia/optimade-python-tools/v0.18.0/optimade-version.json)](https://github.com/Materials-Consortia/OPTIMADE/) | [![GitHub Workflow Status](https://img.shields.io/github/workflow/status/aiidateam/aiida-optimade/aiida-optimade)](https://github.com/aiidateam/aiida-optimade/actions/)<br>[![Codecov](https://img.shields.io/codecov/c/gh/aiidateam/aiida-optimade)](https://codecov.io/gh/aiidateam/aiida-optimade) | [![GitHub last commit](https://img.shields.io/github/last-commit/aiidateam/aiida-optimade)](https://github.com/aiidateam/aiida-optimade) |
@@ -122,7 +129,7 @@ Navigate to `http://localhost:5000/v1/info`
 
 ### With Docker
 
-Adapt `profiles/test_django.json` and `profiles/docker-compose.yml` appropriately.
+Adapt `profiles/test_psql_dos.json` and `profiles/docker-compose.yml` appropriately.
 
 ```shell
 $ docker-compose -f profiles/docker-compose.yml up --build
@@ -157,3 +164,11 @@ To release the new version, go to GitHub release API of the repo create a new re
 The release action will be triggered by newly created release.
 Note, the tag should start with a `v` and be followed by a full semantic version (see [SemVer](https://semver.org)).
 For example: `v2.3.12`.
+
+[AiiDA v2 range]: https://img.shields.io/badge/AiiDA->=2.0.0,<3.0.0-007ec6.svg?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAACMAAAAhCAYAAABTERJSAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAAFhgAABYYBG6Yz4AAAABl0RVh0U29mdHdhcmUAd3d3Lmlua3NjYXBlLm9yZ5vuPBoAAAUbSURBVFiFzZhrbFRVEMd%2Fc%2B5uu6UUbIFC%2FUAUVEQCLbQJBIiBDyiImJiIhmohYNCkqJAQxASLF8tDgYRHBLXRhIcKNtFEhVDgAxBJqgmVh4JEKg3EIn2QYqBlt917xg%2BFss%2ByaDHOtzsz5z%2B%2FuZl7ztmF%2F5HJvxVQN6cPYX8%2FPLnOmsvNAvqfwuib%2FbNIk9cQeQnLcKRL5xLIV%2Fic9eJeunjPYbRs4FjQSpTB3aS1IpRKeeOOewajy%2FKKEO8Q0DuVdKy8IqsbPulxGHUfCBBu%2BwUYGuFuBTK7wQnht6PEbf4tlRomVRjCbXNjQEB0AyrFQOL5ENIJm7dTLZE6DPJCnEtFZVXDLny%2B4Sjv0PmmYu1ZdUek9RiMgoDmJ8V0L7XJqsZ3UW8YsBOwEeHeeFce7jEYXBy0m9m4BbXqSj2%2Bxnkg26MCVrN6DEZcwggtd8pTFx%2Fh3B9B50YLaFOPwXQKUt0tBLegtSomfBlfY13PwijbEnhztGzgJsK5h9W9qeWwBqjvyhB2iBs1Qz0AU974DciRGO8CVN8AJhAeMAdA3KbrKEtvxhsI%2B9emWiJlGBEU680Cfk%2BSsVqXZvcFYGXjF8ABVJ%2BTNfVXehyms1zzn1gmIOxLEB6E31%2FWBe5rnCarmo7elf7dJEeaLh80GasliI5F6Q9cAz1GY1OJVNDxTzQTw7iY%2FHEZRQY7xqJ9RU2LFe%2FYqakdP911ha0XhjjiTVAkDwgatWfCGeYocx8M3glG8g8EXhSrLrHnEFJ5Ymow%2FkhIYv6ttYUW1iFmEqqxdVoUs9FmsDYSqmtmJh3Cl1%2BVtl2s7owDUdocR5bceiyoSivGTT5vzpbzL1uoBpmcAAQgW7ArnKD9ng9rc%2BNgrobSNwpSkkhcRN%2BvmXLjIsDovYHHEfmsYFygPAnIDEQrQPzJYCOaLHLUfIt7Oq0LJn9fxkSgNCb1qEIQ5UKgT%2Fs6gJmVOOroJhQBXVqw118QtWLdyUxEP45sUpSzqP7RDdFYMyB9UReMiF1MzPwoUqHt8hjGFFeP5wZAbZ%2F0%2BcAtAAcji6LeSq%2FMYiAvSsdw3GtrfVSVFUBbIhwRWYR7yOcr%2FBi%2FB1MSJZ16JlgH1AGM3EO2QnmMyrSbTSiACgFBv4yCUapZkt9qwWVL7aeOyHvArJjm8%2Fz9BhdI4XcZgz2%2FvRALosjsk1ODOyMcJn9%2FYI6IrkS5vxMGdUwou2YKfyVqJpn5t9aNs3gbQMbdbkxnGdsr4bTHm2AxWo9yNZK4PXR3uzhAh%2BM0AZejnCrGdy0UvJxl0oMKgWSLR%2B1LH2aE9ViejiFs%2BXn6bTjng3MlIhJ1I1TkuLdg6OcAbD7Xx%2Bc3y9TrWAiSHqVkbZ2v9ilCo6s4AjwZCzFyD9mOL305nV9aonvsQeT2L0gVk4OwOJqXXVRW7naaxswDKVdlYLyMXAnntteYmws2xcVVZzq%2BtHPAooQggmJkc6TLSusOiL4RKgwzzYU1iFQgiUBA1H7E8yPau%2BZl9P7AblVNebtHqTgxLfRqrNvZWjsHZFuqMqKcDWdlFjF7UGvX8Jn24DyEAykJwNcdg0OvJ4p5pQ9tV6SMlP4A0PNh8aYze1ArROyUNTNouy8tNF3Rt0CSXb6bRFl4%2FIfQzNMjaE9WwpYOWQnOdEF%2BTdJNO0iFh7%2BI0kfORzQZb6P2kymS9oTxzBiM9rUqLWr1WE5G6ODhycQd%2FUnNVeMbcH68hYkGycNoUNWc8fxaxfwhDbHpfwM5oeTY7rUX8QAAAABJRU5ErkJggg%3D%3D
+
+[AiiDA v1 range]: https://img.shields.io/badge/AiiDA->=1.6.0,<2.0.0-007ec6.svg?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAACMAAAAhCAYAAABTERJSAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAAFhgAABYYBG6Yz4AAAABl0RVh0U29mdHdhcmUAd3d3Lmlua3NjYXBlLm9yZ5vuPBoAAAUbSURBVFiFzZhrbFRVEMd%2Fc%2B5uu6UUbIFC%2FUAUVEQCLbQJBIiBDyiImJiIhmohYNCkqJAQxASLF8tDgYRHBLXRhIcKNtFEhVDgAxBJqgmVh4JEKg3EIn2QYqBlt917xg%2BFss%2ByaDHOtzsz5z%2B%2FuZl7ztmF%2F5HJvxVQN6cPYX8%2FPLnOmsvNAvqfwuib%2FbNIk9cQeQnLcKRL5xLIV%2Fic9eJeunjPYbRs4FjQSpTB3aS1IpRKeeOOewajy%2FKKEO8Q0DuVdKy8IqsbPulxGHUfCBBu%2BwUYGuFuBTK7wQnht6PEbf4tlRomVRjCbXNjQEB0AyrFQOL5ENIJm7dTLZE6DPJCnEtFZVXDLny%2B4Sjv0PmmYu1ZdUek9RiMgoDmJ8V0L7XJqsZ3UW8YsBOwEeHeeFce7jEYXBy0m9m4BbXqSj2%2Bxnkg26MCVrN6DEZcwggtd8pTFx%2Fh3B9B50YLaFOPwXQKUt0tBLegtSomfBlfY13PwijbEnhztGzgJsK5h9W9qeWwBqjvyhB2iBs1Qz0AU974DciRGO8CVN8AJhAeMAdA3KbrKEtvxhsI%2B9emWiJlGBEU680Cfk%2BSsVqXZvcFYGXjF8ABVJ%2BTNfVXehyms1zzn1gmIOxLEB6E31%2FWBe5rnCarmo7elf7dJEeaLh80GasliI5F6Q9cAz1GY1OJVNDxTzQTw7iY%2FHEZRQY7xqJ9RU2LFe%2FYqakdP911ha0XhjjiTVAkDwgatWfCGeYocx8M3glG8g8EXhSrLrHnEFJ5Ymow%2FkhIYv6ttYUW1iFmEqqxdVoUs9FmsDYSqmtmJh3Cl1%2BVtl2s7owDUdocR5bceiyoSivGTT5vzpbzL1uoBpmcAAQgW7ArnKD9ng9rc%2BNgrobSNwpSkkhcRN%2BvmXLjIsDovYHHEfmsYFygPAnIDEQrQPzJYCOaLHLUfIt7Oq0LJn9fxkSgNCb1qEIQ5UKgT%2Fs6gJmVOOroJhQBXVqw118QtWLdyUxEP45sUpSzqP7RDdFYMyB9UReMiF1MzPwoUqHt8hjGFFeP5wZAbZ%2F0%2BcAtAAcji6LeSq%2FMYiAvSsdw3GtrfVSVFUBbIhwRWYR7yOcr%2FBi%2FB1MSJZ16JlgH1AGM3EO2QnmMyrSbTSiACgFBv4yCUapZkt9qwWVL7aeOyHvArJjm8%2Fz9BhdI4XcZgz2%2FvRALosjsk1ODOyMcJn9%2FYI6IrkS5vxMGdUwou2YKfyVqJpn5t9aNs3gbQMbdbkxnGdsr4bTHm2AxWo9yNZK4PXR3uzhAh%2BM0AZejnCrGdy0UvJxl0oMKgWSLR%2B1LH2aE9ViejiFs%2BXn6bTjng3MlIhJ1I1TkuLdg6OcAbD7Xx%2Bc3y9TrWAiSHqVkbZ2v9ilCo6s4AjwZCzFyD9mOL305nV9aonvsQeT2L0gVk4OwOJqXXVRW7naaxswDKVdlYLyMXAnntteYmws2xcVVZzq%2BtHPAooQggmJkc6TLSusOiL4RKgwzzYU1iFQgiUBA1H7E8yPau%2BZl9P7AblVNebtHqTgxLfRqrNvZWjsHZFuqMqKcDWdlFjF7UGvX8Jn24DyEAykJwNcdg0OvJ4p5pQ9tV6SMlP4A0PNh8aYze1ArROyUNTNouy8tNF3Rt0CSXb6bRFl4%2FIfQzNMjaE9WwpYOWQnOdEF%2BTdJNO0iFh7%2BI0kfORzQZb6P2kymS9oTxzBiM9rUqLWr1WE5G6ODhycQd%2FUnNVeMbcH68hYkGycNoUNWc8fxaxfwhDbHpfwM5oeTY7rUX8QAAAABJRU5ErkJggg%3D%3D
+
+[Python v3.7-v3.9]: https://img.shields.io/badge/python-3.7%20%7C%203.8%20%7C%203.9-blue
+
+[OPTIMADE from OPT]: https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/Materials-Consortia/optimade-python-tools/v0.18.0/optimade-version.json
diff --git a/aiida_optimade/cli/__init__.py b/aiida_optimade/cli/__init__.py
@@ -1,8 +1,2 @@
-# pylint: disable=wrong-import-position
-import click_completion
-
-# Activate the completion of parameter types provided by the click_completion package
-click_completion.init()
-
 # Import to populate sub commands
 from aiida_optimade.cli import cmd_calc, cmd_init, cmd_run  # noqa: E402,F401
diff --git a/aiida_optimade/cli/cmd_calc.py b/aiida_optimade/cli/cmd_calc.py
@@ -40,6 +40,10 @@ def calc(obj: dict, fields: Tuple[str], force_yes: bool, silent: bool):
     from aiida import load_profile
     from aiida.cmdline.utils import echo
 
+    # The default aiida.cmdline loglevel inherit from aiida loglevel is REPORT
+    # Here we use INFO loglevel for the operations
+    echo.CMDLINE_LOGGER.setLevel("INFO")
+
     try:
         profile: str = obj.get("profile").name
     except AttributeError:

diff --git a/aiida_optimade/cli/cmd_init.py b/aiida_optimade/cli/cmd_init.py
@@ -48,6 +48,10 @@ def init(obj: dict, force: bool, silent: bool, mongo: bool, filename: str):
     from aiida import load_profile
     from aiida.cmdline.utils import echo
 
+    # The default aiida.cmdline loglevel inherit from aiida loglevel is REPORT
+    # Here we use INFO loglevel for the operations
+    echo.CMDLINE_LOGGER.setLevel("INFO")
+
     filename: Path = Path(filename) if filename else filename
 
     if mongo and filename:

diff --git a/aiida_optimade/mappers/structures.py b/aiida_optimade/mappers/structures.py
@@ -22,8 +22,8 @@ class StructureMapper(ResourceMapper):
     """Map 'structure' resources from OPTIMADE to AiiDA"""
 
     TRANSLATORS: Dict[str, AiidaEntityTranslator] = {
-        "data.cif.CifData.": CifDataTranslator,
-        "data.structure.StructureData.": StructureDataTranslator,
+        "data.core.cif.CifData.": CifDataTranslator,
+        "data.core.structure.StructureData.": StructureDataTranslator,
     }
     REQUIRED_ATTRIBUTES = set(StructureResourceAttributes.schema().get("required"))
     # This should be REQUIRED_FIELDS, but should be set as such in `optimade`

diff --git a/aiida_optimade/routers/structures.py b/aiida_optimade/routers/structures.py
@@ -24,7 +24,7 @@
 ROUTER = APIRouter(redirect_slashes=True)
 
 STRUCTURES = AiidaCollection(
-    entities=["data.cif.CifData.", "data.structure.StructureData."],
+    entities=["data.core.cif.CifData.", "data.core.structure.StructureData."],
     resource_cls=StructureResource,
     resource_mapper=StructureMapper,
 )

diff --git a/aiida_optimade/translators/entities.py b/aiida_optimade/translators/entities.py
@@ -1,8 +1,9 @@
 from typing import Any, List, Union
+from aiida import orm
 from aiida.orm.nodes import Node
 from aiida.orm.querybuilder import QueryBuilder
 
-from aiida_optimade.common import AiidaEntityNotFound, AiidaError, LOGGER
+from aiida_optimade.common import AiidaEntityNotFound, LOGGER
 
 
 __all__ = ("AiidaEntityTranslator",)
@@ -138,41 +139,12 @@ def _store_attributes_mongo(self) -> None:
 
     def _store_attributes_node_extra(self) -> None:
         """Store new attributes in Node extras"""
-        from aiida.manage.manager import get_manager
-
         optimade = self._get_optimade_extras()
         if optimade:
             optimade.update(self.new_attributes)
         else:
             optimade = self.new_attributes
-        extras = (
-            self._get_unique_node_property("extras")
-            if self._get_unique_node_property("extras")
-            else {}
-        )
-        extras[self.EXTRAS_KEY] = optimade
-
-        profile = get_manager().get_profile()
-        LOGGER.debug(
-            "(%s) Updating Node %s in AiiDA DB!", profile.database_backend, self._pk
-        )
-        if profile.database_backend == "django":
-            from aiida.backends.djsite.db.models import DbNode
-
-            with get_manager().get_backend().transaction():
-                DbNode.objects.filter(pk=self._pk).update(extras=extras)
-        elif profile.database_backend == "sqlalchemy":
-            from aiida.backends.sqlalchemy.models.node import DbNode
-
-            with get_manager().get_backend().transaction() as session:
-                session.query(DbNode).filter(DbNode.id == self._pk).update(
-                    values={"extras": extras}
-                )
-        else:
-            raise AiidaError(
-                f'Unknown AiiDA backend "{profile.database_backend}" for profile'
-                f"{profile}"
-            )
 
-        # For posterity, this is how to do the same, going through AiiDA's API:
-        # self._node.set_extra(self.EXTRAS_KEY, optimade)
+        LOGGER.debug("Updating Node %s in AiiDA DB!", self._pk)
+        node = orm.load_node(self._pk)
+        node.set_extra(self.EXTRAS_KEY, optimade)
diff --git a/profiles/config.j2 b/profiles/config.j2
@@ -1,26 +1,34 @@
 {
     "CONFIG_VERSION": {
-        "CURRENT": 4,
-        "OLDEST_COMPATIBLE": 3
+        "CURRENT": 8,
+        "OLDEST_COMPATIBLE": 8
     },
     "default_profile": "{{ aiida_profile }}",
     "profiles": {
         "{{ aiida_profile }}": {
-            "AIIDADB_ENGINE": "postgresql_psycopg2",
-            "AIIDADB_PASS": "{{ aiida_db_pass }}",
-            "AIIDADB_NAME": "{{ aiida_db_name }}",
-            "AIIDADB_HOST": "{{ aiida_db_host }}",
-            "AIIDADB_USER": "{{ aiida_db_user }}",
-            "AIIDADB_PORT": {{ aiida_db_port }},
-            "AIIDADB_BACKEND": "{{ aiida_db_backend }}",
-            "default_user_email": "aiida@localhost",
-            "AIIDADB_REPOSITORY_URI": "file:///app/.aiida/repository-quicksetup",
-            "broker_protocol": "amqp",
-            "broker_username": "guest",
-            "broker_password": "guest",
-            "broker_host": "127.0.0.1",
-            "broker_port": 5672,
-            "broker_virtual_host": "",
+            "storage": {
+                "backend": "psql_dos",
+                "config": {
+                    "database_engine": "postgresql_psycopg2",
+                    "database_hostname": "{{ aiida_db_host }}",
+                    "database_port": {{ aiida_db_port }},
+                    "database_name": "{{ aiida_db_name }}",
+                    "database_username": "{{ aiida_db_user }}",
+                    "database_password": "{{ aiida_db_pass }}",
+                    "repository_uri": "file:///app/.aiida/repository-quicksetup"
+                }
+            },
+            "process_control": {
+                "backend": "rabbitmq",
+                "config": {
+                    "broker_protocol": "amqp",
+                    "broker_username": "guest",
+                    "broker_password": "guest",
+                    "broker_host": "127.0.0.1",
+                    "broker_port": 5672,
+                    "broker_virtual_host": ""
+                }
+            },
             "options": {}
         }
     }

diff --git a/profiles/docker-compose-mongo.j2 b/profiles/docker-compose-mongo.j2
@@ -11,7 +11,7 @@ services:
       dockerfile: Dockerfile
       args:
         OPTIMADE_TOOLS_VERSION: 0.18.0
-        AIIDA_VERSION: 1.6.8
+        AIIDA_VERSION: 2.0.1
         CONFIG_FILE: aiida_optimade/config.json
     environment:
       AIIDA_PATH: /app

diff --git a/profiles/docker-compose-mongo.yml b/profiles/docker-compose-mongo.yml
@@ -11,18 +11,18 @@ services:
       dockerfile: Dockerfile
       args:
         OPTIMADE_TOOLS_VERSION: 0.18.0
-        AIIDA_VERSION: 1.6.8
+        AIIDA_VERSION: 2.0.1
         CONFIG_FILE: aiida_optimade/config.json
     environment:
       AIIDA_PATH: /app
-      AIIDA_PROFILE: test_django
+      AIIDA_PROFILE: test_psql_dos
       FORCE_INIT:
       USE_MONGO:
       MONGO_FILENAME:
       AIIDA_OPTIMADE_LOG_LEVEL:
       OPTIMADE_MONGO_URI: 'mongodb://mongo:27017'
     volumes:
-     - '/tmp/test_repository_test_django:/app/.aiida/repository-quicksetup'
+     - '/tmp/test_repository_test_psql_dos:/app/.aiida/repository-quicksetup'
      - '.:/profiles'
      - '../.github/mongo:/app/.mongo'
     ports:

diff --git a/profiles/docker-compose.j2 b/profiles/docker-compose.j2
@@ -9,7 +9,7 @@ services:
       dockerfile: Dockerfile
       args:
         OPTIMADE_TOOLS_VERSION: 0.18.0
-        AIIDA_VERSION: 1.6.8
+        AIIDA_VERSION: 2.0.1
         CONFIG_FILE: aiida_optimade/config.json
     environment:
       AIIDA_PATH: /app