database.py

import os
import sqlite3
import json

DB_PATH = "tasks.db"

def init_db():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    # 任务表
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS download_tasks (
            id TEXT PRIMARY KEY,
            original_link TEXT,
            media_info TEXT,  -- JSON string
            save_name TEXT,
            status TEXT,      -- pending, downloading, completed, failed
            progress INTEGER DEFAULT 0,
            speed TEXT DEFAULT '',
            error_msg TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    # 解析历史表 (可选)
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS parse_history (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            link TEXT,
            media_count INTEGER,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    conn.commit()
    conn.close()

def add_task(task_id, link, media_info, save_name):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute('''
        INSERT INTO download_tasks (id, original_link, media_info, save_name, status)
        VALUES (?, ?, ?, ?, 'pending')
    ''', (task_id, link, json.dumps(media_info), save_name))
    conn.commit()
    conn.close()

def update_task_progress(task_id, progress, speed, status='downloading'):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute('''
        UPDATE download_tasks
        SET progress = ?, speed = ?, status = ?
        WHERE id = ?
    ''', (progress, speed, status, task_id))
    conn.commit()
    conn.close()

def update_task_status(task_id, status, error_msg=''):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute('''
        UPDATE download_tasks
        SET status = ?, error_msg = ?
        WHERE id = ?
    ''', (status, error_msg, task_id))
    conn.commit()
    conn.close()

def clear_all_tasks():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute('DELETE FROM download_tasks')
    conn.commit()
    conn.close()

def get_all_tasks():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()
    cursor.execute('SELECT * FROM download_tasks ORDER BY created_at DESC')
    rows = cursor.fetchall()
    conn.close()
    
    tasks = []
    for row in rows:
        task = dict(row)
        task['media_info'] = json.loads(task['media_info'])
        tasks.append(task)
    return tasks

def clear_completed_tasks():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute("DELETE FROM download_tasks WHERE status IN ('completed', 'failed')")
    conn.commit()
    conn.close()

if __name__ == "__main__":
    init_db()
    print("Database initialized.")