db_version_service

src/basic_memory/config.py

@@ -37,7 +37,11 @@ class ProjectConfig(BaseSettings):
     @property
     def database_path(self) -> Path:
         """Get SQLite database path."""
-        return self.home / DATA_DIR_NAME / DATABASE_NAME
+        database_path = self.home / DATA_DIR_NAME / DATABASE_NAME
+        if not database_path.exists():
+            database_path.parent.mkdir(parents=True, exist_ok=True)
+            database_path.touch()
+        return database_path
 
     @field_validator("home")
     @classmethod

src/basic_memory/db.py

@@ -14,8 +14,7 @@ from sqlalchemy.ext.asyncio import (
     async_scoped_session,
 )
 
-from basic_memory.models import Base
-
+from basic_memory.models import Base, SCHEMA_VERSION
 
 # Module level state
 _engine: Optional[AsyncEngine] = None
@@ -69,12 +68,34 @@ async def scoped_session(
         await factory.remove()
 
 
-async def init_db(session: AsyncSession):
+async def init_db() -> str:
     """Initialize database with required tables."""
-    await session.execute(text("PRAGMA foreign_keys=ON"))
-    conn = await session.connection()
-    await conn.run_sync(Base.metadata.create_all)
-    await session.commit()
+
+    logger.info("Initializing database...")
+
+    async with scoped_session(_session_maker) as session:
+        await session.execute(text("PRAGMA foreign_keys=ON"))
+        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
+    
+    logger.info("Dropping tables...")
+    async with scoped_session(_session_maker) as session:
+        conn = await session.connection()
+        await conn.run_sync(Base.metadata.drop_all)
+        await session.commit()
+        
+    # reset global engine and session_maker
+    _engine = None
+    _session_maker = None
 
 
 async def get_or_create_db(
@@ -83,7 +104,7 @@ async def get_or_create_db(
 ) -> tuple[AsyncEngine, async_sessionmaker[AsyncSession]]:
     """Get or create database engine and session maker."""
     global _engine, _session_maker
-    
+
     if _engine is None:
         db_url = DatabaseType.get_db_url(db_path, db_type)
         logger.debug(f"Creating engine for db_url: {db_url}")
@@ -91,9 +112,7 @@ async def get_or_create_db(
         _session_maker = async_sessionmaker(_engine, expire_on_commit=False)
 
         # Initialize database
-        logger.debug("Initializing database...")
-        async with scoped_session(_session_maker) as db_session:
-            await init_db(db_session)
+        await init_db()
 
     return _engine, _session_maker
 
@@ -101,35 +120,58 @@ async def get_or_create_db(
 async def shutdown_db():
     """Clean up database connections."""
     global _engine, _session_maker
-    
+
     if _engine:
         await _engine.dispose()
         _engine = None
         _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(
     db_path: Path,
-    db_type: DatabaseType = DatabaseType.FILESYSTEM,
+    db_type: DatabaseType = DatabaseType.MEMORY,
     init: bool = True,
 ) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]:
     """Create engine and session factory.
-    
+
     Note: This is primarily used for testing where we want a fresh database
     for each test. For production use, use get_or_create_db() instead.
     """
+
+    global _engine, _session_maker
+    
     db_url = DatabaseType.get_db_url(db_path, db_type)
     logger.debug(f"Creating engine for db_url: {db_url}")
-    engine = create_async_engine(db_url, connect_args={"check_same_thread": False})
+    
+    _engine = create_async_engine(db_url, connect_args={"check_same_thread": False})
     try:
-        factory = async_sessionmaker(engine, expire_on_commit=False)
+        _session_maker = async_sessionmaker(_engine, expire_on_commit=False)
 
         if init:
-            logger.debug("Initializing database...")
-            async with scoped_session(factory) as db_session:
-                await init_db(db_session)
+            await init_db()
 
-        yield engine, factory
+        yield _engine, _session_maker
     finally:
-        await engine.dispose()
+        await _engine.dispose()

src/basic_memory/models/__init__.py

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

src/basic_memory/models/base.py

@@ -1,9 +1,21 @@
 """Base model class for SQLAlchemy models."""
-
+from sqlalchemy import String
 from sqlalchemy.ext.asyncio import AsyncAttrs
-from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 
 
 class Base(AsyncAttrs, DeclarativeBase):
     """Base class for all models"""
-    pass
+    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 .service import BaseService
 from .file_service import FileService
 from .entity_service import EntityService
@@ -8,4 +8,5 @@ __all__ = [
     "BaseService",
     "FileService",
     "EntityService",
+    "DbVersionService"
 ]

src/basic_memory/services/db_version_service.py

@@ -0,0 +1,118 @@
+"""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}")