hammerwork 3.0.0

//! Database migration system for Hammerwork.
//!
//! This module provides a migration framework that allows progressive schema updates
//! while maintaining backward compatibility. Migrations are versioned and tracked
//! to ensure each migration runs only once.
//!
//! ## Quick Start
//!
//! The easiest way to run migrations is using the cargo subcommand:
//!
//! ```bash
//! # Build the subcommand
//! cargo build --bin cargo-hammerwork --features postgres
//!
//! # Run migrations
//! cargo hammerwork migrate --database-url postgresql://localhost/hammerwork
//!
//! # Check status
//! cargo hammerwork status --database-url postgresql://localhost/hammerwork
//! ```
//!
//! ## Application Usage
//!
//! Once migrations are complete, your application simply connects to the database:
//!
#![cfg_attr(feature = "postgres", doc = "```rust,no_run")]
#![cfg_attr(not(feature = "postgres"), doc = "```rust,ignore")]
//! use hammerwork::{Job, JobQueue, queue::DatabaseQueue};
//! use serde_json::json;
//! use std::sync::Arc;
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
//!     // Database schema is already set up by cargo hammerwork migrate
//!     let pool = sqlx::PgPool::connect("postgresql://localhost/hammerwork").await?;
//!     let queue = Arc::new(JobQueue::new(pool));
//!     
//!     // Start using immediately
//!     let job = Job::new("queue".to_string(), json!({"task": "work"}));
//!     queue.enqueue(job).await?;
//!     Ok(())
//! }
//! ```

#[cfg(feature = "postgres")]
pub mod postgres;

#[cfg(feature = "mysql")]
pub mod mysql;

#[cfg(any(feature = "postgres", feature = "mysql", test))]
mod split;

use crate::Result;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sqlx::Database;
use std::collections::HashMap;
use tracing::info;

/// Migration identifier and metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub struct Migration {
    /// Unique migration identifier (e.g., "001_initial_schema")
    pub id: String,
    /// Human-readable description of what this migration does
    pub description: String,
    /// Version number for ordering (e.g., 1, 2, 3...)
    pub version: u32,
    /// When this migration was created
    pub created_at: DateTime<Utc>,
}

impl Migration {
    /// A migration `id` (such as `"021_add_my_table"`) at `version`, created now.
    pub fn new(id: impl Into<String>, description: impl Into<String>, version: u32) -> Self {
        Self {
            id: id.into(),
            description: description.into(),
            version,
            created_at: Utc::now(),
        }
    }
}

with_setters! {
    Migration {
        with_created_at(created_at): DateTime<Utc>,
    }
}

/// Migration execution record
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct MigrationRecord {
    pub migration_id: String,
    pub executed_at: DateTime<Utc>,
    pub execution_time_ms: u64,
}

/// Seals [`MigrationRunner`]: only the runners of this crate implement
/// [`private::Sealed`], so only they can implement the trait.
mod private {
    /// The supertrait that seals [`MigrationRunner`](super::MigrationRunner).
    pub trait Sealed {}

    #[cfg(feature = "postgres")]
    impl Sealed for super::postgres::PostgresMigrationRunner {}

    #[cfg(feature = "mysql")]
    impl Sealed for super::mysql::MySqlMigrationRunner {}
}

/// The database-specific side of [`MigrationManager`].
///
/// Sealed: it is implemented by [`PostgresMigrationRunner`](postgres::PostgresMigrationRunner)
/// and [`MySqlMigrationRunner`](mysql::MySqlMigrationRunner) only, because the built-in
/// migrations exist for those two databases only. It stays usable as a bound and as a
/// trait object (`Box<dyn MigrationRunner<DB> + Send + Sync>`).
#[async_trait::async_trait]
pub trait MigrationRunner<DB: Database>: private::Sealed {
    /// Take the lock that serializes migration runs on the database, waiting while
    /// another process holds it. [`MigrationManager::run_migrations`] holds it for the
    /// whole run and releases it with
    /// [`release_migration_lock`](Self::release_migration_lock).
    ///
    /// The runners use a session lock (`pg_advisory_lock` on PostgreSQL, `GET_LOCK` on
    /// MySQL) on a dedicated connection that also runs every statement of the run; if
    /// the run is abandoned the connection is closed, which releases the lock.
    async fn acquire_migration_lock(&self) -> Result<()>;

    /// Release the lock taken by [`acquire_migration_lock`](Self::acquire_migration_lock).
    async fn release_migration_lock(&self) -> Result<()>;

    /// Run a specific migration SQL for this database type
    async fn run_migration(&self, migration: &Migration, sql: &str) -> Result<()>;

