admin/basicmachines-co-basic-memory
1
0
Fork 0
mirror of https://github.com/basicmachines-co/basic-memory synced 2026-06-21 13:47:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: refactor db schema migrate handling

Browse Source
This commit is contained in:
phernandez
2025-02-15 14:36:12 -06:00
parent a491c2b7d4
commit ca632beb6f
8 changed files with 79 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -42,3 +42,4 @@ ENV/
# macOS
.DS_Store
/.coverage.*
src/basic_memory/api/app.py
+1 -23
View File
@@ -10,35 +10,13 @@ import basic_memory
from basic_memory import db
from basic_memory.config import config as app_config
from basic_memory.api.routers import knowledge, search, memory, resource
from alembic import command
from alembic.config import Config
from basic_memory.db import DatabaseType
from basic_memory.repository.search_repository import SearchRepository
async def run_migrations(): # pragma: no cover
"""Run any pending alembic migrations."""
logger.info("Running database migrations...")
try:
config = Config("alembic.ini")
command.upgrade(config, "head")
logger.info("Migrations completed successfully")
_, session_maker = await db.get_or_create_db(
app_config.database_path, DatabaseType.FILESYSTEM
)
await SearchRepository(session_maker).init_search_index()
except Exception as e:
logger.error(f"Error running migrations: {e}")
raise
@asynccontextmanager
async def lifespan(app: FastAPI): # pragma: no cover
"""Lifecycle manager for the FastAPI app."""
logger.info(f"Starting Basic Memory API {basic_memory.__version__}")
await run_migrations()
await db.run_migrations(app_config)
yield
logger.info("Shutting down Basic Memory API")
await db.shutdown_db()
src/basic_memory/cli/app.py
+10
View File
@@ -1,3 +1,13 @@
import asyncio
import typer
from basic_memory import db
from basic_memory.config import config
from basic_memory.utils import setup_logging
setup_logging(log_file=".basic-memory/basic-memory-cli.log") # pragma: no cover
asyncio.run(db.run_migrations(config))
app = typer.Typer()
src/basic_memory/cli/commands/db.py
+1 -1
View File
@@ -22,4 +22,4 @@ def reset(
from basic_memory.cli.commands.sync import sync
logger.info("Rebuilding search index from filesystem...")
asyncio.run(sync()) # pyright: ignore
sync(watch=False) # pyright: ignore
src/basic_memory/cli/commands/status.py
+5 -7
View File
@@ -25,13 +25,11 @@ async def get_file_change_scanner(
db_type=DatabaseType.FILESYSTEM,
) -> FileChangeScanner: # pragma: no cover
"""Get sync service instance."""
async with db.engine_session_factory(db_path=config.database_path, db_type=db_type) as (
engine,
session_maker,
):
entity_repository = EntityRepository(session_maker)
file_change_scanner = FileChangeScanner(entity_repository)
return file_change_scanner
_, session_maker = await db.get_or_create_db(db_path=config.database_path, db_type=db_type)
entity_repository = EntityRepository(session_maker)
file_change_scanner = FileChangeScanner(entity_repository)
return file_change_scanner
def add_files_to_tree(
src/basic_memory/cli/commands/sync.py
+37 -37
View File
@@ -39,50 +39,48 @@ class ValidationIssue:
error: str
async def get_sync_service(db_type=DatabaseType.FILESYSTEM): # pragma: no cover
async def get_sync_service(): # pragma: no cover
"""Get sync service instance with all dependencies."""
async with db.engine_session_factory(db_path=config.database_path, db_type=db_type) as (
engine,
session_maker,
):
entity_parser = EntityParser(config.home)
markdown_processor = MarkdownProcessor(entity_parser)
file_service = FileService(config.home, markdown_processor)
_, session_maker = await db.get_or_create_db(db_path=config.database_path, db_type=db.DatabaseType.FILESYSTEM)
# Initialize repositories
entity_repository = EntityRepository(session_maker)
observation_repository = ObservationRepository(session_maker)
relation_repository = RelationRepository(session_maker)
search_repository = SearchRepository(session_maker)
entity_parser = EntityParser(config.home)
markdown_processor = MarkdownProcessor(entity_parser)
file_service = FileService(config.home, markdown_processor)
# Initialize services
search_service = SearchService(search_repository, entity_repository, file_service)
link_resolver = LinkResolver(entity_repository, search_service)
# Initialize repositories
entity_repository = EntityRepository(session_maker)
observation_repository = ObservationRepository(session_maker)
relation_repository = RelationRepository(session_maker)
search_repository = SearchRepository(session_maker)
# Initialize services
search_service = SearchService(search_repository, entity_repository, file_service)
link_resolver = LinkResolver(entity_repository, search_service)
# Initialize scanner
file_change_scanner = FileChangeScanner(entity_repository)
# Initialize scanner
file_change_scanner = FileChangeScanner(entity_repository)
# Initialize services
entity_service = EntityService(
entity_parser,
entity_repository,
observation_repository,
relation_repository,
file_service,
link_resolver,
)
# Initialize services
entity_service = EntityService(
entity_parser,
entity_repository,
observation_repository,
relation_repository,
file_service,
link_resolver,
)
# Create sync service
sync_service = SyncService(
scanner=file_change_scanner,
entity_service=entity_service,
entity_parser=entity_parser,
entity_repository=entity_repository,
relation_repository=relation_repository,
search_service=search_service,
)
# Create sync service
sync_service = SyncService(
scanner=file_change_scanner,
entity_service=entity_service,
entity_parser=entity_parser,
entity_repository=entity_repository,
relation_repository=relation_repository,
search_service=search_service,
)
return sync_service
return sync_service
def group_issues_by_directory(issues: List[ValidationIssue]) -> Dict[str, List[ValidationIssue]]:
@@ -154,6 +152,8 @@ def display_detailed_sync_results(knowledge: SyncReport):
async def run_sync(verbose: bool = False, watch: bool = False):
"""Run sync operation."""
sync_service = await get_sync_service()
# Start watching if requested
src/basic_memory/db.py
+21 -24
View File
@@ -4,6 +4,10 @@ from enum import Enum, auto
from pathlib import Path
from typing import AsyncGenerator, Optional
from basic_memory.config import ProjectConfig
from alembic import command
from alembic.config import Config
from loguru import logger
from sqlalchemy import text
from sqlalchemy.ext.asyncio import (
@@ -14,8 +18,8 @@ from sqlalchemy.ext.asyncio import (
async_scoped_session,
)
from basic_memory.models import Base
from basic_memory.models.search import CREATE_SEARCH_INDEX
from basic_memory.repository.search_repository import SearchRepository
# Module level state
_engine: Optional[AsyncEngine] = None
@@ -35,7 +39,7 @@ class DatabaseType(Enum):
logger.info("Using in-memory SQLite database")
return "sqlite+aiosqlite://"
return f"sqlite+aiosqlite:///{db_path}"
return f"sqlite+aiosqlite:///{db_path}" # pragma: no cover
def get_scoped_session_factory(
@@ -69,21 +73,6 @@ async def scoped_session(
await factory.remove()
async def init_db() -> None:
"""Initialize database with required tables."""
if _session_maker is None: # pragma: no cover
raise RuntimeError("Database session maker not initialized")
logger.info("Initializing database...")
async with scoped_session(_session_maker) as session:
await session.execute(text("PRAGMA foreign_keys=ON"))
# recreate search index
await session.execute(CREATE_SEARCH_INDEX)
await session.commit()
async def get_or_create_db(
db_path: Path,
@@ -98,9 +87,6 @@ async def get_or_create_db(
_engine = create_async_engine(db_url, connect_args={"check_same_thread": False})
_session_maker = async_sessionmaker(_engine, expire_on_commit=False)
# Initialize database
await init_db()
assert _engine is not None # for type checker
assert _session_maker is not None # for type checker
return _engine, _session_maker
@@ -120,7 +106,6 @@ async def shutdown_db() -> None: # pragma: no cover
async def engine_session_factory(
db_path: Path,
db_type: DatabaseType = DatabaseType.MEMORY,
init: bool = True,
) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]:
"""Create engine and session factory.
@@ -137,9 +122,6 @@ async def engine_session_factory(
try:
_session_maker = async_sessionmaker(_engine, expire_on_commit=False)
if init:
await init_db()
assert _engine is not None # for type checker
assert _session_maker is not None # for type checker
yield _engine, _session_maker
@@ -148,3 +130,18 @@ async def engine_session_factory(
await _engine.dispose()
_engine = None
_session_maker = None
async def run_migrations(app_config: ProjectConfig, database_type=DatabaseType.FILESYSTEM):
"""Run any pending alembic migrations."""
logger.info("Running database migrations...")
try:
config = Config("alembic.ini")
command.upgrade(config, "head")
logger.info("Migrations completed successfully")
_, session_maker = await get_or_create_db(app_config.database_path, database_type)
await SearchRepository(session_maker).init_search_index()
except Exception as e: # pragma: no cover
logger.error(f"Error running migrations: {e}")
raise
tests/conftest.py
+3 -5
View File
@@ -59,12 +59,10 @@ async def engine_factory(
async with db.engine_session_factory(
db_path=test_config.database_path, db_type=DatabaseType.MEMORY
) 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()
# Create all tables for the DB the engine is connected to
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
yield engine, session_maker