From 276886a3a70b2da61e34aa34cd444fefac73131b Mon Sep 17 00:00:00 2001
From: Conatum <cona@thewisewolf.dev>
Date: Wed, 22 Dec 2021 17:26:49 +0200
Subject: [PATCH] sharding (general): Add launch data filters.

Filter cached reaction role messages by shardid.
Filter expiring rented room by shardid.
Filter scanned study badges by shardid.
Filter resumed study sessions by shardid.
Filter resumed workouts by shardid.

Fix a spacing issue in the log printer.
---
 bot/meta/logger.py                                |  2 +-
 bot/modules/guild_admin/reaction_roles/tracker.py |  3 ++-
 bot/modules/renting/rooms.py                      |  3 ++-
 bot/modules/study/badges/badge_tracker.py         | 12 ++++++++----
 bot/modules/study/tracking/session_tracker.py     |  3 ++-
 bot/modules/workout/tracker.py                    |  4 +++-
 6 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/bot/meta/logger.py b/bot/meta/logger.py
index 3e7bd026..a95500e4 100644
--- a/bot/meta/logger.py
+++ b/bot/meta/logger.py
@@ -17,7 +17,7 @@ logger = logging.getLogger()
 log_fmt = logging.Formatter(
     fmt=('[{asctime}][{levelname:^8}]' +
          '[SHARD {}]'.format(sharding.shard_number) +
-         '{message}'),
+         ' {message}'),
     datefmt='%d/%m | %H:%M:%S',
     style='{'
 )
diff --git a/bot/modules/guild_admin/reaction_roles/tracker.py b/bot/modules/guild_admin/reaction_roles/tracker.py
index f18e3c34..17a64960 100644
--- a/bot/modules/guild_admin/reaction_roles/tracker.py
+++ b/bot/modules/guild_admin/reaction_roles/tracker.py
@@ -12,6 +12,7 @@ from discord import PartialEmoji
 from meta import client
 from core import Lion
 from data import Row
+from data.conditions import THIS_SHARD
 from utils.lib import utc_now
 from settings import GuildSettings
 
@@ -584,5 +585,5 @@ def load_reaction_roles(client):
     """
     Load the ReactionRoleMessages.
     """
-    rows = reaction_role_messages.fetch_rows_where()
+    rows = reaction_role_messages.fetch_rows_where(guildid=THIS_SHARD)
     ReactionRoleMessage._messages = {row.messageid: ReactionRoleMessage(row.messageid) for row in rows}
diff --git a/bot/modules/renting/rooms.py b/bot/modules/renting/rooms.py
index 3e1d19c4..a8c29876 100644
--- a/bot/modules/renting/rooms.py
+++ b/bot/modules/renting/rooms.py
@@ -5,6 +5,7 @@ import datetime
 from cmdClient.lib import SafeCancellation
 
 from meta import client
+from data.conditions import THIS_SHARD
 from settings import GuildSettings
 
 from .data import rented, rented_members
@@ -276,7 +277,7 @@ class Room:
 
 @module.launch_task
 async def load_rented_rooms(client):
-    rows = rented.fetch_rows_where()
+    rows = rented.fetch_rows_where(guildid=THIS_SHARD)
     for row in rows:
         Room(row.channelid).schedule()
     client.log(
diff --git a/bot/modules/study/badges/badge_tracker.py b/bot/modules/study/badges/badge_tracker.py
index 2c0d33fb..721f3962 100644
--- a/bot/modules/study/badges/badge_tracker.py
+++ b/bot/modules/study/badges/badge_tracker.py
@@ -6,8 +6,8 @@ import contextlib
 
 import discord
 
-from meta import client
-from data.conditions import GEQ
+from meta import client, sharding
+from data.conditions import GEQ, THIS_SHARD
 from core.data import lions
 from utils.lib import strfdur
 from settings import GuildSettings
@@ -54,12 +54,16 @@ async def update_study_badges(full=False):
 
     # Retrieve member rows with out of date study badges
     if not full and client.appdata.last_study_badge_scan is not None:
+        # TODO: _extra here is a hack to cover for inflexible conditionals
         update_rows = new_study_badges.select_where(
+            guildid=THIS_SHARD,
             _timestamp=GEQ(client.appdata.last_study_badge_scan or 0),
-            _extra="OR session_start IS NOT NULL"
+            _extra="OR session_start IS NOT NULL AND (guildid >> 22) %% {} = {}".format(
+                sharding.shard_count, sharding.shard_number
+            )
         )
     else:
-        update_rows = new_study_badges.select_where()
+        update_rows = new_study_badges.select_where(guildid=THIS_SHARD)
 
     if not update_rows:
         client.appdata.last_study_badge_scan = datetime.datetime.utcnow()
diff --git a/bot/modules/study/tracking/session_tracker.py b/bot/modules/study/tracking/session_tracker.py
index 9cbe53be..8158f96a 100644
--- a/bot/modules/study/tracking/session_tracker.py
+++ b/bot/modules/study/tracking/session_tracker.py
@@ -7,6 +7,7 @@ from collections import defaultdict
 
 from utils.lib import utc_now
 from data import tables
+from data.conditions import THIS_SHARD
 from core import Lion
 from meta import client
 
@@ -398,7 +399,7 @@ async def _init_session_tracker(client):
     ended = 0
 
     # Grab all ongoing sessions from data
-    rows = current_sessions.fetch_rows_where()
+    rows = current_sessions.fetch_rows_where(guildid=THIS_SHARD)
 
     # Iterate through, resume or end as needed
     for row in rows:
diff --git a/bot/modules/workout/tracker.py b/bot/modules/workout/tracker.py
index 79dc9378..be3438df 100644
--- a/bot/modules/workout/tracker.py
+++ b/bot/modules/workout/tracker.py
@@ -7,6 +7,7 @@ from core import Lion
 from settings import GuildSettings
 from meta import client
 from data import NULL, tables
+from data.conditions import THIS_SHARD
 
 from .module import module
 from .data import workout_sessions
@@ -226,7 +227,8 @@ async def load_workouts(client):
     client.objects['current_workouts'] = {}  # (guildid, userid) -> Row
     # Process any incomplete workouts
     workouts = workout_sessions.fetch_rows_where(
-        duration=NULL
+        duration=NULL,
+        guildid=THIS_SHARD
     )
     count = 0
     for workout in workouts: