omni-parse/omni/source_state.py

"""Состояние источников: включён/выключен, пользовательские источники, результаты проверок.

Две таблицы в той же SQLite, что и записи:
  source_flags  (name, enabled, status, items, note, checked)
  source_custom (name, category, kind, url, adapter, tags, limit, detail, added)
"""
import contextlib
import json
import sqlite3
import time

from . import config

SCHEMA = """
CREATE TABLE IF NOT EXISTS source_flags (
  name TEXT PRIMARY KEY,
  enabled INTEGER DEFAULT 1,
  status INTEGER,
  items INTEGER,
  note TEXT,
  checked INTEGER
);
CREATE TABLE IF NOT EXISTS source_custom (
  name TEXT PRIMARY KEY,
  category TEXT, kind TEXT, url TEXT, adapter TEXT,
  tags TEXT, lim INTEGER DEFAULT 20, detail INTEGER DEFAULT 0, added INTEGER
);
"""


def conn():
    c = sqlite3.connect(str(config.DB), timeout=30)
    c.execute("PRAGMA journal_mode=WAL")
    c.executescript(SCHEMA)
    return c


@contextlib.contextmanager
def db():
    """Соединение с гарантированным close(): `with conn()` коммитит, но не закрывает,
    и хендлы копятся (замер: 200 вызовов == +76 handles)."""
    c = conn()
    try:
        yield c
        c.commit()
    except Exception:
        c.rollback()
        raise
    finally:
        c.close()


def flags():
    """{name: {enabled, status, items, note, checked}}"""
    with db() as c:
        rows = c.execute("SELECT name, enabled, status, items, note, checked FROM source_flags").fetchall()
    return {r[0]: {"enabled": bool(r[1]), "status": r[2], "items": r[3], "note": r[4], "checked": r[5]}
            for r in rows}


def set_enabled(name, enabled):
    now = int(time.time())
    with db() as c:
        c.execute("INSERT INTO source_flags(name, enabled, checked) VALUES(?,?,?) "
                  "ON CONFLICT(name) DO UPDATE SET enabled=excluded.enabled",
                  (name, 1 if enabled else 0, now))
    return {"name": name, "enabled": bool(enabled)}


def record_check(name, status, items, note=""):
    now = int(time.time())
    with db() as c:
        c.execute("INSERT INTO source_flags(name, enabled, status, items, note, checked) VALUES(?,1,?,?,?,?) "
                  "ON CONFLICT(name) DO UPDATE SET status=excluded.status, items=excluded.items, "
                  "note=excluded.note, checked=excluded.checked",
                  (name, status, items, str(note)[:160], now))
    return {"name": name, "status": status, "items": items, "note": str(note)[:160], "checked": now}


def custom():
    with db() as c:
        rows = c.execute("SELECT name, category, kind, url, adapter, tags, lim, detail, added "
                         "FROM source_custom ORDER BY added DESC").fetchall()
    out = []
    for r in rows:
        out.append({"name": r[0], "category": r[1], "kind": r[2], "url": r[3], "adapter": r[4],
                    "tags": tuple(json.loads(r[5] or "[]")), "limit": r[6] or 20,
                    "detail": bool(r[7]), "extra": {}, "added": r[8]})
    return out


def add_custom(name, category, url, kind=None, adapter=None, tags=(), limit=20, detail=False):
    name = (name or "").strip()
    url = (url or "").strip()
    if not name or not url:
        raise ValueError("нужны name и url")
    if not url.startswith("http"):
        raise ValueError("url должен начинаться с http")
    kind, adapter = kind or guess_kind(url), adapter or guess_adapter(url, kind)
    with db() as c:
        c.execute("INSERT OR REPLACE INTO source_custom(name, category, kind, url, adapter, tags, lim, detail, added) "
                  "VALUES(?,?,?,?,?,?,?,?,?)",
                  (name, category or "custom", kind, url, adapter,
                   json.dumps(list(tags), ensure_ascii=False), int(limit), 1 if detail else 0, int(time.time())))
    return {"name": name, "category": category or "custom", "kind": kind, "url": url,
            "adapter": adapter, "tags": list(tags), "limit": limit, "detail": bool(detail)}


def remove_custom(name):
    with db() as c:
        c.execute("DELETE FROM source_custom WHERE name=?", (name,))
        c.execute("DELETE FROM source_flags WHERE name=?", (name,))
    return {"removed": name}


def guess_kind(url):
    u = url.lower().rstrip("/")
    path = u.split("://", 1)[-1]
    path = path[path.find("/"):] if "/" in path else ""
    last = path.rsplit("/", 1)[-1]
    if "github.com" in u and u.count("/") <= 4:
        return "github"
    if last in ("rss", "feed", "atom", "rss.xml", "feed.xml", "atom.xml", "index.xml", "rss2"):
        return "rss"
    if u.endswith(".json") or "/latest.json" in u or "api." in u or "/api/" in u:
        return "json"
    if any(x in u for x in (".rss", "rss.xml", "/feed", ".atom", "/rss/", "atom.xml")):
        return "rss"
    return "html"


def guess_adapter(url, kind=None):
    u = url.lower()
    kind = kind or guess_kind(url)
    if "github.com" in u:
        return "github_org"
    if "discourse" in u or u.endswith("/latest.json"):
        return "discourse"
    if kind == "json":
        return "generic_json"
    if u.endswith(".json") or "api." in u:
        return "generic_json"
    if "t.me/s/" in u:
        return "tg_web"
    if kind == "rss":
        return "rss"
    return "html_listing"