WeReadBoard/fetch_data.py

#!/usr/bin/env python3
"""
微信读书数据采集脚本
从微信读书 API 拉取数据并存入本地 SQLite 数据库
"""

import json
import os
import sqlite3
import time
import urllib.request

API_URL = "https://i.weread.qq.com/api/agent/gateway"
API_KEY = os.environ.get("WEREAD_API_KEY", "")
SKILL_VERSION = "1.0.3"
DB_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "weread.db")
PROJECT_DIR = os.path.dirname(os.path.abspath(__file__))


def api_call(api_name, **kwargs):
    """调用微信读书 API"""
    body = {"api_name": api_name, "skill_version": SKILL_VERSION}
    body.update(kwargs)
    req = urllib.request.Request(
        API_URL,
        data=json.dumps(body).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json",
        },
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.loads(resp.read().decode("utf-8"))


def init_db(db_path):
    """初始化数据库表结构"""
    conn = sqlite3.connect(db_path)
    c = conn.cursor()

    # 书架书籍表
    c.execute("""
        CREATE TABLE IF NOT EXISTS books (
            book_id TEXT PRIMARY KEY,
            title TEXT,
            author TEXT,
            cover TEXT,
            category TEXT,
            finish_reading INTEGER DEFAULT 0,
            read_update_time INTEGER,
            secret INTEGER DEFAULT 0,
            is_top INTEGER DEFAULT 0,
            update_time INTEGER,
            deep_link TEXT
        )
    """)

    # 阅读统计概要表
    c.execute("""
        CREATE TABLE IF NOT EXISTS reading_stats (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            mode TEXT,
            period TEXT,
            total_read_time INTEGER,
            read_days INTEGER,
            day_average_read_time REAL,
            regist_time INTEGER,
            raw_json TEXT
        )
    """)

    # 阅读时长时间序列（按年/月/日分桶）
    c.execute("""
        CREATE TABLE IF NOT EXISTS read_times (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            mode TEXT,
            period TEXT,
            bucket_timestamp INTEGER,
            duration INTEGER
        )
    """)

    # 阅读时长最长的书籍排行
    c.execute("""
        CREATE TABLE IF NOT EXISTS read_longest (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            mode TEXT,
            period TEXT,
            book_id TEXT,
            title TEXT,
            author TEXT,
            cover TEXT,
            read_time INTEGER,
            tags TEXT
        )
    """)

    # 偏好分类
    c.execute("""
        CREATE TABLE IF NOT EXISTS prefer_categories (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            mode TEXT,
            period TEXT,
            category_id INTEGER,
            category_title TEXT,
            parent_category_title TEXT,
            val REAL,
            reading_time INTEGER,
            reading_count INTEGER
        )
    """)

    # 偏好作者
    c.execute("""
        CREATE TABLE IF NOT EXISTS prefer_authors (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            mode TEXT,
            period TEXT,
            author_id INTEGER,
            name TEXT,
            count INTEGER,
            read_time TEXT
        )
    """)

    # 24小时阅读时段分布
    c.execute("""
        CREATE TABLE IF NOT EXISTS prefer_time (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            mode TEXT,
            period TEXT,
            hour_index INTEGER,
            duration INTEGER
        )
    """)

    # 笔记本概览
    c.execute("""
        CREATE TABLE IF NOT EXISTS notebooks (
            book_id TEXT PRIMARY KEY,
            title TEXT,
            author TEXT,
            cover TEXT,
            review_count INTEGER DEFAULT 0,
            note_count INTEGER DEFAULT 0,
            bookmark_count INTEGER DEFAULT 0,
            reading_progress INTEGER DEFAULT 0,
            marked_status INTEGER DEFAULT 0,
            sort INTEGER
        )
    """)

    # 阅读统计摘要项
    c.execute("""
        CREATE TABLE IF NOT EXISTS read_stat_items (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            mode TEXT,
            period TEXT,
            stat_name TEXT,
            counts TEXT,
            scheme TEXT
        )
    """)

    # 书单/归档
    c.execute("""
        CREATE TABLE IF NOT EXISTS archives (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT,
            book_ids TEXT
        )
    """)

    # 数据采集元信息
    c.execute("""
        CREATE TABLE IF NOT EXISTS fetch_meta (
            key TEXT PRIMARY KEY,
            value TEXT
        )
    """)

    conn.commit()
    return conn


