Repository: emcie-co/parlant
File: server/src/emcie/server/core/persistence.py (Python)

Code before:
        for i, d in enumerate(self._collections[collection.name]):
            if _matches_filters(filters, d):
                self._collections[collection.name][i] = updated_document
                return updated_document
        if upsert:
            document = await self.insert_one(collection, updated_document)
            return document

        raise ValueError("No document found matching the provided filters.")

    async def delete_one(
        self,
        collection: CollectionDescriptor,
        filters: dict[str, FieldFilter],
    ) -> None:
        self._collection_descriptors[collection.name] = collection
        for i, d in enumerate(self._collections[collection.name]):
            if _matches_filters(filters, d):
                del self._collections[collection.name][i]
                return
        raise ValueError("No document found matching the provided filters.")

    async def list_collections(self) -> Iterable[str]:
        return self._collections.keys()

    async def get_collection(
        self,
        collection_name: str,
    ) -> CollectionDescriptor:
        return self._collection_descriptors[collection_name]


class JSONFileDocumentDatabase(DocumentDatabase):

    def __init__(self, file_path: Path) -> None:
        self.file_path = file_path
        self._lock = asyncio.Lock()
        self.op_counter = 0
        if not self.file_path.exists():
            self.file_path.write_text(json.dumps({}))
        self.transient_db: TransientDocumentDatabase

    async def __aenter__(self) -> JSONFileDocumentDatabase:
        self.transient_db = TransientDocumentDatabase(await self._load_data())
        return self

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: Optional[object],

Review comment (suggestion):
@dor-emcie Can you remind me why `get_collection()` and `list_collections()` are public? (i.e. not prefixed with an underscore)

Code after:
                self._collections[collection.name][i] = updated_document
                return updated_document["id"]
        if upsert:
            document_id = await self.insert_one(collection, updated_document)
            return document_id

        raise ValueError("No document found matching the provided filters.")

    async def delete_one(
        self,
        collection: CollectionDescriptor,
        filters: dict[str, FieldFilter],
    ) -> None:
        for i, d in enumerate(self._collections[collection.name]):
            if _matches_filters(filters, d):
                del self._collections[collection.name][i]
                return
        raise ValueError("No document found matching the provided filters.")


class JSONFileDocumentDatabase(DocumentDatabase):

    def __init__(self, file_path: Path) -> None:
        self._collection_descriptors: dict[str, CollectionDescriptor] = {}
        self.file_path = file_path
        self._lock = asyncio.Lock()
        self.op_counter = 0
        if not self.file_path.exists():
            self.file_path.write_text(json.dumps({}))
        self.transient_db: TransientDocumentDatabase

    async def __aenter__(self) -> JSONFileDocumentDatabase:
        self.transient_db = TransientDocumentDatabase(await self._load_data())
        return self

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: Optional[object],
    ) -> bool:
        await self.flush()
        return False

    async def _process_operation_counter(self) -> None:
        self.op_counter += 1
        if self.op_counter % 5:
            await self.flush()

    async def _load_data(
        self,