mgu/server/database.py

import sqlite3
import logging
import dataclasses

import argon2
import flask_login

from sqlalchemy import create_engine, Column, Integer, ForeignKey, Text, select
from sqlalchemy.orm import sessionmaker, DeclarativeBase
from sqlalchemy.exc import SQLAlchemyError

import mgU

DATABASE = "./main2.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()

# orm garbage

class Base(DeclarativeBase):
    pass

class MessageORM(Base):
    __tablename__ = "messages"
    id = Column(Integer, primary_key=True, autoincrement=True)
    user = Column(Integer, ForeignKey("users.id"), nullable=False)
    thread = Column(Integer, ForeignKey("threads.id"), nullable=False)
    content = Column(Text, nullable=False)
    time = Column(Integer, nullable=False)

class UserORM(Base):
    __tablename__ = "users"
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(Text, nullable=False, unique=True)
    passhash = Column(Text, nullable=False)

class ForumORM(Base):
    __tablename__ = "forums"
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(Text, nullable=False)

class UserForumsORM(Base):
    __tablename__ = "user_forums"
    user_id = Column(Integer, ForeignKey("users.id"), primary_key=True, nullable=False)
    forum_id = Column(Integer, ForeignKey("forums.id"), primary_key=True, nullable=False)

class ThreadORM(Base):
    __tablename__ = "threads"
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(Text, nullable=False)
    forum = Column(Integer, ForeignKey("forums.id"), nullable=False)

# Raw SQL DTOs

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

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

engine = create_engine(f"sqlite:///{DATABASE}", echo=False, future=True)
SessionLocal = sessionmaker(bind=engine)

def insert_message(message: MessageORM) -> None:
    """
    Insert a new message (from mgU.Message dataclass) and return a new mgU.Message with id, or None on failure.
    """
    session = SessionLocal()
    try:
        session.add(message)
        session.commit()
    except SQLAlchemyError:
        session.rollback()
        return None
    finally:
        session.close()

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]:
    session = SessionLocal()
    try:
        q = (
            select(ForumORM.id, ForumORM.name)
            .join_from(ForumORM, UserForumsORM)
            .where(UserForumsORM.user_id == user_id)
        )
        results = session.execute(q)
        return [
            mgU.Forum(
                name=forum.name,
                id=forum.id
            )
            for forum in results
        ]
    finally:
        session.close()

def get_forum_threads(forum_id: int) -> list[mgU.Thread]:
    session = SessionLocal()
    try:
        q = (
            select(ThreadORM.id, ThreadORM.name, ThreadORM.forum)
            .where(ThreadORM.forum == forum_id)
        )

        results = session.execute(q)
        return [
            mgU.Thread(
                name=thread.name,
                forum=thread.forum,
                id=thread.id
            )
            for thread in results
        ]
    finally:
        session.close()

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

    session = SessionLocal()
    try:
        results = session.query(UserORM).where(UserORM.name == username)
    finally:
        session.close()

    if not results:
        return None

    user_row = results[0]

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

    return User(str(user_row.id), str(user_row.name))