Stop Babysitting Your PostgreSQL Database - Automate Maintenance with AI in 45 Minutes

# Create isolated environment
python3 -m venv postgres_ai_env
source postgres_ai_env/bin/activate

# Install required packages
pip install psycopg2-binary pandas scikit-learn schedule prometheus-client
pip install openai anthropic  # Choose your preferred AI provider

# database_monitor.py
import psycopg2
import pandas as pd
from datetime import datetime, timedelta
import json
import logging

class PostgreSQLMonitor:
    def __init__(self, connection_params):
        self.conn_params = connection_params
        self.setup_logging()
        
    def setup_logging(self):
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler('postgres_ai.log'),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger(__name__)
    
    def connect(self):
        """Establish database connection with retry logic"""
        try:
            conn = psycopg2.connect(**self.conn_params)
            conn.autocommit = True
            return conn
        except psycopg2.Error as e:
            self.logger.error(f"Database connection failed: {e}")
            raise
    
    def get_table_stats(self):
        """Collect table statistics for AI analysis"""
        query = """
        SELECT 
            schemaname,
            tablename,
            n_live_tup,
            n_dead_tup,
            n_tup_ins + n_tup_upd + n_tup_del as total_operations,
            last_vacuum,
            last_autovacuum,
            last_analyze,
            last_autoanalyze
        FROM pg_stat_user_tables
        ORDER BY n_dead_tup DESC;
        """
        
        with self.connect() as conn:
            df = pd.read_sql_query(query, conn)
            return df
    
    def get_query_performance(self):
        """Extract query performance data"""
        query = """
        SELECT 
            queryid,
            query,
            calls,
            total_exec_time,
            mean_exec_time,
            rows,
            100.0 * shared_blks_hit / 
                nullif(shared_blks_hit + shared_blks_read, 0) AS hit_ratio
        FROM pg_stat_statements 
        WHERE calls > 10
        ORDER BY total_exec_time DESC
        LIMIT 50;
        """
        
        with self.connect() as conn:
            df = pd.read_sql_query(query, conn)
            return df

# ai_analyzer.py
from sklearn.ensemble import IsolationForest
from sklearn.preprocessing import StandardScaler
import numpy as np
import openai  # or anthropic
from typing import List, Dict, Any

class DatabaseAIAnalyzer:
    def __init__(self, ai_api_key: str):
        self.scaler = StandardScaler()
        self.anomaly_detector = IsolationForest(contamination=0.1, random_state=42)
        openai.api_key = ai_api_key
        
    def analyze_table_health(self, table_stats: pd.DataFrame) -> Dict[str, Any]:
        """Use AI to assess table health and recommend actions"""
        # Prepare features for ML analysis
        features = table_stats[['n_live_tup', 'n_dead_tup', 'total_operations']].fillna(0)
        
        # Detect anomalies in table behavior
        scaled_features = self.scaler.fit_transform(features)
        anomalies = self.anomaly_detector.fit_predict(scaled_features)
        
        # Identify tables needing attention
        problem_tables = table_stats[anomalies == -1].copy()
        
        recommendations = []
        for _, table in problem_tables.iterrows():
            # Calculate bloat ratio
            bloat_ratio = table['n_dead_tup'] / (table['n_live_tup'] + 1)
            
            if bloat_ratio > 0.2:  # More than 20% dead tuples
                recommendations.append({
                    'table': f"{table['schemaname']}.{table['tablename']}",
                    'action': 'VACUUM FULL',
                    'priority': 'HIGH' if bloat_ratio > 0.5 else 'MEDIUM',
                    'reason': f"Dead tuple ratio: {bloat_ratio:.2%}"
                })
        
        return {
            'anomalous_tables': len(problem_tables),
            'recommendations': recommendations,
            'health_score': self._calculate_health_score(table_stats)
        }
    
    def generate_index_suggestions(self, query_data: pd.DataFrame) -> List[Dict]:
        """Use AI to suggest optimal indexes"""
        # Find slow queries without good cache hit ratios
        slow_queries = query_data[
            (query_data['mean_exec_time'] > 100) & 
            (query_data['hit_ratio'] < 95)
        ].copy()
        
        suggestions = []
        for _, query in slow_queries.iterrows():
            # Use OpenAI to analyze query pattern and suggest indexes
            prompt = f"""
            Analyze this PostgreSQL query and suggest optimal indexes:
            
            Query: {query['query'][:500]}
            Execution time: {query['mean_exec_time']:.2f}ms
            Hit ratio: {query['hit_ratio']:.1f}%
            Calls: {query['calls']}
            
            Provide specific CREATE INDEX statements that would improve performance.
            Consider composite indexes and partial indexes where appropriate.
            """
            
            try:
                response = openai.ChatCompletion.create(
                    model="gpt-4",
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=300,
                    temperature=0.1
                )
                
                suggestions.append({
                    'query_id': query['queryid'],
                    'current_performance': f"{query['mean_exec_time']:.2f}ms",
                    'ai_suggestion': response.choices[0].message.content.strip(),
                    'priority': 'HIGH' if query['mean_exec_time'] > 1000 else 'MEDIUM'
                })
            except Exception as e:
                self.logger.error(f"AI analysis failed for query {query['queryid']}: {e}")
        
        return suggestions
    
    def _calculate_health_score(self, stats: pd.DataFrame) -> float:
        """Calculate overall database health score (0-100)"""
        total_live = stats['n_live_tup'].sum()
        total_dead = stats['n_dead_tup'].sum()
        
        if total_live == 0:
            return 100.0
        
        bloat_ratio = total_dead / (total_live + total_dead)
        health_score = max(0, 100 - (bloat_ratio * 200))  # Scale bloat impact
        
        return round(health_score, 1)

