intl_news/dedup/store.py

"""SQLite 指纹存储。

注意: SimHash 是 64 位无符号整数，SQLite INTEGER 是 64 位有符号
(范围 [-2^63, 2^63-1])。直接存可能溢出/转负数，虽然 XOR 仍然
正确但语义混乱。这里统一存为 16 位 hex TEXT，避免符号问题。
"""

import logging
import sqlite3
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any

from dedup.models import Fingerprint

logger = logging.getLogger(__name__)

DEFAULT_DB_PATH = Path("data/dedup/fingerprints.sqlite3")

_SCHEMA_SQL = """
CREATE TABLE IF NOT EXISTS fingerprints (
    url_hash TEXT PRIMARY KEY,
    content_hash TEXT NOT NULL,
    simhash_hex TEXT NOT NULL,
    source_id TEXT NOT NULL,
    url TEXT NOT NULL,
    title TEXT NOT NULL,
    publish_date TEXT,
    ingested_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_content_hash ON fingerprints(content_hash);
CREATE INDEX IF NOT EXISTS idx_publish_date ON fingerprints(publish_date);
CREATE INDEX IF NOT EXISTS idx_source_id ON fingerprints(source_id);
"""


def _to_hex(simhash: int) -> str:
    """64 位无符号整数 → 16 位 hex 字符串。"""
    return f"{simhash:016x}"


def _from_hex(hex_str: str) -> int:
    """16 位 hex 字符串 → 64 位无符号整数。"""
    return int(hex_str, 16)


def _row_to_fp(row: sqlite3.Row) -> Fingerprint:
    """sqlite3.Row → Fingerprint 模型。"""
    return Fingerprint(
        url_hash=row["url_hash"],
        content_hash=row["content_hash"],
        simhash=_from_hex(row["simhash_hex"]),
        source_id=row["source_id"],
        url=row["url"],
        title=row["title"],
        publish_date=row["publish_date"],
        ingested_at=datetime.fromisoformat(row["ingested_at"]),
    )


class FingerprintStore:
    """SQLite 指纹库。线程不安全（每个线程请新建实例）。"""

    def __init__(self, db_path: str | Path = DEFAULT_DB_PATH) -> None:
        self.db_path = Path(db_path)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._conn: sqlite3.Connection = sqlite3.connect(
            self.db_path, isolation_level=None
        )
        self._conn.row_factory = sqlite3.Row
        self._conn.executescript(_SCHEMA_SQL)
        logger.debug("打开指纹库: %s", self.db_path)

    def close(self) -> None:
        self._conn.close()

    def __enter__(self) -> "FingerprintStore":
        return self

    def __exit__(self, *_: Any) -> None:
        self.close()

    # ------------------------------------------------------------------ #
    # 查询
    # ------------------------------------------------------------------ #

    def get_by_url_hash(self, url_hash: str) -> Fingerprint | None:
        """按 url_hash 精确查询。"""
        row = self._conn.execute(
            "SELECT * FROM fingerprints WHERE url_hash = ?", (url_hash,)
        ).fetchone()
        return _row_to_fp(row) if row else None

    def find_by_content_hash(self, content_hash: str) -> Fingerprint | None:
        """返回任一 content_hash 匹配项。"""
        row = self._conn.execute(
            "SELECT * FROM fingerprints WHERE content_hash = ? LIMIT 1",
            (content_hash,),
        ).fetchone()
        return _row_to_fp(row) if row else None

    def candidates_for_simhash(
        self,
        publish_date: str | None,
        window_days: int,
    ) -> list[Fingerprint]:
        """返回 publish_date ± window_days 内的指纹候选。

        publish_date 为 None 时，不限定窗口（返回全部，慎用）。
        """
        if publish_date is None or window_days < 0:
            rows = self._conn.execute("SELECT * FROM fingerprints").fetchall()
            return [_row_to_fp(r) for r in rows]

        try:
            center = datetime.strptime(publish_date, "%Y-%m-%d")
        except ValueError:
            logger.debug("publish_date 不可解析: %r，退化为全表扫描", publish_date)
            rows = self._conn.execute("SELECT * FROM fingerprints").fetchall()
            return [_row_to_fp(r) for r in rows]

        lo = (center - timedelta(days=window_days)).strftime("%Y-%m-%d")
        hi = (center + timedelta(days=window_days)).strftime("%Y-%m-%d")
        rows = self._conn.execute(
            "SELECT * FROM fingerprints "
            "WHERE publish_date IS NULL OR (publish_date >= ? AND publish_date <= ?)",
            (lo, hi),
        ).fetchall()
        return [_row_to_fp(r) for r in rows]

    def count(self) -> int:
        """指纹总数。"""
        return self._conn.execute("SELECT COUNT(*) FROM fingerprints").fetchone()[0]

    def count_by_source(self) -> dict[str, int]:
        """按 source_id 统计。"""
        rows = self._conn.execute(
            "SELECT source_id, COUNT(*) AS n FROM fingerprints GROUP BY source_id"
        ).fetchall()
        return {r["source_id"]: r["n"] for r in rows}

    def date_range(self) -> tuple[str | None, str | None]:
        """指纹库中最早和最晚的 publish_date。"""
        row = self._conn.execute(
            "SELECT MIN(publish_date) AS lo, MAX(publish_date) AS hi FROM fingerprints"
        ).fetchone()
        return (row["lo"], row["hi"]) if row else (None, None)

    # ------------------------------------------------------------------ #
    # 写入
    # ------------------------------------------------------------------ #

    def upsert(self, fp: Fingerprint) -> None:
        """插入或替换指纹。"""
        self._conn.execute(
            "INSERT OR REPLACE INTO fingerprints "
            "(url_hash, content_hash, simhash_hex, source_id, url, title, "
            " publish_date, ingested_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            (
                fp.url_hash,
                fp.content_hash,
                _to_hex(fp.simhash),
                fp.source_id,
                fp.url,
                fp.title,
                fp.publish_date,
                fp.ingested_at.isoformat(),
            ),
        )

    def delete(self, url_hash: str) -> None:
        """删除指定指纹。"""
        self._conn.execute("DELETE FROM fingerprints WHERE url_hash = ?", (url_hash,))

    def clear(self) -> None:
        """清空指纹库，主要用于测试。"""
        self._conn.execute("DELETE FROM fingerprints")