devranta/examples/crawler/database.py

import logging
from typing import List

import aiosqlite

from devranta.api import Comment, Rant, UserProfile


class DatabaseManager:
    def __init__(self, db_path: str):
        self.db_path = db_path
        self._conn: aiosqlite.Connection | None = None

    async def __aenter__(self):
        logging.info(f"Connecting to database at {self.db_path}...")
        self._conn = await aiosqlite.connect(self.db_path)
        await self._conn.execute("PRAGMA journal_mode=WAL;")
        await self._conn.execute("PRAGMA foreign_keys=ON;")
        await self.create_tables()
        logging.info("Database connection successful.")
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self._conn:
            await self._conn.close()
            logging.info("Database connection closed.")

    async def create_tables(self):
        logging.info("Ensuring database tables exist...")
        await self._conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY,
                username TEXT NOT NULL UNIQUE,
                score INTEGER,
                about TEXT,
                location TEXT,
                created_time INTEGER,
                skills TEXT,
                github TEXT,
                website TEXT
            );
            CREATE TABLE IF NOT EXISTS rants (
                id INTEGER PRIMARY KEY,
                user_id INTEGER,
                text TEXT,
                score INTEGER,
                created_time INTEGER,
                num_comments INTEGER
            );
            CREATE TABLE IF NOT EXISTS comments (
                id INTEGER PRIMARY KEY,
                rant_id INTEGER,
                user_id INTEGER,
                body TEXT,
                score INTEGER,
                created_time INTEGER
            );
        """
        )
        await self._conn.commit()
        logging.info("Table schema verified.")

    async def add_rant(self, rant: Rant):
        await self._conn.execute(
            "INSERT OR IGNORE INTO rants (id, user_id, text, score, created_time, num_comments) VALUES (?, ?, ?, ?, ?, ?)",
            (
                rant["id"],
                rant["user_id"],
                rant["text"],
                rant["score"],
                rant["created_time"],
                rant["num_comments"],
            ),
        )
        await self._conn.commit()

    async def add_comment(self, comment: Comment):
        await self._conn.execute(
            "INSERT OR IGNORE INTO comments (id, rant_id, user_id, body, score, created_time) VALUES (?, ?, ?, ?, ?, ?)",
            (
                comment["id"],
                comment["rant_id"],
                comment["user_id"],
                comment["body"],
                comment["score"],
                comment["created_time"],
            ),
        )
        await self._conn.commit()

    async def add_user(self, user: UserProfile, user_id: int):
        await self._conn.execute(
            "INSERT OR IGNORE INTO users (id, username, score, about, location, created_time, skills, github, website) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                user_id,
                user["username"],
                user["score"],
                user["about"],
                user["location"],
                user["created_time"],
                user["skills"],
                user["github"],
                user["website"],
            ),
        )
        await self._conn.commit()

    async def rant_exists(self, rant_id: int) -> bool:
        async with self._conn.execute(
            "SELECT 1 FROM rants WHERE id = ? LIMIT 1", (rant_id,)
        ) as cursor:
            return await cursor.fetchone() is not None

    async def user_exists(self, user_id: int) -> bool:
        async with self._conn.execute(
            "SELECT 1 FROM users WHERE id = ? LIMIT 1", (user_id,)
        ) as cursor:
            return await cursor.fetchone() is not None

    async def get_random_user_ids(self, limit: int) -> List[int]:
        logging.info(
            f"Fetching up to {limit} random user IDs from database for seeding..."
        )
        query = "SELECT id FROM users ORDER BY RANDOM() LIMIT ?"
        async with self._conn.execute(query, (limit,)) as cursor:
            rows = await cursor.fetchall()
            user_ids = [row[0] for row in rows]
            logging.info(f"Found {len(user_ids)} user IDs to seed.")
            return user_ids