2021-02-09 17:17:19 +00:00
import util
2021-07-28 17:13:32 +00:00
import discord . ext . commands as commands
2021-02-09 17:17:19 +00:00
2021-07-28 17:13:32 +00:00
def check_key ( key ) :
if len ( key ) > 128 : raise ValueError ( " Key too long " )
def preprocess_value ( value ) :
value = value . strip ( )
if len ( value ) > 1024 : raise ValueError ( " Value too long " )
return value
class Userdata ( commands . Cog ) :
def __init__ ( self , bot ) :
self . bot = bot
@commands.group ( name = " userdata " , aliases = [ " data " ] , help = """ Store per-user data AND retrieve it later! Note that, due to the nature of storing things, it is necessary to set userdata before getting it.
2021-02-09 18:19:12 +00:00
Data can either be localized to a guild ( guild scope ) or shared between guilds ( global scope ) , but is always tied to a user . """ )
2021-07-28 17:13:32 +00:00
async def userdata ( self , ctx ) : pass
2021-02-09 17:17:19 +00:00
2021-07-28 17:13:32 +00:00
async def get_userdata ( self , user , guild , key ) :
return ( await self . bot . database . execute_fetchone ( " SELECT * FROM user_data WHERE user_id = ? AND guild_id = ? AND key = ? " , ( user , guild , key ) )
or await self . bot . database . execute_fetchone ( " SELECT * FROM user_data WHERE user_id = ? AND guild_id = ' _global ' AND key = ? " , ( user , key ) ) )
async def set_userdata ( self , user , guild , key , value ) :
await self . bot . database . execute ( " INSERT OR REPLACE INTO user_data VALUES (?, ?, ?, ?) " , ( user , guild , key , value ) )
await self . bot . database . commit ( )
2021-02-09 17:17:19 +00:00
@userdata.command ( help = " Get a userdata key. Checks guild first, then global. " )
2021-07-28 17:13:32 +00:00
async def get ( self , ctx , * , key ) :
2021-08-13 10:45:18 +00:00
row = await self . get_userdata ( ctx . author . id , ctx . guild and ctx . guild . id , key )
2021-02-09 17:17:19 +00:00
if not row :
raise ValueError ( " No such key " )
2021-02-15 20:12:34 +00:00
await ctx . send ( row [ " value " ] )
2021-02-09 17:17:19 +00:00
2021-02-09 17:19:12 +00:00
@userdata.command ( name = " list " , brief = " List userdata keys in a given scope matching a query. " )
2021-07-28 17:13:32 +00:00
async def list_cmd ( self , ctx , query = " % " , scope = " guild " , show_values : bool = False ) :
2021-05-30 19:11:01 +00:00
" List userdata keys in a given scope (guild/global) matching your query (LIKE syntax). Can also show the associated values. "
2021-02-09 17:17:19 +00:00
if scope == " global " :
2021-07-28 17:13:32 +00:00
rows = await self . bot . database . execute_fetchall ( " SELECT * FROM user_data WHERE user_id = ? AND guild_id = ' _global ' AND key LIKE ? " , ( ctx . author . id , query ) )
2021-02-09 17:17:19 +00:00
else :
2021-08-13 10:45:18 +00:00
rows = await self . bot . database . execute_fetchall ( " SELECT * FROM user_data WHERE user_id = ? AND guild_id = ? AND key LIKE ? " , ( ctx . author . id , ctx . guild and ctx . guild . id , query ) )
2021-02-09 17:17:19 +00:00
out = [ ]
for row in rows :
if show_values :
out . append ( f " ** { row [ ' key ' ] } **: { row [ ' value ' ] } " )
else :
out . append ( row [ " key " ] )
if len ( out ) == 0 : return await ctx . send ( " No data " )
2021-07-28 17:13:32 +00:00
await ctx . send ( ( " \n " if show_values else " , " ) . join ( out ) [ : 2000 ] ) # TODO: split better
2021-02-09 17:17:19 +00:00
@userdata.command ( name = " set " , help = " Set a userdata key in the guild scope. " )
2021-07-28 17:13:32 +00:00
async def set_cmd ( self , ctx , key , * , value ) :
2021-02-09 17:17:19 +00:00
check_key ( key )
value = preprocess_value ( value )
2021-08-13 10:45:18 +00:00
await self . set_userdata ( ctx . author . id , ctx . guild and ctx . guild . id , key , value )
2021-02-09 17:17:19 +00:00
await ctx . send ( f " ** { key } ** set (scope guild) " )
@userdata.command ( help = " Set a userdata key in the global scope. " )
2021-07-28 17:13:32 +00:00
async def set_global ( self , ctx , key , * , value ) :
2021-02-09 17:17:19 +00:00
check_key ( key )
value = preprocess_value ( value )
2021-07-28 17:13:32 +00:00
await self . set_userdata ( ctx . author . id , " _global " , key , value )
2021-02-09 17:17:19 +00:00
await ctx . send ( f " ** { key } ** set (scope global) " )
@userdata.command ( )
2021-07-28 17:13:32 +00:00
async def inc ( self , ctx , key , by : int = 1 ) :
2021-02-09 17:19:12 +00:00
" Increase the integer value of a userdata key. "
2021-02-09 17:17:19 +00:00
check_key ( key )
2021-07-28 17:13:32 +00:00
row = await self . get_userdata ( ctx . author . id , ctx . guild . id , key )
2021-02-09 17:17:19 +00:00
if not row :
value = 0
2021-08-13 10:45:18 +00:00
guild = ctx . guild and ctx . guild . id
2021-02-09 17:17:19 +00:00
else :
value = int ( row [ " value " ] )
guild = row [ " guild_id " ]
new_value = value + by
2021-07-28 17:13:32 +00:00
await self . set_userdata ( ctx . author . id , guild , key , preprocess_value ( str ( new_value ) ) )
2021-02-09 17:17:19 +00:00
await ctx . send ( f " ** { key } ** set to { new_value } " )
@userdata.command ( )
2021-07-28 17:13:32 +00:00
async def delete ( self , ctx , * keys ) :
2021-02-09 17:19:12 +00:00
" Delete the specified keys (smallest scope first). "
2021-02-09 17:17:19 +00:00
for key in keys :
2021-08-13 10:45:18 +00:00
row = await self . get_userdata ( ctx . author . id , ctx . guild and ctx . guild . id , key )
2021-02-09 17:17:19 +00:00
if not row :
return await ctx . send ( embed = util . error_embed ( f " No such key { key } " ) )
2021-07-28 17:13:32 +00:00
await self . bot . database . execute ( " DELETE FROM user_data WHERE user_id = ? AND guild_id = ? AND key = ? " , ( ctx . author . id , row [ " guild_id " ] , key ) )
await self . bot . database . commit ( )
await ctx . send ( f " ** { key } ** deleted " )
def setup ( bot ) :
bot . add_cog ( Userdata ( bot ) )