Croc/croccybot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

rewrite: Snapshots and Lion.

Browse Source
This commit is contained in:
Conatum
2022-11-20 08:34:18 +02:00
parent 7bd546126b
commit 2eea40f679
9 changed files with 322 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
bot/analytics/cog.py
View File

@@ -51,6 +51,9 @@ class Analytics(LionCog):
elif before.channel and not after.channel:
# Member left channel
action = VoiceAction.LEFT
else:
# Member change state, we don't need to deal with that
return
event = VoiceEvent(
appname=appname,

12
bot/core/cog.py
View File

@@ -2,7 +2,7 @@ from typing import Optional
import discord
from meta import LionBot, LionCog
from meta import LionBot, LionCog, LionContext
from meta.app import shardname, appname
from meta.logger import log_wrap
from utils.lib import utc_now
@@ -10,6 +10,7 @@ from utils.lib import utc_now
from settings.groups import SettingGroup
from .data import CoreData
from .lion import Lions
class CoreCog(LionCog):
@@ -17,6 +18,7 @@ class CoreCog(LionCog):
self.bot = bot
self.data = CoreData()
bot.db.load_registry(self.data)
self.lions = Lions(bot)
self.app_config: Optional[CoreData.AppConfig] = None
self.bot_config: Optional[CoreData.BotConfig] = None
@@ -32,6 +34,12 @@ class CoreCog(LionCog):
self.app_cmd_cache: list[discord.app_commands.AppCommand] = []
self.cmd_name_cache: dict[str, discord.app_commands.AppCommand] = {}
async def bot_check_once(self, ctx: LionContext):
lion = await self.lions.fetch(ctx.guild.id if ctx.guild else 0, ctx.author.id)
await lion.touch_discord_models(ctx.author)
ctx.alion = lion
return True
async def cog_load(self):
# Fetch (and possibly create) core data rows.
conn = await self.bot.db.get_connection()
@@ -48,6 +56,7 @@ class CoreCog(LionCog):
self.bot.add_listener(self.shard_update_guilds, name='on_guild_remove')
self.bot.core = self
await self.bot.add_cog(self.lions)
# Load the app command cache
for guildid in self.bot.testing_guilds:
@@ -56,6 +65,7 @@ class CoreCog(LionCog):
self.cmd_name_cache = {cmd.name: cmd for cmd in self.app_cmd_cache}
async def cog_unload(self):
await self.bot.remove_cog(self.lions.qualified_name)
self.bot.remove_listener(self.shard_update_guilds, name='on_guild_join')
self.bot.remove_listener(self.shard_update_guilds, name='on_guild_leave')
self.bot.core = None

172
bot/core/data.py
View File

@@ -1,6 +1,9 @@
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
@@ -56,22 +59,86 @@ class CoreData(Registry, name="core"):
guild_count = Integer()
class User(RowModel):
"""User model, representing configuration data for a single user."""
"""
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_: TTLCache[tuple[int], 'User'] = TTLCache(5000, ttl=60*5)
_cache_: WeakCache[tuple[int], 'CoreData.User'] = WeakCache(TTLCache(1000, ttl=60*5))
userid = Integer(primary=True)
timezone = Column()
topgg_vote_reminder = Column()
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."""
"""
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_: TTLCache[tuple[int], 'Guild'] = TTLCache(2500, ttl=60*5)
_cache_: WeakCache[tuple[int], 'CoreData.Guild'] = WeakCache(TTLCache(1000, ttl=60*5))
guildid = Integer(primary=True)
@@ -121,15 +188,41 @@ class CoreData(Registry, name="core"):
name = String()
first_joined_at = Timestamp()
left_at = Timestamp()
unranked_rows = Table('unranked_rows')
donator_roles = Table('donator_roles')
class Member(RowModel):
"""Member model, representing configuration data for a single member."""
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_: TTLCache[tuple[int, int], 'Member'] = TTLCache(5000, ttl=60*5)
_cache_: WeakCache[tuple[int, int], 'CoreData.Member'] = WeakCache(TTLCache(5000, ttl=60*5))
guildid = Integer(primary=True)
userid = Integer(primary=True)
@@ -138,16 +231,18 @@ class CoreData(Registry, name="core"):
coins = Integer()
workout_count = Integer()
last_workout_start = Column()
revision_mute_count = Integer()
last_workout_start = Timestamp()
last_study_badgeid = Integer()
video_warned = Bool()
display_name = String()
_timestamp = Column()
first_joined = Timestamp()
last_left = Timestamp()
_timestamp = Timestamp()
@classmethod
async def add_pending(cls, pending: tuple[int, int, int]) -> list['Member']:
async def add_pending(cls, pending: list[tuple[int, int, int]]) -> list['CoreData.Member']:
"""
Safely add pending coins to a list of members.
@@ -156,5 +251,56 @@ class CoreData(Registry, name="core"):
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)

