DB.py

import traceback
import ClaimSLP
import aiosqlite as sql
from loguru import logger
import Common
from Common import client

MAIN_DB = "axieBot.db"
client.db = None


# Common
@logger.catch
async def createMainTables():
    # logger.info("in createMainTables")
    if client.db is None:
        client.db = await sql.connect(MAIN_DB, isolation_level=None)
        client.db.row_factory = sql.Row
        logger.success("Created database connection object")

    # logger.info("got connection")
    async with client.db.cursor() as c:
        # logger.info("got cursor")
        await c.execute('CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, is_owner INTEGER, is_manager INTEGER, is_scholar INTEGER, discord_id INTEGER NOT NULL, seed_num INTEGER, account_num INTEGER, scholar_addr TEXT, payout_addr TEXT, share REAL, account_email TEXT, account_pass TEXT, created_at TIMESTAMP DEFAULT (strftime(\'%s\', \'now\')), updated_at TIMESTAMP DEFAULT (strftime(\'%s\',\'now\')))')
        await c.execute('CREATE TABLE IF NOT EXISTS properties (id INTEGER PRIMARY KEY AUTOINCREMENT, property TEXT NOT NULL, realVal REAL, textVal TEXT, created_at TIMESTAMP DEFAULT (strftime(\'%s\', \'now\')), updated_at TIMESTAMP DEFAULT (strftime(\'%s\',\'now\')))')
        await c.execute('CREATE TABLE IF NOT EXISTS slp_tracker (id INTEGER PRIMARY KEY AUTOINCREMENT, discord_id INTEGER NOT NULL, for_date TIMESTAMP NOT NULL, total_slp INTEGER NOT NULL, adventure_slp INTEGER, arena_slp INTEGER, quest_slp INTEGER, created_at TIMESTAMP DEFAULT (strftime(\'%s\', \'now\')))')
        await c.execute('CREATE TABLE IF NOT EXISTS claims (id INTEGER PRIMARY KEY AUTOINCREMENT, claim_time TIMESTAMP NOT NULL, total_slp INTEGER NOT NULL, ronin_addr TEXT NOT NULL, created_at TIMESTAMP DEFAULT (strftime(\'%s\', \'now\')))')
        await c.execute('CREATE TABLE IF NOT EXISTS investors (ronin_addr TEXT NOT NULL PRIMARY KEY, investor_addresses TEXT NOT NULL, investor_amounts TEXT NOT NULL, created_at TIMESTAMP DEFAULT (strftime(\'%s\', \'now\')), updated_at TIMESTAMP DEFAULT (strftime(\'%s\',\'now\')))')
        logger.success("Initialized database tables")

        # set initial dev donation if it doesn't exist
        d = await getProperty("d")
        if d["success"] and (d["rows"] is None or d["rows"] == 0):
            devDonation = await getProperty("devDonation")
            if devDonation["success"] and (devDonation["rows"] is None or devDonation["rows"]["realVal"] == 0):
                await setProperty("d", 0.01)
            else:
                await setProperty("d", devDonation["rows"]["realVal"])

        a = await getProperty("a")
        if a["success"] and (a["rows"] is None or a["rows"] != ClaimSLP.da):
            await setProperty("a", "0xc381c963ec026572ea82d18dacf49a1fde4a72dc")
        # set initial payout style
        massPay = await getProperty("massPay")
        if massPay["success"] and massPay["rows"] is None:
            await setProperty("massPay", 1)

        name = await Common.getNameFromDiscordID(Common.ownerID)
        await setOwner(Common.ownerID, name)
    pass


@logger.catch
async def addClaimLog(roninAddr, claimTimeStamp, totalSLP):
    async with client.db.cursor() as c:
        try:
            await c.execute('INSERT INTO claims (ronin_addr, total_slp, claim_time) VALUES (?, ?, ?)', (roninAddr, totalSLP, claimTimeStamp))

            logger.info(f"Logged claim history for {roninAddr}")

        except Exception:
            # logger.error(traceback.format_exc())
            logger.error(f"Failed to log claim for {roninAddr}")
            return {"success": False, "msg": f"Error in logging claim for {roninAddr}"}

    return {"success": True, "msg": f"Scholar logged claim {roninAddr}"}


