scripts/setup_timescaledb.py

#!/usr/bin/env python3
"""
setup_timescaledb.py

Initializes TimescaleDB tables and extensions for MorphGuard metrics collection.
This script should be run once to set up the database schema.

By default, it uses the database connection settings from config.py.
You can override these settings with command-line arguments.
"""

import os
import sys
import argparse
import psycopg2
from psycopg2 import sql
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Add project root to Python path for imports
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

def create_api_keys_table(cursor):
    """Create the api_keys table if it doesn't exist."""
    print("Checking api_keys table...")
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS api_keys (
            id SERIAL PRIMARY KEY,
            key_hash TEXT UNIQUE NOT NULL,
            name TEXT,
            created_at TIMESTAMPTZ DEFAULT NOW(),
            is_active BOOLEAN DEFAULT TRUE,
            permissions TEXT[]
        );
    """)
    print("api_keys table checked/created.")

def create_metrics_tables(host, port, dbname, user, password):
    """
    Create TimescaleDB tables and indexes for metrics collection
    """
    conn = None
    try:
        # Connect to the database
        print(f"Connecting to database {dbname} on {host}:{port} as {user}...")
        # Build connection parameters; use socket for localhost to avoid TCP hangs
        # Build connection parameters (always use host so we get MD5 auth)
        conn_params = {
            'host': host,
            'port': port,
            'dbname': dbname,
            'user': user,
            'password': password,
            'connect_timeout': 10,
            'sslmode': 'disable',  # Disable SSL for development
        }
        conn = psycopg2.connect(**conn_params)
        cursor = conn.cursor()
        conn.autocommit = True
        # Prevent DDL from hanging indefinitely: set short lock & statement timeouts
        try:
            cursor.execute("SET lock_timeout = '5s';")
            cursor.execute("SET statement_timeout = '5min';")
        except psycopg2.Error:
            # ignore if not supported
            pass
        
        # Migration/version-control: track schema version to avoid repeated setup
        TARGET_SCHEMA_VERSION = 7  # bump version to add cloud connector metrics tables
        # Ensure versioning table exists
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS schema_version (
            version INTEGER NOT NULL,
            updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
        );
        """)
        # Read current version
        cursor.execute("SELECT version FROM schema_version ORDER BY updated_at DESC LIMIT 1;")
        row = cursor.fetchone()
        current_version = row[0] if row else 0
        if current_version >= TARGET_SCHEMA_VERSION:
            print(f"Schema is already at version {current_version}, skipping setup.")
            return
        # Otherwise proceed with one-time schema creation/migration
        
        # Check if TimescaleDB extension exists
        cursor.execute("SELECT extname FROM pg_extension WHERE extname = 'timescaledb';")
        if cursor.fetchone() is None:
            print("Enabling TimescaleDB extension...")
            cursor.execute("CREATE EXTENSION IF NOT EXISTS timescaledb;")
            print("TimescaleDB extension enabled successfully.")
        else:
            print("TimescaleDB extension already exists.")
        
        # Create system metrics table
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS device_metrics (
            time TIMESTAMPTZ NOT NULL,
            cpu_percent FLOAT,
            memory_percent FLOAT
        );
        """)
        
        # Convert to hypertable
        try:
            cursor.execute("SELECT create_hypertable('device_metrics', 'time', if_not_exists => TRUE);")
            print("Created device_metrics hypertable.")
        except psycopg2.Error as e:
            if "already a hypertable" in str(e):
                print("device_metrics is already a hypertable.")
            else:
                raise
        
        # Create GPU metrics table
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS gpu_metrics (
            time TIMESTAMPTZ NOT NULL,
            memory_used_mb FLOAT,
            utilization FLOAT,
            temperature_c FLOAT
        );
        """)
        
        # Convert to hypertable
        try:
            cursor.execute("SELECT create_hypertable('gpu_metrics', 'time', if_not_exists => TRUE);")
            print("Created gpu_metrics hypertable.")
        except psycopg2.Error as e:
            if "already a hypertable" in str(e):
                print("gpu_metrics is already a hypertable.")
            else:
                raise
        
        # Create detection metrics table
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS detection_metrics (
            time TIMESTAMPTZ NOT NULL,
            model_name TEXT,
            confidence_score FLOAT,
            is_morphed BOOLEAN,
            processing_time_ms FLOAT,
            image_hash TEXT,
            face_count INTEGER,
            forensic_ticks INTEGER,
            reasoning_path TEXT,
            detection_tier TEXT,
            request_id TEXT
        );
        """)
        
        # Convert to hypertable
        try:
            cursor.execute("SELECT create_hypertable('detection_metrics', 'time', if_not_exists => TRUE);")
            print("Created detection_metrics hypertable.")
        except psycopg2.Error as e:
            if "already a hypertable" in str(e):
                print("detection_metrics is already a hypertable.")
            else:
                raise
        
        # Create demorph metrics table
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS demorph_metrics (
            time TIMESTAMPTZ NOT NULL,
            model_name TEXT,
            processing_time_ms FLOAT,
            original_image_hash TEXT,
            result_image_hash TEXT,
            method TEXT,
            success BOOLEAN,
            request_id TEXT
        );
        """)
        
        # Convert to hypertable
        try:
            cursor.execute("SELECT create_hypertable('demorph_metrics', 'time', if_not_exists => TRUE);")
            print("Created demorph_metrics hypertable.")
        except psycopg2.Error as e:
            if "already a hypertable" in str(e):
                print("demorph_metrics is already a hypertable.")
            else:
                raise
        
        # Create liveness metrics table
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS liveness_metrics (
            time TIMESTAMPTZ NOT NULL,
            is_live BOOLEAN,
            confidence FLOAT,
            attack_type TEXT,
            processing_time_ms FLOAT,
            image_hash TEXT,
            face_count INTEGER,
            request_id TEXT
        );
        """)
        
        # Convert to hypertable
        try:
            cursor.execute("SELECT create_hypertable('liveness_metrics', 'time', if_not_exists => TRUE);")
            print("Created liveness_metrics hypertable.")
        except psycopg2.Error as e:
            if "already a hypertable" in str(e):
                print("liveness_metrics is already a hypertable.")
            else:
                raise
                
        # Create training metrics table
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS training_metrics (
            time TIMESTAMPTZ NOT NULL,
            model_name TEXT,
            epoch INTEGER,
            loss FLOAT,
            accuracy FLOAT,
            val_loss FLOAT,
            val_accuracy FLOAT,
            learning_rate FLOAT,
            batch_size INTEGER,
            training_session_id TEXT
        );
        """)
        
        # Convert to hypertable
        try:
            cursor.execute("SELECT create_hypertable('training_metrics', 'time', if_not_exists => TRUE);")
            print("Created training_metrics hypertable.")
        except psycopg2.Error as e:
            if "already a hypertable" in str(e):
                print("training_metrics is already a hypertable.")
            else:
                raise

        # Apply in-place migrations: ensure new columns exist for index creation
        print("Applying in-place migrations to add missing columns...")
        try:
            cursor.execute("ALTER TABLE detection_metrics ADD COLUMN IF NOT EXISTS model_name TEXT;")
            cursor.execute("ALTER TABLE detection_metrics ADD COLUMN IF NOT EXISTS forensic_ticks INTEGER;")
            cursor.execute("ALTER TABLE detection_metrics ADD COLUMN IF NOT EXISTS reasoning_path TEXT;")
            cursor.execute("ALTER TABLE detection_metrics ADD COLUMN IF NOT EXISTS detection_tier TEXT;")
            cursor.execute("ALTER TABLE demorph_metrics ADD COLUMN IF NOT EXISTS model_name TEXT;")
            cursor.execute("ALTER TABLE training_metrics ADD COLUMN IF NOT EXISTS model_name TEXT;")
        except psycopg2.Error as e:
            print(f"Warning: Could not apply migrations: {e}")
        # Create indexes for faster queries
        print("Creating indexes...")
        try:
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_device_metrics_time ON device_metrics (time DESC);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_gpu_metrics_time ON gpu_metrics (time DESC);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_detection_metrics_time_model ON detection_metrics (time DESC, model_name);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_detection_metrics_is_morphed ON detection_metrics (is_morphed, time DESC);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_detection_metrics_request_id ON detection_metrics (request_id);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_demorph_metrics_time_model ON demorph_metrics (time DESC, model_name);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_demorph_metrics_method ON demorph_metrics (method, time DESC);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_demorph_metrics_request_id ON demorph_metrics (request_id);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_liveness_metrics_time ON liveness_metrics (time DESC);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_liveness_metrics_is_live ON liveness_metrics (is_live, time DESC);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_liveness_metrics_attack_type ON liveness_metrics (attack_type);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_training_metrics_model_time ON training_metrics (model_name, time DESC);")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_training_metrics_session_epoch ON training_metrics (training_session_id, epoch);")
        except psycopg2.Error as e:
            print(f"Warning: Could not create indexes: {e}")
        
        # Create retention policies (keep data for 30/90 days)
        print("Setting up retention policies...")
        try:
            # System metrics - retain for 30 days
            cursor.execute("SELECT add_retention_policy('device_metrics', INTERVAL '30 days', if_not_exists => TRUE);")
            cursor.execute("SELECT add_retention_policy('gpu_metrics', INTERVAL '30 days', if_not_exists => TRUE);")
            
            # Model metrics - retain for 90 days
            cursor.execute("SELECT add_retention_policy('detection_metrics', INTERVAL '90 days', if_not_exists => TRUE);")
            cursor.execute("SELECT add_retention_policy('demorph_metrics', INTERVAL '90 days', if_not_exists => TRUE);")
            cursor.execute("SELECT add_retention_policy('liveness_metrics', INTERVAL '90 days', if_not_exists => TRUE);")
            
            # Training metrics - retain for 365 days (keep longer for historical analysis)
            cursor.execute("SELECT add_retention_policy('training_metrics', INTERVAL '365 days', if_not_exists => TRUE);")
            
            # Cloud connector metrics - retain for 90 days
            cursor.execute("SELECT add_retention_policy('cloud_connector_metrics', INTERVAL '90 days', if_not_exists => TRUE);")
            
            # Verification events - retain for 1 year (important for audit trail)
            cursor.execute("SELECT add_retention_policy('verification_events', INTERVAL '365 days', if_not_exists => TRUE);")
        except psycopg2.Error as e:
            print(f"Warning: Could not set retention policies: {e}")
            print("You may need to manually configure retention policies or upgrade TimescaleDB.")
        # Create model interactions table for AI model mapping
        print("Creating model_interactions table...")
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS model_interactions (
            time TIMESTAMPTZ NOT NULL,
            source_model TEXT NOT NULL,
            dest_model TEXT NOT NULL,
            value BIGINT NOT NULL DEFAULT 1,
            details JSONB
        );
        """)
        # Convert to hypertable
        try:
            cursor.execute("SELECT create_hypertable('model_interactions', 'time', if_not_exists => TRUE);")
            print("Created model_interactions hypertable.")
        except psycopg2.Error as e:
            if "already a hypertable" in str(e):
                print("model_interactions is already a hypertable.")
            else:
                raise
        # Create indexes
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_model_interactions_time ON model_interactions (time DESC);")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_model_interactions_source_dest ON model_interactions (source_model, dest_model);")
        
        # Create cloud connector metrics table
        print("Creating cloud_connector_metrics table...")
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS cloud_connector_metrics (
            time TIMESTAMPTZ NOT NULL,
            provider TEXT NOT NULL,
            operation TEXT NOT NULL,
            success BOOLEAN NOT NULL,
            response_time_ms FLOAT,
            cost_estimate FLOAT,
            api_calls_used INTEGER,
            free_tier_remaining INTEGER,
            error_message TEXT,
            request_id TEXT,
            confidence_score FLOAT,
            details JSONB
        );
        """)
        
        # Convert to hypertable
        try:
            cursor.execute("SELECT create_hypertable('cloud_connector_metrics', 'time', if_not_exists => TRUE);")
            print("Created cloud_connector_metrics hypertable.")
        except psycopg2.Error as e:
            if "already a hypertable" in str(e):
                print("cloud_connector_metrics is already a hypertable.")
            else:
                raise
        
        # Create users table (migrating from SQLite)
        print("Creating users table...")
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS users (
            id SERIAL PRIMARY KEY,
            username TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            role TEXT NOT NULL DEFAULT 'user',
            email TEXT,
            created_at TIMESTAMPTZ DEFAULT NOW(),
            last_login TIMESTAMPTZ,
            active BOOLEAN DEFAULT TRUE
        );
        """)
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_users_username ON users (username);")
        
        # Create API Keys table
        create_api_keys_table(cursor)
        
        # Create indexes for cloud metrics
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_cloud_metrics_time ON cloud_connector_metrics (time DESC);")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_cloud_metrics_provider ON cloud_connector_metrics (provider, time DESC);")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_cloud_metrics_operation ON cloud_connector_metrics (operation, time DESC);")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_cloud_metrics_success ON cloud_connector_metrics (success, time DESC);")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_cloud_metrics_request_id ON cloud_connector_metrics (request_id);")
        
        # Create verification events table for blockchain logging
        print("Creating verification_events table...")
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS verification_events (
            time TIMESTAMPTZ NOT NULL,
            verification_type TEXT NOT NULL,
            subject_hash TEXT NOT NULL,
            result TEXT NOT NULL,
            confidence FLOAT,
            provider TEXT,
            blockchain_tx_hash TEXT,
            verification_id TEXT,
            metadata JSONB
        );
        """)
        
        # Convert to hypertable
        try:
            cursor.execute("SELECT create_hypertable('verification_events', 'time', if_not_exists => TRUE);")
            print("Created verification_events hypertable.")
        except psycopg2.Error as e:
            if "already a hypertable" in str(e):
                print("verification_events is already a hypertable.")
            else:
                raise
        
        # Create indexes for verification events
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_verification_events_time ON verification_events (time DESC);")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_verification_events_type ON verification_events (verification_type, time DESC);")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_verification_events_result ON verification_events (result, time DESC);")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_verification_events_blockchain ON verification_events (blockchain_tx_hash);")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_verification_events_verification_id ON verification_events (verification_id);")
        
        # Create inference metrics table for logging inference time
        print("Creating inference_metrics table...")
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS inference_metrics (
            time TIMESTAMPTZ NOT NULL,
            model_name TEXT,
            duration_ms FLOAT,
            confidence FLOAT,
            face_count INTEGER,
            request_id TEXT
        );
        """)
        try:
            cursor.execute("SELECT create_hypertable('inference_metrics', 'time', if_not_exists => TRUE);")
            print("Created inference_metrics hypertable.")
        except psycopg2.Error as e:
            if "already a hypertable" in str(e):
                print("inference_metrics is already a hypertable.")
            else:
                raise
        
        # Create a unified 'metrics' view for Grafana dashboards
        print("Creating unified metrics view...")
        try:
            cursor.execute("""
        CREATE OR REPLACE VIEW metrics AS
        -- Training metrics: accuracy and loss
        SELECT time,
               'accuracy' AS metric_name,
               accuracy AS value,
               CASE WHEN model_name = 'morph_detector' THEN 'detector'
                    WHEN model_name = 'demorpher' THEN 'demorpher'
                    ELSE model_name END AS model_type
        FROM training_metrics
        UNION ALL
        SELECT time,
               'loss',
               loss,
               CASE WHEN model_name = 'morph_detector' THEN 'detector'
                    WHEN model_name = 'demorpher' THEN 'demorpher'
                    ELSE model_name END
        FROM training_metrics
        UNION ALL
        -- GPU utilization
        SELECT time,
               'gpu_usage',
               utilization AS value,
               'system' AS model_type
        FROM gpu_metrics
        UNION ALL
        -- Inference time
        SELECT time,
               'inference_time',
               duration_ms AS value,
               CASE WHEN model_name = 'morph_detector' THEN 'detector' ELSE model_name END
        FROM inference_metrics
        UNION ALL
        -- Cloud connector response times
        SELECT time,
               'cloud_response_time',
               response_time_ms AS value,
               provider AS model_type
        FROM cloud_connector_metrics
        WHERE response_time_ms IS NOT NULL
        UNION ALL
        -- Cloud API usage
        SELECT time,
               'cloud_api_calls',
               api_calls_used AS value,
               provider AS model_type
        FROM cloud_connector_metrics
        WHERE api_calls_used IS NOT NULL
        UNION ALL
        -- Cloud success rate
        SELECT time,
               'cloud_success_rate',
               CASE WHEN success THEN 1.0 ELSE 0.0 END AS value,
               provider AS model_type
        FROM cloud_connector_metrics;
        """)
        except Exception as e:
            print(f"Warning: could not create unified metrics view: {e}")
        
        conn.commit()
        print("Database setup completed (schema version deployed).")
        # Record completion        # Update schema version
        cursor.execute("INSERT INTO schema_version (version) VALUES (%s);", (TARGET_SCHEMA_VERSION,))
        print(f"Database setup completed (schema version {TARGET_SCHEMA_VERSION}).")
        
    except Exception as e:
        print(f"Error: {e}")
        if conn:
            conn.rollback()
    finally:
        if conn:
            conn.close()

def main():
    # Try to import config.py for default database settings
    try:
        sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
        import config
        default_host = getattr(config, 'DB_HOST', 'localhost')
        default_port = getattr(config, 'DB_PORT', 5432)
        default_dbname = getattr(config, 'DB_NAME', 'morphguard')
        default_user = getattr(config, 'DB_USER', 'morphguard')
        default_password = getattr(config, 'DB_PASSWORD', 'morphguard')
    except ImportError:
        # Use defaults if config.py not found
        default_host = 'localhost'
        default_port = 5432
        default_dbname = 'morphguard'
        default_user = 'morphguard'
        default_password = 'morphguard'
    
    # Parse command-line arguments
    parser = argparse.ArgumentParser(description='Set up TimescaleDB tables for MorphGuard')
    parser.add_argument('--host', default=default_host, help=f'Database host (default: {default_host})')
    parser.add_argument('--port', type=int, default=default_port, help=f'Database port (default: {default_port})')
    parser.add_argument('--dbname', default=default_dbname, help=f'Database name (default: {default_dbname})')
    parser.add_argument('--user', default=default_user, help=f'Database user (default: {default_user})')
    parser.add_argument('--password', default=default_password, help=f'Database password (default: {default_password})')
    
    args = parser.parse_args()
    
    # Create tables
    create_metrics_tables(
        host=args.host,
        port=args.port,
        dbname=args.dbname,
        user=args.user,
        password=args.password
    )

def initialize_schema():
    """
    Initialize TimescaleDB schema programmatically using default settings from config.py.
    """
    try:
        import config
        host = getattr(config, 'DB_HOST', 'localhost')
        port = getattr(config, 'DB_PORT', 5432)
        dbname = getattr(config, 'DB_NAME', 'morphguard')
        user = getattr(config, 'DB_USER', 'morphguard')
        password = getattr(config, 'DB_PASSWORD', 'morphguard')
    except ImportError:
        host = 'localhost'
        port = 5432
        dbname = 'morphguard'
        user = 'morphguard'
        password = 'morphguard'
    create_metrics_tables(host=host, port=port, dbname=dbname, user=user, password=password)

if __name__ == "__main__":
    main()