|
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
|