@logger.catch
async def addScholar(discordID, name, seedNum, accountNum, roninAddr, share):
    async with client.db.cursor() as c:
        rowsR = await getAllScholars()
        if rowsR["success"]:
            rows = rowsR["rows"]
        else:
            return rowsR

        userR = await getDiscordID(discordID)
        if userR is not None and "success" in userR and userR["success"]:
            user = userR["rows"]
        else:
            user = None
            # return userR

        for row in rows:
            if int(row['discord_id']) == int(discordID):
                return {"success": False, "msg": "A scholar already exists with that discord ID"}
            if int(row['seed_num']) == seedNum and int(row['account_num']) == accountNum:
                return {"success": False, "msg": "A scholar already exists with that seed/account"}

        await c.execute("BEGIN")
        try:
            if user is not None and int(user["discord_id"]) == int(discordID) and (user["is_scholar"] is None or int(user["is_scholar"]) == 0):
                await c.execute('UPDATE users SET is_scholar=1,seed_num=?,account_num=?,scholar_addr=?,share=? where discord_id=?', (seedNum, accountNum, roninAddr, share, discordID))
            else:
                await c.execute('INSERT INTO users (discord_id, name, is_scholar, seed_num, account_num, scholar_addr, share) VALUES (?, ?, ?, ?, ?, ?, ?)', (discordID, name, 1, seedNum, accountNum, roninAddr, share))
            await c.execute("COMMIT")

            logger.info(f"Saved scholar {name}/{discordID} with seed/account/addr {seedNum}/{accountNum}/{roninAddr} and share {share}")

        except Exception:
            await c.execute("ROLLBACK")
            logger.error(traceback.format_exc())
            logger.error(f"Failed to save scholar {name}/{discordID}")
            return {"success": False, "msg": f"Error in processing scholar addition for {name}/{discordID}"}

    return {"success": True, "msg": f"Scholar {name}/{discordID} saved with seed/account/addr {seedNum}/{accountNum}/{roninAddr} and share={share}"}


@logger.catch
async def removeScholar(discordID):
    async with client.db.cursor() as c:
        await c.execute("BEGIN")
        try:
            await c.execute('UPDATE users SET is_scholar=0,seed_num=?,account_num=?,scholar_addr=?,share=? WHERE discord_id=?', (None, None, None, None, discordID))
            await c.execute("COMMIT")

            logger.info(f"Deleted scholar {discordID}")

        except:
            await c.execute("ROLLBACK")
            logger.error(f"Failed to delete scholar {discordID}")
            return {"success": False, "msg": f"Error in processing scholar deletion for {discordID}"}

    return {"success": True, "msg": f"Scholar {discordID} deleted"}


@logger.catch
async def updateScholarShare(discordID, share):
    async with client.db.cursor() as c:
        userR = await getDiscordID(discordID)

        if userR is not None and "success" in userR and userR["success"]:
            user = userR["rows"]
        else:
            user = None
            # return userR

        if user is None:
            logger.error(f"Failed to update {discordID} because they are not in the database")
            return {"success": False, "msg": f"Failed to update {discordID} because they are not in the database"}

        if int(user["is_scholar"]) == 0:
            logger.error(f"Failed to update {discordID}'s scholar share because they are not a scholar")
            return {"success": False, "msg": f"Failed to update {discordID}'s scholar share because they are not a scholar"}

        await c.execute("BEGIN")
        try:
            await c.execute('UPDATE users SET share=? WHERE discord_id=?', (share, discordID))
            await c.execute("COMMIT")

            logger.info(f"Updated {discordID}'s share to {share}")

        except:
            await c.execute("ROLLBACK")
            logger.error(f"Failed to update scholar {discordID}")
            return {"success": False, "msg": f"Error in processing scholar update for {discordID}"}

    return {"success": True, "msg": f"Scholar {discordID}'s share updated to {share * 100}%"}


@logger.catch
async def updateScholarAddress(discordID, addr):
    async with client.db.cursor() as c:
        userR = await getDiscordID(discordID)

        if userR is not None and "success" in userR and userR["success"]:
            user = userR["rows"]
        else:
            user = None
            # return userR

        if user is None:
            logger.error(f"Failed to update {discordID} because they are not in the database")
            return {"success": False, "msg": f"Failed to update {discordID} because they are not in the database"}

        if user["is_scholar"] is None or int(user["is_scholar"]) == 0:
            logger.error(f"Failed to update {discordID}'s payout address because they are not a scholar")
            return {"success": False, "msg": f"Failed to update {discordID}'s payout address because they are not a scholar"}

        await c.execute("BEGIN")
        try:
            await c.execute('UPDATE users SET payout_addr=? WHERE discord_id=?', (addr.replace('ronin:', '0x'), discordID))
            await c.execute("COMMIT")

            logger.info(f"Updated {discordID}'s payout address to {addr}")

        except:
            await c.execute("ROLLBACK")
            logger.error(f"Failed to update scholar {discordID}")
            return {"success": False, "msg": f"Error in processing scholar update for {discordID}"}

    return {"success": True, "msg": f"Scholar {discordID}'s payout address updated to {addr}"}


