pyworkflow-engine 0.1.21__py3-none-any.whl → 0.1.23__py3-none-any.whl

pyworkflow/storage/migrations/base.py

@@ -0,0 +1,299 @@
+"""
+Base classes for the database migration framework.
+
+Provides Migration dataclass, MigrationRegistry for tracking migrations,
+and MigrationRunner abstract base class for backend-specific implementations.
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from typing import Any
+
+
+@dataclass
+class Migration:
+    """
+    Represents a database schema migration.
+
+    Attributes:
+        version: Integer version number (must be unique and sequential)
+        description: Human-readable description of what the migration does
+        up_sql: SQL to apply the migration (can be None for Python-based migrations)
+        down_sql: SQL to rollback the migration (optional, for future use)
+        up_func: Optional Python function for complex migrations (receives connection)
+    """
+
+    version: int
+    description: str
+    up_sql: str | None = None
+    down_sql: str | None = None
+    up_func: Callable[[Any], Any] | None = None
+
+    def __post_init__(self) -> None:
+        if self.version < 1:
+            raise ValueError("Migration version must be >= 1")
+        if not self.up_sql and not self.up_func:
+            raise ValueError("Migration must have either up_sql or up_func")
+
+
+@dataclass
+class AppliedMigration:
+    """
+    Record of an applied migration.
+
+    Attributes:
+        version: Migration version number
+        applied_at: When the migration was applied
+        description: Description of the migration
+    """
+
+    version: int
+    applied_at: datetime
+    description: str
+
+
+class MigrationRegistry:
+    """
+    Registry for managing and tracking migrations.
+
+    Maintains an ordered list of migrations and provides methods to
+    get pending migrations based on the current database version.
+    """
+
+    def __init__(self) -> None:
+        self._migrations: dict[int, Migration] = {}
+
+    def register(self, migration: Migration) -> None:
+        """
+        Register a migration.
+
+        Args:
+            migration: Migration to register
+
+        Raises:
+            ValueError: If a migration with the same version already exists
+        """
+        if migration.version in self._migrations:
+            raise ValueError(f"Migration version {migration.version} already registered")
+        self._migrations[migration.version] = migration
+
+    def get_all(self) -> list[Migration]:
+        """
+        Get all registered migrations, ordered by version.
+
+        Returns:
+            List of migrations sorted by version ascending
+        """
+        return [self._migrations[v] for v in sorted(self._migrations.keys())]
+
+    def get_pending(self, current_version: int) -> list[Migration]:
+        """
+        Get migrations that need to be applied.
+
+        Args:
+            current_version: Current schema version (0 if fresh database)
+
+        Returns:
+            List of migrations with version > current_version, sorted ascending
+        """
+        return [self._migrations[v] for v in sorted(self._migrations.keys()) if v > current_version]
+
+    def get_latest_version(self) -> int:
+        """
+        Get the latest migration version.
+
+        Returns:
+            Highest registered version, or 0 if no migrations
+        """
+        return max(self._migrations.keys()) if self._migrations else 0
+
+    def get(self, version: int) -> Migration | None:
+        """
+        Get a specific migration by version.
+
+        Args:
+            version: Migration version number
+
+        Returns:
+            Migration if found, None otherwise
+        """
+        return self._migrations.get(version)
+
+
+# Global migration registry for SQL backends
+_global_registry = MigrationRegistry()
+
+
+def get_global_registry() -> MigrationRegistry:
+    """Get the global migration registry."""
+    return _global_registry
+
+
+def register_migration(migration: Migration) -> None:
+    """Register a migration in the global registry."""
+    _global_registry.register(migration)
+
+
+class MigrationRunner(ABC):
+    """
+    Abstract base class for running migrations on a storage backend.
+
+    Subclasses must implement the backend-specific methods for:
+    - Ensuring the schema_versions table exists
+    - Getting the current schema version
+    - Applying individual migrations
+    - Detecting existing schemas (for backward compatibility)
+    """
+
+    def __init__(self, registry: MigrationRegistry | None = None) -> None:
+        """
+        Initialize the migration runner.
+
+        Args:
+            registry: Migration registry to use (defaults to global registry)
+        """
+        self.registry = registry or get_global_registry()
+
+    @abstractmethod
+    async def ensure_schema_versions_table(self) -> None:
+        """
+        Create the schema_versions table if it doesn't exist.
+
+        The table should have:
+        - version: INTEGER PRIMARY KEY
+        - applied_at: TIMESTAMP NOT NULL
+        - description: TEXT
+        """
+        pass
+
+    @abstractmethod
+    async def get_current_version(self) -> int:
+        """
+        Get the current schema version from the database.
+
+        Returns:
+            Current version (highest applied), or 0 if no migrations applied
+        """
+        pass
+
+    @abstractmethod
+    async def apply_migration(self, migration: Migration) -> None:
+        """
+        Apply a single migration.
+
+        This should:
+        1. Execute the migration SQL/function in a transaction
+        2. Record the migration in schema_versions
+        3. Rollback on failure
+
+        Args:
+            migration: Migration to apply
+
+        Raises:
+            Exception: If migration fails
+        """
+        pass
+
+    @abstractmethod
+    async def detect_existing_schema(self) -> bool:
+        """
+        Detect if the database has an existing schema (pre-versioning).
+
+        This is used for backward compatibility with databases created
+        before the migration framework was added. If tables exist but
+        no schema_versions table, we assume it's a V1 schema.
+
+        Returns:
+            True if existing schema detected, False if fresh database
+        """
+        pass
+
+    @abstractmethod
+    async def record_baseline_version(self, version: int, description: str) -> None:
+        """
+        Record a baseline version without running migrations.
+
+        Used when detecting an existing schema to mark it as a known version.
+
+        Args:
+            version: Version number to record
+            description: Description of the baseline
+        """
+        pass
+
+    async def run_migrations(self) -> list[AppliedMigration]:
+        """
+        Run all pending migrations.
+
+        This is the main entry point for the migration runner. It:
+        1. Ensures the schema_versions table exists
+        2. Detects existing schemas and records baseline if needed
+        3. Applies all pending migrations in order
+
+        Returns:
+            List of applied migrations
+
+        Raises:
+            Exception: If any migration fails (partial migrations are rolled back)
+        """
+        # Ensure we have a schema_versions table
+        await self.ensure_schema_versions_table()
+
+        # Get current version
+        current_version = await self.get_current_version()
+
+        # If no migrations recorded, check for existing schema
+        if current_version == 0:
+            has_existing_schema = await self.detect_existing_schema()
+            if has_existing_schema:
+                # Database has tables but no version tracking
+                # Assume it's at V1 (original schema)
+                await self.record_baseline_version(1, "Baseline: original schema")
+                current_version = 1
+
+        # Get and apply pending migrations
+        pending = self.registry.get_pending(current_version)
+        applied: list[AppliedMigration] = []
+
+        for migration in pending:
+            await self.apply_migration(migration)
+            applied.append(
+                AppliedMigration(
+                    version=migration.version,
+                    applied_at=datetime.now(UTC),
+                    description=migration.description,
+                )
+            )
+
+        return applied
+
+
+# =============================================================================
+# Migration Definitions
+# =============================================================================
+
+# Version 1: Baseline schema (represents the original schema before versioning)
+# This is auto-detected for existing databases
+
+_v1_migration = Migration(
+    version=1,
+    description="Baseline: original schema",
+    up_sql="SELECT 1",  # No-op, baseline is detected not applied
+)
+register_migration(_v1_migration)
+
+
+# Version 2: Add step_id column to events table for optimized queries
+# The actual SQL is backend-specific, but we define the Python function here
+# to handle the backfill logic
+
+_v2_migration = Migration(
+    version=2,
+    description="Add step_id column to events table for optimized has_event() queries",
+    # SQL is None because each backend has different syntax
+    # The up_func will be set by each backend's runner
+    up_sql="SELECT 1",  # Placeholder, actual migration is in backend runners
+)
+register_migration(_v2_migration)

pyworkflow/storage/mysql.py

@@ -17,6 +17,7 @@ import aiomysql
 
 from pyworkflow.engine.events import Event, EventType
 from pyworkflow.storage.base import StorageBackend
+from pyworkflow.storage.migrations import Migration, MigrationRegistry, MigrationRunner
 from pyworkflow.storage.schemas import (
     Hook,
     HookStatus,
@@ -31,6 +32,124 @@ from pyworkflow.storage.schemas import (
 )
 
 
+class MySQLMigrationRunner(MigrationRunner):
+    """MySQL-specific migration runner."""
+
+    def __init__(self, pool: aiomysql.Pool, registry: MigrationRegistry | None = None) -> None:
+        super().__init__(registry)
+        self._pool = pool
+
+    async def ensure_schema_versions_table(self) -> None:
+        """Create schema_versions table if it doesn't exist."""
+        async with self._pool.acquire() as conn, conn.cursor() as cur:
+            await cur.execute("""
+                CREATE TABLE IF NOT EXISTS schema_versions (
+                    version INT PRIMARY KEY,
+                    applied_at DATETIME(6) NOT NULL,
+                    description TEXT
+                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
+            """)
+
+    async def get_current_version(self) -> int:
+        """Get the highest applied migration version."""
+        async with self._pool.acquire() as conn, conn.cursor() as cur:
+            await cur.execute("SELECT COALESCE(MAX(version), 0) as version FROM schema_versions")
+            row = await cur.fetchone()
+            return row[0] if row else 0
+
+    async def detect_existing_schema(self) -> bool:
+        """Check if the events table exists (pre-versioning database)."""
+        async with self._pool.acquire() as conn, conn.cursor() as cur:
+            await cur.execute("""
+                SELECT COUNT(*) FROM information_schema.tables
+                WHERE table_schema = DATABASE() AND table_name = 'events'
+            """)
+            row = await cur.fetchone()
+            return row[0] > 0 if row else False
+
+    async def record_baseline_version(self, version: int, description: str) -> None:
+        """Record a baseline version without running migrations."""
+        async with self._pool.acquire() as conn, conn.cursor() as cur:
+            await cur.execute(
+                """
+                INSERT IGNORE INTO schema_versions (version, applied_at, description)
+                VALUES (%s, %s, %s)
+                """,
+                (version, datetime.now(UTC), description),
+            )
+
+    async def apply_migration(self, migration: Migration) -> None:
+        """Apply a migration with MySQL-specific handling."""
+        async with self._pool.acquire() as conn:
+            await conn.begin()
+            try:
+                async with conn.cursor() as cur:
+                    if migration.version == 2:
+                        # V2: Add step_id column to events table
+                        # First check if events table exists (fresh databases won't have it yet)
+                        await cur.execute("""
+                            SELECT COUNT(*) FROM information_schema.tables
+                            WHERE table_schema = DATABASE() AND table_name = 'events'
+                        """)
+                        row = await cur.fetchone()
+                        table_exists = row[0] > 0 if row else False
+
+                        if table_exists:
+                            # Check if column exists first
+                            await cur.execute("""
+                                SELECT COUNT(*) FROM information_schema.columns
+                                WHERE table_schema = DATABASE()
+                                  AND table_name = 'events'
+                                  AND column_name = 'step_id'
+                            """)
+                            row = await cur.fetchone()
+                            if row[0] == 0:
+                                await cur.execute(
+                                    "ALTER TABLE events ADD COLUMN step_id VARCHAR(255)"
+                                )
+
+                            # Create index for optimized has_event() queries
+                            # Check if index exists first
+                            await cur.execute("""
+                                SELECT COUNT(*) FROM information_schema.statistics
+                                WHERE table_schema = DATABASE()
+                                  AND table_name = 'events'
+                                  AND index_name = 'idx_events_run_id_step_id_type'
+                            """)
+                            row = await cur.fetchone()
+                            if row[0] == 0:
+                                await cur.execute("""
+                                    CREATE INDEX idx_events_run_id_step_id_type
+                                    ON events(run_id, step_id, type)
+                                """)
+
+                            # Backfill step_id from JSON data
+                            await cur.execute("""
+                                UPDATE events
+                                SET step_id = JSON_UNQUOTE(JSON_EXTRACT(data, '$.step_id'))
+                                WHERE step_id IS NULL
+                                  AND JSON_EXTRACT(data, '$.step_id') IS NOT NULL
+                            """)
+                        # If table doesn't exist, schema will be created with step_id column
+                    elif migration.up_func:
+                        await migration.up_func(conn)
+                    elif migration.up_sql and migration.up_sql != "SELECT 1":
+                        await cur.execute(migration.up_sql)
+
+                    # Record the migration
+                    await cur.execute(
+                        """
+                        INSERT INTO schema_versions (version, applied_at, description)
+                        VALUES (%s, %s, %s)
+                        """,
+                        (migration.version, datetime.now(UTC), migration.description),
+                    )
+                await conn.commit()
+            except Exception:
+                await conn.rollback()
+                raise
+
+
 class MySQLStorageBackend(StorageBackend):
     """
     MySQL storage backend using aiomysql for async operations.
@@ -101,11 +220,16 @@ class MySQLStorageBackend(StorageBackend):
             self._initialized = False
 
     async def _initialize_schema(self) -> None:
-        """Create database tables if they don't exist."""
+        """Create database tables if they don't exist and run migrations."""
         if not self._pool:
             await self.connect()
 
         pool = self._ensure_connected()