124
bot/core/lion.py Normal file
View File

@@ -0,0 +1,124 @@
from typing import Optional
from cachetools import LRUCache
import discord
from meta import LionCog, LionBot, LionContext
from settings import InteractiveSetting
from utils.lib import utc_now
from data import WeakCache
from .data import CoreData
class Lion:
"""
A Lion is a high level representation of a Member in the LionBot paradigm.
All members interacted with by the application should be available as Lions.
It primarily provides an interface to the User and Member data.
Lion also provides centralised access to various Member properties and methods,
that would normally be served by other cogs.
Many Lion methods may only be used when the required cogs and extensions are loaded.
A Lion may exist without a Bot instance or a Member in cache,
although the functionality available will be more limited.
There is no guarantee that a corresponding discord Member actually exists.
"""
__slots__ = ('data', 'user_data', 'guild_data', '_member', '__weakref__')
def __init__(self, data: CoreData.Member, user_data: CoreData.User, guild_data: CoreData.Guild):
self.data = data
self.user_data = user_data
self.guild_data = guild_data
self._member: Optional[discord.Member] = None
# Data properties
@property
def key(self):
return (self.data.guildid, self.data.userid)
@property
def guildid(self):
return self.data.guildid
@property
def userid(self):
return self.data.userid
@classmethod
def get(cls, guildid, userid):
return cls._cache_.get((guildid, userid), None)
# Setting interfaces
# Each of these return an initialised member setting
@property
def timezone(self):
pass
@property
def locale(self):
pass
# Time utilities
@property
def now(self):
"""
Returns current time-zone aware time for the member.
"""
pass
# Discord data cache
async def touch_discord_models(self, member: discord.Member):
"""
Update the stored discord data from the givem member.
Intended to be used when we get member data from events that may not be available in cache.
"""
# Can we do these in one query?
if member.guild and (self.guild_data.name != member.guild.name):
await self.guild_data.update(name=member.guild.name)
avatar_key = member.avatar.key if member.avatar else None
await self.user_data.update(avatar_hash=avatar_key, name=member.name, last_seen=utc_now())
if member.display_name != self.data.display_name:
await self.data.update(display_name=member.display_name)
class Lions(LionCog):
def __init__(self, bot: LionBot):
self.bot = bot
# Full Lions cache
# Don't expire Lions with strong references
self._cache_: WeakCache[tuple[int, int], 'Lion'] = WeakCache(LRUCache(5000))
self._settings_: dict[str, InteractiveSetting] = {}
async def fetch(self, guildid, userid) -> Lion:
"""
Fetch or create the given Member.
If the guild or user row doesn't exist, also creates it.
Relies on the core cog existing, to retrieve the core data.
"""
# TODO: Find a way to reduce this to one query, while preserving cache
lion = self._cache_.get((guildid, userid))
if lion is None:
if self.bot.core:
data = self.bot.core.data
else:
raise ValueError("Cannot fetch Lion before core module is attached.")
guild = await data.Guild.fetch_or_create(guildid)
user = await data.User.fetch_or_create(userid)
member = await data.Member.fetch_or_create(guildid, userid)
lion = Lion(member, user, guild)
self._cache_[(guildid, userid)] = lion
return lion
def add_model_setting(self, setting: InteractiveSetting):
self._settings_[setting.__class__.__name__] = setting
return setting