def fetch_shelf(conn):
    """拉取书架数据"""
    print("  拉取书架数据...")
    data = api_call("/shelf/sync")
    c = conn.cursor()

    books = data.get("books", [])
    for book in books:
        c.execute(
            """INSERT OR REPLACE INTO books
            (book_id, title, author, cover, category, finish_reading,
             read_update_time, secret, is_top, update_time, deep_link)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                book.get("bookId"),
                book.get("title", ""),
                book.get("author", ""),
                book.get("cover", ""),
                book.get("category", ""),
                book.get("finishReading", 0),
                book.get("readUpdateTime", 0),
                book.get("secret", 0),
                book.get("isTop", 0),
                book.get("updateTime", 0),
                book.get("deepLink", ""),
            ),
        )

    # 保存书单
    archives = data.get("archive", [])
    c.execute("DELETE FROM archives")
    for arch in archives:
        c.execute(
            "INSERT INTO archives (name, book_ids) VALUES (?, ?)",
            (arch.get("name", ""), json.dumps(arch.get("bookIds", []))),
        )

    conn.commit()
    mp_flag = "yes" if data.get("mp") else "no"
    print(f"  书架: {len(books)} 本电子书, {len(data.get('albums', []))} 个专辑, MP={mp_flag}, {len(archives)} 个书单")


def fetch_readdata(conn, mode, period_label):
    """拉取阅读统计数据"""
    print(f"  拉取阅读统计 [{mode}]...")
    params = {"mode": mode}
    if mode != "overall":
        params["baseTime"] = 0
    data = api_call("/readdata/detail", **params)
    c = conn.cursor()

    total_read_time = data.get("totalReadTime", 0)
    read_days = data.get("readDays", 0)
    day_avg = data.get("dayAverageReadTime", 0)
    regist_time = data.get("registTime", 0)

    c.execute(
        """INSERT INTO reading_stats
        (mode, period, total_read_time, read_days, day_average_read_time,
         regist_time, raw_json)
        VALUES (?, ?, ?, ?, ?, ?, ?)""",
        (
            mode,
            period_label,
            total_read_time,
            read_days,
            day_avg,
            regist_time,
            json.dumps(data, ensure_ascii=False),
        ),
    )

    # 保存 readTimes 分桶数据
    read_times = data.get("readTimes", {})
    for ts, duration in read_times.items():
        c.execute(
            "INSERT INTO read_times (mode, period, bucket_timestamp, duration) VALUES (?, ?, ?, ?)",
            (mode, period_label, int(ts), duration),
        )

    # 保存 dailyReadTimes（如果有）
    daily_times = data.get("dailyReadTimes", {})
    for ts, duration in daily_times.items():
        c.execute(
            "INSERT INTO read_times (mode, period, bucket_timestamp, duration) VALUES (?, ?, ?, ?)",
            (mode + "_daily", period_label, int(ts), duration),
        )

    # 保存阅读最久的书籍
    for item in data.get("readLongest", []):
        book = item.get("book", {})
        album = item.get("albumInfo", {})
        c.execute(
            """INSERT INTO read_longest
            (mode, period, book_id, title, author, cover, read_time, tags)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                mode,
                period_label,
                book.get("bookId") or album.get("albumId", ""),
                book.get("title") or album.get("name", ""),
                book.get("author") or album.get("authorName", ""),
                book.get("cover") or album.get("cover", ""),
                item.get("readTime", 0),
                json.dumps(item.get("tags", []), ensure_ascii=False),
            ),
        )

    # 保存偏好分类
    for cat in data.get("preferCategory", []):
        c.execute(
            """INSERT INTO prefer_categories
            (mode, period, category_id, category_title, parent_category_title,
             val, reading_time, reading_count)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                mode,
                period_label,
                cat.get("categoryId", 0),
                cat.get("categoryTitle", ""),
                cat.get("parentCategoryTitle", ""),
                cat.get("val", 0),
                cat.get("readingTime", 0),
                cat.get("readingCount", 0),
            ),
        )

    # 保存偏好作者
    for author in data.get("preferAuthor", []):
        c.execute(
            """INSERT INTO prefer_authors
            (mode, period, author_id, name, count, read_time)
            VALUES (?, ?, ?, ?, ?, ?)""",
            (
                mode,
                period_label,
                author.get("authorId", 0),
                author.get("name", ""),
                author.get("count", 0),
                author.get("readTime", ""),
            ),
        )

    # 保存24小时时段分布
    prefer_time = data.get("preferTime", [])
    # preferTime 从6点开始到次日5点，映射到 hour_index
    for i, duration in enumerate(prefer_time):
        hour = (6 + i) % 24
        c.execute(
            "INSERT INTO prefer_time (mode, period, hour_index, duration) VALUES (?, ?, ?, ?)",
            (mode, period_label, hour, duration),
        )

    # 保存统计摘要项
    for stat in data.get("readStat", []):
        c.execute(
            """INSERT INTO read_stat_items
            (mode, period, stat_name, counts, scheme)
            VALUES (?, ?, ?, ?, ?)""",
            (
                mode,
                period_label,
                stat.get("stat", ""),
                stat.get("counts", ""),
                stat.get("scheme", ""),
            ),
        )

    conn.commit()
    print(f"  [{mode}] 总时长: {total_read_time}秒 ({total_read_time//3600}小时), 阅读天数: {read_days}")


def fetch_notebooks(conn):
    """拉取笔记本概览（分页拉取全部）"""
    print("  拉取笔记本数据...")
    c = conn.cursor()
    c.execute("DELETE FROM notebooks")

    count = 100
    last_sort = None
    total = 0

    while True:
        params = {"count": count}
        if last_sort is not None:
            params["lastSort"] = last_sort
        data = api_call("/user/notebooks", **params)

        books = data.get("books", [])
        for item in books:
            book = item.get("book", {})
            c.execute(
                """INSERT OR REPLACE INTO notebooks
                (book_id, title, author, cover, review_count, note_count,
                 bookmark_count, reading_progress, marked_status, sort)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    item.get("bookId", ""),
                    book.get("title", ""),
                    book.get("author", ""),
                    book.get("cover", ""),
                    item.get("reviewCount", 0),
                    item.get("noteCount", 0),
                    item.get("bookmarkCount", 0),
                    item.get("readingProgress", 0),
                    item.get("markedStatus", 0),
                    item.get("sort", 0),
                ),
            )
            total += 1

        if data.get("hasMore", 0) == 1 and books:
            last_sort = books[-1].get("sort")
        else:
            break
        time.sleep(0.3)

    conn.commit()
    print(f"  笔记本: {total} 本有笔记的书")