+
+        # Run migrations first (handles schema versioning)
+        runner = MySQLMigrationRunner(pool)
+        await runner.run_migrations()
+
         async with pool.acquire() as conn, conn.cursor() as cur:
             # Workflow runs table
             await cur.execute("""
@@ -140,7 +264,7 @@ class MySQLStorageBackend(StorageBackend):
                     ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
                 """)
 
-            # Events table
+            # Events table (includes step_id column added in V2 migration)
             await cur.execute("""
                     CREATE TABLE IF NOT EXISTS events (
                         event_id VARCHAR(255) PRIMARY KEY,
@@ -149,8 +273,10 @@ class MySQLStorageBackend(StorageBackend):
                         type VARCHAR(100) NOT NULL,
                         timestamp DATETIME(6) NOT NULL,
                         data LONGTEXT NOT NULL DEFAULT '{}',
+                        step_id VARCHAR(255),
                         INDEX idx_events_run_id_sequence (run_id, sequence),
-                        INDEX idx_events_type (type),
+                        INDEX idx_events_run_id_type (run_id, type),
+                        INDEX idx_events_run_id_step_id_type (run_id, step_id, type),
                         FOREIGN KEY (run_id) REFERENCES workflow_runs(run_id) ON DELETE CASCADE
                     ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci
                 """)
