mgu/server/database.py

import sqlite3
import logging
import dataclasses

import argon2
import flask_login

import mgU

DATABASE = "./main.db"

ph = argon2.PasswordHasher()

class User(flask_login.UserMixin):

    def __init__(self, id: str, name: str):
        self.name = name
        self.id = id

    def get_id(self) -> str:
        return self.id

# My awful sqlite boilerplate because i hate orms

def load_query(name: str) -> str:
    path = f'server/sql/{name}.sql'
    return open(path).read()

class Opener():
    def __init__(self):
        self.con = sqlite3.connect(DATABASE)

    def __enter__(self) -> tuple[sqlite3.Connection, sqlite3.Cursor]:
        return self.con, self.con.cursor()

    def __exit__(self, type, value, traceback):
        self.con.commit()
        self.con.close()

# Raw SQL DTOs

@dataclasses.dataclass
class UserDTO:
    name: str
    id: int | None = None

@dataclasses.dataclass
class MessageDTO:
    user_id: UserDTO
    content: str
    time: int
    thread_id: int
    id: int | None = None

def create_user(username: str, password: str) -> mgU.User | None:
    """
    Create a new user with the given username and password.
    Returns the created User object with id, or None if creation fails.
    """
    # Hash the password
    passhash = ph.hash(password)

    try:
        with Opener() as (conn, cur):
            q = load_query("insert_user")
            cur.execute(q, [username, passhash])
            user_id = cur.lastrowid

            return mgU.User(name=username, id=user_id)
    except sqlite3.IntegrityError:
        # Username already exists or other constraint violation
        return None

def insert_message(message: mgU.Message | MessageDTO) -> None:
    """
    Insert a new message (from mgU.Message dataclass) and return a new mgU.Message with id, or None on failure.
    """
    if type(message) == mgU.Message:
        thread_id = message.thread
        user_id = message.user.id
    elif type(message) == MessageDTO:
        thread_id = message.thread_id
        user_id = message.user_id


    with Opener() as (conn, cur):
        q = load_query("insert_message")
        cur.execute(q, [
            user_id,
            thread_id,
            message.content,
            message.time
        ])


def get_messages_by_thread(thread_id: int) -> list[mgU.Message]:
    """
    Return a list of mgU.Message objects for a given thread_id, ordered by id ascending.
    """

    with Opener() as (conn, cur):
        q = load_query("get_messages_by_thread")
        cur.execute(q, [thread_id])
        results = cur.fetchall()
    return [
        mgU.Message(
            user=mgU.User(
                name=msg[0],
                id=msg[1]
            ),
            thread=msg[4],
            content=msg[2],
            time=msg[3],
            id=msg[5]
        )
        for msg in results
    ]

def get_user_forums(user_id: int) -> list[mgU.Forum]:
    q = load_query("get_user_forums")
    with Opener() as (conn, cur):
        cur.execute(q, [user_id])
        results = cur.fetchall()
        return [
            mgU.Forum(
                name=forum[0],
                id=forum[1]
            )
            for forum in results
        ]

def get_forum_threads(forum_id: int) -> list[mgU.Thread]:
    q = load_query("get_forum_threads")
    with Opener() as (conn, cur):
        cur.execute(q, [forum_id])
        results = cur.fetchall()
        return [
            mgU.Thread(
                name=thread[0],
                forum=thread[1],
                id=thread[2]
            )
            for thread in results
        ]

# For flask_login
def get_user(username: str, api_key: str) -> User | None:

    q = load_query("get_user_passhash")
    with Opener() as (conn, cur):
        cur.execute(q)
        results = cur.fetchall()

    if not results:
        return None

    row = results[0]

    user_id, user_name, passhash = row[0], row[1], row[2]

    try:
        ph.verify(passhash, api_key)
    except argon2.exceptions.VerifyMismatchError:
        return None

    return User(str(user_id), str(user_name))

# Returns a list of user_ids that are owners of a forum
def get_forum_owners(forum_id: int) -> list[int]:

    q = load_query("get_forum_owners")
    with Opener() as (conn, cur):
        cur.execute(q)
        results = cur.fetchall()

    if not results:
        return []

    user_ids: list[int] = [row[0] for row in results]

    return user_ids

def create_thread(name: str, forum_id: int) -> mgU.Thread | None:
    """
    Create a new thread in the specified forum.
    Returns the created Thread object with id, or None if creation fails.
    """
    try:
        with Opener() as (conn, cur):
            q = load_query("insert_thread")
            cur.execute(q, [name, forum_id])
            thread_id = cur.lastrowid

            return mgU.Thread(name=name, forum=forum_id, id=thread_id)
    except sqlite3.IntegrityError:
        # Forum doesn't exist or other constraint violation
        return None