nas-tools/info-engine/db.py

import json
import os
from datetime import datetime, timezone

import aiosqlite

DB_PATH = os.environ.get("DB_PATH", "/app/data/info_engine.db")


async def init_db():
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute(
            """
            CREATE TABLE IF NOT EXISTS info_items (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                source TEXT NOT NULL,
                title TEXT NOT NULL,
                summary TEXT NOT NULL,
                url TEXT NOT NULL UNIQUE,
                tags TEXT NOT NULL,
                published_at DATETIME,
                collected_at DATETIME NOT NULL,
                content_text TEXT,
                content_html TEXT
            )
            """
        )
        await db.execute("CREATE INDEX IF NOT EXISTS idx_info_items_collected_at ON info_items(collected_at DESC)")
        await db.execute("CREATE INDEX IF NOT EXISTS idx_info_items_published_at ON info_items(published_at DESC)")
        await db.execute("CREATE INDEX IF NOT EXISTS idx_info_items_source ON info_items(source)")
        await db.commit()


async def upsert_items(items: list[dict]) -> int:
    if not items:
        return 0

    inserted = 0
    now = datetime.now(timezone.utc).isoformat()

    async with aiosqlite.connect(DB_PATH) as db:
        for item in items:
            cursor = await db.execute(
                """
                INSERT OR IGNORE INTO info_items
                (source, title, summary, url, tags, published_at, collected_at, content_text, content_html)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    item["source"],
                    item["title"],
                    item.get("summary", ""),
                    item["url"],
                    json.dumps(item.get("tags", []), ensure_ascii=False),
                    item.get("published_at"),
                    now,
                    item.get("content_text"),
                    item.get("content_html"),
                ),
            )
            inserted += cursor.rowcount
        await db.commit()

    return inserted