@@ -449,6 +575,9 @@ class MySQLStorageBackend(StorageBackend):
         """Record an event to the append-only event log."""
         pool = self._ensure_connected()
 
+        # Extract step_id from event data for indexed column
+        step_id = event.data.get("step_id") if event.data else None
+
         async with pool.acquire() as conn:
             # Use transaction for atomic sequence assignment
             await conn.begin()
@@ -464,8 +593,8 @@ class MySQLStorageBackend(StorageBackend):
 
                     await cur.execute(
                         """
-                        INSERT INTO events (event_id, run_id, sequence, type, timestamp, data)
-                        VALUES (%s, %s, %s, %s, %s, %s)
+                        INSERT INTO events (event_id, run_id, sequence, type, timestamp, data, step_id)
+                        VALUES (%s, %s, %s, %s, %s, %s, %s)
                         """,
                         (
                             event.event_id,
@@ -474,6 +603,7 @@ class MySQLStorageBackend(StorageBackend):
                             event.type.value,
                             event.timestamp,
                             json.dumps(event.data),
+                            step_id,
                         ),
                     )
                 await conn.commit()
@@ -545,6 +675,57 @@ class MySQLStorageBackend(StorageBackend):
 
         return self._row_to_event(row)
 
+    async def has_event(
+        self,
+        run_id: str,
+        event_type: str,
+        **filters: str,
+    ) -> bool:
+        """
+        Check if an event exists using optimized indexed queries.
+
+        When step_id is the only filter, uses a direct indexed query (O(1) lookup).
+        For other filters, falls back to loading events of the type and filtering in Python.
+
+        Args:
+            run_id: Workflow run identifier
+            event_type: Event type to check for
+            **filters: Additional filters for event data fields
+
+        Returns:
+            True if a matching event exists, False otherwise
+        """
+        pool = self._ensure_connected()
+
+        # Optimized path: if only filtering by step_id, use indexed column directly
+        if filters.keys() == {"step_id"}:
+            step_id = str(filters["step_id"])
+            async with pool.acquire() as conn, conn.cursor() as cur:
+                await cur.execute(
+                    """
+                    SELECT 1 FROM events
+                    WHERE run_id = %s AND type = %s AND step_id = %s
+                    LIMIT 1
+                    """,
+                    (run_id, event_type, step_id),
+                )
+                row = await cur.fetchone()
+            return row is not None
+
+        # Fallback: load events of type and filter in Python
+        events = await self.get_events(run_id, event_types=[event_type])
+
+        for event in events:
+            match = True
+            for key, value in filters.items():
+                if str(event.data.get(key)) != str(value):
+                    match = False
+                    break
+            if match:
+                return True
+
+        return False
+
     # Step Operations
 
     async def create_step(self, step: StepExecution) -> None: