Data system refactor and core redesign for public.

Redesigned data and core systems to be public-capable.
This commit is contained in:
2021-09-12 11:04:49 +03:00
parent 459a728968
commit 0183b63c55
33 changed files with 1170 additions and 790 deletions

View File

@@ -1,2 +1,4 @@
from . import tables
from .user import User
from . import data # noqa
from .module import module
from .lion import Lion # noqa

81
bot/core/data.py Normal file
View File

@@ -0,0 +1,81 @@
from psycopg2.extras import execute_values
from cachetools import TTLCache
from data import RowTable, Table
meta = RowTable(
'AppData',
('appid', 'last_study_badge_scan'),
'appid',
attach_as='meta',
)
user_config = RowTable(
'user_config',
('userid', 'timezone'),
'userid',
cache=TTLCache(5000, ttl=60*5)
)
@user_config.save_query
def add_pending(pending):
"""
pending:
List of tuples of the form `(userid, pending_coins, pending_time)`.
"""
with lions.conn:
cursor = lions.conn.cursor()
data = execute_values(
cursor,
"""
UPDATE members
SET
coins = coins + t.coin_diff,
tracked_time = tracked_time + t.time_diff
FROM
(VALUES %s)
AS
t (guildid, userid, coin_diff, time_diff)
WHERE
members.guildid = t.guildid
AND
members.userid = t.userid
RETURNING *
""",
pending,
fetch=True
)
return lions._make_rows(*data)
guild_config = RowTable(
'guild_config',
('guildid', 'admin_role', 'mod_role', 'event_log_channel',
'min_workout_length', 'workout_reward',
'max_tasks', 'task_reward', 'task_reward_limit',
'study_hourly_reward', 'study_hourly_live_bonus',
'study_ban_role', 'max_study_bans'),
'guildid',
cache=TTLCache(1000, ttl=60*5)
)
unranked_roles = Table('unranked_roles')
donator_roles = Table('donator_roles')
lions = RowTable(
'members',
('guildid', 'userid',
'tracked_time', 'coins',
'workout_count', 'last_workout_start',
'last_study_badgeid',
'study_ban_count',
),
('guildid', 'userid'),
cache=TTLCache(5000, ttl=60*5),
attach_as='lions'
)

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

@@ -0,0 +1,145 @@
import pytz
from meta import client
from data import tables as tb
from settings import UserSettings
class Lion:
"""
Class representing a guild Member.
Mostly acts as a transparent interface to the corresponding Row,
but also adds some transaction caching logic to `coins` and `tracked_time`.
"""
__slots__ = ('guildid', 'userid', '_pending_coins', '_pending_time', '_member')
# Members with pending transactions
_pending = {} # userid -> User
# Lion cache. Currently lions don't expire
_lions = {} # (guildid, userid) -> Lion
def __init__(self, guildid, userid):
self.guildid = guildid
self.userid = userid
self._pending_coins = 0
self._pending_time = 0
self._member = None
self._lions[self.key] = self
@classmethod
def fetch(cls, guildid, userid):
"""
Fetch a Lion with the given member.
If they don't exist, creates them.
If possible, retrieves the user from the user cache.
"""
key = (guildid, userid)
if key in cls._lions:
return cls._lions[key]
else:
tb.lions.fetch_or_create(key)
return cls(guildid, userid)
@property
def key(self):
return (self.guildid, self.userid)
@property
def member(self):
"""
The discord `Member` corresponding to this user.
May be `None` if the member is no longer in the guild or the caches aren't populated.
Not guaranteed to be `None` if the member is not in the guild.
"""
if self._member is None:
guild = client.get_guild(self.guildid)
if guild:
self._member = guild.get_member(self.userid)
return self._member
@property
def data(self):
"""
The Row corresponding to this user.
"""
return tb.lions.fetch(self.key)
@property
def settings(self):
"""
The UserSettings object for this user.
"""
return UserSettings(self.userid)
@property
def time(self):
"""
Amount of time the user has spent studying, accounting for pending values.
"""
return int(self.data.tracked_time + self._pending_time)
@property
def coins(self):
"""
Number of coins the user has, accounting for the pending value.
"""
return int(self.data.coins + self._pending_coins)
def localize(self, naive_utc_dt):
"""
Localise the provided naive UTC datetime into the user's timezone.
"""
timezone = self.settings.timezone.value
return naive_utc_dt.replace(tzinfo=pytz.UTC).astimezone(timezone)
def addCoins(self, amount, flush=True):
"""
Add coins to the user, optionally store the transaction in pending.
"""
self._pending_coins += amount
self._pending[self.key] = self
if flush:
self.flush()
def addTime(self, amount, flush=True):
"""
Add time to a user (in seconds), optionally storing the transaction in pending.
"""
self._pending_time += amount
self._pending[self.key] = self
if flush:
self.flush()
def flush(self):
"""
Flush any pending transactions to the database.
"""
self.sync(self)
@classmethod
def sync(cls, *lions):
"""
Flush pending transactions to the database.
Also refreshes the Row cache for updated lions.
"""
lions = lions or list(cls._pending.values())
if lions:
# Build userid to pending coin map
pending = [
(lion.guildid, lion.userid, int(lion._pending_coins), int(lion._pending_time))
for lion in lions
]
# Write to database
tb.lions.queries.add_pending(pending)
# Cleanup pending users
for lion in lions:
lion._pending_coins -= int(lion._pending_coins)
lion._pending_time -= int(lion._pending_time)
cls._pending.pop(lion.key, None)

