alma-memory 0.5.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

alma/storage/migrations/runner.py

@@ -0,0 +1,323 @@
+"""
+ALMA Migration Framework - Migration Runner.
+
+Executes migrations and manages schema version tracking.
+"""
+
+import hashlib
+import logging
+from datetime import datetime, timezone
+from typing import Any, Callable, Dict, List, Optional, Protocol, Type
+
+from alma.storage.migrations.base import (
+    Migration,
+    MigrationError,
+    MigrationRegistry,
+    SchemaVersion,
+    get_registry,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class VersionStore(Protocol):
+    """Protocol for schema version storage."""
+
+    def get_current_version(self) -> Optional[str]:
+        """Get the current schema version."""
+        ...
+
+    def get_version_history(self) -> List[SchemaVersion]:
+        """Get all applied versions in order."""
+        ...
+
+    def record_version(self, version: SchemaVersion) -> None:
+        """Record a new version."""
+        ...
+
+    def remove_version(self, version: str) -> None:
+        """Remove a version record (for rollback)."""
+        ...
+
+
+class MigrationRunner:
+    """
+    Executes schema migrations and tracks versions.
+
+    Handles:
+    - Forward migrations (upgrades)
+    - Backward migrations (rollbacks)
+    - Version tracking
+    - Pre/post migration checks
+    - Dry run mode
+    """
+
+    # Current schema version for fresh installations
+    CURRENT_SCHEMA_VERSION = "1.0.0"
+
+    def __init__(
+        self,
+        version_store: VersionStore,
+        registry: Optional[MigrationRegistry] = None,
+        backend: Optional[str] = None,
+    ):
+        """
+        Initialize migration runner.
+
+        Args:
+            version_store: Storage for version tracking
+            registry: Migration registry (uses global if not provided)
+            backend: Backend name for filtering migrations
+        """
+        self.version_store = version_store
+        self.registry = registry or get_registry()
+        self.backend = backend
+        self._hooks: Dict[str, List[Callable]] = {
+            "pre_migrate": [],
+            "post_migrate": [],
+            "pre_rollback": [],
+            "post_rollback": [],
+        }
+
+    def add_hook(self, event: str, callback: Callable) -> None:
+        """
+        Add a hook callback for migration events.
+
+        Args:
+            event: Event name (pre_migrate, post_migrate, etc.)
+            callback: Function to call
+        """
+        if event in self._hooks:
+            self._hooks[event].append(callback)
+
+    def _run_hooks(self, event: str, *args: Any, **kwargs: Any) -> None:
+        """Run all hooks for an event."""
+        for callback in self._hooks.get(event, []):
+            try:
+                callback(*args, **kwargs)
+            except Exception as e:
+                logger.warning(f"Hook {callback.__name__} failed: {e}")
+
+    def get_current_version(self) -> Optional[str]:
+        """Get the current schema version."""
+        return self.version_store.get_current_version()
+
+    def get_pending_migrations(self) -> List[Type[Migration]]:
+        """Get list of migrations that need to be applied."""
+        current = self.get_current_version()
+        return self.registry.get_pending_migrations(current, self.backend)
+
+    def needs_migration(self) -> bool:
+        """Check if there are pending migrations."""
+        return len(self.get_pending_migrations()) > 0
+
+    def migrate(
+        self,
+        connection: Any,
+        target_version: Optional[str] = None,
+        dry_run: bool = False,
+    ) -> List[str]:
+        """
+        Apply pending migrations.
+
+        Args:
+            connection: Database connection or storage instance
+            target_version: Optional target version (applies all if not specified)
+            dry_run: If True, show what would be done without making changes
+
+        Returns:
+            List of applied migration versions
+
+        Raises:
+            MigrationError: If a migration fails
+        """
+        current = self.get_current_version()
+        pending = self.registry.get_pending_migrations(current, self.backend)
+
+        if target_version:
+            target_tuple = SchemaVersion._parse_version(target_version)
+            pending = [
+                m
+                for m in pending
+                if SchemaVersion._parse_version(m.version) <= target_tuple
+            ]
+
+        if not pending:
+            logger.info("No pending migrations")
+            return []
+
+        applied = []
+        logger.info(f"Found {len(pending)} pending migrations")
+
+        for migration_class in pending:
+            version = migration_class.version
+            migration = migration_class()
+
+            if dry_run:
+                logger.info(
+                    f"[DRY RUN] Would apply migration {version}: {migration.description}"
+                )
+                applied.append(version)
+                continue
+
+            logger.info(f"Applying migration {version}: {migration.description}")
+
+            try:
+                # Run pre-check
+                if not migration.pre_check(connection):
+                    raise MigrationError(
+                        f"Pre-check failed for migration {version}",
+                        version=version,
+                    )
+
+                # Run pre-migrate hooks
+                self._run_hooks("pre_migrate", migration, connection)
+
+                # Apply migration
+                migration.upgrade(connection)
+
+                # Run post-check
+                if not migration.post_check(connection):
+                    raise MigrationError(
+                        f"Post-check failed for migration {version}",
+                        version=version,
+                    )
+
+                # Record version
+                schema_version = SchemaVersion(
+                    version=version,
+                    applied_at=datetime.now(timezone.utc),
+                    description=migration.description,
+                    checksum=self._compute_checksum(migration_class),
+                )
+                self.version_store.record_version(schema_version)
+
+                # Run post-migrate hooks
+                self._run_hooks("post_migrate", migration, connection)
+
+                applied.append(version)
+                logger.info(f"Successfully applied migration {version}")
+
+            except MigrationError:
+                raise
+            except Exception as e:
+                raise MigrationError(
+                    f"Migration {version} failed: {str(e)}",
+                    version=version,
+                    cause=e,
+                ) from e
+
+        return applied
+
+    def rollback(
+        self,
+        connection: Any,
+        target_version: str,
+        dry_run: bool = False,
+    ) -> List[str]:
+        """
+        Roll back to a target version.
+
+        Args:
+            connection: Database connection or storage instance
+            target_version: Version to roll back to
+            dry_run: If True, show what would be done without making changes
+
+        Returns:
+            List of rolled back migration versions
+
+        Raises:
+            MigrationError: If a rollback fails
+        """
+        current = self.get_current_version()
+        if current is None:
+            logger.info("No migrations to roll back")
+            return []
+
+        rollback_migrations = self.registry.get_rollback_migrations(
+            current, target_version, self.backend
+        )
+
+        if not rollback_migrations:
+            logger.info("No migrations to roll back")
+            return []
+
+        rolled_back = []
+        logger.info(f"Rolling back {len(rollback_migrations)} migrations")
+
+        for migration_class in rollback_migrations:
+            version = migration_class.version
+            migration = migration_class()
+
+            if dry_run:
+                logger.info(f"[DRY RUN] Would roll back migration {version}")
+                rolled_back.append(version)
+                continue
+
+            logger.info(f"Rolling back migration {version}")
+
+            try:
+                # Run pre-rollback hooks
+                self._run_hooks("pre_rollback", migration, connection)
+
+                # Apply downgrade
+                migration.downgrade(connection)
+
+                # Remove version record
+                self.version_store.remove_version(version)
+
+                # Run post-rollback hooks
+                self._run_hooks("post_rollback", migration, connection)
+
+                rolled_back.append(version)
+                logger.info(f"Successfully rolled back migration {version}")
+
+            except NotImplementedError as e:
+                raise MigrationError(
+                    f"Migration {version} does not support rollback",
+                    version=version,
+                ) from e
+            except Exception as e:
+                raise MigrationError(
+                    f"Rollback of {version} failed: {str(e)}",
+                    version=version,
+                    cause=e,
+                ) from e
+
+        return rolled_back
+
+    def get_status(self) -> Dict[str, Any]:
+        """
+        Get migration status information.
+
+        Returns:
+            Dict with current version, pending migrations, and history
+        """
+        current = self.get_current_version()
+        pending = self.get_pending_migrations()
+        history = self.version_store.get_version_history()
+
+        return {
+            "current_version": current,
+            "target_version": self.CURRENT_SCHEMA_VERSION,
+            "pending_count": len(pending),
+            "pending_versions": [m.version for m in pending],
+            "applied_count": len(history),
+            "history": [
+                {
+                    "version": v.version,
+                    "applied_at": v.applied_at.isoformat(),
+                    "description": v.description,
+                }
+                for v in history
+            ],
+            "needs_migration": len(pending) > 0,
+        }
+
+    @staticmethod
+    def _compute_checksum(migration_class: Type[Migration]) -> str:
+        """Compute checksum of migration for integrity tracking."""
+        import inspect
+
+        source = inspect.getsource(migration_class)
+        return hashlib.sha256(source.encode()).hexdigest()[:16]

alma/storage/migrations/version_stores.py

@@ -0,0 +1,337 @@
+"""
+ALMA Migration Framework - Version Stores.
+
+Implementations of version tracking for different storage backends.
+"""
+
+import json
+import logging
+import sqlite3
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+from typing import Any, List, Optional
+
+from alma.storage.migrations.base import SchemaVersion
+
+logger = logging.getLogger(__name__)
+
+
+class SQLiteVersionStore:
+    """
+    Version store using SQLite for tracking schema versions.
+
+    Creates a `_schema_versions` table to track applied migrations.
+    """
+
+    def __init__(self, db_path: Path):
+        """
+        Initialize SQLite version store.
+
+        Args:
+            db_path: Path to SQLite database file
+        """
+        self.db_path = Path(db_path)
+        self._init_version_table()
+
+    @contextmanager
+    def _get_connection(self):
+        """Get database connection with context manager."""
+        conn = sqlite3.connect(self.db_path)
+        conn.row_factory = sqlite3.Row
+        try:
+            yield conn
+            conn.commit()
+        except Exception:
+            conn.rollback()
+            raise
+        finally:
+            conn.close()
+
+    def _init_version_table(self) -> None:
+        """Create schema versions table if it doesn't exist."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                CREATE TABLE IF NOT EXISTS _schema_versions (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    version TEXT NOT NULL UNIQUE,
+                    applied_at TEXT NOT NULL,
+                    description TEXT,
+                    checksum TEXT
+                )
+            """)
+            cursor.execute(
+                "CREATE INDEX IF NOT EXISTS idx_schema_version "
+                "ON _schema_versions(version)"
+            )
+
+    def get_current_version(self) -> Optional[str]:
+        """Get the current (latest) schema version."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                SELECT version FROM _schema_versions
+                ORDER BY id DESC LIMIT 1
+            """)
+            row = cursor.fetchone()
+            return row["version"] if row else None
+
+    def get_version_history(self) -> List[SchemaVersion]:
+        """Get all applied versions in chronological order."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                SELECT version, applied_at, description, checksum
+                FROM _schema_versions
+                ORDER BY id ASC
+            """)
+            rows = cursor.fetchall()
+
+        return [
+            SchemaVersion(
+                version=row["version"],
+                applied_at=datetime.fromisoformat(row["applied_at"]),
+                description=row["description"] or "",
+                checksum=row["checksum"],
+            )
+            for row in rows
+        ]
+
+    def record_version(self, version: SchemaVersion) -> None:
+        """Record a new schema version."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                """
+                INSERT INTO _schema_versions (version, applied_at, description, checksum)
+                VALUES (?, ?, ?, ?)
+                """,
+                (
+                    version.version,
+                    version.applied_at.isoformat(),
+                    version.description,
+                    version.checksum,
+                ),
+            )
+        logger.debug(f"Recorded schema version {version.version}")
+
+    def remove_version(self, version: str) -> None:
+        """Remove a version record (for rollback)."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute(
+                "DELETE FROM _schema_versions WHERE version = ?",
+                (version,),
+            )
+        logger.debug(f"Removed schema version {version}")
+
+    def has_version_table(self) -> bool:
+        """Check if the version table exists."""
+        with self._get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                SELECT name FROM sqlite_master
+                WHERE type='table' AND name='_schema_versions'
+            """)
+            return cursor.fetchone() is not None
+
+
+class PostgreSQLVersionStore:
+    """
+    Version store using PostgreSQL for tracking schema versions.
+
+    Creates an `_schema_versions` table in the configured schema.
+    """
+
+    def __init__(self, pool: Any, schema: str = "public"):
+        """
+        Initialize PostgreSQL version store.
+
+        Args:
+            pool: psycopg connection pool
+            schema: Database schema name
+        """
+        self._pool = pool
+        self.schema = schema
+        self._init_version_table()
+
+    @contextmanager
+    def _get_connection(self):
+        """Get database connection from pool."""
+        with self._pool.connection() as conn:
+            yield conn
+
+    def _init_version_table(self) -> None:
+        """Create schema versions table if it doesn't exist."""
+        with self._get_connection() as conn:
+            conn.execute(f"""
+                CREATE TABLE IF NOT EXISTS {self.schema}._schema_versions (
+                    id SERIAL PRIMARY KEY,
+                    version TEXT NOT NULL UNIQUE,
+                    applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                    description TEXT,
+                    checksum TEXT
+                )
+            """)
+            conn.execute(f"""
+                CREATE INDEX IF NOT EXISTS idx_schema_version
+                ON {self.schema}._schema_versions(version)
+            """)
+            conn.commit()
+
+    def get_current_version(self) -> Optional[str]:
+        """Get the current (latest) schema version."""
+        with self._get_connection() as conn:
+            cursor = conn.execute(f"""
+                SELECT version FROM {self.schema}._schema_versions
+                ORDER BY id DESC LIMIT 1
+            """)
+            row = cursor.fetchone()
+            return row["version"] if row else None
+
+    def get_version_history(self) -> List[SchemaVersion]:
+        """Get all applied versions in chronological order."""
+        with self._get_connection() as conn:
+            cursor = conn.execute(f"""
+                SELECT version, applied_at, description, checksum
+                FROM {self.schema}._schema_versions
+                ORDER BY id ASC
+            """)
+            rows = cursor.fetchall()
+
+        return [
+            SchemaVersion(
+                version=row["version"],
+                applied_at=row["applied_at"]
+                if isinstance(row["applied_at"], datetime)
+                else datetime.fromisoformat(str(row["applied_at"])),
+                description=row["description"] or "",
+                checksum=row["checksum"],
+            )
+            for row in rows
+        ]
+
+    def record_version(self, version: SchemaVersion) -> None:
+        """Record a new schema version."""
+        with self._get_connection() as conn:
+            conn.execute(
+                f"""
+                INSERT INTO {self.schema}._schema_versions
+                (version, applied_at, description, checksum)
+                VALUES (%s, %s, %s, %s)
+                """,
+                (
+                    version.version,
+                    version.applied_at,
+                    version.description,
+                    version.checksum,
+                ),
+            )
+            conn.commit()
+        logger.debug(f"Recorded schema version {version.version}")
+
+    def remove_version(self, version: str) -> None:
+        """Remove a version record (for rollback)."""
+        with self._get_connection() as conn:
+            conn.execute(
+                f"DELETE FROM {self.schema}._schema_versions WHERE version = %s",
+                (version,),
+            )
+            conn.commit()
+        logger.debug(f"Removed schema version {version}")
+
+    def has_version_table(self) -> bool:
+        """Check if the version table exists."""
+        with self._get_connection() as conn:
+            cursor = conn.execute(
+                """
+                SELECT EXISTS (
+                    SELECT FROM information_schema.tables
+                    WHERE table_schema = %s
+                    AND table_name = '_schema_versions'
+                )
+            """,
+                (self.schema,),
+            )
+            row = cursor.fetchone()
+            return row["exists"] if row else False
+
+
+class FileBasedVersionStore:
+    """
+    Version store using a JSON file for tracking schema versions.
+
+    Useful for file-based storage backends.
+    """
+
+    def __init__(self, storage_dir: Path):
+        """
+        Initialize file-based version store.
+
+        Args:
+            storage_dir: Directory to store version file
+        """
+        self.storage_dir = Path(storage_dir)
+        self.version_file = self.storage_dir / "_schema_versions.json"
+        self._ensure_file_exists()
+
+    def _ensure_file_exists(self) -> None:
+        """Create version file if it doesn't exist."""
+        self.storage_dir.mkdir(parents=True, exist_ok=True)
+        if not self.version_file.exists():
+            self._write_versions([])
+
+    def _read_versions(self) -> List[dict]:
+        """Read versions from file."""
+        try:
+            with open(self.version_file, "r") as f:
+                return json.load(f)
+        except (json.JSONDecodeError, FileNotFoundError):
+            return []
+
+    def _write_versions(self, versions: List[dict]) -> None:
+        """Write versions to file."""
+        with open(self.version_file, "w") as f:
+            json.dump(versions, f, indent=2, default=str)
+
+    def get_current_version(self) -> Optional[str]:
+        """Get the current (latest) schema version."""
+        versions = self._read_versions()
+        if not versions:
+            return None
+        return versions[-1]["version"]
+
+    def get_version_history(self) -> List[SchemaVersion]:
+        """Get all applied versions in chronological order."""
+        versions = self._read_versions()
+        return [
+            SchemaVersion(
+                version=v["version"],
+                applied_at=datetime.fromisoformat(v["applied_at"]),
+                description=v.get("description", ""),
+                checksum=v.get("checksum"),
+            )
+            for v in versions
+        ]
+
+    def record_version(self, version: SchemaVersion) -> None:
+        """Record a new schema version."""
+        versions = self._read_versions()
+        versions.append(
+            {
+                "version": version.version,
+                "applied_at": version.applied_at.isoformat(),
+                "description": version.description,
+                "checksum": version.checksum,
+            }
+        )
+        self._write_versions(versions)
+        logger.debug(f"Recorded schema version {version.version}")
+
+    def remove_version(self, version: str) -> None:
+        """Remove a version record (for rollback)."""
+        versions = self._read_versions()
+        versions = [v for v in versions if v["version"] != version]
+        self._write_versions(versions)
+        logger.debug(f"Removed schema version {version}")

alma/storage/migrations/versions/__init__.py

@@ -0,0 +1,11 @@
+"""
+ALMA Schema Migrations - Version Definitions.
+
+Each version file contains migrations for that schema version.
+Migrations are automatically registered when imported.
+"""
+
+# Import all version modules to register migrations
+from alma.storage.migrations.versions import v1_0_0
+
+__all__ = ["v1_0_0"]