Introduce 'aiida-optimade calc' CLI command

.github/workflows/ci.yml

@@ -1,6 +1,11 @@
 name: CI
 
-on: [push, pull_request]
+on:
+  pull_request:
+  push:
+    branches:
+      - master
+      - 'push-action/**'
 
 jobs:
 
@@ -85,7 +90,7 @@ jobs:
         AIIDA_TEST_BACKEND: ${{ matrix.backend }}
         AIIDA_TEST_PROFILE: test_${{ matrix.backend }}
         AIIDA_PROFILE: test_${{ matrix.backend }}
-      run: pytest --cov=./aiida_optimade/ --cov-report=xml
+      run: pytest -v --cov=./aiida_optimade/ --cov-report=xml
 
     - name: Upload coverage to Codecov
       if: matrix.python-version == 3.8

aiida_optimade/cli/__init__.py

@@ -5,4 +5,4 @@
 click_completion.init()
 
 # Import to populate sub commands
-from aiida_optimade.cli import cmd_init, cmd_run  # noqa: E402,F401
+from aiida_optimade.cli import cmd_calc, cmd_init, cmd_run  # noqa: E402,F401

aiida_optimade/cli/cmd_calc.py

@@ -0,0 +1,116 @@
+# pylint: disable=protected-access,too-many-locals
+from typing import Tuple
+
+import click
+
+from aiida_optimade.cli.cmd_aiida_optimade import cli
+from aiida_optimade.common.logger import LOGGER
+
+
+@cli.command()
+@click.argument(
+    "fields",
+    type=click.STRING,
+    required=True,
+    nargs=-1,
+)
+@click.option(
+    "-q",
+    "--silent",
+    is_flag=True,
+    default=False,
+    show_default=True,
+    help=(
+        "Do not ask for confirmation when (re-)calculating the OPTIMADE field(s) in "
+        "the AiiDA database."
+    ),
+)
+@click.pass_obj
+def calc(obj: dict, fields: Tuple[str], silent: bool):
+    """Calculate OPTIMADE fields in the AiiDA database."""
+    from aiida import load_profile
+
+    try:
+        profile: str = obj.get("profile").name
+    except AttributeError:
+        profile = None
+    profile = load_profile(profile).name
+
+    try:
+        from aiida_optimade.routers.structures import STRUCTURES
+
+        if not silent:
+            click.confirm(
+                "Are you sure you want to (re-)calculate the field(s): "
+                f"{', '.join(fields)}?",
+                default=True,
+                abort=True,
+                show_default=True,
+            )
+
+        # Remove OPTIMADE fields in OPTIMADE-specific extra
+        extras_key = STRUCTURES.resource_mapper.PROJECT_PREFIX.split(".")[1]
+        query_kwargs = {
+            "filters": {
+                "and": [
+                    {"extras": {"has_key": extras_key}},
+                    {
+                        f"extras.{extras_key}": {
+                            "or": [{"has_key": field} for field in fields]
+                        }
+                    },
+                ]
+            },
+            "project": ["*", "extras.optimade"],
+        }
+
+        number_of_nodes = STRUCTURES.count(**query_kwargs)
+        if number_of_nodes:
+            click.echo(
+                f"Fields found for {number_of_nodes} Nodes. "
+                "The fields will now be removed for these Nodes. "
+                "Note: This may take several minutes!"
+            )
+
+            all_calculated_nodes = STRUCTURES._find_all(**query_kwargs)
+            for node, optimade in all_calculated_nodes:
+                for field in fields:
+                    optimade.pop(field, None)
+                node.set_extra("optimade", optimade)
+                del node
+            del all_calculated_nodes
+
+            click.echo(
+                f"Done removing {', '.join(fields)} from {number_of_nodes} Nodes."
+            )
+
+        click.echo(
+            f"{'Re-c' if number_of_nodes else 'C'}alcuating field(s) {fields} in "
+            f"{profile!r}. Note: This may take several minutes!"
+        )
+
+        STRUCTURES._filter_fields = set()
+        STRUCTURES._alias_filter({field: "" for field in fields})
+        updated_pks = STRUCTURES._check_and_calculate_entities()
+    except click.Abort:
+        click.echo("Aborted!")
+        return
+    except Exception as exc:  # pylint: disable=broad-except
+        from traceback import print_exc
+
+        LOGGER.error("Full exception from 'aiida-optimade calc' CLI:\n%s", print_exc())
+        click.echo(
+            f"An exception happened while trying to initialize {profile!r}:\n{exc!r}"
+        )
+        return
+
+    if updated_pks:
+        click.echo(
+            f"Success! {profile!r} has had {len(fields)} fields calculated for "
+            f"{len(updated_pks)} Nodes for use with AiiDA-OPTIMADE."
+        )
+    else:
+        click.echo(
+            f"No StructureData Nodes found to calculate fields {', '.join(fields)} for "
+            f"{profile!r}."
+        )

