niming_backend/utils/dbhelper.py

from typing import Tuple, Dict, List

from flask import make_response, Response, jsonify
from sqlalchemy.orm import sessionmaker, aliased
from sqlalchemy import desc, func, literal, and_

from utils import pgclass
from utils.misc import error, error_proto
from protobuf_files import niming_pb2

class db:
    _engine = None

    @classmethod
    def __init__(cls, engine):
        cls._engine = engine

    @classmethod
    def getsession(cls):
        Session = sessionmaker(bind=cls._engine)
        return Session()


# role (general) (owner) (admin)
# 獲取單一文章
def solo_article_fetcher(role:str, key) -> Tuple[Dict | bytes, int]: # admin, owner, general
    table = pgclass.SQLarticle  # main
    table2 = aliased(table)     # comment
    with db.getsession() as session:
        # query
        res = session.query(table.id, 
                            table.content, 
                            table.reference,
                            table.file_list,
                            table.hash,
                            table.igid,
                            table.mark,
                            table.ip,
                            func.coalesce(func.array_agg(table2.id), literal([])).label("comments"))
        
        if role == "owner":
            res = res.join(table2, table2.reference == table.id, isouter=True) \
                    .filter(table.hash == key[0], table.id == key[1])
        elif role == "admin":
            res = res.join(table2, table2.reference == table.id, isouter=True) \
                    .filter(table.id == key)
        elif role == "general":
            res = res.join(table2, and_(table2.reference == table.id, table2.mark == "visible"), isouter=True) \
                    .filter(table.id == key, table.mark == "visible")
        
        res = res.group_by(table.id, table.content, table.reference, table.file_list,
                           table.hash, table.igid, table.mark, table.ip).first()
        if res is None: 
            return error_proto("fetch", "Post not found"), 404

        # mapping
        one = {
            "id": res[0],
            "content":res[1],
            "igid":res[5],
            "mark":res[6],
            "reference":res[2],
            "files_id":res[3],
            "comments":res[8]
        }
        
        if role == "admin": 
            one["ip"] = res[7]
        if role == "owner" or role == "admin": 
            one["hash"] = res[4]
        
    return one, 200


# 獲取文章列表
def multi_article_fetcher(role:str, page:str, count:int) -> Tuple[bytes, int]: # general, admin
    # checker
    if page is None or not page.isdigit(): 
        return error_proto("fetch", "Arguments error"), 400
    page = int(page)*count

    table = pgclass.SQLarticle
    resfn = niming_pb2.FetchResponse(
        status = niming_pb2.Status.Success
    )

    with db.getsession() as session:
        # query
        res = session.query(table)
        if role == "general":
            res = res.filter(table.mark == "visible", table.reference == None)
        elif role == "admin":
            res = res.filter(table.reference == None)
        res = res.order_by(desc(table.id)).offset(page).limit(count).all()
        
        # mapping
        for r in res:
            one = niming_pb2.FetchResponse.Message(
                id = r.id,
                content = r.content,
                files_id = r.file_list,
                igid = r.igid,
                mark = r.mark,
                ref = r.reference
            )
            if role == "admin": # 如果是管理員 多給ip 跟 hash
                one.hash = r.hash
                one.ip = r.ip
            resfn.posts.append(one)
    
    return resfn.SerializeToString(), 200


# 刪除文章
def solo_article_remover(role:str, hash:str=None, id:int=None) -> Tuple[Dict | bytes, int]: # admin, owner
    key = None
    if role == "admin": key = id
    elif role == "owner": key = (hash, id)
    
    table = pgclass.SQLarticle
    with db.getsession() as session:
        # 獲取本體
        if role == "admin":
            res = session.query(table).filter(table.id == key).first()
        elif role == "owner":
            res = session.query(table).filter(table.hash == key[0], table.id == key[1]).first()
        if res is None: # 檢查本體是否存在
            return error_proto("fetch", "Post not found!"), 404
        # 刪本體
        session.delete(res)
        session.commit()
    
    return {"id":res.id, "mark":res.mark}, 200


# 獲取檔案
def solo_file_fetcher(role:str, id:int) -> Tuple[Response, int]: # general, admin
    table = pgclass.SQLarticle
    ftab  = pgclass.SQLfile

    with db.getsession() as session:
        fres = session.query(ftab).filter(ftab.id == id).first()
        if fres is None: # 檢查檔案是否存在
            return error("File not found"), 404
        
        if role == "general":
            article = session.query(table).filter(table.hash == fres.reference, table.mark == 'visible').first()
        elif role == "admin":
            article = session.query(table).filter(table.hash == fres.reference).first()
        if article is None: # 檢查文章本體是否存在/可以閱覽
            return error("File not found"), 404
    
    resp = make_response(fres.binary)
    resp.headers.set("Content-Type", fres.type)
    resp.headers.set("Content-Disposition", f"attachment; filename=file{fres.id}")
    return resp, 200