add FKs to db connections

src/basic_memory/db.py

@@ -1,90 +1,101 @@
 """Database configuration and initialization for basic-memory."""
+
+from contextlib import asynccontextmanager
 from enum import Enum
 from pathlib import Path
 from typing import AsyncGenerator
-from contextlib import asynccontextmanager
 
 from loguru import logger
-from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncEngine, AsyncSession
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import (
+    create_async_engine,
+    async_sessionmaker,
+    AsyncEngine,
+    AsyncSession,
+)
 from sqlalchemy.pool import StaticPool
 
 from basic_memory.models import Base
 
+
 class DatabaseType(Enum):
     """Types of database configurations."""
-    MEMORY = "memory"      # In-memory SQLite for testing
-    FILESYSTEM = "file"    # File-based SQLite for projects
 
-def get_database_url(project_path: Path, db_type: DatabaseType, ) -> str:
+    MEMORY = "memory"  # In-memory SQLite for testing
+    FILESYSTEM = "file"  # File-based SQLite for projects
+
+
+def get_database_url(
+    project_path: Path,
+    db_type: DatabaseType,
+) -> str:
     """
     Get database URL based on type and optional project path.
-    
+
     Args:
         db_type: Type of database to configure
         project_path: Project directory for file-based DBs (required if type is FILESYSTEM)
-        
+
     Returns:
         Database URL string
-    
+
     Raises:
         ValueError: If project_path is required but not provided
     """
     match db_type:
         case DatabaseType.MEMORY:
             return "sqlite+aiosqlite:///:memory:"
-            
+
         case DatabaseType.FILESYSTEM:
             if not project_path:
                 raise ValueError("project_path required for filesystem database")
-                
+
             # Ensure data directory exists
             data_dir = project_path / "data"
             data_dir.mkdir(parents=True, exist_ok=True)
-            
+
             db_path = data_dir / "memory.db"
             return f"sqlite+aiosqlite:///{db_path}"
 
+
 async def init_database(url: str, echo: bool = False) -> tuple[AsyncEngine, async_sessionmaker]:
     """
     Initialize database with schema.
-    
+
     Args:
         url: Database URL
         echo: Whether to echo SQL statements
-        
+
     Returns:
         Configured async engine and session factory
     """
     # Configure engine based on URL
     connect_args = {"check_same_thread": False}
+
     if url == "sqlite+aiosqlite:///:memory:":
         engine = create_async_engine(
             url,
             echo=echo,
             poolclass=StaticPool,  # Single connection for in-memory
-            connect_args=connect_args
+            connect_args=connect_args,
         )
     else:
-        engine = create_async_engine(
-            url,
-            echo=echo,
-            connect_args=connect_args
-        )
+        engine = create_async_engine(url, echo=echo, connect_args=connect_args)
 
     # Create tables
     async with engine.begin() as conn:
         await conn.run_sync(Base.metadata.create_all)
 
     # Create session factory for this engine
-    session_factory = async_sessionmaker(
-        engine,
-        expire_on_commit=False
-    )
-        
+    session_factory = async_sessionmaker(engine, expire_on_commit=False)
+
     return engine, session_factory
 
+
 @asynccontextmanager
-async def engine_session_factory(project_path: Path, db_type=DatabaseType.FILESYSTEM) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]:
+async def engine_session_factory(
+    project_path: Path, db_type=DatabaseType.FILESYSTEM
+) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]:
     """Get database engine and session factory with proper lifecycle management."""
     url = get_database_url(project_path, db_type=db_type)
     engine, session_factory = await init_database(url)
