Repository: Materials-Consortia/optimade-python-tools
Pull request #2175 (2024-10-30 11:20:48)

Quality of life changes around JSONL & MongoDB
- Better error handling and logging on JSONL insert
- Parallel inserts from MongoDB
- Convenience methods and config for creating indexes: nothing changes by default, but indexes can now be made when inserting from JSONL data.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/pull/2175?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia) Report
Attention: Patch coverage is `48.00000%` with `13 lines` in your changes missing coverage. Please review.
> Project coverage is 90.65%. Comparing base [(`ccc8b64`)](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/commit/ccc8b64fb844fac9b93a137438e4e32074b0770b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia) to head [(`7ff086e`)](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/commit/7ff086e316f4ecc3752559c156db61f7069c1c97?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia).

| [Files with missing lines](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/pull/2175?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia) | Patch % | Lines |
|---|---|---|
| [optimade/utils.py](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/pull/2175?src=pr&el=tree&filepath=optimade%2Futils.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia#diff-b3B0aW1hZGUvdXRpbHMucHk=) | 50.00% | [9 Missing :warning: ](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/pull/2175?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia) |
| [optimade/server/entry\_collections/mongo.py](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/pull/2175?src=pr&el=tree&filepath=optimade%2Fserver%2Fentry_collections%2Fmongo.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia#diff-b3B0aW1hZGUvc2VydmVyL2VudHJ5X2NvbGxlY3Rpb25zL21vbmdvLnB5) | 40.00% | [3 Missing :warning: ](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/pull/2175?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia) |
| [...made/server/entry\_collections/entry\_collections.py](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/pull/2175?src=pr&el=tree&filepath=optimade%2Fserver%2Fentry_collections%2Fentry_collections.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia#diff-b3B0aW1hZGUvc2VydmVyL2VudHJ5X2NvbGxlY3Rpb25zL2VudHJ5X2NvbGxlY3Rpb25zLnB5) | 50.00% | [1 Missing :warning: ](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/pull/2175?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #2175      +/-   ##
==========================================
- Coverage   90.80%   90.65%   -0.16%     
==========================================
  Files          75       75              
  Lines        4819     4836      +17     
==========================================
+ Hits         4376     4384       +8     
- Misses        443      452       +9     
```

| [Flag](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/pull/2175/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia) | Coverage Δ | |
|---|---|---|
| [project](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/pull/2175/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia) | `90.65% <48.00%> (-0.16%)` | :arrow_down: |
| [validator](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/pull/2175/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia) | `90.65% <48.00%> (-0.16%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/Materials-Consortia/optimade-python-tools/pull/2175?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Materials-Consortia).

ml-evs: I'm going to merge, can potentially backport these changes to 1.1.x but not super important, otherwise will wait for 1.2.x release

Patch:
diff --git a/optimade/server/config.py b/optimade/server/config.py
index 29ba11a1..7b56469c 100644
--- a/optimade/server/config.py
+++ b/optimade/server/config.py
@@ -180,6 +180,16 @@ class ServerConfig(BaseSettings):
         ),
     ] = None
 
+    create_default_index: Annotated[
+        bool,
+        Field(
+            description=(
+                "Whether to create a set of default indices "
+                "for supporting databases after inserting JSONL data."
+            )
+        ),
+    ] = False
+
     use_real_mongo: Annotated[
         bool | None,
         Field(description="DEPRECATED: force usage of MongoDB over any other backend."),
diff --git a/optimade/server/entry_collections/elasticsearch.py b/optimade/server/entry_collections/elasticsearch.py
index 7b0fffc7..5bc52016 100644
--- a/optimade/server/entry_collections/elasticsearch.py
+++ b/optimade/server/entry_collections/elasticsearch.py
@@ -48,14 +48,19 @@ class ElasticCollection(EntryCollection):
         self.client = client if client else CLIENT
         self.name = name
 
-        # If we are creating a new collection from scratch, also create the index,
-        # otherwise assume it has already been created externally
-        if CONFIG.insert_test_data:
-            self.create_optimade_index()
-
     def count(self, *args, **kwargs) -> int:
         raise NotImplementedError
 
+    def create_default_index(self) -> None:
+        """Create the default index for the collection.
+
+        For Elastic, the default is to create a search index
+        over all relevant OPTIMADE fields based on the configured
+        mapper.
+
+        """
+        return self.create_optimade_index()
+
     def create_optimade_index(self) -> None:
         """Load or create an index that can handle aliased OPTIMADE fields and attach it
         to the current client.
@@ -111,7 +116,7 @@ class ElasticCollection(EntryCollection):
         """Returns the total number of entries in the collection."""
         return Search(using=self.client, index=self.name).execute().hits.total.value
 
-    def insert(self, data: list[EntryResource]) -> None:
+    def insert(self, data: list[EntryResource | dict]) -> None:
         """Add the given entries to the underlying database.
 
         Warning:
diff --git a/optimade/server/entry_collections/entry_collections.py b/optimade/server/entry_collections/entry_collections.py
index 308842ce..caa4593f 100644
--- a/optimade/server/entry_collections/entry_collections.py
+++ b/optimade/server/entry_collections/entry_collections.py
@@ -119,11 +119,16 @@ class EntryCollection(ABC):
         """Returns the total number of entries in the collection."""
 
     @abstractmethod
-    def insert(self, data: list[EntryResource]) -> None:
+    def insert(self, data: list[EntryResource | dict]) -> None:
         """Add the given entries to the underlying database.
 
+        Warning:
+            No validation is performed on the incoming data, this data
+            should have been mapped to the appropriate format before
+            insertion.
+
         Arguments:
-            data: The entry resource objects to add to the database.
+            data: The entries to add to the database.
 
         """
 
@@ -274,6 +279,30 @@ class EntryCollection(ABC):
 
         return self._all_fields
 
+    def create_index(self, field: str, unique: bool = False) -> None:
+        """Create an index on the given field, as stored in the database.
+
+        Arguments:
+            field: The database field to index (i.e., if different from the OPTIMADE field,
+                the mapper should be used to convert between the two).
+            unique: Whether or not the index should be unique.
+
+        """
+        raise NotImplementedError
+
+    def create_default_index(self) -> None:
+        """Create the default index for the collection.
+
+        For example, a database backend could override this method to
+        create a unique index on the `id` field, so that it can be called
+        on server startup.
+
+        This method should use a mapper to convert any OPTIMADE field names
+        to the corresponding stored names in the database.
+
+        """
+        raise NotImplementedError
+
     def get_attribute_fields(self) -> set[str]:
         """Get the set of attribute fields
 
diff --git a/optimade/server/entry_collections/mongo.py b/optimade/server/entry_collections/mongo.py
index 77d2c694..6740f6a5 100644
--- a/optimade/server/entry_collections/mongo.py
+++ b/optimade/server/entry_collections/mongo.py
@@ -91,17 +91,43 @@ class MongoCollection(EntryCollection):
             except ExecutionTimeout:
                 return None
 
-    def insert(self, data: list[EntryResource]) -> None:
+    def insert(self, data: list[EntryResource | dict]) -> None:
         """Add the given entries to the underlying database.
 
         Warning:
-            No validation is performed on the incoming data.
+            No validation is performed on the incoming data, this data
+            should have been mapped to the appropriate format before
+            insertion.
 
         Arguments:
-            data: The entry resource objects to add to the database.
+            data: The entries to add to the database.
 
         """
-        self.collection.insert_many(data)
+        self.collection.insert_many(data, ordered=False)
+
+    def create_index(self, field: str, unique: bool = False) -> None:
+        """Create an index on the given field, as stored in the database.
+
+        If any error is raised during index creation, this method should faithfully
+        return it, except for the simple case where an identical index already exists.
+
+        Arguments:
+            field: The database field to index (i.e., if different from the OPTIMADE field,
+                the mapper should be used to convert between the two).
+            unique: Whether or not the index should be unique.
+
+        """
+        self.collection.create_index(field, unique=unique, background=True)
+
+    def create_default_index(self) -> None:
+        """Create the default index for the collection.
+
+        For MongoDB, the default is to create a unique index
+        on the `id` field. This method should obey any configured
+        mappers.
+
+        """
+        self.create_index(self.resource_mapper.get_backend_field("id"), unique=True)
 
     def handle_query_params(
         self, params: EntryListingQueryParams | SingleEntryQueryParams
diff --git a/optimade/server/main.py b/optimade/server/main.py
index f563f365..74f770ea 100644
--- a/optimade/server/main.py
+++ b/optimade/server/main.py
@@ -123,7 +123,7 @@ if CONFIG.insert_test_data or CONFIG.insert_from_jsonl:
                 f"Requested JSONL file does not exist: {jsonl_path}. Please specify an absolute group."
             )
 
-        insert_from_jsonl(jsonl_path)
+        insert_from_jsonl(jsonl_path, create_default_index=CONFIG.create_default_index)
 
         LOGGER.debug("Inserted data from JSONL file: %s", jsonl_path)
         if CONFIG.insert_test_data:
diff --git a/optimade/utils.py b/optimade/utils.py
index 45829726..306bb327 100644
--- a/optimade/utils.py
+++ b/optimade/utils.py
@@ -7,7 +7,6 @@ import contextlib
 import json
 from collections.abc import Container, Iterable
 from pathlib import Path
-from traceback import print_exc
 from typing import TYPE_CHECKING, Optional
 
 from requests.exceptions import SSLError
@@ -25,17 +24,19 @@ PROVIDER_LIST_URLS = (
 )
 
 
-def insert_from_jsonl(jsonl_path: Path) -> None:
+def insert_from_jsonl(jsonl_path: Path, create_default_index: bool = False) -> None:
     """Insert OPTIMADE JSON lines data into the database.
 
     Arguments:
         jsonl_path: Path to the JSON lines file.
+        create_default_index: Whether to create a default index on the `id` field.
 
     """
     from collections import defaultdict
 
     import bson.json_util
 
+    from optimade.server.logger import LOGGER
     from optimade.server.routers import ENTRY_COLLECTIONS
 
     batch = defaultdict(list)
@@ -50,6 +51,16 @@ def insert_from_jsonl(jsonl_path: Path) -> None:
             )
         jsonl_path = _jsonl_path
 
+    # If the chosen database backend supports it, make the default indices
+    if create_default_index:
+        for entry_type in ENTRY_COLLECTIONS:
+            try:
+                ENTRY_COLLECTIONS[entry_type].create_default_index()
+            except NotImplementedError:
+                pass
+
+    bad_rows: int = 0
+    good_rows: int = 0
     with open(jsonl_path) as handle:
         header = handle.readline()
         header_jsonl = json.loads(header)
@@ -57,12 +68,19 @@ def insert_from_jsonl(jsonl_path: Path) -> None:
             "x-optimade"
         ), "No x-optimade header, not sure if this is a JSONL file"
 
-        for json_str in handle:
+        for line_no, json_str in enumerate(handle):
             try:
-                entry = bson.json_util.loads(json_str)
+                if json_str.strip():
+                    entry = bson.json_util.loads(json_str)
+                else:
+                    LOGGER.warning("Could not read any data from L%s", line_no)
+                    bad_rows += 1
+                    continue
             except json.JSONDecodeError:
-                print(f"Could not read entry as JSON: {json_str}")
-                print_exc()
+                from optimade.server.logger import LOGGER
+
+                LOGGER.warning("Could not read entry L%s JSON: '%s'", line_no, json_str)
+                bad_rows += 1
                 continue
             try:
                 id = entry.get("id", None)
@@ -76,19 +94,26 @@ def insert_from_jsonl(jsonl_path: Path) -> None:
                 # Append the data to the batch
                 batch[_type].append(inp_data)
             except Exception as exc:
-                print(f"Error with entry {entry}: {exc}")
-                print_exc()
+                LOGGER.warning(f"Error with entry at L{line_no} -- {entry} -- {exc}")
+                bad_rows += 1
                 continue
 
             if len(batch[_type]) >= batch_size:
                 ENTRY_COLLECTIONS[_type].insert(batch[_type])
                 batch[_type] = []
 
+            good_rows += 1
+
         # Insert any remaining data
         for entry_type in batch:
             ENTRY_COLLECTIONS[entry_type].insert(batch[entry_type])
             batch[entry_type] = []
 
+        if bad_rows:
+            LOGGER.warning("Could not read %d rows from the JSONL file", bad_rows)
+
+        LOGGER.info("Inserted %d rows from the JSONL file", good_rows)
+
 
 def mongo_id_for_database(database_id: str, database_type: str) -> str:
     """Produce a MongoDB ObjectId for a database"""


Test patch:
diff --git a/tests/server/entry_collections/test_indexes.py b/tests/server/entry_collections/test_indexes.py
new file mode 100644
index 00000000..006bf2e8
--- /dev/null
+++ b/tests/server/entry_collections/test_indexes.py
@@ -0,0 +1,40 @@
+import pytest
+from bson import ObjectId
+
+from optimade.server.config import CONFIG
+
+
+@pytest.mark.skipif(
+    CONFIG.database_backend.value not in ("mongomock", "mongodb"),
+    reason="Skipping index test when testing the elasticsearch backend.",
+)
+def test_indexes_are_created_where_appropriate(client):
+    """Test that with the test config, default indices are made by
+    supported backends. This is tested by checking that we cannot insert
+    an entry with the same underlying ID as the test data, and that this
+    returns the appopriate database-specific error.
+
+    """
+    import pymongo.errors
+
+    from optimade.server.query_params import EntryListingQueryParams
+    from optimade.server.routers import ENTRY_COLLECTIONS
+
+    # get one structure with and try to reinsert it
+    for _type in ENTRY_COLLECTIONS:
+        result, _, _, _, _ = ENTRY_COLLECTIONS[_type].find(
+            EntryListingQueryParams(page_limit=1)
+        )
+        assert result is not None
+        if isinstance(result, list):
+            result = result[0]
+
+        # The ID is mapped to the test data ID (e.g., 'task_id'), so the index is actually on that
+        id_field = ENTRY_COLLECTIONS[_type].resource_mapper.get_backend_field("id")
+
+        # Take the raw database result, extract the OPTIMADE ID and try to insert a canary
+        # document containing just that ID, plus a fake MongoDB ID to avoid '_id' clashes
+        canary = {id_field: result["id"], "_id": ObjectId(24 * "0")}
+        # Match either for "Duplicate" (mongomock) or "duplicate" (mongodb)
+        with pytest.raises(pymongo.errors.BulkWriteError, match="uplicate"):
+            ENTRY_COLLECTIONS[_type].insert([canary])  # type: ignore
diff --git a/tests/test_config.json b/tests/test_config.json
index b011371a..43787f4e 100644
--- a/tests/test_config.json
+++ b/tests/test_config.json
@@ -17,6 +17,7 @@
     "mongo_count_timeout": 0,
     "index_base_url": "http://localhost:5001",
     "insert_from_jsonl": "optimade/server/data/test_data.jsonl",
+    "create_default_index": true,
     "provider_fields": {
         "structures": [
             "band_gap",