# maintenance_executor.py
import schedule
import time
from datetime import datetime, timedelta
from typing import List, Dict

class MaintenanceExecutor:
    def __init__(self, monitor: PostgreSQLMonitor, analyzer: DatabaseAIAnalyzer):
        self.monitor = monitor
        self.analyzer = analyzer
        self.maintenance_log = []
        
    def safe_vacuum(self, table_name: str, vacuum_type: str = 'VACUUM') -> bool:
        """Execute VACUUM with safety checks"""
        try:
            # Check if table is currently being heavily used
            activity_query = f"""
            SELECT count(*) as active_queries 
            FROM pg_stat_activity 
            WHERE query ILIKE '%{table_name}%' 
            AND state = 'active';
            """
            
            with self.monitor.connect() as conn:
                with conn.cursor() as cur:
                    cur.execute(activity_query)
                    active_count = cur.fetchone()[0]
                    
                    # Skip if table is busy (more than 2 active queries)
                    if active_count > 2:
                        self.monitor.logger.warning(f"Skipping {vacuum_type} on {table_name} - table busy")
                        return False
                    
                    # Execute maintenance
                    start_time = datetime.now()
                    cur.execute(f"{vacuum_type} {table_name};")
                    duration = (datetime.now() - start_time).total_seconds()
                    
                    self.maintenance_log.append({
                        'timestamp': start_time,
                        'action': vacuum_type,
                        'table': table_name,
                        'duration_seconds': duration,
                        'status': 'SUCCESS'
                    })
                    
                    self.monitor.logger.info(f"{vacuum_type} completed on {table_name} in {duration:.2f}s")
                    return True
                    
        except Exception as e:
            self.monitor.logger.error(f"{vacuum_type} failed on {table_name}: {e}")
            self.maintenance_log.append({
                'timestamp': datetime.now(),
                'action': vacuum_type,
                'table': table_name,
                'status': 'FAILED',
                'error': str(e)
            })
            return False
    
    def create_suggested_index(self, index_sql: str, estimated_benefit: str) -> bool:
        """Safely create AI-suggested indexes"""
        try:
            # Extract index name for logging
            index_name = index_sql.split('INDEX ')[1].split(' ON')[0]
            
            with self.monitor.connect() as conn:
                with conn.cursor() as cur:
                    # Create index concurrently to avoid blocking
                    concurrent_sql = index_sql.replace('CREATE INDEX', 'CREATE INDEX CONCURRENTLY')
                    
                    start_time = datetime.now()
                    cur.execute(concurrent_sql)
                    duration = (datetime.now() - start_time).total_seconds()
                    
                    self.maintenance_log.append({
                        'timestamp': start_time,
                        'action': 'CREATE_INDEX',
                        'target': index_name,
                        'duration_seconds': duration,
                        'estimated_benefit': estimated_benefit,
                        'status': 'SUCCESS'
                    })
                    
                    self.monitor.logger.info(f"Index {index_name} created successfully in {duration:.2f}s")
                    return True
                    
        except Exception as e:
            self.monitor.logger.error(f"Index creation failed: {e}")
            return False
    
    def run_maintenance_cycle(self):
        """Execute complete AI-driven maintenance cycle"""
        self.monitor.logger.info("Starting AI maintenance cycle")
        
        # Collect current database state
        table_stats = self.monitor.get_table_stats()
        query_stats = self.monitor.get_query_performance()
        
        # Get AI analysis
        health_analysis = self.analyzer.analyze_table_health(table_stats)
        index_suggestions = self.analyzer.generate_index_suggestions(query_stats)
        
        self.monitor.logger.info(f"Database health score: {health_analysis['health_score']}/100")
        
        # Execute high-priority maintenance
        for rec in health_analysis['recommendations']:
            if rec['priority'] == 'HIGH':
                self.safe_vacuum(rec['table'], rec['action'])
                time.sleep(30)  # Brief pause between operations
        
        # Implement top index suggestions
        for suggestion in index_suggestions[:3]:  # Limit to top 3 suggestions
            if suggestion['priority'] == 'HIGH':
                # Extract CREATE INDEX statement from AI response
                lines = suggestion['ai_suggestion'].split('\n')
                for line in lines:
                    if line.strip().startswith('CREATE INDEX'):
                        self.create_suggested_index(line.strip(), suggestion['current_performance'])
                        time.sleep(60)  # Longer pause for index creation
                        break
        
        self.monitor.logger.info("AI maintenance cycle completed")
        
        # Generate summary report
        return self.generate_maintenance_report()
    
    def generate_maintenance_report(self) -> Dict:
        """Generate summary of maintenance actions"""
        recent_actions = [
            action for action in self.maintenance_log 
            if action['timestamp'] > datetime.now() - timedelta(hours=24)
        ]
        
        return {
            'timestamp': datetime.now(),
            'actions_performed': len(recent_actions),
            'successful_actions': len([a for a in recent_actions if a['status'] == 'SUCCESS']),
            'total_maintenance_time': sum(
                a.get('duration_seconds', 0) for a in recent_actions
            ),
            'recent_actions': recent_actions[-10:]  # Last 10 actions
        }