36
bot/core/module.py Normal file
View File

@@ -0,0 +1,36 @@
import logging
import asyncio
from meta import client, conf
from LionModule import LionModule
from .lion import Lion
module = LionModule("Core")
async def _lion_sync_loop():
while True:
while not client.is_ready():
await asyncio.sleep(1)
client.log(
"Running lion data sync.",
context="CORE",
level=logging.DEBUG,
post=False
)
Lion.sync()
await asyncio.sleep(conf.bot.getint("lion_sync_period"))
@module.launch_task
async def launch_lion_sync_loop(client):
asyncio.create_task(_lion_sync_loop())
@module.unload_task
async def final_lion_sync(client):
Lion.sync()

View File

@@ -1,31 +0,0 @@
from psycopg2.extras import execute_values
from cachetools import TTLCache
from data import RowTable, Table
users = RowTable(
'lions',
('userid', 'tracked_time', 'coins'),
'userid',
cache=TTLCache(5000, ttl=60*5)
)
@users.save_query
def add_coins(userid_coins):
with users.conn:
cursor = users.conn.cursor()
data = execute_values(
cursor,
"""
UPDATE lions
SET coins = coins + t.diff
FROM (VALUES %s) AS t (userid, diff)
WHERE lions.userid = t.userid
RETURNING *
""",
userid_coins,
fetch=True
)
return users._make_rows(*data)

View File

@@ -1,105 +0,0 @@
from . import tables as tb
from meta import conf, client
class User:
"""
Class representing a "Lion", i.e. a member of the managed guild.
Mostly acts as a transparent interface to the corresponding Row,
but also adds some transaction caching logic to `coins`.
"""
__slots__ = ('userid', '_pending_coins', '_member')
# Users with pending transactions
_pending = {} # userid -> User
# User cache. Currently users don't expire
_users = {} # userid -> User
def __init__(self, userid):
self.userid = userid
self._pending_coins = 0
self._users[self.userid] = self
@classmethod
def fetch(cls, userid):
"""
Fetch a User with the given userid.
If they don't exist, creates them.
If possible, retrieves the user from the user cache.
"""
if userid in cls._users:
return cls._users[userid]
else:
tb.users.fetch_or_create(userid)
return cls(userid)
@property
def member(self):
"""
The discord `Member` corresponding to this user.
May be `None` if the member is no longer in the guild or the caches aren't populated.
Not guaranteed to be `None` if the member is not in the guild.
"""
if self._member is None:
self._member = client.get_guild(conf.meta.getint('managed_guild_id')).get_member(self.userid)
@property
def data(self):
"""
The Row corresponding to this user.
"""
return tb.users.fetch(self.userid)
@property
def time(self):
"""
Amount of time the user has spent.. studying?
"""
return self.data.tracked_time
@property
def coins(self):
"""
Number of coins the user has, accounting for the pending value.
"""
return self.data.coins + self._pending_coins
def addCoins(self, amount, flush=True):
"""
Add coins to the user, optionally store the transaction in pending.
"""
self._pending_coins += amount
if self._pending_coins != 0:
self._pending[self.userid] = self
else:
self._pending.pop(self.userid, None)
if flush:
self.flush()
def flush(self):
"""
Flush any pending transactions to the database.
"""
self.sync(self)
@classmethod
def sync(cls, *users):
"""
Flush pending transactions to the database.
Also refreshes the Row cache for updated users.
"""
users = users or list(cls._pending.values())
if users:
# Build userid to pending coin map
userid_coins = [(user.userid, user._pending_coins) for user in users]
# Write to database
tb.users.queries.add_coins(userid_coins)
# Cleanup pending users
for user in users:
user._pending_coins = 0
cls._pending.pop(user.userid, None)