Groove_Coaster_2_Server/api/database.py

import sqlalchemy
from sqlalchemy import Table, Column, Integer, String, DateTime, Boolean, JSON
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy import select, update

from config import REDIS_ADDRESS, USE_REDIS_CACHE

import os
import databases
import datetime
import time

if USE_REDIS_CACHE:
    import redis.asyncio as aioredis

redis = None

DB_NAME = "player.db"
DB_PATH = os.path.join(os.getcwd(), DB_NAME)
DATABASE_URL = f"sqlite+aiosqlite:///{DB_PATH}"

CACHE_DB_NAME = "cache.db"
CACHE_DB_PATH = os.path.join(os.getcwd(), CACHE_DB_NAME)
CACHE_DATABASE_URL = f"sqlite+aiosqlite:///{CACHE_DB_PATH}"

cache_database = databases.Database(CACHE_DATABASE_URL)
cache_metadata = sqlalchemy.MetaData()

database = databases.Database(DATABASE_URL)
metadata = sqlalchemy.MetaData()

user = Table(
    "user",
    metadata,
    Column("id", Integer, primary_key=True, autoincrement=True),
    Column("username", String(20), unique=True, nullable=False),
    Column("password_hash", String(255), nullable=False),
    Column("device_id", String(512)),
    Column("data", String, nullable=True),
    Column("save_id", String, nullable=True),
    Column("crc", Integer, nullable=True),
    Column("timestamp", DateTime, default=None),
    Column("coin_mp", Integer, default=1),
)

daily_reward = Table(
    "daily_reward",
    metadata,
    Column("id", Integer, primary_key=True, autoincrement=True),
    Column("device_id", String(512)),
    Column("timestamp", DateTime, default=datetime.datetime.utcnow),
    Column("my_stage", String),
    Column("my_avatar", String),
    Column("item", String),
    Column("day", Integer),
    Column("coin", Integer),
    Column("lvl", Integer),
    Column("title", Integer),
    Column("avatar", Integer),
)

result = Table(
    "result",
    metadata,
    Column("rid", Integer, primary_key=True, autoincrement=True),
    Column("vid", String(512), nullable=False),
    Column("tid", String(512), nullable=False),
    Column("sid", Integer, nullable=False),
    Column("stts", String(64)),
    Column("id", Integer),
    Column("mode", Integer),
    Column("avatar", Integer),
    Column("score", Integer),
    Column("high_score", String(128)),
    Column("play_rslt", String(128)),
    Column("item", Integer),
    Column("os", String(16)),
    Column("os_ver", String(16)),
    Column("ver", String(16)),
    Column("mike", Integer),
)

whitelist = Table(
    "whitelist",
    metadata,
    Column("id", String(512), primary_key=True)
)
blacklist = Table(
    "blacklist",
    metadata,
    Column("id", String(512), primary_key=True),
    Column("reason", String(256))
)

batch_token = Table(
    "batch_token",
    metadata,
    Column("id", Integer, primary_key=True, autoincrement=True),
    Column("token", String(512), nullable=False, unique=True),
    Column("sid", Integer, nullable=False),
    Column("verification_name", String(64), nullable=False, default=False),
    Column("verification_id", String(64), nullable=False),
    Column("expire_at", Integer, default=int(time.time()) + 1800),
)

admins = Table(
    "admins",
    metadata,
    Column("id", Integer, primary_key=True, autoincrement=True),
    Column("username", String(32), unique=True, nullable=False),
    Column("password", String(64), nullable=False),
    Column("token", String(256), unique=True, nullable=True)
)

ranking_cache = Table(
    "ranking_cache",
    cache_metadata,
    Column("id", Integer, primary_key=True, autoincrement=True),
    Column("key", String(16), nullable=False),
    Column("value", JSON, nullable=False),
    Column("expire_at", Integer)
)

async def init_db():
    global redis
    if not os.path.exists(DB_PATH):
        print("[DB] Creating new database:", DB_PATH)

    if not os.path.exists(CACHE_DB_PATH):
        print("[DB] Creating new cache database:", CACHE_DB_PATH)
    
    cache_engine = create_async_engine(CACHE_DATABASE_URL, echo=False)
    async with cache_engine.begin() as conn:
        await conn.run_sync(cache_metadata.create_all)
    await cache_engine.dispose()
    
    engine = create_async_engine(DATABASE_URL, echo=False)
    
    async with engine.begin() as conn:
        await conn.run_sync(metadata.create_all)
    
    await engine.dispose()
    print("[DB] Database initialized successfully.")
    await ensure_user_columns()
    if USE_REDIS_CACHE:
        print("[DB] Connecting to Redis at", REDIS_ADDRESS)
        redis = await aioredis.from_url("redis://" + REDIS_ADDRESS)

async def get_user_data(uid, data_field):
    query = select(user.c[data_field]).where(user.c.device_id == uid[b'vid'][0].decode())
    result = await database.fetch_one(query)
    return result[data_field] if result else None

async def set_user_data(uid, data_field, new_data):
    query = (
        update(user)
        .where(user.c.device_id == uid[b'vid'][0].decode())
        .values({data_field: new_data})
    )
    async with database.transaction():
        await database.execute(query)
        
async def check_whitelist(uid):
    query = select(whitelist.c.id).where(whitelist.c.id == uid[b'vid'][0].decode())
    result = await database.fetch_one(query)
    return result is not None

async def check_blacklist(uid):
    device_id = uid[b'vid'][0].decode()
    user_data = await get_user_data(uid, "username")
    username = user_data[0] if user_data else None

    query = select(blacklist.c.id).where(
        (blacklist.c.id == device_id) | (blacklist.c.id == username)
    )
    result = await database.fetch_one(query)
    return result is None

async def ensure_user_columns():
    import aiosqlite

    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute("PRAGMA table_info(user);") as cursor:
            columns = [row[1] async for row in cursor]

        alter_needed = False
        if "save_id" not in columns:
            await db.execute("ALTER TABLE user ADD COLUMN save_id TEXT;")
            alter_needed = True
        if "coin_mp" not in columns:
            await db.execute("ALTER TABLE user ADD COLUMN coin_mp INTEGER DEFAULT 1;")
            alter_needed = True
        if alter_needed:
            await db.commit()
            print("[DB] Added missing columns to user table.")