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