# Schedule automated maintenance
def setup_maintenance_schedule(executor: MaintenanceExecutor):
    """Configure automated maintenance schedule"""
    # Run light maintenance every 4 hours
    schedule.every(4).hours.do(executor.run_maintenance_cycle)
    
    # Run comprehensive analysis daily at 2 AM
    schedule.every().day.at("02:00").do(executor.run_maintenance_cycle)
    
    # Weekly deep maintenance on Sundays at 3 AM
    schedule.every().sunday.at("03:00").do(
        lambda: executor.run_maintenance_cycle()
    )
    
    return schedule

# monitoring_dashboard.py
from flask import Flask, render_template, jsonify
import json
from datetime import datetime, timedelta
import threading

app = Flask(__name__)

class MaintenanceDashboard:
    def __init__(self, monitor: PostgreSQLMonitor, executor: MaintenanceExecutor):
        self.monitor = monitor
        self.executor = executor
        
    @app.route('/')
    def dashboard(self):
        return render_template('dashboard.html')
    
    @app.route('/api/health')
    def get_health_metrics(self):
        """API endpoint for current database health"""
        try:
            table_stats = self.monitor.get_table_stats()
            query_stats = self.monitor.get_query_performance()
            
            health_data = {
                'timestamp': datetime.now().isoformat(),
                'total_tables': len(table_stats),
                'tables_needing_vacuum': len(
                    table_stats[table_stats['n_dead_tup'] > table_stats['n_live_tup'] * 0.2]
                ),
                'avg_query_time': query_stats['mean_exec_time'].mean() if not query_stats.empty else 0,
                'slow_queries': len(query_stats[query_stats['mean_exec_time'] > 1000]),
                'maintenance_actions_24h': len([
                    a for a in self.executor.maintenance_log 
                    if a['timestamp'] > datetime.now() - timedelta(hours=24)
                ])
            }
            
            return jsonify(health_data)
            
        except Exception as e:
            return jsonify({'error': str(e)}), 500
    
    @app.route('/api/recent-actions')
    def get_recent_actions(self):
        """API endpoint for recent maintenance actions"""
        recent = [
            action for action in self.executor.maintenance_log[-20:]
            if isinstance(action.get('timestamp'), datetime)
        ]
        
        # Convert datetime objects for JSON serialization
        for action in recent:
            action['timestamp'] = action['timestamp'].isoformat()
        
        return jsonify(recent)

def run_monitoring_server(dashboard: MaintenanceDashboard, port: int = 5000):
    """Run the monitoring dashboard server"""
    app.run(host='0.0.0.0', port=port, debug=False)