@@ -94,20 +105,25 @@ async def engine_session_factory(project_path: Path, db_type=DatabaseType.FILESY
     finally:
         await engine.dispose()
 
+
 @asynccontextmanager
-async def session(session_factory: async_sessionmaker[AsyncSession]) -> AsyncGenerator[AsyncSession, None]:
+async def session(
+    session_factory: async_sessionmaker[AsyncSession],
+) -> AsyncGenerator[AsyncSession, None]:
     """
     Get database session with proper lifecycle management.
-    
+
     Args:
         session_factory: Async session factory to create session from
-        
+
     Yields:
         AsyncSession configured for engine
     """
     # Create and yield session
     session = session_factory()
     try:
+        # Ensure foreign keys enabled for this session
+        await session.execute(text("PRAGMA foreign_keys=ON"))
         yield session
         await session.commit()
     except Exception:
@@ -116,11 +132,12 @@ async def session(session_factory: async_sessionmaker[AsyncSession]) -> AsyncGen
     finally:
         await session.close()
 
+
 async def dispose_database(engine: AsyncEngine):
     """
     Clean up database engine.
-    
+
     Args:
         engine: Engine to dispose
     """
-    await engine.dispose()
+    await engine.dispose()

src/basic_memory/models.py

@@ -5,6 +5,8 @@ from sqlalchemy import String, DateTime, ForeignKey, Text, Integer, text, Unique
 from sqlalchemy.orm import Mapped, mapped_column, relationship, DeclarativeBase
 from sqlalchemy.ext.asyncio import AsyncAttrs
 
+from basic_memory.utils import normalize_entity_id
+
 
 class Base(AsyncAttrs, DeclarativeBase):
     """Base class for all models"""
@@ -68,8 +70,8 @@ class Entity(Base):
     @classmethod
     def generate_id(cls, entity_type: str, name: str) -> str:
         """Generate a filesystem path-based ID for this entity."""
-        # Normalize name for filesystem (handle spaces, special chars etc)
-        safe_name = name.lower().replace(" ", "_")
+        # Use common normalization for filesystem safety
+        safe_name = normalize_entity_id(name)
         return f"{entity_type}/{safe_name}"
 
     def get_file_path(self) -> str:

src/basic_memory/schemas.py

@@ -3,11 +3,16 @@
 from typing import List, Optional, Annotated, TypeAlias
 
 from annotated_types import Len
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, BeforeValidator
+
+from basic_memory.utils import normalize_entity_id
 
 # Base Models
 Observation: TypeAlias = str
 
+# Custom field types with validation
+EntityId = Annotated[str, BeforeValidator(normalize_entity_id)]
+
 
 class Relation(BaseModel):
     """
@@ -15,8 +20,8 @@ class Relation(BaseModel):
     Relations are always stored in active voice (e.g. "created", "teaches", etc.)
     """
 
-    from_id: str
-    to_id: str
+    from_id: EntityId
+    to_id: EntityId
     relation_type: str
     context: Optional[str] = None
 
@@ -28,7 +33,7 @@ class Entity(BaseModel):
     associated observations.
     """
 
-    id: Optional[str] = None
+    id: Optional[EntityId] = None
     name: str
     entity_type: str
     description: Optional[str] = None
@@ -47,7 +52,7 @@ class Entity(BaseModel):
 class AddObservationsRequest(BaseModel):
     """Schema for adding observations to an entity."""
 
-    entity_id: str
+    entity_id: EntityId
     context: Optional[str] = None
     observations: List[Observation]
 
@@ -67,7 +72,7 @@ class SearchNodesRequest(BaseModel):
 class OpenNodesRequest(BaseModel):
     """Request schema for open_nodes tool."""
 
-    names: Annotated[List[str], Len(min_length=1)]
+    names: Annotated[List[EntityId], Len(min_length=1)]
 
 
 class CreateRelationsRequest(BaseModel):
@@ -82,7 +87,7 @@ class CreateRelationsRequest(BaseModel):
 class DeleteEntityRequest(BaseModel):
     """Request schema for delete_entities tool."""
 
-    entity_ids: List[str]
+    entity_ids: List[EntityId]
 
 
 class DeleteRelationsRequest(BaseModel):
@@ -94,7 +99,7 @@ class DeleteRelationsRequest(BaseModel):
 class DeleteObservationsRequest(BaseModel):
     """Request schema for delete_observations tool."""
 
-    entity_id: str
+    entity_id: EntityId
     deletions: List[Observation]
 
 
@@ -118,7 +123,7 @@ class ObservationResponse(SQLAlchemyModel):
 class ObservationsResponse(SQLAlchemyModel):
     """Schema for bulk observation operation results."""
 
-    entity_id: str
+    entity_id: EntityId
     observations: List[ObservationResponse]
 
 
@@ -129,7 +134,7 @@ class RelationResponse(Relation, SQLAlchemyModel):
 class EntityResponse(SQLAlchemyModel):
     """Schema for entity data returned from the service."""
 
-    id: str
+    id: EntityId
     name: str
     entity_type: str
     description: Optional[str] = None
@@ -159,7 +164,7 @@ class OpenNodesResponse(SQLAlchemyModel):
 class AddObservationsResponse(SQLAlchemyModel):
     """Response for add_observations tool."""
 
-    entity_id: str
+    entity_id: EntityId
     observations: List[ObservationResponse]
 
 

src/basic_memory/utils.py

@@ -0,0 +1,13 @@
+"""Utility functions for basic-memory."""
+
+def normalize_entity_id(entity_id: str) -> str:
+    """
+    Normalize an entity ID by converting to lowercase and replacing spaces with underscores.
+    
+    Args:
+        entity_id: Raw entity ID to normalize
+        
+    Returns:
+        Normalized entity ID suitable for filesystem and database use
+    """
+    return entity_id.lower().replace(" ", "_")