finish EnityService

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.

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:

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