PyGetGPT/dao/db/user.py

import sqlite3, time
import configUtil

def getConnection() -> sqlite3.Connection:
    return sqlite3.connect(configUtil.ConfigUtil("config.ini").get("database","path"))


def init():
    conn = getConnection()
    cursor = conn.cursor()
    cursor.execute(
        '''
        CREATE TABLE User (
            uid TEXT,
            email TEXT,
            password_hash TEXT,
            created_at INT,
            surplus INT);
        '''
        )
    conn.commit()
    cursor.close()
    conn.close()
    print("数据库初始化完成")


def addUser(uid: str, email: str, password_hash: str) -> bool:
    conn = getConnection()
    cursor = conn.cursor()
    try:
        cursor.execute(
            "INSERT INTO User (uid, email, password_hash, created_at, surplus) VALUES (?, ?, ?, ?, ?);",
            [uid, email, password_hash, int(time.time()), 0]
        )
        conn.commit()
        return True
    except sqlite3.IntegrityError:
        return False
    

def checkUser(uid: str,password_hash) -> bool:
    """检查用户与密码是否合法 可输入邮箱或uid"""
    conn = getConnection()
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM User WHERE ( uid = ? AND password_hash = ? ) OR ( email = ? AND password_hash = ? )",[uid,password_hash,uid,password_hash])
    result = cursor.fetchone()
    return result != None


def getUser(uid: str) -> dict:
    """获取用户信息"""
    conn = getConnection()
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM User WHERE uid = ?",[uid])
    result = cursor.fetchone()
    return {"uid":result[0],"email":result[1],"password_hash":result[2],"created_at":result[3],"surplus":result[4]}


def updateUserSurplus(uid: str, surplus: int) -> bool:
    """更新用户剩余额度"""
    conn = getConnection()
    cursor = conn.cursor()
    try:
        cursor.execute("UPDATE User SET surplus = ? WHERE uid = ?",[surplus,uid])
        conn.commit()
        cursor.close()
        conn.close()
    except: return False
    return True

def getUserSurplus(uid: str) -> int:
    """获取用户剩余额度"""
    return getUser(uid)["surplus"]


def updateUserPasswd(uid: str, password_hash: str) -> bool:
    """更新用户密码"""
    conn = getConnection()
    cursor = conn.cursor()
    try:
        cursor.execute("UPDATE User SET password_hash = ? WHERE uid = ?",[password_hash,uid])
        conn.commit()
        cursor.close()
        conn.close()
    except: return False
    return True