diff --git a/fiftyone/core/collections.py b/fiftyone/core/collections.py index 1ae796033b..6145c8d8b9 100644 --- a/fiftyone/core/collections.py +++ b/fiftyone/core/collections.py @@ -19,7 +19,7 @@ import warnings from bson import ObjectId -from pymongo import InsertOne, UpdateOne, UpdateMany +from pymongo import InsertOne, UpdateOne, UpdateMany, WriteConcern import eta.core.serial as etas import eta.core.utils as etau @@ -9126,7 +9126,9 @@ def get_index_information(self, include_stats=False): return index_info - def create_index(self, field_or_spec, unique=False, **kwargs): + def create_index( + self, field_or_spec, unique=False, acknowledged=True, **kwargs + ): """Creates an index on the given field or with the given specification, if necessary. @@ -9160,6 +9162,8 @@ def create_index(self, field_or_spec, unique=False, **kwargs): :meth:`pymongo:pymongo.collection.Collection.create_index` for supported values unique (False): whether to add a uniqueness constraint to the index + acknowledged (True): whether to wait and acknowledge index + creation **kwargs: optional keyword arguments for :meth:`pymongo:pymongo.collection.Collection.create_index` @@ -9238,10 +9242,17 @@ def create_index(self, field_or_spec, unique=False, **kwargs): # Satisfactory index already exists return index_name + # Setting `w=0` sets `acknowledged=False` in pymongo + write_concern = WriteConcern(w=0) if not acknowledged else None + if is_frame_index: - coll = self._dataset._frame_collection + coll = self._dataset._get_frame_collection( + write_concern=write_concern + ) else: - coll = self._dataset._sample_collection + coll = self._dataset._get_sample_collection( + write_concern=write_concern + ) name = coll.create_index(index_spec, unique=unique, **kwargs) diff --git a/fiftyone/core/dataset.py b/fiftyone/core/dataset.py index e4601ade8b..061286d070 100644 --- a/fiftyone/core/dataset.py +++ b/fiftyone/core/dataset.py @@ -21,7 +21,15 @@ import cachetools from deprecated import deprecated import mongoengine.errors as moe -from pymongo import DeleteMany, InsertOne, ReplaceOne, UpdateMany, UpdateOne +from pymongo import ( + DeleteMany, + InsertOne, + ReplaceOne, + UpdateMany, + UpdateOne, + WriteConcern, +) +from pymongo.collection import Collection from pymongo.errors import CursorNotFound, BulkWriteError import eta.core.serial as etas @@ -322,6 +330,7 @@ def __init__( self._run_cache = cachetools.LRUCache(5) self._deleted = False + self._write_concern = None if not _virtual: self._update_last_loaded_at() @@ -1172,14 +1181,22 @@ def stats( def _sample_collstats(self): conn = foo.get_db_conn() - return conn.command("collstats", self._sample_collection_name) + return conn.command( + "collstats", + self._sample_collection_name, + write_concern=self._write_concern, + ) def _frame_collstats(self): if self._frame_collection_name is None: return None conn = foo.get_db_conn() - return conn.command("collstats", self._frame_collection_name) + return conn.command( + "collstats", + self._frame_collection_name, + write_concern=self._write_concern, + ) def first(self): """Returns the first sample in the dataset. @@ -7023,7 +7040,12 @@ def _sample_collection_name(self): @property def _sample_collection(self): - return foo.get_db_conn()[self._sample_collection_name] + return self._get_sample_collection(write_concern=self._write_concern) + + def _get_sample_collection(self, write_concern=None) -> Collection: + return foo.get_db_conn().get_collection( + self._sample_collection_name, write_concern=write_concern + ) @property def _frame_collection_name(self): @@ -7031,10 +7053,15 @@ def _frame_collection_name(self): @property def _frame_collection(self): + return self._get_frame_collection(write_concern=self._write_concern) + + def _get_frame_collection(self, write_concern=None) -> Collection: if self._frame_collection_name is None: return None - return foo.get_db_conn()[self._frame_collection_name] + return foo.get_db_conn().get_collection( + self._frame_collection_name, write_concern=write_concern + ) @property def _frame_indexes(self):