from itertools import chain from psycopg import sql from cachetools import TTLCache from data import Table, Registry, Column, RowModel from data.models import WeakCache from data.columns import Integer, String, Bool, Timestamp class CoreData(Registry, name="core"): class AppConfig(RowModel): """ Schema ------ CREATE TABLE app_config( appname TEXT PRIMARY KEY, created_at TIMESTAMPTZ NOT NULL DEFAULT now() ); """ _tablename_ = 'app_config' appname = String(primary=True) created_at = Timestamp() class BotConfig(RowModel): """ Schema ------ CREATE TABLE bot_config( appname TEXT PRIMARY KEY REFERENCES app_config(appname) ON DELETE CASCADE, default_skin TEXT ); """ _tablename_ = 'bot_config' appname = String(primary=True) default_skin = String() class Shard(RowModel): """ Schema ------ CREATE TABLE shard_data( shardname TEXT PRIMARY KEY, appname TEXT REFERENCES bot_config(appname) ON DELETE CASCADE, shard_id INTEGER NOT NULL, shard_count INTEGER NOT NULL, last_login TIMESTAMPTZ, guild_count INTEGER ); """ _tablename_ = 'shard_data' shardname = String(primary=True) appname = String() shard_id = Integer() shard_count = Integer() last_login = Timestamp() guild_count = Integer() class User(RowModel): """ User model, representing configuration data for a single user. Schema ------ CREATE TABLE user_config( userid BIGINT PRIMARY KEY, timezone TEXT, topgg_vote_reminder BOOLEAN, avatar_hash TEXT, name TEXT, API_timestamp BIGINT, gems INTEGER DEFAULT 0, first_seen TIMESTAMPTZ DEFAULT now(), last_seen TIMESTAMPTZ ); """ _tablename_ = "user_config" _cache_: WeakCache[tuple[int], 'CoreData.User'] = WeakCache(TTLCache(1000, ttl=60*5)) userid = Integer(primary=True) timezone = String() topgg_vote_reminder = Bool() avatar_hash = String() name = String() API_timestamp = Integer() gems = Integer() first_seen = Timestamp() last_seen = Timestamp() class Guild(RowModel): """ Guild model, representing configuration data for a single guild. Schema ------ CREATE TABLE guild_config( guildid BIGINT PRIMARY KEY, admin_role BIGINT, mod_role BIGINT, event_log_channel BIGINT, mod_log_channel BIGINT, alert_channel BIGINT, studyban_role BIGINT, min_workout_length INTEGER, workout_reward INTEGER, max_tasks INTEGER, task_reward INTEGER, task_reward_limit INTEGER, study_hourly_reward INTEGER, study_hourly_live_bonus INTEGER, renting_price INTEGER, renting_category BIGINT, renting_cap INTEGER, renting_role BIGINT, renting_sync_perms BOOLEAN, accountability_category BIGINT, accountability_lobby BIGINT, accountability_bonus INTEGER, accountability_reward INTEGER, accountability_price INTEGER, video_studyban BOOLEAN, video_grace_period INTEGER, greeting_channel BIGINT, greeting_message TEXT, returning_message TEXT, starting_funds INTEGER, persist_roles BOOLEAN, daily_study_cap INTEGER, pomodoro_channel BIGINT, name TEXT, first_joined_at TIMESTAMPTZ DEFAULT now(), left_at TIMESTAMPTZ ); """ _tablename_ = "guild_config" _cache_: WeakCache[tuple[int], 'CoreData.Guild'] = WeakCache(TTLCache(1000, ttl=60*5)) guildid = Integer(primary=True) admin_role = Integer() mod_role = Integer() event_log_channel = Integer() mod_log_channel = Integer() alert_channel = Integer() studyban_role = Integer() max_study_bans = Integer() min_workout_length = Integer() workout_reward = Integer() max_tasks = Integer() task_reward = Integer() task_reward_limit = Integer() study_hourly_reward = Integer() study_hourly_live_bonus = Integer() daily_study_cap = Integer() renting_price = Integer() renting_category = Integer() renting_cap = Integer() renting_role = Integer() renting_sync_perms = Bool() accountability_category = Integer() accountability_lobby = Integer() accountability_bonus = Integer() accountability_reward = Integer() accountability_price = Integer() video_studyban = Bool() video_grace_period = Integer() greeting_channel = Integer() greeting_message = String() returning_message = String() starting_funds = Integer() persist_roles = Bool() pomodoro_channel = Integer() name = String() first_joined_at = Timestamp() left_at = Timestamp() unranked_rows = Table('unranked_rows') donator_roles = Table('donator_roles') member_ranks = Table('member_ranks') class Member(RowModel): """ Member model, representing configuration data for a single member. Schema ------ CREATE TABLE members( guildid BIGINT, userid BIGINT, tracked_time INTEGER DEFAULT 0, coins INTEGER DEFAULT 0, workout_count INTEGER DEFAULT 0, revision_mute_count INTEGER DEFAULT 0, last_workout_start TIMESTAMP, last_study_badgeid INTEGER REFERENCES study_badges ON DELETE SET NULL, video_warned BOOLEAN DEFAULT FALSE, display_name TEXT, first_joined TIMESTAMPTZ DEFAULT now(), last_left TIMESTAMPTZ, _timestamp TIMESTAMP DEFAULT (now() at time zone 'utc'), PRIMARY KEY(guildid, userid) ); CREATE INDEX member_timestamps ON members (_timestamp); """ _tablename_ = 'members' _cache_: WeakCache[tuple[int, int], 'CoreData.Member'] = WeakCache(TTLCache(5000, ttl=60*5)) guildid = Integer(primary=True) userid = Integer(primary=True) tracked_time = Integer() coins = Integer() workout_count = Integer() revision_mute_count = Integer() last_workout_start = Timestamp() last_study_badgeid = Integer() video_warned = Bool() display_name = String() first_joined = Timestamp() last_left = Timestamp() _timestamp = Timestamp() @classmethod async def add_pending(cls, pending: list[tuple[int, int, int]]) -> list['CoreData.Member']: """ Safely add pending coins to a list of members. Arguments --------- pending: List of tuples of the form `(guildid, userid, pending_coins)`. """ query = sql.SQL(""" UPDATE members SET coins = LEAST(coins + t.coin_diff, 2147483647) FROM (VALUES {}) AS t (guildid, userid, coin_diff) WHERE members.guildid = t.guildid AND members.userid = t.userid RETURNING * """).format( sql.SQL(', ').join( sql.SQL("({}, {}, {})").format(sql.Placeholder(), sql.Placeholder(), sql.Placeholder()) for _ in pending ) ) # TODO: Replace with copy syntax/query? conn = await cls.table.connector.get_connection() async with conn.cursor() as cursor: await cursor.execute( query, tuple(chain(*pending)) ) rows = await cursor.fetchall() return cls._make_rows(*rows) @classmethod async def get_member_rank(cls, guildid, userid, untracked): """ Get the time and coin ranking for the given member, ignoring the provided untracked members. """ conn = await cls.table.connector.get_connection() async with conn.cursor() as curs: await curs.execute( """ SELECT time_rank, coin_rank FROM ( SELECT userid, row_number() OVER (ORDER BY total_tracked_time DESC, userid ASC) AS time_rank, row_number() OVER (ORDER BY total_coins DESC, userid ASC) AS coin_rank FROM members_totals WHERE guildid=%s AND userid NOT IN %s ) AS guild_ranks WHERE userid=%s """, (guildid, tuple(untracked), userid) ) return (await curs.fetchone()) or (None, None)