Skip to content

Commit ca5fa55

Browse files
phernandezclaude
andcommitted
refactor: add shared initialization service
- Created shared initialization service for database migrations and file sync - Refactored CLI, API, and MCP server to use the shared service - Added tests for the initialization service - Improves logging and error handling for startup processes 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: phernandez <paul@basicmachines.co>
1 parent 7144035 commit ca5fa55

File tree

6 files changed

+267
-83
lines changed

6 files changed

+267
-83
lines changed

src/basic_memory/api/app.py

Lines changed: 3 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
"""FastAPI application for basic-memory knowledge graph API."""
22

3-
import asyncio
43
from contextlib import asynccontextmanager
54

65
from fastapi import FastAPI, HTTPException
@@ -10,47 +9,14 @@
109
from basic_memory import db
1110
from basic_memory.api.routers import knowledge, memory, project_info, resource, search
1211
from basic_memory.config import config as project_config
13-
from basic_memory.config import config_manager
14-
from basic_memory.sync import SyncService, WatchService
15-
16-
17-
async def run_background_sync(
18-
sync_service: SyncService, watch_service: WatchService
19-
): # pragma: no cover
20-
logger.info(f"Starting watch service to sync file changes in dir: {project_config.home}")
21-
# full sync
22-
await sync_service.sync(project_config.home, show_progress=False)
23-
24-
# watch changes
25-
await watch_service.run()
12+
from basic_memory.services.initialization import initialize_app
2613

2714

2815
@asynccontextmanager
2916
async def lifespan(app: FastAPI): # pragma: no cover
3017
"""Lifecycle manager for the FastAPI app."""
31-
await db.run_migrations(project_config)
32-
33-
# app config
34-
basic_memory_config = config_manager.load_config()
35-
logger.info(f"Sync changes enabled: {basic_memory_config.sync_changes}")
36-
logger.info(
37-
f"Update permalinks on move enabled: {basic_memory_config.update_permalinks_on_move}"
38-
)
39-
40-
watch_task = None
41-
if basic_memory_config.sync_changes:
42-
# import after migrations have run
43-
from basic_memory.cli.commands.sync import get_sync_service
44-
45-
sync_service = await get_sync_service()
46-
watch_service = WatchService(
47-
sync_service=sync_service,
48-
file_service=sync_service.entity_service.file_service,
49-
config=project_config,
50-
)
51-
watch_task = asyncio.create_task(run_background_sync(sync_service, watch_service))
52-
else:
53-
logger.info("Sync changes disabled. Skipping watch service.")
18+
# Initialize database and file sync services
19+
sync_service, watch_service, watch_task = await initialize_app(project_config)
5420

5521
# proceed with startup
5622
yield

src/basic_memory/cli/commands/mcp.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
"""MCP server command."""
22

3+
import asyncio
34
from loguru import logger
45

56
import basic_memory
67
from basic_memory.cli.app import app
78
from basic_memory.config import config, config_manager
9+
from basic_memory.services.initialization import initialize_app
810

911
# Import mcp instance
1012
from basic_memory.mcp.server import mcp as mcp_server # pragma: no cover
@@ -25,6 +27,18 @@ def mcp(): # pragma: no cover
2527
logger.info(f"Starting Basic Memory MCP server {basic_memory.__version__}")
2628
logger.info(f"Project: {project_name}")
2729
logger.info(f"Project directory: {home_dir}")
30+
31+
# Run initialization before starting MCP server
32+
# Note: Although we already run migrations via ensure_migrations in the CLI callback,
33+
# we do a full initialization here to ensure file sync is properly set up too
34+
try:
35+
loop = asyncio.new_event_loop()
36+
asyncio.set_event_loop(loop)
37+
loop.run_until_complete(initialize_app(config))
38+
except Exception as e:
39+
logger.error(f"Error during app initialization: {e}")
40+
logger.info("Continuing with server startup despite initialization error")
41+
2842
logger.info(f"Sync changes enabled: {basic_memory_config.sync_changes}")
2943
logger.info(
3044
f"Update permalinks on move enabled: {basic_memory_config.update_permalinks_on_move}"

src/basic_memory/cli/main.py

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,19 +21,13 @@
2121
tool,
2222
)
2323
from basic_memory.config import config
24-
from basic_memory.db import run_migrations as db_run_migrations
24+
from basic_memory.services.initialization import ensure_initialization
2525

