reset_upload_data.py

"""
업로드된 파일과 데이터베이스 정보를 삭제하는 스크립트
"""
import os
import sys
import sqlite3
from pathlib import Path

# UTF-8 인코딩 설정
if sys.platform == 'win32':
    sys.stdout.reconfigure(encoding='utf-8')
    sys.stderr.reconfigure(encoding='utf-8')

def reset_upload_data():
    """업로드된 파일과 DB 데이터 삭제"""
    try:
        # 1. 데이터베이스 경로 확인
        db_path = os.path.join('instance', 'finance_analysis.db')
        if not os.path.exists(db_path):
            # 상대 경로로도 시도
            db_path = 'finance_analysis.db'
            if not os.path.exists(db_path):
                print(f"데이터베이스 파일을 찾을 수 없습니다: {db_path}")
                return False
        
        print(f"데이터베이스 경로: {db_path}")
        
        # 2. 데이터베이스 연결
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        # 3. 업로드된 파일 목록 조회 (삭제 전에 파일 경로 저장)
        cursor.execute("SELECT id, file_path FROM uploaded_file")
        files_to_delete = cursor.fetchall()
        
        print(f"\n삭제할 파일 수: {len(files_to_delete)}개")
        
        # 4. 파일 시스템에서 파일 삭제
        deleted_files = 0
        failed_files = []
        
        for file_id, file_path in files_to_delete:
            if file_path and os.path.exists(file_path):
                try:
                    os.remove(file_path)
                    deleted_files += 1
                    print(f"  ✓ 파일 삭제: {os.path.basename(file_path)}")
                except Exception as e:
                    failed_files.append((file_path, str(e)))
                    print(f"  ✗ 파일 삭제 실패: {os.path.basename(file_path)} - {str(e)}")
        
        # 5. 데이터베이스에서 데이터 삭제
        print(f"\n데이터베이스 데이터 삭제 중...")
        
        # document_chunk 테이블 삭제 (RAG용 청크)
        cursor.execute("DELETE FROM document_chunk")
        chunk_count = cursor.rowcount
        print(f"  ✓ document_chunk: {chunk_count}개 청크 삭제")
        
        # uploaded_file 테이블 삭제
        cursor.execute("DELETE FROM uploaded_file")
        file_count = cursor.rowcount
        print(f"  ✓ uploaded_file: {file_count}개 파일 정보 삭제")
        
        # chat_message 테이블 삭제 (선택사항 - 채팅 기록도 삭제)
        cursor.execute("DELETE FROM chat_message")
        message_count = cursor.rowcount
        print(f"  ✓ chat_message: {message_count}개 메시지 삭제")
        
        # chat_session 테이블 삭제 (선택사항 - 채팅 세션도 삭제)
        cursor.execute("DELETE FROM chat_session")
        session_count = cursor.rowcount
        print(f"  ✓ chat_session: {session_count}개 세션 삭제")
        
        # 변경사항 커밋
        conn.commit()
        conn.close()
        
        print(f"\n{'='*60}")
        print(f"삭제 완료!")
        print(f"  - 파일 시스템: {deleted_files}개 파일 삭제")
        print(f"  - 데이터베이스: {file_count}개 파일 정보, {chunk_count}개 청크 삭제")
        if failed_files:
            print(f"\n삭제 실패한 파일 ({len(failed_files)}개):")
            for file_path, error in failed_files:
                print(f"  - {file_path}: {error}")
        print(f"{'='*60}\n")
        
        return True
        
    except Exception as e:
        print(f"\n오류 발생: {str(e)}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == '__main__':
    import sys
    
    print("="*60)
    print("업로드 데이터 초기화 스크립트")
    print("="*60)
    print("\n다음 데이터가 삭제됩니다:")
    print("  - 업로드된 모든 파일 (uploads 폴더)")
    print("  - 데이터베이스의 uploaded_file 테이블")
    print("  - 데이터베이스의 document_chunk 테이블 (RAG 청크)")
    print("  - 데이터베이스의 chat_message 테이블 (채팅 기록)")
    print("  - 데이터베이스의 chat_session 테이블 (채팅 세션)")
    print("\n사용자 정보는 유지됩니다.")
    print("="*60)
    
    # 명령줄 인자로 --yes가 있으면 확인 없이 실행
    auto_confirm = '--yes' in sys.argv or '-y' in sys.argv
    
    if not auto_confirm:
        try:
            confirm = input("\n정말 삭제하시겠습니까? (yes/no): ")
        except (EOFError, KeyboardInterrupt):
            print("\n입력이 취소되었습니다. 자동 실행 모드로 진행합니다.")
            confirm = 'yes'
    else:
        confirm = 'yes'
        print("\n자동 실행 모드: 확인 없이 삭제를 진행합니다.")
    
    if confirm.lower() in ['yes', 'y', '예', 'ㅇ']:
        success = reset_upload_data()
        if success:
            print("\n초기화가 완료되었습니다. 이제 새로 파일을 업로드할 수 있습니다.")
        else:
            print("\n초기화 중 오류가 발생했습니다.")
    else:
        print("\n취소되었습니다.")