Passed
Push — main ( 55bb28...245da3 )
by Bartosz
02:31 queued 01:14
created

g_roles()   A

Complexity

Conditions 1

Size

Total Lines 2
Code Lines 2

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 1
eloc 2
nop 1
dl 0
loc 2
rs 10
c 0
b 0
f 0
1
from discord.ext import commands, tasks
2
from discord_slash import cog_ext
3
from modules.get_settings import get_settings
4
from sqlalchemy import create_engine, Table, Column, Integer, String, MetaData, ForeignKey, \
5
    BigInteger, update, select, DateTime, delete
6
from sqlalchemy.dialects.mysql import insert
7
import cogs.cogbase as cogbase
8
from datetime import datetime
9
import pandas as pd
10
11
metadata_obj = MetaData()
12
member = Table('member', metadata_obj,
13
               Column('id', BigInteger, primary_key=True),
14
               Column('name', String(50), nullable=False),
15
               Column('display_name', String(50), nullable=False)
16
               )
17
fk_member_id = "member.id"
18
19
spots = Table('spots', metadata_obj,
20
              Column('member_id', BigInteger, ForeignKey(fk_member_id), primary_key=True),
21
              Column('legendary', Integer, default=0),
22
              Column('rare', Integer, default=0),
23
              Column('common', Integer, default=0),
24
              Column('event1', Integer, default=0),
25
              Column('event2', Integer, default=0)
26
              )
27
28
spots_temp = Table('spots_temp', metadata_obj,
29
                   Column('member_id', BigInteger, ForeignKey(fk_member_id), primary_key=True),
30
                   Column('legendary', Integer, default=0),
31
                   Column('rare', Integer, default=0),
32
                   Column('common', Integer, default=0),
33
                   Column('event1', Integer, default=0),
34
                   Column('event2', Integer, default=0)
35
                   )
36
37
warn = Table('warn', metadata_obj,
38
             Column('id', Integer, primary_key=True),
39
             Column('member_id', BigInteger, ForeignKey(fk_member_id)),
40
             Column('reason', String(120), nullable=False),
41
             Column('date', DateTime, nullable=False),
42
             )
43
44
coords = Table('coords', metadata_obj,
45
               Column('id', Integer, primary_key=True),
46
               Column('coords', String(100), nullable=False),
47
               Column('monster_type', String(20), nullable=False)
48
               )
49
50
51
# noinspection PyPropertyAccess
52
class DatabaseCog(cogbase.BaseCog):
53
    user = get_settings("DB_U")
54
    password = get_settings("DB_P")
55
    engine = create_engine(f"mysql+mysqldb://{user}:{password}@localhost/server_database")
56
    metadata_obj.create_all(engine)
57
    conn = engine.connect()
58
59
    def __init__(self, base):
60
        super().__init__(base)
61
62
        # Connect to database
63
        self.engine = DatabaseCog.engine
64
        metadata_obj.create_all(self.engine)
65
        self.conn = DatabaseCog.conn
66
        self.db_update_loop.start()
67
68
    # ----- BASE DATABASE OPERATIONS -----
69
70
    # Add or update member in member table
71
    def db_add_update_member(self, _member):
72
        stmt = insert(member).values(
73
            id=_member.id, name=_member.name,
74
            display_name=_member.display_name)
75
        do_update_stmt = stmt.on_duplicate_key_update(
76
            name=stmt.inserted.name, display_name=stmt.inserted.display_name
77
        )
78
        self.conn.execute(do_update_stmt)
79
80
    # Add or update spots in spots table
81
    def db_add_update_spots(self, spots_table, guild_member):
82
        stmt = insert(spots_table).values(
83
            member_id=guild_member.id)
84
        do_update_stmt = stmt.on_duplicate_key_update(member_id=stmt.inserted.member_id)
85
        self.conn.execute(do_update_stmt)
86
87
    # Add or refresh all guild members and spots to database
88
    async def db_update(self):
89
        guild = self.bot.get_guild(self.bot.guild[0])
90
91
        print(f"[{self.__class__.__name__}]: Refreshing member and spots tables")
92
        for guild_member in guild.members:
93
            # Member tables
94
            self.db_add_update_member(guild_member)
95
            # Spots tables
96
            self.db_add_update_spots(spots, guild_member)
97
            self.db_add_update_spots(spots_temp, guild_member)
98
        print(f"[{self.__class__.__name__}]: Member and spots tables refreshed")
99
100
    @tasks.loop(hours=12)
101
    async def db_update_loop(self):
102
        await self.db_update()
103
104
    @db_update_loop.before_loop
105
    async def before_db_update_loop(self):
106
        print(f'[{self.__class__.__name__}]: Waiting until Bot is ready')
107
        await self.bot.wait_until_ready()
108
109
    # Add member to database on member join
110
    @commands.Cog.listener()
111
    async def on_member_join(self, _member):
112
        self.db_add_update_member(_member)
113
        self.db_add_update_spots(spots, _member)
114
        self.db_add_update_spots(spots_temp, _member)
115
116
    # Command for loading data from previous bot
117
    @cog_ext.cog_slash(name="updateSpotsWithOld", guild_ids=cogbase.GUILD_IDS,
118
                       description="Change N-Word channel name",
119
                       permissions=cogbase.PERMISSION_ADMINS)