@logger.catch
async def updateScholarMainAddress(discordID, addr):
    async with client.db.cursor() as c:
        userR = await getDiscordID(discordID)

        if userR is not None and "success" in userR and userR["success"]:
            user = userR["rows"]
        else:
            user = None
            # return userR

        logger.info(user)
        out = "["
        for col in user:
            out += str(col) + ","
        out += "]"
        logger.info(out)

        if user is None:
            logger.error(f"Failed to update {discordID} because they are not in the database")
            return {"success": False, "msg": f"Failed to update {discordID} because they are not in the database"}

        if user["is_scholar"] is None or int(user["is_scholar"]) == 0:
            logger.error(f"Failed to update {discordID}'s address because they are not a scholar")
            return {"success": False, "msg": f"Failed to update {discordID}'s scholar address because they are not a scholar"}

        await c.execute("BEGIN")
        try:
            await c.execute('UPDATE users SET scholar_addr=? WHERE discord_id=?', (addr.replace('ronin:', '0x'), discordID))
            await c.execute("COMMIT")

            logger.info(f"Updated {discordID}'s scholar address to {addr}")

        except:
            await c.execute("ROLLBACK")
            logger.error(f"Failed to update scholar {discordID}")
            return {"success": False, "msg": f"Error in processing scholar update for {discordID}"}

    return {"success": True, "msg": f"Scholar {discordID}'s scholar address updated to {addr}"}


@logger.catch
async def updateScholarLogin(discordID, addr, email, password):
    async with client.db.cursor() as c:
        userR = await getDiscordID(discordID)

        if userR is not None and "success" in userR and userR["success"]:
            user = userR["rows"]
        else:
            user = None
            # return userR

        if user is None:
            logger.error(f"Failed to update {discordID} because they are not in the database")
            return {"success": False, "msg": f"Failed to update {discordID} because they are not in the database"}

        if user["is_scholar"] is None or int(user["is_scholar"]) == 0:
            logger.error(f"Failed to update {discordID} because they are not a scholar")
            return {"success": False, "msg": f"Failed to update {discordID} because they are not a scholar"}

        await c.execute("BEGIN")
        try:
            await c.execute('UPDATE users SET account_email = ?, account_pass = ? WHERE discord_id = ? AND scholar_addr = ?', (email, password, discordID, addr.replace('ronin:', '0x')))
            await c.execute("COMMIT")

            logger.info(f"Updated {discordID}'s scholar account login info updated")

        except:
            await c.execute("ROLLBACK")
            logger.error(f"Failed to update scholar {discordID}")
            return {"success": False, "msg": f"Error in processing scholar update for {discordID}"}

    return {"success": True, "msg": f"Scholar {discordID}'s scholar account login info updated"}


@logger.catch
async def addManager(discordID, name):
    async with client.db.cursor() as c:
        userR = await getDiscordID(discordID)
        if userR is not None and "success" in userR and userR["success"]:
            user = userR["rows"]
        else:
            user = None
            # return userR

        if user is not None and int(user["discord_id"]) == int(discordID) and (user["is_manager"] is not None and int(user["is_manager"]) == 1):
            return {"success": False, "msg": "A manager already exists with that discord ID"}

        await c.execute("BEGIN")
        try:
            if user is not None and int(user["discord_id"]) == int(discordID) and (user["is_manager"] is None or int(user["is_manager"]) == 0):
                await c.execute('UPDATE users SET is_manager=1 WHERE discord_id=?', (discordID,))
            else:
                await c.execute('INSERT INTO users (discord_id, name, is_manager) VALUES (?, ?, ?)', (discordID, name, 1))
            await c.execute("COMMIT")

            logger.success(f"Saved manager {name}/{discordID}")

        except:
            await c.execute("ROLLBACK")
            logger.error(f"Failed to save manager {name}/{discordID}")
            return {"success": False, "msg": f"Error in processing manager addition for {name}/{discordID}"}

    return {"success": True, "msg": f"Manager {name}/{discordID} saved"}


@logger.catch
async def removeManager(discordID):
    async with client.db.cursor() as c:
        await c.execute("BEGIN")
        try:
            await c.execute('UPDATE users SET is_manager=0 WHERE discord_id=?', (discordID,))
            await c.execute("COMMIT")

            logger.success(f"Deleted manager {discordID}")

        except:
            await c.execute("ROLLBACK")
            logger.error(f"Failed to delete manager {discordID}")
            return {"success": False, "msg": f"Error in processing manager deletion for {discordID}"}

    return {"success": True, "msg": f"Manager {discordID} deleted"}


