aboutsummaryrefslogtreecommitdiff
path: root/code/database.py
blob: 34bc954dd871641e6392136f85fe41393b0b54f5 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
import aiosqlite
import random
from typing import Tuple, List

from global_variables import CONNECTION


async def init_database():
    cur = CONNECTION.cursor()
    cur.execute("CREATE TABLE IF NOT EXISTS guildData (guild_id BIGINT, user_id BIGINT, exp BIGINT, PRIMARY KEY (guild_id, user_id))")
    cur.execute("CREATE TABLE IF NOT EXISTS mute (guild_id BIGINT, role_id BIGINT, PRIMARY KEY (guild_id, role_id))")
    cur.execute("CREATE TABLE IF NOT EXISTS tempmute (guild_id BIGINT, user_id BIGINT, role_id BIGINT, time TIMESTAMP, PRIMARY KEY (guild_id, user_id))")
    cur.execute("CREATE TABLE IF NOT EXISTS warnings (warn_id BIGINT, guild_id BIGINT, user_id BIGINT, warning TEXT, warn_time DATE, warned_by BIGINT, PRIMARY KEY (warn_id))")
    cur.execute("CREATE TABLE IF NOT EXISTS economy (user_id BIGINT NOT NULL PRIMARY KEY, money BIGINT NOT NULL DEFAULT 0)")
    cur.execute("CREATE TABLE IF NOT EXISTS profile (user_id BIGINT, rank_name TEXT, rank_int BIGINT, UNIQUE (user_id, rank_name, rank_int))")
    CONNECTION.commit()

    cur = await aiosqlite.connect("code/count/count.db")
    await cur.execute("CREATE TABLE IF NOT EXISTS count (count INTEGER)")
    await cur.commit()
    await cur.close()


Entry = Tuple[int, int]

class Database:
    def __init__(self, bot):
        self.bot = bot

    async def get_entry(self, user_id: int) -> Entry:
        cur = CONNECTION.cursor()
        cur.execute("SELECT * FROM economy WHERE user_id = %s", (user_id,))
        result = cur.fetchone()
        if result:
            return result
        return await self.new_entry(user_id)

    async def new_entry(self, user_id: int) -> Entry:
        try:
            cur = CONNECTION.cursor()
            cur.execute("INSERT INTO economy(user_id, money) VALUES(%s, %s)", (user_id, 0))
            CONNECTION.commit()
            return await self.get_entry(user_id)
        except:
            return await self.get_entry(user_id)

    async def remove_entry(self, user_id: int) -> None:
        cur = CONNECTION.cursor()
        cur.execute("DELETE FROM economy WHERE user_id = %s", (user_id,))
        CONNECTION.commit()

    async def set_money(self, user_id: int, money: int) -> Entry:
        cur = CONNECTION.cursor()
        cur.execute("UPDATE economy SET money = %s WHERE user_id = %s", (money, user_id))
        CONNECTION.commit()
        return await self.get_entry(user_id)

    async def add_money(self, user_id: int, money_to_add: int) -> Entry:
        money = (await self.get_entry(user_id))[1]
        total = money + money_to_add
        if total < 0:
            total = 0
        await self.set_money(user_id, total)
        return await self.get_entry(user_id)

    async def random_entry(self) -> Entry:
        cur = CONNECTION.cursor()
        cur.execute("SELECT * FROM economy")
        return random.choice(cur.fetchall())

    async def top_entries(self, n: int=0) -> List[Entry]:
        cur = CONNECTION.cursor()
        cur.execute("SELECT * FROM economy ORDER BY money DESC")
        return cur.fetchmany(n) if n else cur.fetchall()