def main():
    if not API_KEY:
        print("错误: 请先设置 WEREAD_API_KEY 环境变量")
        print("  export WEREAD_API_KEY=wrk-xxxxxxxx")
        return

    print("=" * 50)
    print("微信读书数据采集")
    print("=" * 50)

    # 初始化数据库
    print(f"\n[1/4] 初始化数据库: {DB_PATH}")
    conn = init_db(DB_PATH)

    # 清空旧数据（保留表结构）
    for table in [
        "reading_stats", "read_times", "read_longest", "prefer_categories",
        "prefer_authors", "prefer_time", "read_stat_items",
    ]:
        conn.execute(f"DELETE FROM {table}")
    conn.commit()

    # 拉取书架
    print("\n[2/4] 拉取书架数据")
    fetch_shelf(conn)

    # 拉取阅读统计（overall + 当前年）
    print("\n[3/4] 拉取阅读统计数据")
    fetch_readdata(conn, "overall", "overall")
    current_year = time.localtime().tm_year
    fetch_readdata(conn, "annually", str(current_year))

    # 拉取笔记本
    print("\n[4/4] 拉取笔记本数据")
    fetch_notebooks(conn)

    # 记录采集时间
    c = conn.cursor()
    c.execute(
        "INSERT OR REPLACE INTO fetch_meta (key, value) VALUES (?, ?)",
        ("last_fetch_time", time.strftime("%Y-%m-%d %H:%M:%S")),
    )
    conn.commit()

    # 统计
    c.execute("SELECT COUNT(*) FROM books")
    book_count = c.fetchone()[0]
    c.execute("SELECT COUNT(*) FROM notebooks")
    notebook_count = c.fetchone()[0]
    c.execute("SELECT SUM(total_read_time) FROM reading_stats WHERE mode='overall'")
    total_time = c.fetchone()[0] or 0

    print("\n" + "=" * 50)
    print("采集完成!")
    print(f"  书架书籍: {book_count} 本")
    print(f"  笔记本: {notebook_count} 本")
    print(f"  总阅读时长: {total_time // 3600} 小时 {(total_time % 3600) // 60} 分钟")
    print(f"  数据库: {DB_PATH}")
    print("=" * 50)

    conn.close()


if __name__ == "__main__":
    main()