add db check at startup

This commit is contained in:
phernandez
2025-02-07 16:03:31 -06:00
parent e2882dbf84
commit ffbd091ad3
13 changed files with 360 additions and 269 deletions
+1
View File
@@ -27,6 +27,7 @@ dependencies = [
"dateparser>=1.2.0",
"watchfiles>=1.0.4",
"fastapi[standard]>=0.115.8",
"alembic>=1.14.1",
]
[project.optional-dependencies]
+2 -2
View File
@@ -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,
)
+1 -24
View File
@@ -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(
+7 -9
View File
@@ -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",
]
-11
View File
@@ -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}')"
+2 -2
View File
@@ -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"
]
@@ -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}")
@@ -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}")
+1 -13
View File
@@ -22,7 +22,7 @@ from basic_memory.repository.relation_repository import RelationRepository
from basic_memory.repository.search_repository import SearchRepository
from basic_memory.services import (
EntityService,
DbVersionService,
DatabaseService,
)
from basic_memory.services.file_service import FileService
from basic_memory.services.link_resolver import LinkResolver
@@ -393,15 +393,3 @@ def watch_service(sync_service, file_service, test_config):
file_service=file_service,
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
)
+158
View File
@@ -0,0 +1,158 @@
"""Tests for DatabaseService."""
from datetime import datetime, timedelta
from pathlib import Path
from typing import AsyncGenerator
import pytest
import pytest_asyncio
from sqlalchemy import Column, String, Table, text
from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, AsyncSession
from basic_memory import db
from basic_memory.config import ProjectConfig
from basic_memory.db import DatabaseType
from basic_memory.models import Base
from basic_memory.services.database_service import DatabaseService
from basic_memory.sync import SyncService
@pytest_asyncio.fixture(scope="function")
async def engine_factory(
test_config,
) -> AsyncGenerator[tuple[AsyncEngine, async_sessionmaker[AsyncSession]], None]:
"""Special version of the engine factory fixture that uses a FILESYSTEM db_type"""
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_asyncio.fixture
async def database_service(
test_config: ProjectConfig,
sync_service: SyncService,
) -> DatabaseService:
"""Create DatabaseManagementService instance for testing."""
return DatabaseService(
config=test_config,
db_type = DatabaseType.FILESYSTEM
)
@pytest.mark.asyncio
async def test_check_db_initializes_new_db(
database_service: DatabaseService,
):
"""Test that check_db initializes new database."""
# Ensure DB doesn't exist
if Path(database_service.db_path).exists():
Path(database_service.db_path).unlink()
# Check DB - should initialize
assert await database_service.check_db()
@pytest.mark.asyncio
async def test_check_db_rebuilds_on_schema_mismatch(
database_service: DatabaseService,
session_maker,
):
"""Test that check_db rebuilds DB when schema doesn't match."""
# Initialize DB first
assert await database_service.check_db()
# Alter an existing table to remove a column
async with db.scoped_session(session_maker) as session:
conn = await session.connection()
# Create temp table
await conn.execute(text("""
CREATE TABLE entity_temp (
id INTEGER PRIMARY KEY,
title TEXT,
entity_type TEXT,
content_type TEXT,
permalink TEXT,
file_path TEXT,
checksum TEXT,
created_at TIMESTAMP,
updated_at TIMESTAMP
-- Deliberately omit entity_metadata column
)
"""))
# Drop original table
await conn.execute(text("DROP TABLE entity"))
# Rename temp table
await conn.execute(text("ALTER TABLE entity_temp RENAME TO entity"))
await session.commit()
# Check DB - should detect missing column and rebuild
assert await database_service.check_db()
# Verify entity_metadata column exists now
async with db.scoped_session(session_maker) as session:
result = await session.execute(text("""
SELECT sql FROM sqlite_master
WHERE type='table' AND name='entity'
"""))
create_sql = result.scalar()
assert 'entity_metadata' in create_sql.lower()
@pytest.mark.asyncio
async def test_backup_creates_timestamped_file(
database_service: DatabaseService,
):
"""Test that backup creates properly named backup file."""
if database_service.db_type == db.DatabaseType.MEMORY:
return
# Create dummy DB file
database_service.db_path.parent.mkdir(parents=True, exist_ok=True)
database_service.db_path.write_text("test content")
# Create backup
backup_path = await database_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
@pytest.mark.asyncio
async def test_cleanup_backups_keeps_recent(
database_service: DatabaseService,
):
"""Test that cleanup_backups keeps N most recent backups."""
if database_service.db_type == db.DatabaseType.MEMORY:
return
# Create backup directory
backup_dir = database_service.db_path.parent
backup_dir.mkdir(parents=True, exist_ok=True)
# Create some test backup files with different timestamps
backup_times = [
datetime.now() - timedelta(days=i)
for i in range(7) # Create 7 backups
]
for dt in backup_times:
timestamp = dt.strftime("%Y%m%d_%H%M%S")
backup_path = database_service.db_path.with_suffix(f".{timestamp}.backup")
backup_path.write_text("test backup")
# Set mtime to match our timestamp
backup_path.touch()
ts = dt.timestamp()
# Clean up keeping 5 most recent
await database_service.cleanup_backups(keep_count=5)
# Check that we have exactly 5 backups left
backup_pattern = "*.backup"
remaining = list(backup_dir.glob(backup_pattern))
assert len(remaining) == 5
-88
View File
@@ -1,88 +0,0 @@
"""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
+2 -2
View File
@@ -777,8 +777,8 @@ test content
await sync_service.sync(test_config.home)
# Check permalinks
file_one_content, _ = await file_service.read_file(two_file)
assert "permalink: one-1" in file_one_content
file_two_content, _ = await file_service.read_file(two_file)
assert "permalink: one-1" in file_two_content
# Run another time
await sync_service.sync(test_config.home)
Generated
+28
View File
@@ -13,6 +13,20 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/00/c4/c93eb22025a2de6b83263dfe3d7df2e19138e345bca6f18dba7394120930/aiosqlite-0.20.0-py3-none-any.whl", hash = "sha256:36a1deaca0cac40ebe32aac9977a6e2bbc7f5189f23f4a54d5908986729e5bd6", size = 15564 },
]
[[package]]
name = "alembic"
version = "1.14.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "mako" },
{ name = "sqlalchemy" },
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/99/09/f844822e4e847a3f0bd41797f93c4674cd4d2462a3f6c459aa528cdf786e/alembic-1.14.1.tar.gz", hash = "sha256:496e888245a53adf1498fcab31713a469c65836f8de76e01399aa1c3e90dd213", size = 1918219 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/54/7e/ac0991d1745f7d755fc1cd381b3990a45b404b4d008fc75e2a983516fbfe/alembic-1.14.1-py3-none-any.whl", hash = "sha256:1acdd7a3a478e208b0503cd73614d5e4c6efafa4e73518bb60e4f2846a37b1c5", size = 233565 },
]
[[package]]
name = "annotated-types"
version = "0.7.0"
@@ -51,6 +65,7 @@ version = "0.0.1"
source = { editable = "." }
dependencies = [
{ name = "aiosqlite" },
{ name = "alembic" },
{ name = "dateparser" },
{ name = "fastapi", extra = ["standard"] },
{ name = "greenlet" },
@@ -87,6 +102,7 @@ dev = [
[package.metadata]
requires-dist = [
{ name = "aiosqlite", specifier = ">=0.20.0" },
{ name = "alembic", specifier = ">=1.14.1" },
{ name = "dateparser", specifier = ">=1.2.0" },
{ name = "fastapi", extras = ["standard"], specifier = ">=0.115.8" },
{ name = "greenlet", specifier = ">=3.1.1" },
@@ -430,6 +446,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/0c/29/0348de65b8cc732daa3e33e67806420b2ae89bdce2b04af740289c5c6c8c/loguru-0.7.3-py3-none-any.whl", hash = "sha256:31a33c10c8e1e10422bfd431aeb5d351c7cf7fa671e3c4df004162264b28220c", size = 61595 },
]
[[package]]
name = "mako"
version = "1.3.9"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "markupsafe" },
]
sdist = { url = "https://files.pythonhosted.org/packages/62/4f/ddb1965901bc388958db9f0c991255b2c469349a741ae8c9cd8a562d70a6/mako-1.3.9.tar.gz", hash = "sha256:b5d65ff3462870feec922dbccf38f6efb44e5714d7b593a656be86663d8600ac", size = 392195 }
wheels = [
{ url = "https://files.pythonhosted.org/packages/cd/83/de0a49e7de540513f53ab5d2e105321dedeb08a8f5850f0208decf4390ec/Mako-1.3.9-py3-none-any.whl", hash = "sha256:95920acccb578427a9aa38e37a186b1e43156c87260d7ba18ca63aa4c7cbd3a1", size = 78456 },
]
[[package]]
name = "markdown-it-py"
version = "3.0.0"