    /// Check if the migration tracking table exists
    async fn migration_table_exists(&self) -> Result<bool>;

    /// Create the migration tracking table
    async fn create_migration_table(&self) -> Result<()>;

    /// Get list of migrations that have already been executed
    async fn get_executed_migrations(&self) -> Result<Vec<MigrationRecord>>;

    /// Record that a migration has been executed
    async fn record_migration(&self, migration: &Migration, execution_time_ms: u64) -> Result<()>;
}

/// Main migration manager
pub struct MigrationManager<DB: Database> {
    runner: Box<dyn MigrationRunner<DB> + Send + Sync>,
    migrations: HashMap<String, (Migration, String, String)>, // id -> (metadata, postgres_sql, mysql_sql)
}

impl<DB: Database> MigrationManager<DB> {
    /// Create a new migration manager
    pub fn new(runner: Box<dyn MigrationRunner<DB> + Send + Sync>) -> Self {
        let mut manager = Self {
            runner,
            migrations: HashMap::new(),
        };

        // Register all built-in migrations
        manager.register_builtin_migrations();
        manager
    }

    /// Register a migration with both PostgreSQL and MySQL SQL
    pub fn register_migration(
        &mut self,
        migration: Migration,
        postgres_sql: String,
        mysql_sql: String,
    ) {
        self.migrations
            .insert(migration.id.clone(), (migration, postgres_sql, mysql_sql));
    }

    /// Run all pending migrations.
    ///
    /// Safe to call from several processes at once (for example every replica of a
    /// service starting with `auto_migrate`, or the CLI while the application starts):
    /// the run holds the runner's migration lock
    /// ([`MigrationRunner::acquire_migration_lock`]) from before it reads which
    /// migrations were applied until the last one is recorded, so concurrent runs wait
    /// for each other and each migration runs once. The lock is released on every path,
    /// including errors.
    pub async fn run_migrations(&self) -> Result<()> {
        info!("Starting migration process...");

        self.runner.acquire_migration_lock().await?;
        let result = self.run_pending_migrations().await;
        let released = self.runner.release_migration_lock().await;
        if let Err(e) = &released {
            tracing::warn!("Failed to release the migration lock: {e}");
        }
        result?;
        released
    }

    /// Apply the migrations not yet recorded, in version order. The caller holds the
    /// migration lock, so the list of applied migrations is read after any concurrent
    /// run finished.
    async fn run_pending_migrations(&self) -> Result<()> {
        // Ensure migration table exists
        if !self.runner.migration_table_exists().await? {
            info!("Creating migration tracking table...");
            self.runner.create_migration_table().await?;
        }

        // Get executed migrations
        let executed = self.runner.get_executed_migrations().await?;
        let executed_ids: std::collections::HashSet<String> =
            executed.iter().map(|r| r.migration_id.clone()).collect();

        // Sort migrations by version
        let mut pending_migrations: Vec<_> = self
            .migrations
            .values()
            .filter(|(migration, _, _)| !executed_ids.contains(&migration.id))
            .collect();
        pending_migrations.sort_by_key(|(migration, _, _)| migration.version);

        if pending_migrations.is_empty() {
            info!("No pending migrations to run");
            return Ok(());
        }

        info!("Found {} pending migrations", pending_migrations.len());

        // Run each pending migration
        for (migration, postgres_sql, mysql_sql) in pending_migrations {
            info!(
                "Running migration: {} - {}",
                migration.id, migration.description
            );

            let start_time = std::time::Instant::now();

            // Choose SQL based on database type
            let sql = if std::any::type_name::<DB>().contains("Postgres") {
                postgres_sql
            } else {
                mysql_sql
            };

            // Execute the migration
            self.runner.run_migration(migration, sql).await?;

            let execution_time_ms = start_time.elapsed().as_millis() as u64;

            // Record successful execution
            self.runner
                .record_migration(migration, execution_time_ms)
                .await?;

            info!(
                "Completed migration {} in {}ms",
                migration.id, execution_time_ms
            );
        }

        info!("All migrations completed successfully");
        Ok(())
    }

    /// The migrations this version of the library registers, in version order, whether or
    /// not they were applied. Reads no database state.
    pub fn registered_migrations(&self) -> Vec<Migration> {
        let mut migrations: Vec<_> = self
            .migrations
            .values()
            .map(|(migration, _, _)| migration.clone())
            .collect();
        migrations.sort_by_key(|migration| migration.version);
        migrations
    }

