Add option to create unique ID index when inserting from JSONL

optimade/server/entry_collections/entry_collections.py

@@ -274,11 +274,11 @@ def all_fields(self) -> set[str]:
 
         return self._all_fields
 
-    def create_index(self, fields: str | set[str], unique: bool = False) -> None:
-        """Create an index on the given fields.
+    def create_index(self, field: str, unique: bool = False) -> None:
+        """Create an index on the given field.
 
         Arguments:
-            fields: The fields, or single field, to index.
+            field: The field to index.
             unique: Whether or not the index should be unique.
 
         """

optimade/server/entry_collections/mongo.py

@@ -106,18 +106,15 @@ def insert(self, data: list[EntryResource]) -> None:
         except Exception:
             pass
 
-    def create_index(self, fields: str | set[str], unique: bool = False) -> None:
-        """Create an index on the given fields.
+    def create_index(self, field: str, unique: bool = False) -> None:
+        """Create an index on the given field.
 
         Arguments:
-            fields: The fields, or single field, to index.
+            field: The field to index.
             unique: Whether or not the index should be unique.
 
         """
-        if isinstance(fields, str):
-            fields = {fields}
-
-        self.collection.create_indexes(fields, unique=unique, background=True)
+        self.collection.create_index(field, unique=unique, background=True)
 
     def handle_query_params(
         self, params: EntryListingQueryParams | SingleEntryQueryParams

optimade/utils.py

@@ -24,11 +24,12 @@
 )
 
 
-def insert_from_jsonl(jsonl_path: Path) -> None:
+def insert_from_jsonl(jsonl_path: Path, make_default_index: bool = False) -> None:
     """Insert OPTIMADE JSON lines data into the database.
 
     Arguments:
         jsonl_path: Path to the JSON lines file.
+        make_default_index: Whether to create a default index on the `id` field.
 
     """
     from collections import defaultdict
@@ -100,6 +101,13 @@ def insert_from_jsonl(jsonl_path: Path) -> None:
             ENTRY_COLLECTIONS[entry_type].insert(batch[entry_type])
             batch[entry_type] = []
 
+        if make_default_index:
+            for entry_type in ENTRY_COLLECTIONS:
+                try:
+                    ENTRY_COLLECTIONS[entry_type].create_index("id", unique=True)
+                except NotImplementedError:
+                    pass
+
         if bad_rows:
             LOGGER.warning("Could not read %d rows from the JSONL file", bad_rows)