add db check at startup

src/basic_memory/api/app.py

@@ -9,7 +9,7 @@ from loguru import logger
 from basic_memory import db
 from basic_memory.api.routers import knowledge, search, memory, resource
 from basic_memory.config import config
-from basic_memory.services import DbVersionService
+from basic_memory.services import DatabaseService
 
 
 @asynccontextmanager
@@ -28,7 +28,7 @@ async def check_db(app: FastAPI):
     logger.info("Checking database state")
 
     # Initialize DB management service
-    db_service = DbVersionService(
+    db_service = DatabaseService(
         config=config,
     )
 

src/basic_memory/db.py

@@ -68,7 +68,7 @@ async def scoped_session(
         await factory.remove()
 
 
-async def init_db() -> str:
+async def init_db() -> None:
     """Initialize database with required tables."""
 
     logger.info("Initializing database...")
@@ -78,11 +78,8 @@ async def init_db() -> str:
         conn = await session.connection()
         await conn.run_sync(Base.metadata.create_all)
         
-        version = await set_schema_version(session, SCHEMA_VERSION)
         await session.commit()
 
-    return version
-
 async def drop_db():
     """Drop all database tables."""
     global _engine, _session_maker
@@ -127,26 +124,6 @@ async def shutdown_db():
         _session_maker = None
 
 
-async def get_schema_version(session: AsyncSession) -> Optional[str]:
-    """Get current schema version from DB."""
-    try:
-        result = await session.execute(text("SELECT version FROM schema_version LIMIT 1"))
-        row = result.first()
-        return row[0] if row else None
-    except Exception as e:
-        logger.error(f"Error getting schema version: {e}")
-        return None
-
-
-async def set_schema_version(session: AsyncSession, version: str):
-    """Set schema version in DB."""
-    await session.execute(text("DELETE FROM schema_version"))
-    await session.execute(
-        text("INSERT INTO schema_version (version) VALUES (:version)"), {"version": version}
-    )    
-    await session.commit()
-    logger.info(f"Set schema version to {version}")
-
 
 @asynccontextmanager
 async def engine_session_factory(

src/basic_memory/models/__init__.py

@@ -1,17 +1,15 @@
 """Models package for basic-memory."""
+
 import basic_memory
-from basic_memory.models.base import Base, SchemaVersion
+from basic_memory.models.base import Base
 from basic_memory.models.knowledge import Entity, Observation, Relation, ObservationCategory
 
 SCHEMA_VERSION = basic_memory.__version__ + "-" + "003"
 
 __all__ = [
-    'Base',
-    'Entity',
-    'Observation', 
-    'ObservationCategory', 
-    'Relation',
-    'SchemaVersion'
+    "Base",
+    "Entity",
+    "Observation",
+    "ObservationCategory",
+    "Relation",
 ]
-
-

src/basic_memory/models/base.py

@@ -8,14 +8,3 @@ class Base(AsyncAttrs, DeclarativeBase):
     """Base class for all models"""
     pass
 
-
-class SchemaVersion(Base):
-    """Track database schema version."""
-
-    __tablename__ = "schema_version"
-
-    # Only one row will exist
-    version: Mapped[str] = mapped_column(String, primary_key=True)
-
-    def __repr__(self) -> str:
-        return f"SchemaVersion(version='{self.version}')"

src/basic_memory/services/__init__.py

@@ -1,5 +1,5 @@
 """Services package."""
-from .db_version_service import DbVersionService
+from .database_service import DatabaseService
 from .service import BaseService
 from .file_service import FileService
 from .entity_service import EntityService
@@ -8,5 +8,5 @@ __all__ = [
     "BaseService",
     "FileService",
     "EntityService",
-    "DbVersionService"
+    "DatabaseService"
 ]

src/basic_memory/services/database_service.py

@@ -0,0 +1,158 @@
+"""Service for managing database lifecycle and schema validation."""
+
+from datetime import datetime
+from pathlib import Path
+from typing import Optional, Tuple, List
+
+from alembic.runtime.migration import MigrationContext
+from alembic.autogenerate import compare_metadata
+from loguru import logger
+from sqlalchemy import MetaData
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from basic_memory import db
+from basic_memory.config import ProjectConfig
+from basic_memory.models import Base
+
+
+async def check_schema_matches_models(session: AsyncSession) -> Tuple[bool, List[str]]:
+    """Check if database schema matches SQLAlchemy models.
+    
+    Returns:
+        tuple[bool, list[str]]: (matches, list of differences)
+    """
+    # Get current DB schema via migration context
+    conn = await session.connection()
+
+    def _compare_schemas(connection):
+        context = MigrationContext.configure(connection)
+        return compare_metadata(context, Base.metadata)
+
+    # Run comparison in sync context
+    differences = await conn.run_sync(_compare_schemas)
+    
+    if not differences:
+        return True, []
+        
+    # Format differences into readable messages
+    diff_messages = []
+    for diff in differences:
+        if diff[0] == 'add_table':
+            diff_messages.append(f"Missing table: {diff[1].name}")
+        elif diff[0] == 'remove_table':
+            diff_messages.append(f"Extra table: {diff[1].name}")
+        elif diff[0] == 'add_column':
+            diff_messages.append(f"Missing column: {diff[3]} in table {diff[2]}")
+        elif diff[0] == 'remove_column':
+            diff_messages.append(f"Extra column: {diff[3]} in table {diff[2]}")
+        elif diff[0] == 'modify_type':
+            diff_messages.append(f"Column type mismatch: {diff[3]} in table {diff[2]}")
+            
+    return False, diff_messages
+
+
+class DatabaseService:
+    """Manages database lifecycle including schema validation and backups."""
+
+    def __init__(
+        self,
+        config: ProjectConfig,
+        db_type: db.DatabaseType = db.DatabaseType.FILESYSTEM,
+    ):
+        self.config = config
+        self.db_path = Path(config.database_path)
+        self.db_type = db_type
+
+    async def create_backup(self) -> Optional[Path]:
+        """Create backup of existing database file.
+
+        Returns:
+            Optional[Path]: Path to backup file if created, None if no DB exists
+        """
+        if self.db_type == db.DatabaseType.MEMORY:
+            return None  # Skip backups for in-memory DB
+
+        if not self.db_path.exists():
+            return None
+
+        # Create backup with timestamp
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        backup_path = self.db_path.with_suffix(f".{timestamp}.backup")
+
+        try:
+            self.db_path.rename(backup_path)
+            logger.info(f"Created database backup: {backup_path}")
+            
+            # make a new empty file
+            self.db_path.touch()
+            return backup_path
+        except Exception as e:
+            logger.error(f"Failed to create database backup: {e}")
+            return None
+
+    async def initialize_db(self):
+        """Initialize database with current schema."""
+        logger.info("Initializing database...")
+
+        if self.db_type == db.DatabaseType.FILESYSTEM:
+            await self.create_backup()
+
+        # Drop existing tables if any
+        await db.drop_db()
+
+        # Create tables with current schema
+        await db.get_or_create_db(
+            db_path=self.db_path,
+            db_type=self.db_type
+        )
+
+        logger.info("Database initialized with current schema")
+
+    async def check_db(self) -> bool:
+        """Check database state and rebuild if schema doesn't match models.
+        
+        Returns:
+            bool: True if DB is ready for use, False if initialization failed
+        """
+        try:
+            _, session_maker = await db.get_or_create_db(
+                db_path=self.db_path,
+                db_type=self.db_type
+            )
+            async with db.scoped_session(session_maker) as db_session:
+                # Check actual schema matches
+                matches, differences = await check_schema_matches_models(db_session)
+                if not matches:
+                    logger.warning("Database schema does not match models:")
+                    for diff in differences:
+                        logger.warning(f"  {diff}")
+                    logger.info("Rebuilding database to match current models...")
+                    await self.initialize_db()
+                    return True
+                    
+                logger.info("Database schema matches models")
+                return True
+
+        except Exception as e:
+            logger.error(f"Database initialization failed: {e}")
+            return False
+
+    async def cleanup_backups(self, keep_count: int = 5):
+        """Clean up old database backups, keeping the N most recent."""
+        if self.db_type == db.DatabaseType.MEMORY:
+            return  # Skip cleanup for in-memory DB
+
+        backup_pattern = "*.backup"  # Use relative pattern
+        backups = sorted(
+            self.db_path.parent.glob(backup_pattern),
+            key=lambda p: p.stat().st_mtime,
+            reverse=True,
+        )
+
+        # Remove old backups
+        for backup in backups[keep_count:]:
+            try:
+                backup.unlink()
+                logger.debug(f"Removed old backup: {backup}")
+            except Exception as e:
+                logger.error(f"Failed to remove backup {backup}: {e}")

src/basic_memory/services/db_version_service.py

@@ -1,118 +0,0 @@
-"""Service for managing database lifecycle and schema updates."""
-
-from datetime import datetime
-from pathlib import Path
-from typing import Optional
-
-from loguru import logger
-
-from basic_memory import db
-from basic_memory.config import ProjectConfig
-from basic_memory.db import DatabaseType
-from basic_memory.models import SCHEMA_VERSION
-
-
-class DbVersionService:
-    """Manages database lifecycle including initialization, backups, and schema updates."""
-
-    def __init__(
-        self,
-        config: ProjectConfig,
-        db_type: DatabaseType = DatabaseType.FILESYSTEM
-    ):
-        self.config = config
-        self.db_path = Path(config.database_path)
-        self.db_type = db_type
-
-    async def create_backup(self) -> Optional[Path]:
-        """Create backup of existing database file.
-
-        Returns:
-            Optional[Path]: Path to backup file if created, None if no DB exists
-        """
-        if self.db_type == db.DatabaseType.MEMORY:
-            return None
-        
-        if not self.db_path.exists():
-            return None
-
-        # Create backup with timestamp
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        backup_path = self.db_path.with_suffix(f".{timestamp}.backup")
-
-        try:
-            self.db_path.rename(backup_path)
-            logger.info(f"Created database backup: {backup_path}")
-            
-            # make a new empty file
-            self.db_path.touch()
-            return backup_path
-        except Exception as e:
-            logger.error(f"Failed to create database backup: {e}")
-            return None
-
-    async def initialize_db(self):
-        """Initialize database for first use."""
-
-        logger.info("Initializing database...")
-
-        if self.db_type == db.DatabaseType.FILESYSTEM:
-            await self.create_backup()
-            
-        # Drop existing tables if any
-        await db.drop_db()
-
-        # Create tables with current schema
-        await db.get_or_create_db(db_path=self.db_path)
-
-        logger.info(f"Database initialized with schema version {SCHEMA_VERSION}")
-
-    async def check_db(self) -> bool:
-        """Check database state and initialize/update if needed.
-
-        Returns:
-            bool: True if DB is ready for use, False if initialization failed
-        """
-        try:
-            _, session_maker = await db.get_or_create_db(db_path=self.db_path)
-            async with db.scoped_session(session_maker) as db_session:
-                db_version = await db.get_schema_version(db_session)
-
-                if db_version is None:
-                    logger.info("No existing database found, initializing...")
-                    await self.initialize_db()
-                elif db_version != SCHEMA_VERSION:
-                    logger.info(
-                        f"Schema version mismatch (DB: {db_version}, Current: {SCHEMA_VERSION}), rebuilding..."
-                    )
-                    await self.initialize_db()
-                else:
-                    logger.info(f"Database schema version {db_version} matches current version")
-
-            return True
-
-        except Exception as e:
-            logger.error(f"Database initialization failed: {e}")
-            return False
-
-    async def cleanup_backups(self, keep_count: int = 5):
-        """Clean up old database backups, keeping the N most recent."""
-
-        # Skip cleanup for in-memory DB
-        if self.db_type == db.DatabaseType.MEMORY:
-            return  
-
-        backup_pattern = "*.backup"  # Use relative pattern
-        backups = sorted(
-            self.db_path.parent.glob(backup_pattern),  # Use parent dir for glob
-            key=lambda p: p.stat().st_mtime,
-            reverse=True,
-        )
-
-        # Remove old backups
-        for backup in backups[keep_count:]:
-            try:
-                backup.unlink()
-                logger.debug(f"Removed old backup: {backup}")
-            except Exception as e:
-                logger.error(f"Failed to remove backup {backup}: {e}")