|
import logging
|
|
import aiosqlite
|
|
from typing import List
|
|
from devranta.api import Rant, Comment, 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
|
|
|