2626

2727
# Helper function to run database migrations
2828
def ensure_migrations(): # pragma: no cover
2929
"""Ensure database migrations are run before executing commands."""
30-
try:
31-
logger.info("Running database migrations on startup...")
32-
asyncio.run(db_run_migrations(config))
33-
except Exception as e:
34-
logger.error(f"Error running migrations: {e}")
35-
# Continue execution even if migrations fail
36-
# The actual command might still work or will fail with a more specific error
30+
ensure_initialization(config)
3731

3832

3933
# Version command
@@ -77,8 +71,8 @@ def main(
7771

7872

7973
if __name__ == "__main__": # pragma: no cover
80-
# Run database migrations
81-
asyncio.run(db_run_migrations(config))
74+
# Run initialization
75+
ensure_initialization(config)
8276

8377
# start the app
8478
app()
Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
"""Shared initialization service for Basic Memory.
2+
3+
This module provides shared initialization functions used by both CLI and API
4+
to ensure consistent application startup across all entry points.
5+
"""
6+
7+
import asyncio
8+
from typing import Optional, Tuple
9+
10+
from loguru import logger
11+
12+
from basic_memory import db
13+
from basic_memory.config import ProjectConfig, config_manager
14+
from basic_memory.sync import SyncService, WatchService
15+
16+
# Import this inside functions to avoid circular imports
17+
# from basic_memory.cli.commands.sync import get_sync_service
18+
19+
20+
async def initialize_database(app_config: ProjectConfig) -> None:
21+
"""Run database migrations to ensure schema is up to date.
22+
23+
Args:
24+
app_config: The Basic Memory project configuration
25+
"""
26+
try:
27+
logger.info("Running database migrations...")
28+
await db.run_migrations(app_config)
29+
logger.info("Migrations completed successfully")
30+
except Exception as e:
31+
logger.error(f"Error running migrations: {e}")
32+
# Allow application to continue - it might still work
33+
# depending on what the error was, and will fail with a
34+
# more specific error if the database is actually unusable
35+
36+
37+
async def initialize_file_sync(
38+
app_config: ProjectConfig,
39+
) -> Tuple[Optional[SyncService], Optional[WatchService], Optional[asyncio.Task]]:
40+
"""Initialize file synchronization services.
41+
42+
Args:
43+
app_config: The Basic Memory project configuration
44+
45+
Returns:
46+
Tuple of (sync_service, watch_service, watch_task) if sync is enabled,
47+
or (None, None, None) if sync is disabled
48+
"""
49+
# Load app configuration
50+
basic_memory_config = config_manager.load_config()
51+
logger.info(f"Sync changes enabled: {basic_memory_config.sync_changes}")
52+
logger.info(
53+
f"Update permalinks on move enabled: {basic_memory_config.update_permalinks_on_move}"
54+
)
55+
56+
if not basic_memory_config.sync_changes:
57+
logger.info("Sync changes disabled. Skipping watch service.")
58+
return None, None, None
59+
60+
try:
61+
# Import here to avoid circular imports
62+
from basic_memory.cli.commands.sync import get_sync_service
63+
64+
# Initialize sync service
65+
sync_service = await get_sync_service()
66+
67+
# Initialize watch service
68+
watch_service = WatchService(
69+
sync_service=sync_service,
70+
file_service=sync_service.entity_service.file_service,
71+
config=app_config,
72+
)
73+
74+
# Start background sync task
75+
logger.info(f"Starting watch service to sync file changes in dir: {app_config.home}")
76+
77+
# Create the background task for running sync
78+
async def run_background_sync():
79+
# Run initial full sync
80+
await sync_service.sync(app_config.home, show_progress=False)
81+
# Start watching for changes
82+
await watch_service.run()
83+
84+
watch_task = asyncio.create_task(run_background_sync())
85+
86+
return sync_service, watch_service, watch_task
87+
except Exception as e:
88+
logger.error(f"Error initializing file sync: {e}")
89+
return None, None, None
90+
91+
92+
async def initialize_app(
93+
app_config: ProjectConfig,
94+
) -> Tuple[Optional[SyncService], Optional[WatchService], Optional[asyncio.Task]]:
95+
"""Initialize the Basic Memory application.
96+
97+
This function handles all initialization steps needed for both API and CLI:
98+
- Running database migrations
99+
- Setting up file synchronization
100+
101+
Args:
102+
app_config: The Basic Memory project configuration
103+
104+
Returns:
105+
Tuple of (sync_service, watch_service, watch_task) if sync is enabled,
106+
or (None, None, None) if sync is disabled or initialization failed
107+
"""
108+
# Initialize database first
109+
await initialize_database(app_config)
110+
111+
# Initialize file sync services
112+
return await initialize_file_sync(app_config)
113+
114+
115+
def ensure_initialization(app_config: ProjectConfig) -> None:
116+
"""Ensure initialization runs in a synchronous context.
117+
118+
This is a wrapper for the async initialize_app function that can be
119+
called from synchronous code like CLI entry points.
120+
121+
Args:
122+
app_config: The Basic Memory project configuration
123+
"""
124+
try:
125+
asyncio.run(initialize_database(app_config))
126+
except Exception as e:
127+
logger.error(f"Error during initialization: {e}")
128+
# Continue execution even if initialization fails
129+
# The command might still work, or will fail with a
130+
# more specific error message

tests/cli/test_cli_tools.py

Lines changed: 25 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@
33
These tests use real MCP tools with the test environment instead of mocks.
44
"""
55

6-
# Import the ensure_migrations function from main.py for testing
7-
from basic_memory.cli.main import ensure_migrations
6+
# Import for testing
7+
import asyncio
88

99
import io
1010
from datetime import datetime, timedelta
@@ -415,38 +415,27 @@ def test_continue_conversation_no_results(cli_env):
415415
assert "The supplied query did not return any information" in result.stdout
416416

417417

418-
@patch("basic_memory.cli.main.db_run_migrations")
419-
def test_ensure_migrations_runs_migrations(mock_run_migrations, test_config, monkeypatch):
420-
"""Test that ensure_migrations runs migrations."""
421-
422-
# Configure mock
423-
async def mock_async_success(*args, **kwargs):
424-
return True
425-
426-
mock_run_migrations.return_value = mock_async_success()
427-
418+
@patch("basic_memory.services.initialization.initialize_database")
419+
def test_ensure_migrations_functionality(mock_initialize_database, test_config, monkeypatch):
420+
"""Test the database initialization functionality."""
421+
from basic_memory.services.initialization import ensure_initialization
422+
428423
# Call the function
429-
ensure_migrations()
430-
431-
# Check that run_migrations was called
432-
mock_run_migrations.assert_called_once()
433-
434-
435-
@patch("basic_memory.cli.main.db_run_migrations")
436-
@patch("basic_memory.cli.main.logger")
437-
def test_ensure_migrations_handles_errors(
438-
mock_logger, mock_run_migrations, test_config, monkeypatch
439-
):
440-
"""Test that ensure_migrations handles errors gracefully."""
441-
442-
# Configure mock to raise an exception when awaited
443-
async def mock_async_error(*args, **kwargs):
444-
raise Exception("Test error")
445-
446-
mock_run_migrations.side_effect = mock_async_error
447-
448-
# Call the function
449-
ensure_migrations()
450-
451-
# Check that error was logged
452-
mock_logger.error.assert_called_once()
424+
ensure_initialization(test_config)
425+
426+
# The underlying asyncio.run should call our mocked function
427+
mock_initialize_database.assert_called_once()
428+
429+
430+
@patch("basic_memory.services.initialization.initialize_database")
431+
def test_ensure_migrations_handles_errors(mock_initialize_database, test_config, monkeypatch):
432+
"""Test that initialization handles errors gracefully."""
433+
from basic_memory.services.initialization import ensure_initialization
434+
435+
# Configure mock to raise an exception
436+
mock_initialize_database.side_effect = Exception("Test error")
437+
438+
# Call the function - should not raise exception
439+
ensure_initialization(test_config)
440+
441+
# We're just making sure it doesn't crash by calling it

0 commit comments

Comments
 (0)