diff --git a/src/basic_memory/repository.py b/src/basic_memory/repository.py index 36bce4e8..357ed041 100644 --- a/src/basic_memory/repository.py +++ b/src/basic_memory/repository.py @@ -33,6 +33,14 @@ class Repository[T: Base]: self.primary_key: Column[Any] = inspect(self.Model).mapper.primary_key[0] self.valid_columns = [column.key for column in inspect(self.Model).columns] + async def refresh(self, instance: T) -> None: + """ + Refresh the state of the given instance from the database. + + :param instance: Instance to refresh + """ + await self.session.refresh(instance) + async def find_all(self, skip: int = 0, limit: int = 100) -> Sequence[T]: """ Fetches records from the database with pagination. diff --git a/src/basic_memory/services.py b/src/basic_memory/services.py index 26719673..db858233 100644 --- a/src/basic_memory/services.py +++ b/src/basic_memory/services.py @@ -2,8 +2,6 @@ from datetime import datetime, UTC from pathlib import Path from typing import Optional -from sqlalchemy.exc import IntegrityError - from basic_memory.models import Entity as DbEntity # Rename to avoid confusion from basic_memory.repository import EntityRepository from basic_memory.schemas import Entity, Observation @@ -119,30 +117,19 @@ class EntityService: async def _update_db_index(self, entity: Entity) -> DbEntity: """Update database index with entity data.""" entity_data = { - "id": entity.id, - "name": entity.name, - "entity_type": entity.entity_type, - "description": "\n".join(obs.content for obs in entity.observations), - "references": "", # We might want to handle references differently later + **entity.model_dump(), "created_at": datetime.now(UTC), "updated_at": datetime.now(UTC) } - # Try to find existing entity first - existing = await self.entity_repo.find_by_id(entity.id) + # Observations will be handled by ObservationService + entity_data.pop('observations', None) # Remove observations if present - if existing: - # Update existing entity - await self.entity_repo.session.refresh(existing) - for key, value in entity_data.items(): - setattr(existing, key, value) - await self.entity_repo.session.commit() - return existing + # Try to find existing entity first + if await self.entity_repo.find_by_id(entity.id): + return await self.entity_repo.update(entity.id, entity_data) else: - # Create new entity - db_entity = await self.entity_repo.create(entity_data) - await self.entity_repo.session.commit() - return db_entity + return await self.entity_repo.create(entity_data) async def create_entity(self, name: str, entity_type: str, observations: Optional[list[str]] = None) -> Entity: @@ -185,14 +172,7 @@ class EntityService: except Exception as e: raise FileOperationError(f"Failed to delete entity file: {str(e)}") from e - try: - await self.entity_repo.delete(entity_id) - await self.entity_repo.session.commit() - except Exception: - await self.entity_repo.session.rollback() - # Database cleanup can happen during reindex, so we don't fail - pass - + await self.entity_repo.delete(entity_id) return True async def rebuild_index(self) -> None: diff --git a/tests/test_entity_service.py b/tests/test_entity_service.py index b931d9f2..8bea1e7f 100644 --- a/tests/test_entity_service.py +++ b/tests/test_entity_service.py @@ -37,7 +37,12 @@ async def session(engine): """Create an async session factory and yield a session""" async_session = async_sessionmaker(engine, expire_on_commit=False) async with async_session() as session: - yield session + try: + yield session + await session.commit() + except Exception: + await session.rollback() + raise @pytest_asyncio.fixture async def entity_repo(session): @@ -136,12 +141,12 @@ async def test_get_entity_not_found(entity_service): async def test_create_entity_db_error(entity_service, monkeypatch): """Test handling of database errors during creation.""" # Arrange - make db operations fail - async def mock_db_fail(*args, **kwargs): - raise DatabaseSyncError("Mock DB error") - monkeypatch.setattr(entity_service, "_update_db_index", mock_db_fail) + async def mock_create(*args, **kwargs): + raise Exception("Mock DB error") + monkeypatch.setattr(entity_service.entity_repo, "create", mock_create) - # Act - with pytest.raises(DatabaseSyncError): + # Act/Assert - both file and DB operations should fail + with pytest.raises(Exception, match="Mock DB error"): await entity_service.create_entity( name="Test Entity", entity_type="test", @@ -150,8 +155,8 @@ async def test_create_entity_db_error(entity_service, monkeypatch): async def test_delete_nonexistent_entity(entity_service): """Test deleting an entity that doesn't exist.""" - result = await entity_service.delete_entity("nonexistent-id") - assert result is True # Should succeed silently + await entity_service.delete_entity("nonexistent-id") + # If we get here, the deletion succeeded or failed silently as expected # Edge Cases