From bf2aace1c3b97f90fd78647051e357e5b7aea026 Mon Sep 17 00:00:00 2001 From: Angelos Angelopoulos Date: Thu, 24 Oct 2024 15:58:29 -0400 Subject: [PATCH] Change .env MongoDB key names, optimize MongoDB document existence check, optimize task existence check --- docker/.env.example | 4 ++-- docker/docker-compose.yml | 8 ++++---- eos/persistence/abstract_async_repository.py | 2 +- eos/persistence/mongodb_async_repository.py | 7 +++---- eos/tasks/task_manager.py | 2 +- tests/test_mongodb_async_repository.py | 7 +------ 6 files changed, 12 insertions(+), 18 deletions(-) diff --git a/docker/.env.example b/docker/.env.example index 2b57dda..e04a4de 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -2,8 +2,8 @@ COMPOSE_PROJECT_NAME=eos # MongoDB admin credentials -EOS_MONGO_INITDB_ROOT_USERNAME= -EOS_MONGO_INITDB_ROOT_PASSWORD= +EOS_MONGODB_ROOT_USER= +EOS_MONGODB_ROOT_PASSWORD= # MinIO admin credentials EOS_MINIO_ROOT_USER= diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index d535c39..5ca6cdc 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -4,15 +4,15 @@ services: context: . dockerfile: mongodb/Dockerfile args: - - MONGO_INITDB_ROOT_USERNAME=${EOS_MONGO_INITDB_ROOT_USERNAME} - - MONGO_INITDB_ROOT_PASSWORD=${EOS_MONGO_INITDB_ROOT_PASSWORD} + - MONGO_INITDB_ROOT_USERNAME=${EOS_MONGODB_ROOT_USER} + - MONGO_INITDB_ROOT_PASSWORD=${EOS_MONGODB_ROOT_PASSWORD} image: eos-mongodb/latest container_name: eos-mongodb hostname: eos-mongodb restart: unless-stopped environment: - MONGO_INITDB_ROOT_USERNAME: ${EOS_MONGO_INITDB_ROOT_USERNAME} - MONGO_INITDB_ROOT_PASSWORD: ${EOS_MONGO_INITDB_ROOT_PASSWORD} + MONGO_INITDB_ROOT_USERNAME: ${EOS_MONGODB_ROOT_USER} + MONGO_INITDB_ROOT_PASSWORD: ${EOS_MONGODB_ROOT_PASSWORD} ports: - "27017:27017" networks: diff --git a/eos/persistence/abstract_async_repository.py b/eos/persistence/abstract_async_repository.py index 518c8b8..f8c5a10 100644 --- a/eos/persistence/abstract_async_repository.py +++ b/eos/persistence/abstract_async_repository.py @@ -16,7 +16,7 @@ async def count(self, **query: dict) -> int: pass @abstractmethod - async def exists(self, count: int = 1, **query: dict) -> bool: + async def exists(self, **query: dict) -> bool: pass @abstractmethod diff --git a/eos/persistence/mongodb_async_repository.py b/eos/persistence/mongodb_async_repository.py index 207d80a..3cb80ad 100644 --- a/eos/persistence/mongodb_async_repository.py +++ b/eos/persistence/mongodb_async_repository.py @@ -47,16 +47,15 @@ async def count(self, session: AgnosticClientSession | None = None, **kwargs) -> """ return await self._collection.count_documents(session=session, filter=kwargs) - async def exists(self, count: int = 1, session: AgnosticClientSession | None = None, **kwargs) -> bool: + async def exists(self, session: AgnosticClientSession | None = None, **kwargs) -> bool: """ - Check if the number of entities that match the query exist in the collection. + Check if an entity exists in the collection. - :param count: The number of entities to check for. :param session: The optional session to use for the operation. :param kwargs: Query parameters. :return: Whether the entity exists. """ - return await self.count(session=session, **kwargs) >= count + return await self._collection.find_one(kwargs, {"_id": 1}, session=session) is not None async def get_one(self, session: AgnosticClientSession | None = None, **kwargs) -> dict[str, Any]: """ diff --git a/eos/tasks/task_manager.py b/eos/tasks/task_manager.py index e2350a4..611cba8 100644 --- a/eos/tasks/task_manager.py +++ b/eos/tasks/task_manager.py @@ -62,7 +62,7 @@ async def create_task( :param containers: The input containers for the task. :param metadata: Additional metadata to be stored with the task. """ - if await self._tasks.get_one(experiment_id=experiment_id, id=task_id): + if await self._tasks.exists(experiment_id=experiment_id, id=task_id): raise EosTaskExistsError(f"Cannot create task '{task_id}' as a task with that ID already exists.") task_spec = self._configuration_manager.task_specs.get_spec_by_type(task_type) diff --git a/tests/test_mongodb_async_repository.py b/tests/test_mongodb_async_repository.py index 1bd0593..96a3046 100644 --- a/tests/test_mongodb_async_repository.py +++ b/tests/test_mongodb_async_repository.py @@ -68,13 +68,8 @@ async def test_count_and_exists(self, repository): ] await asyncio.gather(*[repository.create(entity) for entity in entities]) - count, exists, exists_more = await asyncio.gather( - repository.count(), repository.exists(count=2), repository.exists(count=3) - ) - + count = await repository.count() assert count == 2 - assert exists - assert not exists_more @pytest.mark.asyncio async def test_delete_many(self, repository):