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

alma/storage/migrations/__init__.py

@@ -0,0 +1,21 @@
+"""
+ALMA Schema Migration Framework.
+
+Provides version tracking and migration capabilities for storage backends.
+"""
+
+from alma.storage.migrations.base import (
+    Migration,
+    MigrationError,
+    MigrationRegistry,
+    SchemaVersion,
+)
+from alma.storage.migrations.runner import MigrationRunner
+
+__all__ = [
+    "Migration",
+    "MigrationError",
+    "MigrationRegistry",
+    "MigrationRunner",
+    "SchemaVersion",
+]

alma/storage/migrations/base.py

@@ -0,0 +1,321 @@
+"""
+ALMA Migration Framework - Base Classes.
+
+Provides abstract migration classes and version tracking utilities.
+"""
+
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Callable, Dict, List, Optional, Type
+
+logger = logging.getLogger(__name__)
+
+
+class MigrationError(Exception):
+    """Exception raised when a migration fails."""
+
+    def __init__(
+        self,
+        message: str,
+        version: Optional[str] = None,
+        cause: Optional[Exception] = None,
+    ):
+        self.version = version
+        self.cause = cause
+        super().__init__(message)
+
+
+@dataclass
+class SchemaVersion:
+    """
+    Represents a schema version record.
+
+    Attributes:
+        version: Semantic version string (e.g., "1.0.0")
+        applied_at: When the migration was applied
+        description: Human-readable description of changes
+        checksum: Optional hash for integrity verification
+    """
+
+    version: str
+    applied_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    description: str = ""
+    checksum: Optional[str] = None
+
+    def __lt__(self, other: "SchemaVersion") -> bool:
+        """Compare versions for sorting."""
+        return self._parse_version(self.version) < self._parse_version(other.version)
+
+    @staticmethod
+    def _parse_version(version: str) -> tuple:
+        """Parse version string into comparable tuple."""
+        parts = version.split(".")
+        result = []
+        for part in parts:
+            try:
+                result.append(int(part))
+            except ValueError:
+                result.append(part)
+        return tuple(result)
+
+
+class Migration(ABC):
+    """
+    Abstract base class for schema migrations.
+
+    Subclasses must implement upgrade() and optionally downgrade().
+
+    Example:
+        class AddTagsColumn(Migration):
+            version = "1.1.0"
+            description = "Add tags column to heuristics table"
+
+            def upgrade(self, connection):
+                connection.execute(
+                    "ALTER TABLE heuristics ADD COLUMN tags TEXT"
+                )
+
+            def downgrade(self, connection):
+                connection.execute(
+                    "ALTER TABLE heuristics DROP COLUMN tags"
+                )
+    """
+
+    # These must be set by subclasses
+    version: str = ""
+    description: str = ""
+    # Optional: previous version this migration depends on
+    depends_on: Optional[str] = None
+
+    @abstractmethod
+    def upgrade(self, connection: Any) -> None:
+        """
+        Apply the migration.
+
+        Args:
+            connection: Database connection or storage instance
+        """
+        pass
+
+    def downgrade(self, connection: Any) -> None:
+        """
+        Revert the migration (optional).
+
+        Args:
+            connection: Database connection or storage instance
+
+        Raises:
+            NotImplementedError: If downgrade is not supported
+        """
+        raise NotImplementedError(
+            f"Downgrade not implemented for migration {self.version}"
+        )
+
+    def pre_check(self, connection: Any) -> bool:
+        """
+        Optional pre-migration check.
+
+        Override to verify prerequisites before migration.
+
+        Args:
+            connection: Database connection
+
+        Returns:
+            True if migration can proceed, False otherwise
+        """
+        return True
+
+    def post_check(self, connection: Any) -> bool:
+        """
+        Optional post-migration verification.
+
+        Override to verify migration was successful.
+
+        Args:
+            connection: Database connection
+
+        Returns:
+            True if migration was successful
+        """
+        return True
+
+
+class MigrationRegistry:
+    """
+    Registry for available migrations.
+
+    Manages migration discovery, ordering, and execution planning.
+    """
+
+    def __init__(self) -> None:
+        self._migrations: Dict[str, Type[Migration]] = {}
+        self._backend_migrations: Dict[str, Dict[str, Type[Migration]]] = {}
+
+    def register(
+        self, migration_class: Type[Migration], backend: Optional[str] = None
+    ) -> Type[Migration]:
+        """
+        Register a migration class.
+
+        Args:
+            migration_class: The migration class to register
+            backend: Optional backend name (e.g., "sqlite", "postgresql")
+
+        Returns:
+            The migration class (for use as decorator)
+        """
+        version = migration_class.version
+        if not version:
+            raise ValueError(
+                f"Migration {migration_class.__name__} must have a version"
+            )
+
+        if backend:
+            if backend not in self._backend_migrations:
+                self._backend_migrations[backend] = {}
+            self._backend_migrations[backend][version] = migration_class
+            logger.debug(f"Registered migration {version} for backend {backend}")
+        else:
+            self._migrations[version] = migration_class
+            logger.debug(f"Registered global migration {version}")
+
+        return migration_class
+
+    def get_migration(
+        self, version: str, backend: Optional[str] = None
+    ) -> Optional[Type[Migration]]:
+        """
+        Get a migration class by version.
+
+        Args:
+            version: Version string to look up
+            backend: Optional backend name
+
+        Returns:
+            Migration class or None if not found
+        """
+        if backend and backend in self._backend_migrations:
+            migration = self._backend_migrations[backend].get(version)
+            if migration:
+                return migration
+        return self._migrations.get(version)
+
+    def get_all_migrations(
+        self, backend: Optional[str] = None
+    ) -> List[Type[Migration]]:
+        """
+        Get all migrations in version order.
+
+        Args:
+            backend: Optional backend name to filter migrations
+
+        Returns:
+            List of migration classes sorted by version
+        """
+        migrations = dict(self._migrations)
+        if backend and backend in self._backend_migrations:
+            migrations.update(self._backend_migrations[backend])
+
+        return [
+            cls
+            for _, cls in sorted(
+                migrations.items(),
+                key=lambda x: SchemaVersion._parse_version(x[0]),
+            )
+        ]
+
+    def get_pending_migrations(
+        self,
+        current_version: Optional[str],
+        backend: Optional[str] = None,
+    ) -> List[Type[Migration]]:
+        """
+        Get migrations that need to be applied.
+
+        Args:
+            current_version: Current schema version (None if fresh install)
+            backend: Optional backend name
+
+        Returns:
+            List of migration classes that need to be applied
+        """
+        all_migrations = self.get_all_migrations(backend)
+
+        if current_version is None:
+            return all_migrations
+
+        current = SchemaVersion._parse_version(current_version)
+        return [
+            m
+            for m in all_migrations
+            if SchemaVersion._parse_version(m.version) > current
+        ]
+
+    def get_rollback_migrations(
+        self,
+        current_version: str,
+        target_version: str,
+        backend: Optional[str] = None,
+    ) -> List[Type[Migration]]:
+        """
+        Get migrations that need to be rolled back.
+
+        Args:
+            current_version: Current schema version
+            target_version: Target version to roll back to
+            backend: Optional backend name
+
+        Returns:
+            List of migration classes to roll back (in reverse order)
+        """
+        all_migrations = self.get_all_migrations(backend)
+
+        current = SchemaVersion._parse_version(current_version)
+        target = SchemaVersion._parse_version(target_version)
+
+        rollback = [
+            m
+            for m in all_migrations
+            if target < SchemaVersion._parse_version(m.version) <= current
+        ]
+
+        # Return in reverse order for rollback
+        return list(reversed(rollback))
+
+
+# Global registry instance
+_global_registry = MigrationRegistry()
+
+
+def register_migration(
+    backend: Optional[str] = None,
+) -> Callable[[Type[Migration]], Type[Migration]]:
+    """
+    Decorator to register a migration class.
+
+    Args:
+        backend: Optional backend name
+
+    Example:
+        @register_migration()
+        class MyMigration(Migration):
+            version = "1.0.0"
+            ...
+
+        @register_migration(backend="postgresql")
+        class PostgresSpecificMigration(Migration):
+            version = "1.0.1"
+            ...
+    """
+
+    def decorator(cls: Type[Migration]) -> Type[Migration]:
+        return _global_registry.register(cls, backend)
+
+    return decorator
+
+
+def get_registry() -> MigrationRegistry:
+    """Get the global migration registry."""
+    return _global_registry

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]