TK-BOM-Project/backend/scripts/optimize_database.py

"""
데이터베이스 성능 최적화 스크립트
인덱스 생성, 쿼리 최적화, 통계 업데이트
"""

import os
import psycopg2
from psycopg2 import sql
from ..utils.logger import get_logger

logger = get_logger(__name__)

# 환경 변수 로드
DB_HOST = os.getenv("DB_HOST", "postgres")
DB_PORT = os.getenv("DB_PORT", "5432")
DB_NAME = os.getenv("DB_NAME", "tk_mp_bom")
DB_USER = os.getenv("DB_USER", "tkmp_user")
DB_PASSWORD = os.getenv("DB_PASSWORD", "tkmp_password_2025")


def optimize_database():
    """데이터베이스 성능 최적화"""
    
    try:
        conn = psycopg2.connect(
            host=DB_HOST,
            port=DB_PORT,
            database=DB_NAME,
            user=DB_USER,
            password=DB_PASSWORD
        )
        
        with conn.cursor() as cursor:
            # 1. 핵심 인덱스 생성
            print("🔧 핵심 인덱스 생성 중...")
            
            indexes = [
                # materials 테이블 인덱스
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_materials_file_id ON materials(file_id);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_materials_category ON materials(classified_category);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_materials_drawing_name ON materials(drawing_name);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_materials_line_no ON materials(line_no);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_materials_revision_status ON materials(revision_status);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_materials_purchase_confirmed ON materials(purchase_confirmed);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_materials_material_hash ON materials(material_hash);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_materials_main_nom ON materials(main_nom);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_materials_material_grade ON materials(material_grade);",
                
                # files 테이블 인덱스
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_files_job_no ON files(job_no);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_files_revision ON files(revision);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_files_uploaded_by ON files(uploaded_by);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_files_upload_date ON files(upload_date);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_files_is_active ON files(is_active);",
                
                # 복합 인덱스 (자주 함께 사용되는 컬럼들)
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_materials_file_category ON materials(file_id, classified_category);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_materials_drawing_line ON materials(drawing_name, line_no);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_files_job_revision ON files(job_no, revision);",
                
                # material_purchase_tracking 인덱스
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_mpt_material_hash ON material_purchase_tracking(material_hash);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_mpt_job_revision ON material_purchase_tracking(job_no, revision);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_mpt_purchase_status ON material_purchase_tracking(purchase_status);",
                
                # 상세 테이블들 인덱스
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_pipe_details_material_id ON pipe_details(material_id);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_fitting_details_material_id ON fitting_details(material_id);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_flange_details_material_id ON flange_details(material_id);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_valve_details_material_id ON valve_details(material_id);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_bolt_details_material_id ON bolt_details(material_id);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gasket_details_material_id ON gasket_details(material_id);",
                
                # 사용자 관련 인덱스
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_users_username ON users(username);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_users_status ON users(status);",
                "CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_users_role ON users(role);",
            ]
            
            for index_sql in indexes:
                try:
                    cursor.execute(index_sql)
                    conn.commit()
                    print(f"✅ 인덱스 생성: {index_sql.split('idx_')[1].split(' ')[0] if 'idx_' in index_sql else 'unknown'}")
                except Exception as e:
                    print(f"⚠️ 인덱스 생성 실패: {e}")
                    conn.rollback()
            
            # 2. 통계 업데이트
            print("📊 테이블 통계 업데이트 중...")
            
            tables_to_analyze = [
                'materials', 'files', 'projects', 'users',
                'material_purchase_tracking', 'pipe_details',
                'fitting_details', 'flange_details', 'valve_details',
                'bolt_details', 'gasket_details'
            ]
            
            for table in tables_to_analyze:
                try:
                    cursor.execute(f"ANALYZE {table};")
                    conn.commit()
                    print(f"✅ 통계 업데이트: {table}")
                except Exception as e:
                    print(f"⚠️ 통계 업데이트 실패 ({table}): {e}")
                    conn.rollback()
            
            # 3. VACUUM 실행 (선택적)
            print("🧹 데이터베이스 정리 중...")
            try:
                # VACUUM은 트랜잭션 외부에서 실행해야 함
                conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
                cursor.execute("VACUUM ANALYZE;")
                print("✅ VACUUM ANALYZE 완료")
            except Exception as e:
                print(f"⚠️ VACUUM 실패: {e}")
            
            print("✅ 데이터베이스 최적화 완료")
            
    except Exception as e:
        print(f"❌ 데이터베이스 최적화 실패: {e}")
        return False
    finally:
        if conn:
            conn.close()
    
    return True


def get_database_stats():
    """데이터베이스 통계 조회"""
    
    try:
        conn = psycopg2.connect(
            host=DB_HOST,
            port=DB_PORT,
            database=DB_NAME,
            user=DB_USER,
            password=DB_PASSWORD
        )
        
        with conn.cursor() as cursor:
            # 테이블별 레코드 수
            print("📊 테이블별 레코드 수:")
            
            tables = ['materials', 'files', 'projects', 'users', 'material_purchase_tracking']
            
            for table in tables:
                try:
                    cursor.execute(f"SELECT COUNT(*) FROM {table};")
                    count = cursor.fetchone()[0]
                    print(f"  {table}: {count:,}개")
                except Exception as e:
                    print(f"  {table}: 조회 실패 ({e})")
            
            # 인덱스 사용률 확인
            print("\n🔍 인덱스 사용률:")
            cursor.execute("""
                SELECT 
                    schemaname,
                    tablename,
                    indexname,
                    idx_tup_read,
                    idx_tup_fetch
                FROM pg_stat_user_indexes 
                WHERE idx_tup_read > 0
                ORDER BY idx_tup_read DESC
                LIMIT 10;
            """)
            
            for row in cursor.fetchall():
                print(f"  {row[1]}.{row[2]}: {row[3]:,} reads, {row[4]:,} fetches")
            
    except Exception as e:
        print(f"❌ 통계 조회 실패: {e}")
    finally:
        if conn:
            conn.close()


if __name__ == "__main__":
    print("🚀 데이터베이스 성능 최적화 시작")
    
    # 현재 통계 확인
    get_database_stats()
    
    # 최적화 실행
    if optimize_database():
        print("✅ 최적화 완료")
        
        # 최적화 후 통계 확인
        print("\n📊 최적화 후 통계:")
        get_database_stats()
    else:
        print("❌ 최적화 실패")