From ef79b8476efe7293ffb84e99b1dcd3aee527721d Mon Sep 17 00:00:00 2001 From: Matthew Evans Date: Sat, 16 Nov 2024 23:45:41 +0000 Subject: [PATCH] Create unique ID index when inserting from JSONL --- .../server/entry_collections/entry_collections.py | 6 +++--- optimade/server/entry_collections/mongo.py | 11 ++++------- optimade/utils.py | 3 +++ 3 files changed, 10 insertions(+), 10 deletions(-) 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..4f302e038 100644 --- a/optimade/utils.py +++ b/optimade/utils.py @@ -100,6 +100,9 @@ def insert_from_jsonl(jsonl_path: Path) -> None: ENTRY_COLLECTIONS[entry_type].insert(batch[entry_type]) batch[entry_type] = [] + for entry_type in ENTRY_COLLECTIONS: + ENTRY_COLLECTIONS[entry_type].create_index("id", unique=True) + if bad_rows: LOGGER.warning("Could not read %d rows from the JSONL file", bad_rows)