-
Notifications
You must be signed in to change notification settings - Fork 30
/
database_handler.py
370 lines (314 loc) · 16.6 KB
/
database_handler.py
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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
import logging
import aiosqlite
from pathlib import Path
from datetime import datetime
from typing import Tuple, List, Union
from helpers import misc
from helpers import licence_helper
from helpers.errors import DefaultGuildRoleNotSet, DatabaseMissingData
logger = logging.getLogger(__name__)
class DatabaseHandler:
DB_PATH = "databases/"
DB_EXTENSION = ".sqlite3"
@classmethod
async def create_instance(cls, db_name: str = "main"):
""""
Can't use await in __init__ so we create a factory pattern.
To correctly create this object you need to call :
await DatabaseHandler.create_instance()
"""
self = DatabaseHandler()
self.db_name = db_name
self.connection = await self._get_connection()
logger.info("Connection to database established.")
return self
def __init__(self):
self.db_name = None
self.connection = None
async def _get_connection(self) -> aiosqlite.core.Connection:
"""
Returns a connection to the db, if db doesn't exist create new
:return: aiosqlite.core.Connection
"""
path = DatabaseHandler._construct_path(self.db_name)
if Path(path).is_file():
conn = await aiosqlite.connect(path)
return conn
else:
logger.warning("Database not found! Creating fresh ...")
misc.check_create_directory(DatabaseHandler.DB_PATH)
return await DatabaseHandler._create_database(path)
@staticmethod
def _construct_path(db_name: str) -> str:
return DatabaseHandler.DB_PATH + db_name + DatabaseHandler.DB_EXTENSION
@staticmethod
async def _create_database(path: str) -> aiosqlite.core.Connection:
"""
:param path: path where database will be created, including file name and extension
:return: aiosqlite.core.Connection
"""
conn = await aiosqlite.connect(path)
await conn.execute("CREATE TABLE GUILDS "
"("
"GUILD_ID TEXT PRIMARY KEY, "
"PREFIX TEXT CHECK(PREFIX IS NULL OR LENGTH(PREFIX) <= 5), "
"ENABLE_LOG_CHANNEL TINYINT DEFAULT 0, "
"LOG_CHANNEL_ID TEXT, "
"DEFAULT_LICENSE_ROLE_ID TEXT, "
"DEFAULT_LICENSE_DURATION_HOURS UNSIGNED BIG INT DEFAULT 720"
")"
)
await conn.execute("CREATE TABLE LICENSED_MEMBERS "
"("
"MEMBER_ID TEXT, "
"GUILD_ID TEXT, "
"EXPIRATION_DATE DATE, "
"LICENSED_ROLE_ID TEXT, "
"UNIQUE(MEMBER_ID, LICENSED_ROLE_ID)"
")"
)
await conn.execute("CREATE TABLE GUILD_LICENSES "
"("
"LICENSE TEXT PRIMARY KEY, "
"GUILD_ID TEXT, "
"LICENSED_ROLE_ID TEXT, "
"LICENSE_DURATION_HOURS UNSIGNED BIG INT"
")"
)
await conn.commit()
logger.info("Database successfully created!")
return conn
async def update_database(self, query: str, *args):
await self.connection.execute(query, args)
await self.connection.commit()
# TABLE GUILDS #######################################################################
async def setup_new_guild(self, guild_id: int, default_prefix: str):
insert_guild_query = "INSERT INTO GUILDS(GUILD_ID, PREFIX) VALUES(?,?)"
await self.update_database(insert_guild_query, guild_id, default_prefix)
async def get_guild_prefix(self, guild_id: int) -> str:
query = "SELECT PREFIX FROM GUILDS WHERE GUILD_ID=?"
async with self.connection.execute(query, (guild_id,)) as cursor:
row = await cursor.fetchone()
return row[0]
async def get_all_guild_ids(self):
"""
:return: a tuple of all guild ids (ints)
"""
query = "SELECT GUILD_ID FROM GUILDS"
async with self.connection.execute(query) as cursor:
results = await cursor.fetchall()
return tuple(int(guild_id[0]) for guild_id in results)
async def change_guild_prefix(self, guild_id: int, prefix: str):
"""
:param guild_id: int id of guild to change the prefix to in the database
:param prefix: str prefix to change to. Length is limited by table GUILDS layout (max 5 chars)
:raise: IntegrityError if the prefix has too many chars (max 5)
"""
query = "UPDATE GUILDS SET PREFIX=? WHERE GUILD_ID=?"
await self.update_database(query, prefix, guild_id)
async def change_default_guild_role(self, guild_id: int, role_id: int):
query = "UPDATE GUILDS SET DEFAULT_LICENSE_ROLE_ID=? WHERE GUILD_ID=?"
await self.update_database(query, role_id, guild_id)
async def change_default_license_expiration(self, guild_id: int, expiration_hours: int):
query = "UPDATE GUILDS SET DEFAULT_LICENSE_DURATION_HOURS=? WHERE GUILD_ID=?"
await self.update_database(query, expiration_hours, guild_id)
async def get_default_guild_license_role_id(self, guild_id: int) -> int:
"""
Gets the default license role id from specific guild.
This role will be used as link when no role argument is passed in 'generate' command
:return: int default license role id
:raise: DefaultGuildRoleNotSet if it's None
"""
query = "SELECT DEFAULT_LICENSE_ROLE_ID FROM GUILDS WHERE GUILD_ID=?"
async with self.connection.execute(query, (guild_id,)) as cursor:
row = await cursor.fetchone()
try:
return int(row[0])
except TypeError:
raise DefaultGuildRoleNotSet("Default guild license not set!\n\n"
"For more information call command:\n"
"{prefix}help default_role\n\n"
"If still in doubt call:\n"
"{prefix}help")
async def get_default_guild_license_duration_hours(self, guild_id: int) -> int:
"""
Gets the default license duration from specific guild.
:return: int representing hours of license duration
"""
query = "SELECT DEFAULT_LICENSE_DURATION_HOURS FROM GUILDS WHERE GUILD_ID=?"
async with self.connection.execute(query, (guild_id,)) as cursor:
row = await cursor.fetchone()
#
if not row:
# License duration has default value.
# So if this is None it means the guild is not found in database.
raise DatabaseMissingData(f"Guild {guild_id} not found in database!")
return int(row[0])
async def get_guild_info(self, guild_id: int) -> Tuple[str, str, int]:
"""
:param guild_id:
:return: tuple(str prefix, str role_id, int expiration hours)
"""
query = "SELECT * FROM GUILDS WHERE GUILD_ID=?"
async with self.connection.execute(query, (guild_id,)) as cursor:
row = await cursor.fetchone()
# ('guild_id', 'prefix', 0, None, 'role_id', hours)
return row[1], row[4], row[5]
# TABLE LICENSED_MEMBERS #############################################################
async def add_new_licensed_member(self, member_id: int, guild_id: int,
expiration_date: datetime, licensed_role_id: int):
query = "INSERT INTO LICENSED_MEMBERS(MEMBER_ID, GUILD_ID, EXPIRATION_DATE, LICENSED_ROLE_ID) VALUES(?,?,?,?)"
await self.update_database(query, member_id, guild_id, expiration_date, licensed_role_id)
async def delete_licensed_member(self, member_id: int, licensed_role_id: int):
"""
Called when member licensed role has expired
The member row is deleted from table LICENSED_MEMBERS.
MEMBER_ID and LICENSED_ROLE_ID are unique so we only need that to differentiate
"""
delete_query = "DELETE FROM LICENSED_MEMBERS WHERE MEMBER_ID=? AND LICENSED_ROLE_ID=?"
await self.update_database(delete_query, member_id, licensed_role_id)
async def get_member_license_expiration_date(self, member_id: int, licensed_role_id: int) -> str:
query = "SELECT EXPIRATION_DATE FROM LICENSED_MEMBERS WHERE MEMBER_ID=? AND LICENSED_ROLE_ID=?"
async with self.connection.execute(query, (member_id, licensed_role_id)) as cursor:
row = await cursor.fetchone()
if row is not None:
return row[0]
else:
raise DatabaseMissingData(f"ID {member_id} doesn't exists in database table LICENSED_MEMBERS.")
async def get_member_data(self, guild_id: int, member_id: int) -> List[Tuple]:
"""
Return type:
[(), ()...]
Note that returned LICENSED_ROLE_ID is string
"""
query = "SELECT LICENSED_ROLE_ID, EXPIRATION_DATE FROM LICENSED_MEMBERS WHERE GUILD_ID=? AND MEMBER_ID=?"
async with self.connection.execute(query, (guild_id, member_id)) as cursor:
results = await cursor.fetchall()
if results is not None:
return results
else:
raise DatabaseMissingData(f"No active licenses for member {member_id} in guild {guild_id}.")
async def get_guild_licensed_roles_total_count(self, guild_id: int) -> int:
query = "SELECT COUNT(*) FROM LICENSED_MEMBERS WHERE GUILD_ID=?"
async with self.connection.execute(query, (guild_id,)) as cursor:
result = await cursor.fetchone()
return result[0]
async def get_licensed_roles_total_count(self) -> int:
query = "SELECT COUNT(*) FROM LICENSED_MEMBERS"
async with self.connection.execute(query) as cursor:
result = await cursor.fetchone()
return result[0]
# TABLE GUILD_LICENSES ###############################################################
async def get_license_data(self, license: str) -> Union[Tuple[int, int], None]:
"""
Returns licensed role id that the param license is linked to
:param license: license the role is linked to
:return: tuple(int guild id, int license role id)
"""
query = "SELECT GUILD_ID, LICENSED_ROLE_ID FROM GUILD_LICENSES WHERE LICENSE=?"
async with self.connection.execute(query, (license,)) as cursor:
row = await cursor.fetchone()
# TODO: Temporal quick fix. Refactor
if row is None:
return None
else:
return int(row[0]), int(row[1])
async def get_license_duration_hours(self, license):
"""
:param license: It's unique so no need to pass any additional argument
:return: int representing license duration in hours
"""
query = "SELECT LICENSE_DURATION_HOURS FROM GUILD_LICENSES WHERE LICENSE=?"
async with self.connection.execute(query, (license,)) as cursor:
row = await cursor.fetchone()
return int(row[0])
async def generate_guild_licenses(self, number: int, guild_id: int,
license_role_id: int, license_duration: int) -> list:
"""
:param number: int larger than 0, number of licenses to generate
:param guild_id: int guild id. Needed to differentiate guilds even though licenses are unique
for example to avoid members activating a valid license from one guild into
another guild (where linked roles don't exist).
:param license_role_id: role to link to the license
:param license_duration: int representing license duration in hours
:return: list of all generated licenses
"""
licenses = licence_helper.generate_multiple(number)
query = """INSERT INTO GUILD_LICENSES(LICENSE, GUILD_ID, LICENSED_ROLE_ID, LICENSE_DURATION_HOURS)
VALUES(?,?,?,?)"""
for license in licenses:
await self.connection.execute(query, (license, guild_id, license_role_id, license_duration))
await self.connection.commit()
return licenses
async def delete_license(self, license: str):
"""
Called for example when member has redeemed license.
The license is deleted from table GUILD_LICENSES.
Guild ID is not needed since licenses are unique.
But if deleting from withing the guild better to call is_valid_license beforehand
to prevent deletion of other guild licenses in other guilds.
:param license: license to delete
"""
delete_query = "DELETE FROM GUILD_LICENSES WHERE LICENSE=?"
await self.update_database(delete_query, license)
async def get_guild_licenses(self, number: int, guild_id: int, license_role_id: int) -> list:
"""
Returns list of licenses that are linked to license_role_id role and their duration time.
:param number: int larger than 0, max number of licenses to return
:param guild_id: int guild id. Needed to differentiate guilds even though licenses are unique
for example to avoid members activating a valid license from one guild into
another guild (where linked roles don't exist).
:param license_role_id: we get only those licenses that are linked to this role id
:return: List of tuples in format [('license', license_duration_int_hours)]
"""
query = """SELECT LICENSE, LICENSE_DURATION_HOURS FROM GUILD_LICENSES
WHERE GUILD_ID=? AND LICENSED_ROLE_ID=? LIMIT ?"""
async with self.connection.execute(query, (guild_id, license_role_id, number)) as cursor:
return await cursor.fetchall()
async def get_guild_license_total_count(self, guild_id: int) -> int:
query = "SELECT COUNT(*) FROM GUILD_LICENSES WHERE GUILD_ID=?"
async with self.connection.execute(query, (guild_id,)) as cursor:
result = await cursor.fetchone()
return result[0]
async def get_stored_license_total_count(self) -> int:
query = "SELECT COUNT(*) FROM GUILD_LICENSES"
async with self.connection.execute(query) as cursor:
result = await cursor.fetchone()
return result[0]
async def is_valid_license(self, license: str, guild_id: int) -> bool:
"""
:param license: License to check
:param guild_id: int guild id. Needed to differentiate guilds even though licenses are unique
for example to avoid members activating a valid license from one guild into
another guild (where linked roles don't exist).
:return: True if license is valid, False otherwise
"""
query = "SELECT LICENSE FROM GUILD_LICENSES WHERE LICENSE=? AND GUILD_ID=?"
async with self.connection.execute(query, (license, guild_id)) as cursor:
row = await cursor.fetchone()
if row is not None:
return True
return False
async def get_random_licenses(self, guild_id: int, amount: int):
query = """SELECT LICENSE, LICENSED_ROLE_ID, LICENSE_DURATION_HOURS FROM GUILD_LICENSES
WHERE GUILD_ID=? ORDER BY RANDOM() LIMIT ?"""
async with self.connection.execute(query, (guild_id, amount)) as cursor:
return await cursor.fetchall()
async def remove_all_stored_guild_licenses(self, guild_id: int):
query = "DELETE FROM GUILD_LICENSES WHERE GUILD_ID=?"
await self.update_database(query, guild_id)
# ALL TABLES #########################################################################
async def remove_all_guild_data(self, guild_id: int, guild_table_too=False):
queries = ["DELETE FROM LICENSED_MEMBERS WHERE GUILD_ID=?",
"DELETE FROM GUILD_LICENSES WHERE GUILD_ID=?"]
if guild_table_too:
queries.append("DELETE FROM GUILDS WHERE GUILD_ID=?")
for query in queries:
await self.connection.execute(query, (guild_id,))
await self.connection.commit()
async def remove_all_guild_role_data(self, role_id: int):
queries = ["DELETE FROM LICENSED_MEMBERS WHERE LICENSED_ROLE_ID=?",
"DELETE FROM GUILD_LICENSES WHERE LICENSED_ROLE_ID=?"]
for query in queries:
await self.connection.execute(query, (role_id,))
await self.connection.commit()