76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
|
import logging
|
||
|
|
import sqlite3
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
DEFAULT_DB_DIR = Path.home() / ".forcepilot" / "wecom"
|
||
|
|
DEFAULT_TTL_SECONDS = 7200
|
||
|
|
|
||
|
|
|
||
|
|
class PersistentDeduplicator:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
db_path: str | None = None,
|
||
|
|
ttl: int = DEFAULT_TTL_SECONDS,
|
||
|
|
):
|
||
|
|
if db_path is None:
|
||
|
|
DEFAULT_DB_DIR.mkdir(parents=True, exist_ok=True)
|
||
|
|
db_path = str(DEFAULT_DB_DIR / "dedupe.db")
|
||
|
|
|
||
|
|
self._db_path = db_path
|
||
|
|
self._ttl = ttl
|
||
|
|
self._lock = threading.Lock()
|
||
|
|
self._conn: sqlite3.Connection | None = None
|
||
|
|
self._init_db()
|
||
|
|
|
||
|
|
def _init_db(self):
|
||
|
|
with self._lock:
|
||
|
|
conn = sqlite3.connect(self._db_path)
|
||
|
|
conn.execute("CREATE TABLE IF NOT EXISTS dedupe ( msg_id TEXT PRIMARY KEY, created_at REAL NOT NULL)")
|
||
|
|
conn.execute("CREATE INDEX IF NOT EXISTS idx_dedupe_created ON dedupe(created_at)")
|
||
|
|
conn.commit()
|
||
|
|
self._conn = conn
|
||
|
|
|
||
|
|
def is_duplicate(self, msg_id: str) -> bool:
|
||
|
|
if not msg_id:
|
||
|
|
return False
|
||
|
|
|
||
|
|
self._evict_expired()
|
||
|
|
|
||
|
|
with self._lock:
|
||
|
|
try:
|
||
|
|
cur = self._conn.execute("SELECT 1 FROM dedupe WHERE msg_id = ?", (msg_id,))
|
||
|
|
if cur.fetchone():
|
||
|
|
return True
|
||
|
|
|
||
|
|
self._conn.execute(
|
||
|
|
"INSERT OR REPLACE INTO dedupe(msg_id, created_at) VALUES (?, ?)",
|
||
|
|
(msg_id, time.time()),
|
||
|
|
)
|
||
|
|
self._conn.commit()
|
||
|
|
return False
|
||
|
|
except Exception:
|
||
|
|
logger.exception("PersistentDeduplicator error for msg_id=%s", msg_id)
|
||
|
|
return False
|
||
|
|
|
||
|
|
def _evict_expired(self):
|
||
|
|
cutoff = time.time() - self._ttl
|
||
|
|
with self._lock:
|
||
|
|
try:
|
||
|
|
self._conn.execute("DELETE FROM dedupe WHERE created_at < ?", (cutoff,))
|
||
|
|
self._conn.commit()
|
||
|
|
except Exception:
|
||
|
|
logger.exception("PersistentDeduplicator evict error")
|
||
|
|
|
||
|
|
def close(self):
|
||
|
|
with self._lock:
|
||
|
|
if self._conn:
|
||
|
|
try:
|
||
|
|
self._conn.close()
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
self._conn = None
|