database.py

#!/usr/bin/env python3
"""
FhirFlame PostgreSQL Database Manager
Handles persistent storage for job tracking, processing history, and system metrics
Uses the existing PostgreSQL database from the Langfuse infrastructure
"""

import sqlite3
import json
import time
import os
from datetime import datetime
from typing import Dict, List, Any, Optional

try:
    import psycopg2
    import psycopg2.extras
except ImportError:
    psycopg2 = None

class DatabaseManager:
    """
    PostgreSQL database manager for FhirFlame job tracking and processing history
    Connects to the existing langfuse-db PostgreSQL instance
    """
    
    def __init__(self):
        self.db_config = {
            'host': os.getenv('DB_HOST', 'langfuse-db'),
            'port': int(os.getenv('DB_PORT', 5432)),
            'database': os.getenv('DB_NAME', 'langfuse'),
            'user': os.getenv('DB_USER', 'langfuse'),
            'password': os.getenv('DB_PASSWORD', 'langfuse'),
            'connect_timeout': int(os.getenv('DB_CONNECT_TIMEOUT', '3'))
        }
        backend = os.getenv('DB_BACKEND', 'auto').lower()
        self._use_sqlite = backend == 'sqlite' or (
            backend == 'auto' and bool(os.getenv('SPACE_ID'))
        )
        self.init_database()

    def get_connection(self):
        """Get PostgreSQL connection with proper configuration, fallback to SQLite"""
        if self._use_sqlite:
            return self._get_sqlite_connection()

        if psycopg2 is None:
            print("ℹ️ psycopg2 is unavailable; using SQLite")
            self._use_sqlite = True
            return self._get_sqlite_connection()

        try:
            return psycopg2.connect(**self.db_config)
        except Exception as e:
            print(f"❌ Database connection failed: {e}")
            self._use_sqlite = True
            return self._get_sqlite_connection()

    def _get_sqlite_connection(self):
        sqlite_path = os.getenv('SQLITE_DB_PATH', 'fhirflame_fallback.db')
        conn = sqlite3.connect(sqlite_path)
        conn.row_factory = sqlite3.Row
        return conn

    @staticmethod
    def _is_sqlite(conn) -> bool:
        return isinstance(conn, sqlite3.Connection)

    def _table(self, conn, name: str) -> str:
        return name if self._is_sqlite(conn) else f"fhirflame.{name}"

    def _cursor(self, conn, dict_rows: bool = False):
        if self._is_sqlite(conn):
            return conn.cursor()
        if dict_rows:
            return conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        return conn.cursor()

    def _sql(self, conn, query: str) -> str:
        return query.replace('%s', '?') if self._is_sqlite(conn) else query

    def init_database(self):
        """Initialize database schema with proper tables and indexes"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()

            if not self._is_sqlite(conn):
                cursor.execute('CREATE SCHEMA IF NOT EXISTS fhirflame')

            jobs_table = self._table(conn, 'jobs')
            batch_jobs_table = self._table(conn, 'batch_jobs')
            json_type = 'TEXT' if self._is_sqlite(conn) else 'JSONB'
            
            # Create jobs table with comprehensive tracking
            cursor.execute(f'''
                CREATE TABLE IF NOT EXISTS {jobs_table} (
                    id VARCHAR(255) PRIMARY KEY,
                    job_type VARCHAR(50) NOT NULL,
                    name TEXT NOT NULL,
                    text_input TEXT,
                    status VARCHAR(20) NOT NULL DEFAULT 'pending',
                    provider_used VARCHAR(50),
                    success BOOLEAN,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    completed_at TIMESTAMP,
                    processing_time VARCHAR(50),
                    entities_found INTEGER,
                    error_message TEXT,
                    result_data {json_type},
                    file_path TEXT,
                    batch_id VARCHAR(255),
                    workflow_type VARCHAR(50)
                )
            ''')
            
            # Create batch jobs table
            cursor.execute(f'''
                CREATE TABLE IF NOT EXISTS {batch_jobs_table} (
                    id VARCHAR(255) PRIMARY KEY,
                    workflow_type VARCHAR(50) NOT NULL,
                    status VARCHAR(20) NOT NULL DEFAULT 'pending',
                    batch_size INTEGER DEFAULT 0,
                    processed_count INTEGER DEFAULT 0,
                    success_count INTEGER DEFAULT 0,
                    failed_count INTEGER DEFAULT 0,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    completed_at TIMESTAMP
                )
            ''')
            
            # Create indexes for performance
            cursor.execute(f'CREATE INDEX IF NOT EXISTS idx_fhirflame_jobs_status ON {jobs_table}(status)')
            cursor.execute(f'CREATE INDEX IF NOT EXISTS idx_fhirflame_jobs_created_at ON {jobs_table}(created_at)')
            cursor.execute(f'CREATE INDEX IF NOT EXISTS idx_fhirflame_jobs_job_type ON {jobs_table}(job_type)')
            cursor.execute(f'CREATE INDEX IF NOT EXISTS idx_fhirflame_batch_jobs_status ON {batch_jobs_table}(status)')
            
            if not self._is_sqlite(conn):
                cursor.execute('''
                    CREATE OR REPLACE FUNCTION fhirflame.update_updated_at_column()
                    RETURNS TRIGGER AS $$
                    BEGIN
                        NEW.updated_at = CURRENT_TIMESTAMP;
                        RETURN NEW;
                    END;
                    $$ language 'plpgsql'
                ''')
                cursor.execute('''
                    DROP TRIGGER IF EXISTS update_fhirflame_jobs_updated_at ON fhirflame.jobs
                ''')
                cursor.execute('''
                    CREATE TRIGGER update_fhirflame_jobs_updated_at
                    BEFORE UPDATE ON fhirflame.jobs
                    FOR EACH ROW
                    EXECUTE FUNCTION fhirflame.update_updated_at_column()
                ''')
            
            conn.commit()
            cursor.close()
            conn.close()
            backend = 'SQLite' if self._use_sqlite else 'PostgreSQL'
            print(f"✅ {backend} database initialized")
            
        except Exception as e:
            print(f"❌ Database initialization failed: {e}")
            # Don't raise - allow app to continue with in-memory fallback
    
    def add_job(self, job_data: Dict[str, Any]) -> bool:
        """Add a new job to the configured database."""
        try:
            conn = self.get_connection()
            cursor = self._cursor(conn)
            jobs_table = self._table(conn, 'jobs')
            
            # Ensure required fields
            job_id = job_data.get('id', f"job_{int(time.time())}")
            job_type = job_data.get('job_type', 'text')
            name = job_data.get('name', 'Unknown Job')
            status = job_data.get('status', 'pending')
            
            query = f'''
                INSERT INTO {jobs_table} (
                    id, job_type, name, text_input, status, provider_used,
                    success, processing_time, entities_found, error_message,
                    result_data, file_path, batch_id, workflow_type
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                ON CONFLICT (id) DO UPDATE SET
                    status = EXCLUDED.status,
                    updated_at = CURRENT_TIMESTAMP
            '''
            cursor.execute(self._sql(conn, query), (
                job_id,
                job_type,
                name,
                job_data.get('text_input'),
                status,
                job_data.get('provider_used'),
                job_data.get('success'),
                job_data.get('processing_time'),
                job_data.get('entities_found'),
                job_data.get('error_message'),
                json.dumps(job_data.get('result_data')) if job_data.get('result_data') else None,
                job_data.get('file_path'),
                job_data.get('batch_id'),
                job_data.get('workflow_type')
            ))
            
            conn.commit()
            cursor.close()
            conn.close()
            print(f"✅ Job added to database: {job_id}")
            return True
            
        except Exception as e:
            print(f"❌ Failed to add job to PostgreSQL database: {e}")
            return False
    
    def update_job(self, job_id: str, updates: Dict[str, Any]) -> bool:
        """Update an existing job in the configured database."""
        try:
            conn = self.get_connection()
            cursor = self._cursor(conn)
            jobs_table = self._table(conn, 'jobs')
            
            # Build update query dynamically
            update_fields = []
            values = []
            
            for field, value in updates.items():
                if field in ['status', 'provider_used', 'success', 'processing_time', 
                           'entities_found', 'error_message', 'result_data', 'completed_at']:
                    update_fields.append(f"{field} = %s")
                    if field == 'result_data' and value is not None:
                        values.append(json.dumps(value))
                    else:
                        values.append(value)
            
            if update_fields:
                values.append(job_id)
                
                update_fields.append("updated_at = CURRENT_TIMESTAMP")
                query = f"UPDATE {jobs_table} SET {', '.join(update_fields)} WHERE id = %s"
                cursor.execute(self._sql(conn, query), values)
                
                conn.commit()
                cursor.close()
                conn.close()
                print(f"✅ Job updated in database: {job_id}")
                return True

            cursor.close()
            conn.close()
            return False
            
        except Exception as e:
            print(f"❌ Failed to update job in PostgreSQL database: {e}")
            return False
    
    def get_job(self, job_id: str) -> Optional[Dict[str, Any]]:
        """Get a specific job from the configured database."""
        try:
            conn = self.get_connection()
            cursor = self._cursor(conn, dict_rows=True)
            jobs_table = self._table(conn, 'jobs')
            
            cursor.execute(
                self._sql(conn, f"SELECT * FROM {jobs_table} WHERE id = %s"),
                (job_id,)
            )
            row = cursor.fetchone()
            cursor.close()
            conn.close()
            
            if row:
                job_data = dict(row)
                if job_data.get('result_data'):
                    try:
                        job_data['result_data'] = json.loads(job_data['result_data'])
                    except:
                        pass
                return job_data
            return None
            
        except Exception as e:
            print(f"❌ Failed to get job from PostgreSQL database: {e}")
            return None
    
    def get_jobs_history(self, limit: int = 50) -> List[Dict[str, Any]]:
        """Get recent jobs for UI display"""
        try:
            conn = self.get_connection()
            cursor = self._cursor(conn, dict_rows=True)
            jobs_table = self._table(conn, 'jobs')
            
            query = f'''
                SELECT * FROM {jobs_table}
                ORDER BY created_at DESC 
                LIMIT %s
            '''
            cursor.execute(self._sql(conn, query), (limit,))
            
            rows = cursor.fetchall()
            cursor.close()
            conn.close()
            
            jobs = []
            for row in rows:
                job_data = dict(row)
                if job_data.get('result_data'):
                    try:
                        job_data['result_data'] = json.loads(job_data['result_data'])
                    except:
                        pass
                jobs.append(job_data)
            
            print(f"✅ Retrieved {len(jobs)} jobs from database")
            return jobs
            
        except Exception as e:
            print(f"❌ Failed to get jobs history from PostgreSQL: {e}")
            return []
    
    def get_dashboard_metrics(self) -> Dict[str, int]:
        """Get dashboard metrics from the configured database."""
        try:
            conn = self.get_connection()
            cursor = self._cursor(conn)
            jobs_table = self._table(conn, 'jobs')
            
            # Get total jobs
            cursor.execute(f"SELECT COUNT(*) FROM {jobs_table}")
            total_jobs = cursor.fetchone()[0]
            
            # Get completed jobs
            cursor.execute(f"SELECT COUNT(*) FROM {jobs_table} WHERE status = 'completed'")
            completed_jobs = cursor.fetchone()[0]
            
            # Get successful jobs
            cursor.execute(f"SELECT COUNT(*) FROM {jobs_table} WHERE success = true")
            successful_jobs = cursor.fetchone()[0]
            
            # Get failed jobs
            cursor.execute(f"SELECT COUNT(*) FROM {jobs_table} WHERE success = false")
            failed_jobs = cursor.fetchone()[0]
            
            # Get active jobs
            cursor.execute(f"SELECT COUNT(*) FROM {jobs_table} WHERE status IN ('pending', 'processing')")
            active_jobs = cursor.fetchone()[0]
            
            cursor.close()
            conn.close()
            
            metrics = {
                'total_jobs': total_jobs,
                'completed_jobs': completed_jobs,
                'successful_jobs': successful_jobs,
                'failed_jobs': failed_jobs,
                'active_jobs': active_jobs
            }
            
            print(f"✅ Retrieved dashboard metrics: {metrics}")
            return metrics
            
        except Exception as e:
            print(f"❌ Failed to get dashboard metrics from PostgreSQL: {e}")
            return {
                'total_jobs': 0,
                'completed_jobs': 0,
                'successful_jobs': 0,
                'failed_jobs': 0,
                'active_jobs': 0
            }
    
    def add_batch_job(self, batch_data: Dict[str, Any]) -> bool:
        """Add a batch job to the configured database."""
        try:
            conn = self.get_connection()
            cursor = self._cursor(conn)
            batch_jobs_table = self._table(conn, 'batch_jobs')
            
            batch_id = batch_data.get('id', f"batch_{int(time.time())}")
            
            query = f'''
                INSERT INTO {batch_jobs_table} (
                    id, workflow_type, status, batch_size, processed_count,
                    success_count, failed_count
                ) VALUES (%s, %s, %s, %s, %s, %s, %s)
                ON CONFLICT (id) DO UPDATE SET
                    status = EXCLUDED.status,
                    processed_count = EXCLUDED.processed_count,
                    success_count = EXCLUDED.success_count,
                    failed_count = EXCLUDED.failed_count,
                    updated_at = CURRENT_TIMESTAMP
            '''
            cursor.execute(self._sql(conn, query), (
                batch_id,
                batch_data.get('workflow_type', 'unknown'),
                batch_data.get('status', 'pending'),
                batch_data.get('batch_size', 0),
                batch_data.get('processed_count', 0),
                batch_data.get('success_count', 0),
                batch_data.get('failed_count', 0)
            ))
            
            conn.commit()
            cursor.close()
            conn.close()
            print(f"✅ Batch job added to database: {batch_id}")
            return True
            
        except Exception as e:
            print(f"❌ Failed to add batch job to PostgreSQL database: {e}")
            return False

# Global database instance
db_manager = DatabaseManager()

def get_db_connection():
    """Backward compatibility function"""
    return db_manager.get_connection()
def clear_all_jobs():
    """Clear all jobs from the database - utility function for UI"""
    try:
        db_manager = DatabaseManager()
        conn = db_manager.get_connection()
        cursor = conn.cursor()
        
        # Clear both regular jobs and batch jobs
        cursor.execute(f"DELETE FROM {db_manager._table(conn, 'jobs')}")
        cursor.execute(f"DELETE FROM {db_manager._table(conn, 'batch_jobs')}")
        
        conn.commit()
        cursor.close()
        conn.close()
        
        print("✅ All jobs cleared from database")
        return True
        
    except Exception as e:
        print(f"❌ Failed to clear database: {e}")
        return False