src/app/storage.py

from fastapi import APIRouter
import pgsoft.pgfile as pgfile
from pgsoft.pghash.md5 import md5
from pgsoft.pgdate.date_utils import beijing
from time import sleep
import json
import os


router = APIRouter(prefix="/file", tags=["File Service"])
admin_token_md5 = "01112732d503cd3f0898d185a11c8b92"
dataset_id = "pgsoft/game"
tempdir = "game"
db_token = os.environ.get("db_token")
if db_token:
    print(db_token[:5])


@router.get("/download")
def download_file(game: str, token: str, gamecode: str):
    if md5(token) != admin_token_md5:
        print(f"[{beijing()}][download file] failed")
        return {"status": "Failure", "detail": "Invalid token"}
    game = game.strip().lower()
    filename = gamecode.strip() + ".json"
    remotepath = "/".join([game, filename[:2], filename])
    res = pgfile.download(
        dataset_id,
        remotepath=remotepath,
        repo_type="dataset",
        localdir=tempdir,
        token=db_token,
    )
    if not res:
        print(f"[{beijing()}][download file] failed")
        return {"status": "Failure", "detail": "File not found or server error"}
    with open(res, "r") as f:
        outp = json.load(f)
    print(f"[{beijing()}][download file] OK")
    return {"status": "OK", "result": outp}


@router.post("/upload")
def upload_file(game: str, token: str, content: dict):
    if md5(token) != admin_token_md5:
        print(f"[{beijing()}][upload file] failed")
        return {"status": "Failure", "detail": "Invalid token"}
    game = game.strip().lower()
    needed_keys = ["game-file", "device-id"]
    for key in needed_keys:
        if key not in content:
            print(f"[{beijing()}][upload file] failed, missed key {key}")
            return {"status": "Failure", "detail": f"missed key {key}"}

    if not isinstance(content["device-id"], str):
        print(f'[{beijing()}][upload file] failed, "device-id" is not a str')
        return {"status": "Failure", "detail": '"device-id" is not a str'}
    if not isinstance(content["game-file"], dict):
        print(f'[{beijing()}][upload file] failed, "game-file" is not a dict')
        return {"status": "Failure", "detail": '"game-file" is not a dict'}

    obj = {
        "upload-time": beijing().__str__(),
        "game-file": content["game-file"],
    }
    maxtry = 5
    for retry in range(maxtry):
        md5code = md5(obj)
        remotepath = "/".join([game, md5code[:2], md5code + ".json"])
        if not pgfile.api.file_exists(
            repo_id=dataset_id,
            filename=remotepath,
            repo_type="dataset",
            token=db_token,
        ):
            break
        sleep(0.1)
        obj["upload-time"] = beijing().__str__()
        maxtry -= 1
    if not maxtry and pgfile.api.file_exists(
        repo_id=dataset_id,
        filename=remotepath,
        repo_type="dataset",
        token=db_token,
    ):
        print(f"[{beijing()}][upload file] failed, timeout, please retry")
        return {"status": "Failure", "detail": "timeout, please retry"}
    filedir = os.path.join(tempdir, game, md5code[:2])
    if not os.path.exists(filedir):
        os.makedirs(filedir)
    filepath = os.path.join(filedir, md5code + ".json")
    content_indented = json.dumps(content, indent=4)
    with open(filepath, "w") as f:
        f.write(content_indented)
    res = pgfile.upload(
        filepath,
        remotepath,
        dataset_id,
        "dataset",
        db_token,
        f"Updated at {beijing()}",
    )
    if not res:
        print(f"[{beijing()}][upload file] failed")
        return {"status": "Failure", "detail": "server error"}
    print(f"[{beijing()}][upload file] OK")
    return {"status": "OK", "result": md5code}


@router.get("/list")
def list_files(game: str, token: str):
    if md5(token) != admin_token_md5:
        print(f"[{beijing()}][list files] failed")
        return {"status": "Failure", "detail": "Invalid token"}
    game = game.strip().lower()
    games = pgfile.list_files(
        repo_id=dataset_id,
        repo_type="dataset",
        token=db_token,
    )
    if games is None:
        print(f"[{beijing()}][list files] failed")
        return {"status": "Failure", "detail": "server error"}
    games = {
        item.split(".")[0][-32:]: item
        for item in games
        if item.endswith(".json") and item.startswith(game)
    }
    print(f"[{beijing()}][list files] OK")
    return {"status": "OK", "result": games}