aiida_optimade/cli/cmd_run.py

@@ -78,4 +78,5 @@ def run(obj: dict, log_level: str, debug: bool, host: str, port: int, reload: bo
         host=host,
         port=port,
         log_level=log_level,
+        debug=debug,
     )

aiida_optimade/entry_collections.py

@@ -320,7 +320,7 @@ def _check_and_calculate_entities(self) -> List[int]:
             key for key in self.resource_mapper.PROJECT_PREFIX.split(".") if key
         ]
         filter_fields = [
-            {"!has_key": field for field in self._get_extras_filter_fields()}
+            {"!has_key": field} for field in self._get_extras_filter_fields()
         ]
         necessary_entities_qb = self._find_all(
             filters={

aiida_optimade/main.py

@@ -59,7 +59,7 @@
 # Load AiiDA profile
 PROFILE_NAME = os.getenv("AIIDA_PROFILE")
 load_profile(PROFILE_NAME)
-LOGGER.debug("AiiDA Profile: %s", PROFILE_NAME)
+LOGGER.info("AiiDA Profile: %s", PROFILE_NAME)
 
 # Load links in mongomock
 LINKS_DATA = Path(__file__).parent.joinpath("data/links.json").resolve()

aiida_optimade/mappers/structures.py

@@ -31,17 +31,18 @@ def build_attributes(cls, retrieved_attributes: dict, entry_pk: int) -> dict:
         """
         import json
 
-        res = {}
         float_fields_stored_as_strings = {"elements_ratios"}
 
         # Add existing attributes
         missing_attributes = cls.ALL_ATTRIBUTES.copy()
-        for existing_attribute, value in retrieved_attributes.items():
-            if existing_attribute in float_fields_stored_as_strings and value:
-                value = json.loads(str(value))
-            res[existing_attribute] = value
-            if existing_attribute in missing_attributes:
-                missing_attributes.remove(existing_attribute)
+        existing_attributes = set(retrieved_attributes.keys())
+        missing_attributes.difference_update(existing_attributes)
+        for field in float_fields_stored_as_strings:
+            if field in existing_attributes and retrieved_attributes.get(field):
+                retrieved_attributes[field] = json.loads(
+                    str(retrieved_attributes[field])
+                )
+        res = retrieved_attributes.copy()
 
         # Create and add new attributes
         if missing_attributes:

tests/cli/test_calc.py

@@ -0,0 +1,142 @@
+# pylint: disable=unused-argument,too-many-locals
+def test_calc_all_new(run_cli_command, aiida_profile, top_dir):
+    """Test `aiida-optimade -p profile_name calc` works for non-existant fields.
+
+    By "non-existant" the meaning is calculating fields that don't already exist for
+    any Nodes.
+    """
+    from aiida import orm
+    from aiida.tools.importexport import import_data
+
+    from aiida_optimade.cli import cmd_calc
+    from aiida_optimade.translators.entities import AiidaEntityTranslator
+
+    # Clear database and get initialized_nodes.aiida
+    aiida_profile.reset_db()
+    archive = top_dir.joinpath("tests/cli/static/initialized_nodes.aiida")
+    import_data(archive, silent=True)
+
+    fields = ["elements", "chemical_formula_hill"]
+
+    extras_key = AiidaEntityTranslator.EXTRAS_KEY
+    original_data = (
+        orm.QueryBuilder()
+        .append(
+            orm.StructureData,
+            filters={
+                f"extras.{extras_key}": {"or": [{"has_key": field} for field in fields]}
+            },
+            project=["*", f"extras.{extras_key}"],
+        )
+        .all()
+    )
+
+    # Remove these fields
+    for node, optimade in original_data:
+        for field in fields:
+            optimade.pop(field, None)
+        node.set_extra(extras_key, optimade)
+        del node
+    del original_data
+
+    n_structure_data = (
+        orm.QueryBuilder()
+        .append(
+            orm.StructureData,
+            filters={
+                f"extras.{extras_key}": {
+                    "or": [{"!has_key": field} for field in fields]
+                }
+            },
+        )
+        .count()
+    )
+
+    options = ["--silent"] + fields
+    result = run_cli_command(cmd_calc.calc, options)
+
+    assert (
+        f"Fields found for {n_structure_data} Nodes." not in result.stdout
+    ), result.stdout
+    assert (
+        "The fields will now be removed for these Nodes." not in result.stdout
+    ), result.stdout
+
+    assert "Success!" in result.stdout, result.stdout
+    assert f"calculated for {n_structure_data} Nodes" in result.stdout, result.stdout
+
+    n_updated_structure_data = (
+        orm.QueryBuilder()
+        .append(
+            orm.StructureData,
+            filters={
+                f"extras.{extras_key}": {"or": [{"has_key": field} for field in fields]}
+            },
+        )
+        .count()
+    )
+
+    assert n_structure_data == n_updated_structure_data
+
+    # Repopulate database with the "proper" test data
+    aiida_profile.reset_db()
+    original_data = top_dir.joinpath("tests/static/test_structuredata.aiida")
+    import_data(original_data, silent=True)
+
+
+def test_calc(run_cli_command, aiida_profile, top_dir):
+    """Test `aiida-optimade -p profile_name calc` works."""
+    from aiida import orm
+    from aiida.tools.importexport import import_data
+
+    from aiida_optimade.cli import cmd_calc
+    from aiida_optimade.translators.entities import AiidaEntityTranslator
+
+    # Clear database and get initialized_nodes.aiida
+    aiida_profile.reset_db()
+    archive = top_dir.joinpath("tests/cli/static/initialized_nodes.aiida")
+    import_data(archive, silent=True)
+
+    fields = ["elements", "chemical_formula_hill"]
+
+    extras_key = AiidaEntityTranslator.EXTRAS_KEY
+
+    n_structure_data = (
+        orm.QueryBuilder()
+        .append(
+            orm.StructureData,
+            filters={
+                f"extras.{extras_key}": {"or": [{"has_key": field} for field in fields]}
+            },
+        )
+        .count()
+    )
+
+    options = ["--silent"] + fields
+    result = run_cli_command(cmd_calc.calc, options)
+
+    assert f"Fields found for {n_structure_data} Nodes." in result.stdout, result.stdout
+    assert (
+        "The fields will now be removed for these Nodes." in result.stdout
+    ), result.stdout
+
+    assert "Success!" in result.stdout, result.stdout
+    assert f"calculated for {n_structure_data} Nodes" in result.stdout, result.stdout
+
+    n_updated_structure_data = (
+        orm.QueryBuilder()
+        .append(
+            orm.StructureData,
+            filters={
+                f"extras.{extras_key}": {"or": [{"has_key": field} for field in fields]}
+            },
+        )
+        .count()
+    )
+
+    assert n_structure_data == n_updated_structure_data
+
+    # Repopulate database with the "proper" test data
+    aiida_profile.reset_db()
+    original_data = top_dir.joinpath("tests/static/test_structuredata.aiida")
+    import_data(original_data, silent=True)

tests/cli/test_run.py

@@ -55,11 +55,17 @@ def test_run(run_server):  # pylint: disable=unused-argument
 
 
 def test_log_level_debug(run_and_terminate_server):
-    """Test passing log level "debug" to `aiida-optimade run`"""
+    """Test passing log level "debug" to `aiida-optimade run`
+
+    In the latest versions of uvicorn, setting the log-level to "debug"
+    is not enough to create debug log messages in stdout.
+    One needs to also be in debug mode, i.e., either set `reload=True`
+    or set `debug=True`.
+    """
     options = ["--log-level", "debug"]
     output, errors = run_and_terminate_server(command="run", options=options)
     assert "DEBUG MODE" in output, f"output: {output!r}, errors: {errors!r}"
-    assert "DEBUG:" in output, f"output: {output!r}, errors: {errors!r}"
+    assert "DEBUG:" not in output, f"output: {output!r}, errors: {errors!r}"
 
 
 def test_log_level_warning(run_and_terminate_server):
@@ -121,6 +127,5 @@ def test_env_var_is_set(run_and_terminate_server):
     if fixture_profile == "test_profile":
         # This is for local tests only
         fixture_profile = "optimade_sqla"
-    options = ["--log-level", "debug"]
-    output, errors = run_and_terminate_server(command="run", options=options)
+    output, errors = run_and_terminate_server(command="run")
     assert fixture_profile in output, f"output: {output!r}, errors: {errors!r}"