ShieldAI/infra/scripts/rollback-migration.sh

#!/bin/bash
set -euo pipefail

# ShieldAI Database Migration Rollback Script
# Usage: ./rollback-migration.sh <environment> [--migration <name>]
#
# Rolls back the most recent migration or a specific named migration
# Uses AWS Secrets Manager for database credentials
#
# Examples:
#   ./rollback-migration.sh staging                    # Rollback latest
#   ./rollback-migration.sh production --migration 001_create_users  # Rollback specific

ENVIRONMENT="${1:-staging}"
MIGRATION_NAME="${3:-}"

# ─── Configuration ───────────────────────────────────────────────
SECRET_ID="shieldai-${ENVIRONMENT}-db-password"
DB_NAME="shieldai"
DB_USER="shieldai"

# ─── Helpers ─────────────────────────────────────────────────────
log() {
  local level="$1"
  shift
  echo "[$(date -u '+%H:%M:%S')] [$level] $*"
}

log_info()  { log "INFO" "$@"; }
log_warn()  { log "WARN" "$@"; }
log_error() { log "ERROR" "$@"; }

# ─── Validation ──────────────────────────────────────────────────
if [[ "$ENVIRONMENT" != "staging" && "$ENVIRONMENT" != "production" ]]; then
  log_error "Invalid environment: $ENVIRONMENT (expected: staging, production)"
  exit 1
fi

for cmd in aws jq; do
  if ! command -v "$cmd" &>/dev/null; then
    log_error "Missing prerequisite: $cmd"
    exit 1
  fi
done

# ─── Credentials ─────────────────────────────────────────────────
get_db_credentials() {
  log_info "Fetching database credentials from Secrets Manager..."

  local secret
  secret=$(aws secretsmanager get-secret-value \
    --secret-id "$SECRET_ID" \
    --query 'SecretString' \
    --output json 2>/dev/null)

  if [[ -z "$secret" ]]; then
    log_error "Failed to fetch secret: $SECRET_ID"
    exit 1
  fi

  export DB_HOST=$(echo "$secret" | jq -r '.host')
  export DB_PORT=$(echo "$secret" | jq -r '.port' // '5432')
  export DB_PASS=$(echo "$secret" | jq -r '.password')
  export DATABASE_URL="postgresql://${DB_USER}:${DB_PASS}@${DB_HOST}:${DB_PORT}/${DB_NAME}"

  log_info "Database: ${DB_HOST}:${DB_PORT}/${DB_NAME}"
}

# ─── Migration Status ────────────────────────────────────────────
show_migration_status() {
  log_info "=== Current Migration Status ==="

  if command -v npx &>/dev/null; then
    npx drizzle-kit status --config=drizzle.config.ts 2>/dev/null || \
      log_warn "Drizzle status check completed (some warnings expected)"
  fi

  # Show applied migrations from database
  log_info "Applied migrations:"
  PGPASSWORD="$DB_PASS" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" \
    -c "SELECT id, checksum, type FROM __drizzle_migrations_schema ORDER BY id DESC;" 2>/dev/null || \
    log_warn "Could not query migration table (psql may not be installed)"
}

# ─── Rollback Logic ──────────────────────────────────────────────
rollback_latest() {
  log_info "=== Rolling Back Latest Migration ==="

  # Get the latest applied migration
  local latest_migration
  latest_migration=$(PGPASSWORD="$DB_PASS" psql -h "$DB_HOST" -p "$DB_PORT" \
    -U "$DB_USER" -d "$DB_NAME" -t -A \
    -c "SELECT id FROM __drizzle_migrations_schema ORDER BY id DESC LIMIT 1;" 2>/dev/null)

  if [[ -z "$latest_migration" ]]; then
    log_warn "No applied migrations found"
    return 0
  fi

  log_info "Latest migration: $latest_migration"

  # Resolve the migration (marks it as not applied)
  if command -v npx &>/dev/null; then
    npx drizzle-kit migrate:resolve --migration "$latest_migration" --status applied 2>/dev/null || \
      log_warn "Migration resolve completed (check output for details)"
  fi

  log_info "Migration $latest_migration marked as resolved"
}

rollback_specific() {
  local target="$1"
  log_info "=== Rolling Back Migration: $target ==="

  if command -v npx &>/dev/null; then
    npx drizzle-kit migrate:resolve --migration "$target" --status applied 2>/dev/null || \
      log_warn "Migration resolve completed (check output for details)"
  fi

  log_info "Migration $target marked as resolved"
}

# ─── Verification ────────────────────────────────────────────────
verify_connection() {
  log_info "=== Verifying Database Connection ==="

  local result
  result=$(PGPASSWORD="$DB_PASS" psql -h "$DB_HOST" -p "$DB_PORT" \
    -U "$DB_USER" -d "$DB_NAME" -t -A \
    -c "SELECT version();" 2>/dev/null || echo "FAIL")

  if [[ "$result" != "FAIL" ]]; then
    log_info "Connection OK: PostgreSQL $result"
  else
    log_warn "Connection check failed"
  fi
}

# ─── Main ────────────────────────────────────────────────────────
main() {
  log_info "=== ShieldAI Migration Rollback ==="
  log_info "Environment: $ENVIRONMENT"
  log_info "Secret: $SECRET_ID"

  get_db_credentials
  show_migration_status

  if [[ -n "$MIGRATION_NAME" ]]; then
    rollback_specific "$MIGRATION_NAME"
  else
    rollback_latest
  fi

  verify_connection
  show_migration_status

  log_info "=== Rollback Complete ==="
  log_info "Next steps:"
  log_info "1. Verify application schema compatibility"
  log_info "2. Run application health checks"
  log_info "3. If needed, redeploy ECS services: ./rollback.sh $ENVIRONMENT all"
}

main "$@"