diff --git a/optimade/server/entry_collections/entry_collections.py b/optimade/server/entry_collections/entry_collections.py index ba84498fa..c987f4d78 100644 --- a/optimade/server/entry_collections/entry_collections.py +++ b/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. """ diff --git a/optimade/server/entry_collections/mongo.py b/optimade/server/entry_collections/mongo.py index 59d3f6cce..baa1602b7 100644 --- a/optimade/server/entry_collections/mongo.py +++ b/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 diff --git a/optimade/utils.py b/optimade/utils.py index 2cf53c2c8..a3dfa48fd 100644 --- a/optimade/utils.py +++ b/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)