120
    async def db_update_spots_old(self, ctx):
121
        import json
122
        with open('server_files/old_base_test.json', 'r', encoding='utf-8-sig') as fp:
123
            old_db = json.load(fp)
124
125
        for mem_id in old_db:
126
            stmt = insert(spots).values(
127
                member_id=mem_id, legendary=old_db[mem_id]["type_1"],
128
                rare=old_db[mem_id]["type_0"])
129
            do_update_stmt = stmt.on_duplicate_key_update(legendary=stmt.inserted.legendary,
130
                                                          rare=stmt.inserted.rare,
131
                                                          common=stmt.inserted.common)
132
            self.conn.execute(do_update_stmt)
133
            stmt = insert(spots_temp).values(
134
                member_id=mem_id, legendary=old_db[mem_id]["type_1"],
135
                rare=old_db[mem_id]["type_0"])
136
            do_update_stmt = stmt.on_duplicate_key_update(legendary=stmt.inserted.legendary,
137
                                                          rare=stmt.inserted.rare,
138
                                                          common=stmt.inserted.common)
139
            self.conn.execute(do_update_stmt)
140
        await ctx.send(f"Spot tables updated with old data", delete_after=3.0)
141
        print(f'[{self.__class__.__name__}]: Spot tables updated with old data')
142
143
    # ----- SPOTTING OPERATIONS -----
144
145
    # Update spots tables
146
    @classmethod
147
    async def db_count_spot(cls, _id: int, monster_type: str):
148
        # Get member nr of spots for certain monster type
149
        stmt = select(spots.c.member_id, spots.c.legendary, spots.c.rare, spots.c.common).where(
150
            spots.c.member_id == _id)
151
        result = cls.conn.execute(stmt)
152
        counter = []
153
        for nr_of_kills in result.columns(monster_type):
154
            counter = nr_of_kills[0]
155
        stmt = update(spots).where(spots.c.member_id == _id).values({f"{monster_type}": counter + 1})
156
        cls.conn.execute(stmt)
157
158
        stmt = select(spots_temp.c.member_id, spots_temp.c.legendary, spots_temp.c.rare, spots_temp.c.common).where(
159
            spots_temp.c.member_id == _id)
160
        result = cls.conn.execute(stmt)
161
        for nr_of_kills in result.columns(monster_type):
162
            counter = nr_of_kills[0]
163
        stmt = update(spots_temp).where(spots_temp.c.member_id == _id).values({f"{monster_type}": counter + 1})
164
        cls.conn.execute(stmt)
165
166
    # Save coords from spotting channels to database
167
    @classmethod
168
    async def db_save_coords(cls, _coords: str, _monster_type):
169
        stmt = insert(coords).values(coords=_coords, monster_type=_monster_type)
170
        cls.conn.execute(stmt)
171
172
    # Clear data from spots_temp table(for events etc)
173
    @classmethod
174
    async def db_clear_spots_temp_table(cls):
175
        stmt = delete(spots_temp)
176
        cls.conn.execute(stmt)
177
178
    # ----- LEADERBOARD OPERATIONS -----
179
180
    # Return all members' spots
181
    @classmethod
182
    async def db_get_spots_df(cls):
183
        stmt = select(spots.c.member_id, member.c.display_name, spots.c.legendary, spots.c.rare,
184
                      spots.c.common).select_from(member).join(spots, member.c.id == spots.c.member_id)
185
        df = pd.read_sql(stmt, cls.conn)
186
        return df
187
188
    # ----- WARN OPERATIONS -----
189
190
    # Add member's warn to database
191
    @classmethod
192
    async def db_add_warn(cls, _member: int, _reason: str):
193
        stmt = insert(warn).values(member_id=_member, reason=_reason, date=datetime.now())
194
        cls.conn.execute(stmt)
195
196
    # Get member's warns from database
197
    @classmethod
198
    async def db_get_warns(cls, _member: int):
199
        stmt = select(warn.c.reason, warn.c.date).select_from(member).join(warn, member.c.id == warn.c.member_id).where(
200
            member.c.id == _member)
201
        result = cls.conn.execute(stmt)
202
        date_warn = []
203
        counter = 0
204
        for warns in result.columns("reason", "date"):
205
            reason_with_date = [warns[1], warns[0]]
206
            date_warn.append(reason_with_date)
207
            counter += 1
208
        warns_list = [': \t'.join([str(elem) for elem in sublist]) for sublist in date_warn]
209
        return warns_list, counter
210
211
    # Remove all member's warns
212
    @classmethod
213
    async def db_remove_warns(cls, _member: int):
214
        stmt = delete(warn).where(warn.c.member_id == _member)
215
        cls.conn.execute(stmt)
216
217
    # ----- MEMBER OPERATIONS -----
218
219
    # Return member's spots
220
    @classmethod
221
    async def db_get_member_stats(cls, _member: int):
222
        stmt = select(member.c.display_name, spots.c.legendary, spots.c.rare, spots.c.common).select_from(member).join(
223
            spots,
224
            member.c.id == spots.c.member_id).where(spots.c.member_id == _member)
225
        df = pd.read_sql(stmt, cls.conn)
226
        return df
227
228
229
def setup(bot: commands.Bot):
230
    bot.add_cog(DatabaseCog(bot))
231