hbr-crawler/test_db_connection.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
資料庫連線測試腳本
用於測試資料庫連線、建立資料庫和資料表
"""
import os
import sys
import logging
from pathlib import Path

# 設定日誌
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 加入專案路徑
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from dotenv import load_dotenv
from hbr_crawler.hbr_crawler.database import DatabaseManager, get_database_manager

# 載入 .env 檔案
load_dotenv()

# 資料庫連線資訊（從環境變數讀取）
DB_CONFIG = {
    'host': os.environ.get('DB_HOST'),
    'port': int(os.environ.get('DB_PORT', 3306)),
    'user': os.environ.get('DB_USER'),
    'password': os.environ.get('DB_PASSWORD'),
    'database': os.environ.get('DB_NAME')
}

# 檢查必要的環境變數
missing_vars = [k for k, v in DB_CONFIG.items() if v is None and k != 'port']
if missing_vars:
    print(f"錯誤: 缺少必要的環境變數: {', '.join(['DB_' + k.upper() for k in missing_vars])}")
    print("請在 .env 檔案中設定這些值，參考 .env.example")
    sys.exit(1)


def test_basic_connection():
    """測試基本連線（不指定資料庫）"""
    print("\n" + "="*50)
    print("測試 1: 基本資料庫連線（不指定資料庫）")
    print("="*50)

    db_manager = DatabaseManager(
        host=DB_CONFIG['host'],
        port=DB_CONFIG['port'],
        user=DB_CONFIG['user'],
        password=DB_CONFIG['password'],
        database=None
    )

    if db_manager.test_connection():
        print("✓ 基本連線測試成功")
        return True
    else:
        print("✗ 基本連線測試失敗")
        return False


def create_database():
    """建立 HBR_scraper 資料庫（如果需要）"""
    print("\n" + "="*50)
    print("測試 2: 檢查資料庫連線（使用現有資料庫）")
    print("="*50)

    db_manager = DatabaseManager(
        host=DB_CONFIG['host'],
        port=DB_CONFIG['port'],
        user=DB_CONFIG['user'],
        password=DB_CONFIG['password'],
        database=None
    )

    # 嘗試建立資料庫（可能需要管理員權限）
    try:
        if db_manager.create_database('HBR_scraper'):
            print("✓ 資料庫建立成功（或已存在）")
            return True
        else:
            print("✗ 資料庫建立失敗")
            print("提示：如果沒有 CREATE DATABASE 權限，請請管理員協助建立資料庫")
            return False
    except Exception as e:
        print(f"✗ 資料庫建立失敗: {e}")
        print("提示：如果沒有 CREATE DATABASE 權限，請請管理員協助建立資料庫")
        print("或者使用現有的資料庫，修改 DB_NAME 設定")
        return False


def test_database_connection():
    """測試連接到指定資料庫"""
    print("\n" + "="*50)
    print(f"測試 3: 連接到 {DB_CONFIG['database']} 資料庫")
    print("="*50)

    db_manager = DatabaseManager(
        host=DB_CONFIG['host'],
        port=DB_CONFIG['port'],
        user=DB_CONFIG['user'],
        password=DB_CONFIG['password'],
        database=DB_CONFIG['database']
    )

    if db_manager.test_connection(DB_CONFIG['database']):
        print("✓ 資料庫連線測試成功")
        return True
    else:
        print("✗ 資料庫連線測試失敗")
        return False


def create_tables():
    """建立資料表"""
    print("\n" + "="*50)
    print("測試 4: 建立資料表")
    print("="*50)

    db_manager = DatabaseManager(
        host=DB_CONFIG['host'],
        port=DB_CONFIG['port'],
        user=DB_CONFIG['user'],
        password=DB_CONFIG['password'],
        database=DB_CONFIG['database']
    )

    sql_file = project_root / 'create_tables.sql'

    if not sql_file.exists():
        print(f"✗ SQL 檔案不存在: {sql_file}")
        return False

    if db_manager.execute_sql_file(str(sql_file), DB_CONFIG['database']):
        print("✓ 資料表建立成功")
        return True
    else:
        print("✗ 資料表建立失敗")
        return False


def verify_tables():
    """驗證資料表是否建立成功"""
    print("\n" + "="*50)
    print("測試 5: 驗證資料表")
    print("="*50)

    db_manager = DatabaseManager(
        host=DB_CONFIG['host'],
        port=DB_CONFIG['port'],
        user=DB_CONFIG['user'],
        password=DB_CONFIG['password'],
        database=DB_CONFIG['database']
    )

    expected_tables = ['articles', 'tags', 'article_tags']

    try:
        tables = db_manager.execute_query(
            "SHOW TABLES",
            database=DB_CONFIG['database']
        )

        # 取得資料表名稱列表
        table_names = [list(table.values())[0] for table in tables]

        print(f"找到 {len(table_names)} 個資料表: {', '.join(table_names)}")

        for table in expected_tables:
            if table in table_names:
                print(f"✓ 資料表 {table} 存在")
            else:
                print(f"✗ 資料表 {table} 不存在")

        return all(table in table_names for table in expected_tables)
    except Exception as e:
        print(f"✗ 驗證資料表失敗: {e}")
        return False


def main():
    """主函數"""
    print("\n" + "="*60)
    print("HBR 爬蟲系統 - 資料庫連線測試")
    print("="*60)

    results = []

    # 執行測試
    results.append(("基本連線", test_basic_connection()))
    results.append(("建立資料庫", create_database()))
    results.append(("資料庫連線", test_database_connection()))
    results.append(("建立資料表", create_tables()))
    results.append(("驗證資料表", verify_tables()))

    # 顯示測試結果摘要
    print("\n" + "="*60)
    print("測試結果摘要")
    print("="*60)

    for test_name, result in results:
        status = "✓ 通過" if result else "✗ 失敗"
        print(f"{test_name}: {status}")

    all_passed = all(result for _, result in results)

    if all_passed:
        print("\n✓ 所有測試通過！資料庫設定完成。")
        return 0
    else:
        print("\n✗ 部分測試失敗，請檢查錯誤訊息。")
        return 1


if __name__ == '__main__':
    sys.exit(main())