# Simple HTML template for dashboard (save as templates/dashboard.html)
dashboard_html = """
<!DOCTYPE html>
<html>
<head>
    <title>PostgreSQL AI Maintenance Dashboard</title>
    <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
    <style>
        body { font-family: Arial, sans-serif; margin: 20px; }
        .metric { display: inline-block; margin: 10px; padding: 15px; 
                 border: 1px solid #ddd; border-radius: 5px; }
        .metric h3 { margin: 0 0 10px 0; color: #333; }
        .metric .value { font-size: 24px; font-weight: bold; color: #007cba; }
        .status-ok { color: #28a745; }
        .status-warning { color: #ffc107; }
        .status-error { color: #dc3545; }
        table { width: 100%; border-collapse: collapse; margin-top: 20px; }
        th, td { padding: 10px; border: 1px solid #ddd; text-align: left; }
        th { background-color: #f8f9fa; }
    </style>
</head>
<body>
    <h1>PostgreSQL AI Maintenance Dashboard</h1>
    
    <div id="metrics">
        <div class="metric">
            <h3>Database Health</h3>
            <div class="value" id="health-score">--</div>
        </div>
        <div class="metric">
            <h3>Tables Needing Attention</h3>
            <div class="value" id="tables-attention">--</div>
        </div>
        <div class="metric">
            <h3>Avg Query Time</h3>
            <div class="value" id="avg-query-time">--ms</div>
        </div>
        <div class="metric">
            <h3>Actions (24h)</h3>
            <div class="value" id="actions-24h">--</div>
        </div>
    </div>
    
    <h2>Recent Maintenance Actions</h2>
    <table id="actions-table">
        <thead>
            <tr>
                <th>Timestamp</th>
                <th>Action</th>
                <th>Target</th>
                <th>Status</th>
                <th>Duration</th>
            </tr>
        </thead>
        <tbody id="actions-body">
        </tbody>
    </table>
    
    <script>
        function updateDashboard() {
            fetch('/api/health')
                .then(response => response.json())
                .then(data => {
                    document.getElementById('health-score').textContent = 
                        data.tables_needing_vacuum === 0 ? '100' : '85';
                    document.getElementById('tables-attention').textContent = data.tables_needing_vacuum;
                    document.getElementById('avg-query-time').textContent = 
                        Math.round(data.avg_query_time) + 'ms';
                    document.getElementById('actions-24h').textContent = data.maintenance_actions_24h;
                });
            
            fetch('/api/recent-actions')
                .then(response => response.json())
                .then(actions => {
                    const tbody = document.getElementById('actions-body');
                    tbody.innerHTML = '';
                    actions.forEach(action => {
                        const row = tbody.insertRow();
                        row.innerHTML = `
                            <td>${new Date(action.timestamp).toLocaleString()}</td>
                            <td>${action.action}</td>
                            <td>${action.table || action.target || 'N/A'}</td>
                            <td class="status-${action.status.toLowerCase()}">${action.status}</td>
                            <td>${action.duration_seconds ? action.duration_seconds.toFixed(2) + 's' : 'N/A'}</td>
                        `;
                    });
                });
        }
        
        // Update dashboard every 30 seconds
        updateDashboard();
        setInterval(updateDashboard, 30000);
    </script>
</body>
</html>
"""

# main.py
import os
import time
import threading
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Configuration
DATABASE_CONFIG = {
    'host': os.getenv('DB_HOST', 'localhost'),
    'database': os.getenv('DB_NAME', 'your_database'),
    'user': os.getenv('DB_USER', 'postgres'),
    'password': os.getenv('DB_PASSWORD'),
    'port': os.getenv('DB_PORT', 5432)
}

OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')

def main():
    """Main application entry point"""
    print("🚀 Starting PostgreSQL AI Maintenance System")
    
    # Initialize components
    monitor = PostgreSQLMonitor(DATABASE_CONFIG)
    analyzer = DatabaseAIAnalyzer(OPENAI_API_KEY)
    executor = MaintenanceExecutor(monitor, analyzer)
    dashboard = MaintenanceDashboard(monitor, executor)
    
    # Set up maintenance schedule
    maintenance_schedule = setup_maintenance_schedule(executor)
    
    # Start monitoring dashboard in separate thread
    dashboard_thread = threading.Thread(
        target=run_monitoring_server,
        args=(dashboard, 5000),
        daemon=True
    )
    dashboard_thread.start()
    
    print("✅ Dashboard started at http://localhost:5000")
    print("✅ AI maintenance scheduler active")
    print("✅ System monitoring enabled")
    
    # Run scheduled maintenance
    try:
        while True:
            maintenance_schedule.run_pending()
            time.sleep(60)  # Check every minute
            
    except KeyboardInterrupt:
        print("\n🛑 Shutting down AI maintenance system")
        print("💾 Maintenance logs saved")

if __name__ == "__main__":
    main()

# Database Configuration
DB_HOST=localhost
DB_NAME=your_production_db
DB_USER=maintenance_user
DB_PASSWORD=secure_password
DB_PORT=5432

# AI Configuration
OPENAI_API_KEY=sk-your-key-here

# Optional: Slack webhook for alerts
SLACK_WEBHOOK_URL=https://hooks.slack.com/your-webhook