    /// Get status of all migrations
    pub async fn get_migration_status(&self) -> Result<Vec<(Migration, bool)>> {
        let executed = self.runner.get_executed_migrations().await?;
        let executed_ids: std::collections::HashSet<String> =
            executed.iter().map(|r| r.migration_id.clone()).collect();

        let mut status: Vec<_> = self
            .migrations
            .values()
            .map(|(migration, _, _)| (migration.clone(), executed_ids.contains(&migration.id)))
            .collect();
        status.sort_by_key(|(migration, _)| migration.version);

        Ok(status)
    }

    /// Register all built-in migrations
    fn register_builtin_migrations(&mut self) {
        // Migration 001: Initial schema with basic job table
        self.register_migration(
            Migration {
                id: "001_initial_schema".to_string(),
                description: "Create initial hammerwork_jobs table".to_string(),
                version: 1,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-01-01T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("001_initial_schema.postgres.sql").to_string(),
            include_str!("001_initial_schema.mysql.sql").to_string(),
        );

        // Migration 002: Add priority system
        self.register_migration(
            Migration {
                id: "002_add_priority".to_string(),
                description: "Add priority field and indexes for job prioritization".to_string(),
                version: 2,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-02-01T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("002_add_priority.postgres.sql").to_string(),
            include_str!("002_add_priority.mysql.sql").to_string(),
        );

        // Migration 003: Add timeout functionality
        self.register_migration(
            Migration {
                id: "003_add_timeouts".to_string(),
                description: "Add timeout_seconds and timed_out_at fields".to_string(),
                version: 3,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-03-01T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("003_add_timeouts.postgres.sql").to_string(),
            include_str!("003_add_timeouts.mysql.sql").to_string(),
        );

        // Migration 004: Add cron scheduling
        self.register_migration(
            Migration {
                id: "004_add_cron".to_string(),
                description: "Add cron scheduling fields and indexes".to_string(),
                version: 4,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-04-01T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("004_add_cron.postgres.sql").to_string(),
            include_str!("004_add_cron.mysql.sql").to_string(),
        );

        // Migration 005: Add batch processing
        self.register_migration(
            Migration {
                id: "005_add_batches".to_string(),
                description: "Add batch processing table and job batch_id field".to_string(),
                version: 5,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-05-01T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("005_add_batches.postgres.sql").to_string(),
            include_str!("005_add_batches.mysql.sql").to_string(),
        );

        // Migration 006: Add result storage
        self.register_migration(
            Migration {
                id: "006_add_result_storage".to_string(),
                description: "Add result storage fields for job execution results".to_string(),
                version: 6,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-06-01T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("006_add_result_storage.postgres.sql").to_string(),
            include_str!("006_add_result_storage.mysql.sql").to_string(),
        );

        // Migration 007: Add job dependencies for workflow support
        self.register_migration(
            Migration {
                id: "007_add_dependencies".to_string(),
                description: "Add job dependencies and workflow support".to_string(),
                version: 7,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-07-01T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("007_add_dependencies.postgres.sql").to_string(),
            include_str!("007_add_dependencies.mysql.sql").to_string(),
        );

        // Migration 008: Add result configuration storage
        self.register_migration(
            Migration {
                id: "008_add_result_config".to_string(),
                description: "Add result configuration storage fields".to_string(),
                version: 8,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-08-01T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("008_add_result_config.postgres.sql").to_string(),
            include_str!("008_add_result_config.mysql.sql").to_string(),
        );

        // Migration 009: Add job tracing and correlation support
        self.register_migration(
            Migration {
                id: "009_add_tracing".to_string(),
                description: "Add distributed tracing and correlation fields".to_string(),
                version: 9,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-09-01T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("009_add_tracing.postgres.sql").to_string(),
            include_str!("009_add_tracing.mysql.sql").to_string(),
        );

        // Migration 010: Add job archival support and archive table
        self.register_migration(
            Migration {
                id: "010_add_archival".to_string(),
                description: "Add job archival support and archive table".to_string(),
                version: 10,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-10-01T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("010_add_archival.postgres.sql").to_string(),
            include_str!("010_add_archival.mysql.sql").to_string(),
        );

        // Migration 011: Add encryption support for job payloads
        self.register_migration(
            Migration {
                id: "011_add_encryption".to_string(),
                description: "Add job payload encryption and key management".to_string(),
                version: 11,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-11-01T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("011_add_encryption.postgres.sql").to_string(),
            include_str!("011_add_encryption.mysql.sql").to_string(),
        );

        // Migration 012: Optimize dependencies using native PostgreSQL arrays
        self.register_migration(
            Migration {
                id: "012_optimize_dependencies".to_string(),
                description: "Optimize job dependencies using native PostgreSQL UUID arrays"
                    .to_string(),
                version: 12,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-12-01T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("012_optimize_dependencies.postgres.sql").to_string(),
            include_str!("012_optimize_dependencies.mysql.sql").to_string(),
        );

        // Migration 013: Add encryption key audit trails
        self.register_migration(
            Migration {
                id: "013_add_key_audit".to_string(),
                description: "Add encryption key audit trail and operation logging".to_string(),
                version: 13,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-12-15T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("013_add_key_audit.postgres.sql").to_string(),
            include_str!("013_add_key_audit.mysql.sql").to_string(),
        );

        // Migration 014: Add queue pause functionality
        self.register_migration(
            Migration {
                id: "014_add_queue_pause".to_string(),
                description: "Add queue pause and resume functionality for operational control"
                    .to_string(),
                version: 14,
                created_at: chrono::DateTime::parse_from_rfc3339("2025-12-30T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("014_add_queue_pause.postgres.sql").to_string(),
            include_str!("014_add_queue_pause.mysql.sql").to_string(),
        );

        // Migration 015: Add job leases for stale job recovery
        self.register_migration(
            Migration {
                id: "015_add_job_leases".to_string(),
                description: "Add job heartbeat and lease columns for stale job recovery"
                    .to_string(),
                version: 15,
                created_at: chrono::DateTime::parse_from_rfc3339("2026-10-08T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("015_add_job_leases.postgres.sql").to_string(),
            include_str!("015_add_job_leases.mysql.sql").to_string(),
        );

        // Migration 016: Unique encryption keys per (key_id, key_version)
        self.register_migration(
            Migration {
                id: "016_versioned_encryption_keys".to_string(),
                description: "Keep every encryption key version and allow 'Update' audit records"
                    .to_string(),
                version: 16,
                created_at: chrono::DateTime::parse_from_rfc3339("2026-10-08T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("016_versioned_encryption_keys.postgres.sql").to_string(),
            include_str!("016_versioned_encryption_keys.mysql.sql").to_string(),
        );

        // Migration 017: KMS-wrapped data keys for aws:// and gcp:// key sources
        self.register_migration(
            Migration {
                id: "017_add_kms_data_keys".to_string(),
                description: "Persist KMS-wrapped data keys so AWS/GCP KMS keys survive restarts"
                    .to_string(),
                version: 17,
                created_at: chrono::DateTime::parse_from_rfc3339("2026-10-08T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("017_add_kms_data_keys.postgres.sql").to_string(),
            include_str!("017_add_kms_data_keys.mysql.sql").to_string(),
        );

        // Migration 018: Persist per-job retry strategies; index batch progress checks
        self.register_migration(
            Migration {
                id: "018_add_job_retry_strategy".to_string(),
                description: "Store each job's retry strategy and index jobs by batch and status"
                    .to_string(),
                version: 18,
                created_at: chrono::DateTime::parse_from_rfc3339("2026-10-08T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("018_add_job_retry_strategy.postgres.sql").to_string(),
            include_str!("018_add_job_retry_strategy.mysql.sql").to_string(),
        );

        // Migration 019: Keep workflow membership in the archive
        self.register_migration(
            Migration {
                id: "019_add_archive_workflow".to_string(),
                description: "Store workflow_id and workflow_name of archived jobs".to_string(),
                version: 19,
                created_at: chrono::DateTime::parse_from_rfc3339("2026-10-11T00:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("019_add_archive_workflow.postgres.sql").to_string(),
            include_str!("019_add_archive_workflow.mysql.sql").to_string(),
        );

        // Migration 020: Widen archived retry strategies (MySQL), index archived jobs by
        // workflow, keep microseconds in MySQL archive timestamps
        self.register_migration(
            Migration {
                id: "020_archive_retry_strategy_and_workflow_index".to_string(),
                description: "Widen archived retry_strategy and index archived jobs by workflow"
                    .to_string(),
                version: 20,
                created_at: chrono::DateTime::parse_from_rfc3339("2026-10-11T12:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("020_archive_retry_strategy_and_workflow_index.postgres.sql").to_string(),
            include_str!("020_archive_retry_strategy_and_workflow_index.mysql.sql").to_string(),
        );

        // Migration 021: Store MySQL times as DATETIME(6) (lifts the year-2038 limit);
        // a no-op on PostgreSQL
        self.register_migration(
            Migration {
                id: "021_mysql_datetime".to_string(),
                description: "Convert MySQL TIMESTAMP columns to DATETIME(6) (year-2038 limit)"
                    .to_string(),
                version: 21,
                created_at: chrono::DateTime::parse_from_rfc3339("2026-10-11T18:00:00Z")
                    .unwrap()
                    .with_timezone(&Utc),
            },
            include_str!("021_mysql_datetime.postgres.sql").to_string(),
            include_str!("021_mysql_datetime.mysql.sql").to_string(),
        );
    }
}