diff --git a/src/basic_memory/config.py b/src/basic_memory/config.py index ad02c498..3997c59c 100644 --- a/src/basic_memory/config.py +++ b/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 diff --git a/src/basic_memory/db.py b/src/basic_memory/db.py index 9ec02ce5..ce727927 100644 --- a/src/basic_memory/db.py +++ b/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() \ No newline at end of file + await _engine.dispose() diff --git a/src/basic_memory/models/__init__.py b/src/basic_memory/models/__init__.py index 8345aa1a..d7b4e312 100644 --- a/src/basic_memory/models/__init__.py +++ b/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' -] \ No newline at end of file + 'Relation', + 'SchemaVersion' +] + + diff --git a/src/basic_memory/models/base.py b/src/basic_memory/models/base.py index 98106e99..5b0d8c7d 100644 --- a/src/basic_memory/models/base.py +++ b/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 \ No newline at end of file + 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}')" \ No newline at end of file diff --git a/src/basic_memory/services/__init__.py b/src/basic_memory/services/__init__.py index 3e46e87e..db21de08 100644 --- a/src/basic_memory/services/__init__.py +++ b/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" ] diff --git a/src/basic_memory/services/db_version_service.py b/src/basic_memory/services/db_version_service.py new file mode 100644 index 00000000..86fbb1be --- /dev/null +++ b/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}") diff --git a/tests/conftest.py b/tests/conftest.py index 58c96da5..f20586b2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,7 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession, AsyncEngine, async_sessionmaker from basic_memory import db from basic_memory.config import ProjectConfig -from basic_memory.db import DatabaseType +from basic_memory.db import DatabaseType, init_db from basic_memory.markdown import EntityParser from basic_memory.markdown.markdown_processor import MarkdownProcessor from basic_memory.models import Base @@ -21,7 +21,8 @@ from basic_memory.repository.observation_repository import ObservationRepository from basic_memory.repository.relation_repository import RelationRepository from basic_memory.repository.search_repository import SearchRepository from basic_memory.services import ( - EntityService, +EntityService, +DbVersionService, ) from basic_memory.services.file_service import FileService from basic_memory.services.link_resolver import LinkResolver @@ -393,3 +394,14 @@ def watch_service(sync_service, file_service, test_config): config=test_config ) + +@pytest_asyncio.fixture +async def db_version_service( + test_config: ProjectConfig, + sync_service: SyncService, +) -> DbVersionService: + """Create DatabaseManagementService instance for testing.""" + return DbVersionService( + config=test_config, + db_type = DatabaseType.FILESYSTEM + ) \ No newline at end of file diff --git a/tests/services/test_db_version_service.py b/tests/services/test_db_version_service.py new file mode 100644 index 00000000..1b2fecc3 --- /dev/null +++ b/tests/services/test_db_version_service.py @@ -0,0 +1,88 @@ +"""Tests for DatabaseManagementService.""" + +from pathlib import Path +from datetime import datetime, timedelta +from typing import AsyncGenerator + +import pytest +import pytest_asyncio +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, AsyncSession + +from basic_memory import db +from basic_memory.db import DatabaseType +from basic_memory.models import Base +from basic_memory.services.db_version_service import DbVersionService + + +# special version of the engine factory with a filesystem db type +@pytest_asyncio.fixture(scope="function") +async def engine_factory( + test_config, +) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]: + """Create engine and session factory using in-memory SQLite database.""" + async with db.engine_session_factory( + db_path=test_config.database_path, db_type=DatabaseType.FILESYSTEM + ) as (engine, session_maker): + # Initialize database + async with db.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) + + yield engine, session_maker + +@pytest.mark.asyncio +async def test_check_db_initializes_new_db(db_version_service: DbVersionService, session_maker): + """Test that check_db initializes new database.""" + # Ensure DB doesn't exist + if Path(db_version_service.db_path).exists(): + Path(db_version_service.db_path).unlink() + + # Check DB - should initialize + assert await db_version_service.check_db() + + # Verify schema version was set + async with db.scoped_session(session_maker) as session: + version = await db.get_schema_version(session) + assert version == db.SCHEMA_VERSION + + +@pytest.mark.asyncio +async def test_check_db_rebuilds_on_version_mismatch( + db_version_service: DbVersionService, session_maker +): + """Test that check_db rebuilds DB when schema version doesn't match.""" + # Initialize DB first + assert await db_version_service.check_db() + + # Set old version + async with db.scoped_session(session_maker) as session: + await db.set_schema_version(session, "000") + + # Check DB - should rebuild + assert await db_version_service.check_db() + + # Verify version was updated + async with db.scoped_session(session_maker) as session: + version = await db.get_schema_version(session) + assert version == db.SCHEMA_VERSION + + +@pytest.mark.asyncio +async def test_backup_creates_timestamped_file( + db_version_service: DbVersionService, +): + """Test that backup creates properly named backup file.""" + # Create dummy DB file + db_version_service.db_path.write_text("test content") + + # Create backup + backup_path = await db_version_service.create_backup() + + assert backup_path is not None + assert backup_path.exists() + assert backup_path.suffix == ".backup" + assert datetime.now().strftime("%Y%m%d") in backup_path.name + +