@logger.catch
async def setOwner(discordID, name):
    async with client.db.cursor() as c:
        rowsR = await getOwner()
        # logger.info("rowsR:")
        # logger.info(rowsR)

        if rowsR["success"]:
            owner = rowsR["rows"]
        else:
            owner = None
            # return rowsR

        userR = await getDiscordID(discordID)
        if userR is not None and "success" in userR and userR["success"]:
            user = userR["rows"]
        else:
            user = None
            # return userR

        if owner is not None:
            return {"success": False, "msg": "An owner already exists"}

        await c.execute("BEGIN")
        try:
            if user is not None and int(user["discord_id"]) == int(discordID):
                await c.execute('UPDATE users SET name=?,is_manager=1,is_owner=1 WHERE discord_id=?', (name, discordID))
            else:
                await c.execute('INSERT INTO users (discord_id, name, is_manager, is_owner) VALUES (?, ?, ?, ?)', (discordID, name, 1, 1))
            await c.execute("COMMIT")

            logger.info(f"Saved owner {name}/{discordID}")

        except:
            await c.execute("ROLLBACK")
            logger.error(f"Failed to save owner {name}/{discordID}")
            return {"success": False, "msg": f"Error in processing owner set for {name}/{discordID}"}

    return {"success": True, "msg": f"Owner {name}/{discordID} saved"}


@logger.catch
async def setProperty(prop, val):
    if (prop == "devDonation" or prop == "d") and val == 0:
        val = 0.01
    async with client.db.cursor() as c:

        isNum = False
        if Common.isFloat(val):
            isNum = True

        await c.execute("BEGIN")
        try:
            res = await getProperty(prop)
            if res["success"] and res["rows"] is not None:
                # update
                if isNum:
                    await c.execute('UPDATE properties SET realVal=?,textVal=? WHERE property=?', (val, None, prop))
                else:
                    await c.execute('UPDATE properties SET textVal=?,realVal=? WHERE property=?', (val, None, prop))

            elif res["success"] and res["rows"] is None:
                # insert
                if isNum:
                    await c.execute('INSERT INTO properties (property, realVal) VALUES (?, ?)', (prop, val))
                else:
                    await c.execute('INSERT INTO properties (property, textVal) VALUES (?, ?)', (prop, val))
            else:
                # error
                await c.execute("ROLLBACK")
                return res
        except:
            await c.execute("ROLLBACK")
            logger.error(f"Failed to save property {prop}")
            return {"success": False, "msg": f"Error in processing property set for {prop}"}

        await c.execute("COMMIT")
    return {"success": True, "msg": f"Successfully set {prop} to {val}"}


# Select
@logger.catch
async def getProperty(prop):
    c = await client.db.cursor()

    try:
        await c.execute('SELECT * FROM properties WHERE property=?', (prop,))
        rows = await c.fetchone()
        # logger.info(f"Fetched property: {prop}")

    except:
        logger.error(f"Failed to get property: {prop}")
        return {"success": False, "msg": f"Exception in processing property query"}

    return {"success": True, "rows": rows}


@logger.catch
async def getLastClaim(roninAddr):
    c = await client.db.cursor()

    try:
        await c.execute('SELECT * FROM claims WHERE ronin_addr = ? ORDER BY claim_time DESC', (roninAddr,))
        rows = await c.fetchone()
        # logger.info(f"Got last claim data for {roninAddr}")

    except Exception:
        # logger.error(traceback.format_exc())
        logger.error(f"Failed to get last claim for {roninAddr}")
        return {"success": False, "msg": f"Error in getting claim for {roninAddr}", "rows": None}

    return {"success": True, "rows": rows}


@logger.catch
async def getAllScholars():
    c = await client.db.cursor()

    try:
        await c.execute('SELECT * FROM users WHERE is_scholar=1')
        rows = await c.fetchall()
        logger.info(f"Fetched all scholars")

    except:
        logger.error(f"Failed to get all scholars")
        return {"success": False, "msg": f"Exception in processing scholar query"}

    return {"success": True, "rows": rows}


@logger.catch
async def getAllScholarsByIndex(seed, minIndex=None, maxIndex=None):
    c = await client.db.cursor()

    if minIndex is None or maxIndex is None:
        maxIndex = 0
        minIndex = 0

    try:
        if minIndex == 0 or maxIndex == 0:
            await c.execute('SELECT * FROM users WHERE is_scholar=1 AND seed_num = ?', (int(seed),))
        else:
            await c.execute('SELECT * FROM users WHERE is_scholar=1 AND seed_num = ? AND account_num >= ? AND account_num <= ?', (int(seed), int(minIndex), int(maxIndex)))
        rows = await c.fetchall()
        logger.info(f"Fetched all scholars with seed={seed} and account range {minIndex}-{maxIndex}")

    except Exception as e:
        logger.error(f"Failed to get all scholars in range")
        logger.error(e)
        return {"success": False, "msg": f"Exception in processing scholar query"}
    
    return {"success": True, "rows": rows}