11
bot/data/columns.py
View File

@@ -119,10 +119,9 @@ class Column(ColumnExpr, Generic[T]):
def __set_name__(self, owner, name):
# Only allow setting the owner once
if self.owner is None:
self.name = self.name or name
self.owner = owner
self.expr = sql.Identifier(self.owner._schema_, self.owner._tablename_, self.name)
self.name = self.name or name
self.owner = owner
self.expr = sql.Identifier(self.owner._schema_, self.owner._tablename_, self.name)
@overload
def __get__(self: 'Column[T]', obj: None, objtype: "None | Type['RowModel']") -> 'Column[T]':
@@ -136,10 +135,8 @@ class Column(ColumnExpr, Generic[T]):
# Get value from row data or session
if obj is None:
return self
elif obj is self.owner:
return obj.data[self.name]
else:
return self
return obj.data[self.name]
class Integer(Column[int]):

6
bot/data/models.py
View File

@@ -69,7 +69,11 @@ class RowTable(Table, Generic[RowT]):
).where(*args, **kwargs)
class WeakCache(MutableMapping):
WK = TypeVar('WK')
WV = TypeVar('WV')
class WeakCache(Generic[WK, WV], MutableMapping[WK, WV]):
def __init__(self, ref_cache):
self.ref_cache = ref_cache
self.weak_cache = WeakValueDictionary()

2
bot/meta/LionBot.py
View File

@@ -110,7 +110,7 @@ class LionBot(Bot):
if isinstance(ctx.command, HybridCommand) and ctx.command.app_command:
cmd_str = ctx.command.app_command.to_dict()
try:
raise exception from None
raise exception
except (HybridCommandError, CommandInvokeError, appCommandInvokeError):
try:
if isinstance(exception.original, (HybridCommandError, CommandInvokeError, appCommandInvokeError)):

3
bot/meta/LionContext.py
View File

@@ -8,6 +8,7 @@ from discord.ext.commands import Context
if TYPE_CHECKING:
from .LionBot import LionBot
from core.lion import Lion
logger = logging.getLogger(__name__)
@@ -44,6 +45,8 @@ class LionContext(Context['LionBot']):
Extends Context to add Lion-specific methods and attributes.
Also adds several contextual wrapped utilities for simpler user during command invocation.
"""
alion: 'Lion'
def __repr__(self):
parts = {}
if self.interaction is not None:

18
data/migration/v12-13/migration.sql
View File

@@ -1,11 +1,10 @@
-- Add gem support
ALTER TABLE user_config ADD COLUMN name TEXT;
INSERT INTO VersionHistory (version, author) VALUES (13, 'v12-v13 migration');
-- Add first_joined_at to guild table
-- Add left_at to guild table
ALTER TABLE guild_config ADD COLUMN first_joined_at TIMESTAMPTZ;
ALTER TABLE guild_config ADD COLUMN first_joined_at TIMESTAMPTZ DEFAULT now();
ALTER TABLE guild_config ADD COLUMN left_at TIMESTAMPTZ;
ALTER TABLE members ADD COLUMN first_joined TIMESTAMPTZ DEFAULT now();
ALTER TABLE members ADD COLUMN last_left TIMESTAMPTZ;
ALTER TABLE user_config ADD COLUMN first_seen TIMESTAMPTZ DEFAULT now();
ALTER TABLE user_config ADD COLUMN last_seen TIMESTAMPTZ;
-- Bot config data
@@ -114,3 +113,10 @@ CREATE TABLE analytics.gui_renders(
cardname TEXT NOT NULL,
duration INTEGER NOT NULL
) INHERITS (analytics.events);
-- TODO: Correct foreign keys for member table
-- TODO: Add name to user
-- TODO: Add first_joined and last_left time to member
-- TODO: Add first_seen and last_seen time to User
INSERT INTO VersionHistory (version, author) VALUES (13, 'v12-v13 migration');