gt-ai-os-community/scripts/migrations/020_migrate_env_api_keys.py

#!/usr/bin/env python3
"""
Migration 020: Import GROQ_API_KEY from environment to database

Migrates API keys from .env file to encrypted database storage for test-company tenant.
This is part of the move away from environment variables for API keys (#158, #219).

Idempotency: Checks if key already exists before importing
Target: test-company tenant only (as specified in requirements)

Usage:
    python scripts/migrations/020_migrate_env_api_keys.py

Environment variables required:
    - GROQ_API_KEY: The Groq API key to migrate (optional - skips if not set)
    - API_KEY_ENCRYPTION_KEY: Fernet encryption key (auto-generated if not set)
    - CONTROL_PANEL_DB_HOST: Database host (default: localhost)
    - CONTROL_PANEL_DB_PORT: Database port (default: 5432)
    - CONTROL_PANEL_DB_NAME: Database name (default: gt2_admin)
    - CONTROL_PANEL_DB_USER: Database user (default: postgres)
    - ADMIN_POSTGRES_PASSWORD: Database password
"""
import os
import sys
import json
import logging
from datetime import datetime

try:
    from cryptography.fernet import Fernet
    import psycopg2
except ImportError as e:
    print(f"Missing required package: {e}")
    print("Run: pip install cryptography psycopg2-binary")
    sys.exit(1)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Migration constants
TARGET_TENANT_DOMAIN = "test-company"
PROVIDER = "groq"
MIGRATION_ID = "020"


def get_db_connection():
    """Get database connection using environment variables or defaults"""
    try:
        conn = psycopg2.connect(
            host=os.getenv("CONTROL_PANEL_DB_HOST", "localhost"),
            port=os.getenv("CONTROL_PANEL_DB_PORT", "5432"),
            database=os.getenv("CONTROL_PANEL_DB_NAME", "gt2_admin"),
            user=os.getenv("CONTROL_PANEL_DB_USER", "postgres"),
            password=os.getenv("ADMIN_POSTGRES_PASSWORD", "dev_password_change_in_prod")
        )
        return conn
    except psycopg2.Error as e:
        logger.error(f"Database connection failed: {e}")
        raise


def get_encryption_key() -> str:
    """Get or generate Fernet encryption key"""
    key = os.getenv("API_KEY_ENCRYPTION_KEY")
    if not key:
        # Generate a new key - in production this should be persisted
        key = Fernet.generate_key().decode()
        logger.warning("Generated new API_KEY_ENCRYPTION_KEY - add to .env for persistence:")
        logger.warning(f"  API_KEY_ENCRYPTION_KEY={key}")
    return key


def check_env_key_exists() -> str | None:
    """Check if GROQ_API_KEY environment variable exists and is valid"""
    groq_key = os.getenv("GROQ_API_KEY")

    # Skip placeholder values
    placeholder_values = [
        "gsk_your_actual_groq_api_key_here",
        "gsk_placeholder",
        "",
        None
    ]

    if groq_key in placeholder_values:
        logger.info("GROQ_API_KEY not set or is placeholder - skipping migration")
        return None

    # Validate format
    if not groq_key.startswith("gsk_"):
        logger.warning(f"GROQ_API_KEY has invalid format (should start with 'gsk_')")
        return None

    return groq_key


def get_tenant_id(conn, domain: str) -> int | None:
    """Get tenant ID by domain"""
    with conn.cursor() as cur:
        cur.execute(
            "SELECT id FROM tenants WHERE domain = %s AND deleted_at IS NULL",
            (domain,)
        )
        row = cur.fetchone()
        return row[0] if row else None


def check_db_key_exists(conn, tenant_id: int) -> bool:
    """Check if Groq key already exists in database for tenant"""
    with conn.cursor() as cur:
        cur.execute(
            "SELECT api_keys FROM tenants WHERE id = %s",
            (tenant_id,)
        )
        row = cur.fetchone()
        if row and row[0]:
            api_keys = row[0] if isinstance(row[0], dict) else json.loads(row[0])
            if PROVIDER in api_keys and api_keys[PROVIDER].get("key"):
                return True
    return False


def migrate_api_key(conn, tenant_id: int, api_key: str, encryption_key: str) -> bool:
    """Encrypt and store API key in database"""
    try:
        cipher = Fernet(encryption_key.encode())
        encrypted_key = cipher.encrypt(api_key.encode()).decode()

        api_keys_data = {
            PROVIDER: {
                "key": encrypted_key,
                "secret": None,
                "enabled": True,
                "metadata": {
                    "migrated_from": "environment",
                    "migration_id": MIGRATION_ID,
                    "migration_date": datetime.utcnow().isoformat()
                },
                "updated_at": datetime.utcnow().isoformat(),
                "updated_by": f"migration-{MIGRATION_ID}"
            }
        }

        with conn.cursor() as cur:
            cur.execute(
                """
                UPDATE tenants
                SET api_keys = %s::jsonb,
                    api_key_encryption_version = 'v1',
                    updated_at = NOW()
                WHERE id = %s
                """,
                (json.dumps(api_keys_data), tenant_id)
            )
        conn.commit()
        return True

    except Exception as e:
        conn.rollback()
        logger.error(f"Failed to migrate API key: {e}")
        return False


def run_migration() -> bool:
    """Main migration logic"""
    logger.info(f"=== Migration {MIGRATION_ID}: Import GROQ_API_KEY from environment ===")

    # Step 1: Check if env var exists
    groq_key = check_env_key_exists()
    if not groq_key:
        logger.info("Migration skipped: No valid GROQ_API_KEY in environment")
        return True  # Not an error - just nothing to migrate

    logger.info(f"Found GROQ_API_KEY in environment (length: {len(groq_key)})")

    # Step 2: Connect to database
    try:
        conn = get_db_connection()
        logger.info("Connected to database")
    except Exception as e:
        logger.error(f"Failed to connect to database: {e}")
        return False

    try:
        # Step 3: Get tenant ID
        tenant_id = get_tenant_id(conn, TARGET_TENANT_DOMAIN)
        if not tenant_id:
            logger.warning(f"Tenant '{TARGET_TENANT_DOMAIN}' not found - skipping migration")
            logger.info("This is expected for fresh installs before tenant creation")
            return True

        logger.info(f"Found tenant '{TARGET_TENANT_DOMAIN}' with ID: {tenant_id}")

        # Step 4: Check if DB key already exists (idempotency)
        if check_db_key_exists(conn, tenant_id):
            logger.info("Migration already complete - Groq key exists in database")
            return True

        # Step 5: Get/generate encryption key
        encryption_key = get_encryption_key()

        # Step 6: Migrate the key
        logger.info(f"Migrating GROQ_API_KEY to database for tenant {tenant_id}...")
        if migrate_api_key(conn, tenant_id, groq_key, encryption_key):
            logger.info(f"=== Migration {MIGRATION_ID} completed successfully ===")
            logger.info("The GROQ_API_KEY env var can now be removed from docker-compose.yml")
            return True
        else:
            logger.error(f"Migration {MIGRATION_ID} failed")
            return False

    except Exception as e:
        logger.error(f"Migration failed with error: {e}")
        return False
    finally:
        conn.close()


if __name__ == "__main__":
    success = run_migration()
    sys.exit(0 if success else 1)