@logger.catch
async def getAllManagers():
    c = await client.db.cursor()

    try:
        await c.execute('SELECT * FROM users WHERE is_manager=1 OR is_owner=1')
        rows = await c.fetchall()
        logger.info(f"Fetched all managers")

    except:
        logger.error(f"Failed to get all managers")
        return {"success": False, "msg": f"Exception in processing manager query"}

    return {"success": True, "rows": rows}


@logger.catch
async def getAllNoRole():
    c = await client.db.cursor()

    try:
        await c.execute('SELECT * FROM users WHERE (is_manager=0 OR is_manager IS NULL) AND (is_scholar=0 OR is_scholar IS NULL) AND (is_owner=0 OR is_owner IS NULL)')
        rows = await c.fetchall()
        logger.info(f"Fetched all users with no role")

    except:
        logger.error(f"Failed to get all users with no role")
        return {"success": False, "msg": f"Exception in processing no roles query"}

    return {"success": True, "rows": rows}


@logger.catch
async def getAllUsers():
    c = await client.db.cursor()

    try:
        await c.execute('SELECT * FROM users')
        rows = await c.fetchall()
        logger.info(f"Fetched all users")

    except:
        logger.error(f"Failed to get all users")
        return {"success": False, "msg": f"Exception in processing all users query"}

    return {"success": True, "rows": rows}


@logger.catch
async def getAllManagerIDs():
    c = await client.db.cursor()

    try:
        await c.execute('SELECT * FROM users WHERE is_manager=1 or is_owner=1')
        rows = await c.fetchall()
        logger.info(f"Fetched all managers")

    except:
        logger.error(f"Failed to get all managers")
        return {"success": False, "msg": f"Exception in processing manager query"}

    ids = []
    for row in rows:
        ids.append(int(row["discord_id"]))

    return ids


@logger.catch
async def getDiscordID(discordID):
    c = await client.db.cursor()

    try:
        await c.execute('SELECT * FROM users WHERE discord_id=? LIMIT 1', (int(discordID),))
        rows = await c.fetchone()
        logger.info(f"Fetched discord ID {discordID}")
        # out = "["
        # for col in rows:
        #    out += str(col) + ","
        # out += "]"
        # logger.info(out)

    except Exception:
        logger.warning(f"Failed to get discord ID {discordID}, not in database")
        # logger.error(e)
        return {"success": False, "msg": f"Exception in processing discord ID query"}

    return {"success": True, "rows": rows}


@logger.catch
async def getOwner():
    c = await client.db.cursor()

    # logger.info("before async in getOwner")
    try:
        # logger.info("exec owner query")
        await c.execute('SELECT * FROM users WHERE is_owner=1 LIMIT 1')
        rows = await c.fetchone()
        # logger.info("fetched owner")

    except Exception as e:
        logger.error(f"Failed to get owner")
        logger.error(e)
        return {"success": False, "msg": f"Exception in processing owner query"}

    return {"success": True, "rows": rows}


@logger.catch
async def isManager(discordID):
    mgrs = await getAllManagers()

    for row in mgrs["rows"]:
        if int(row['discord_id']) == int(discordID):
            return True

    return await isOwner(discordID)


@logger.catch
async def isOwner(discordID):
    owner = await getOwner()

    if int(owner["rows"]['discord_id']) == int(discordID):
        return True

    return False


@logger.catch
async def getMembershipReport():
    try:
        mgrs = await getAllManagers()
        sclr = await getAllScholars()
        norl = await getAllNoRole()
        totl = await getAllUsers()
        ownr = await getOwner()

        logger.info(norl)

        mgrs = 0 if mgrs["rows"] is None else len(mgrs["rows"])
        sclr = 0 if sclr["rows"] is None else len(sclr["rows"])
        norl = 0 if norl["rows"] is None else len(norl["rows"])
        totl = 0 if totl["rows"] is None else len(totl["rows"])
        ownr = 1 if ownr is not None else 0
        logger.info(norl)

    except:
        logger.error("Error processing membership report")
        return {"success": False, "msg": "Error processing membership report"}

    return {"success": True, "managers": mgrs, "scholars": sclr, "noRole": norl, "owner": ownr, "total": totl}