Support new guild specific avatars

3 files changed, 44 insertions, 8 deletions

diff --git a/alembic/versions/d6c9452c3940_add_discord_py_2_0_avatar_support.py b/alembic/versions/d6c9452c3940_add_discord_py_2_0_avatar_support.py
new file mode 100644
index 0000000..1fcfe9b
--- /dev/null
+++ b/alembic/versions/d6c9452c3940_add_discord_py_2_0_avatar_support.py
@@ -0,0 +1,28 @@
+"""Add Discord.py 2.0 avatar support
+
+Revision ID: d6c9452c3940
+Revises: 911049796159
+Create Date: 2021-10-31 12:48:03.323113
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = 'd6c9452c3940'
+down_revision = '911049796159'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('users', sa.Column('guild_avatar_hash', sa.String(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('users', 'guild_avatar_hash')
+    # ### end Alembic commands ###
diff --git a/metricity/bot.py b/metricity/bot.py
index e19bca1..a33c40b 100644
--- a/metricity/bot.py
+++ b/metricity/bot.py
@@ -157,7 +157,8 @@ async def on_guild_available(guild: Guild) -> None:
         users.append({
             "id": str(user.id),
             "name": user.name,
-            "avatar_hash": user.avatar,
+            "avatar_hash": getattr(user.avatar, "key", None),
+            "guild_avatar_hash": getattr(user.guild_avatar, "key", None),
             "joined_at": user.joined_at,
             "created_at": user.created_at,
             "is_staff": BotConfig.staff_role_id in [role.id for role in user.roles],
@@ -193,7 +194,8 @@ async def on_member_join(member: Member) -> None:
         await db_user.update(
             id=str(member.id),
             name=member.name,
-            avatar_hash=member.avatar,
+            avatar_hash=getattr(member.avatar, "key", None),
+            guild_avatar_hash=getattr(member.guild_avatar, "key", None),
             joined_at=member.joined_at,
             created_at=member.created_at,
             is_staff=BotConfig.staff_role_id in [role.id for role in member.roles],
@@ -206,7 +208,8 @@ async def on_member_join(member: Member) -> None:
             await User.create(
                 id=str(member.id),
                 name=member.name,
-                avatar_hash=member.avatar,
+                avatar_hash=getattr(member.avatar, "key", None),
+                guild_avatar_hash=getattr(member.guild_avatar, "key", None),
                 joined_at=member.joined_at,
                 created_at=member.created_at,
                 is_staff=BotConfig.staff_role_id in [role.id for role in member.roles],
@@ -250,7 +253,8 @@ async def on_member_update(before: Member, member: Member) -> None:
     if db_user := await User.get(str(member.id)):
         if (
             db_user.name != member.name or
-            db_user.avatar_hash != member.avatar or
+            db_user.avatar_hash != getattr(member.avatar, "key", None) or
+            db_user.guild_avatar_hash != getattr(member.guild_avatar, "key", None) or
             BotConfig.staff_role_id in
             [role.id for role in member.roles] != db_user.is_staff
             or db_user.pending is not member.pending
@@ -258,7 +262,8 @@ async def on_member_update(before: Member, member: Member) -> None:
             await db_user.update(
                 id=str(member.id),
                 name=member.name,
-                avatar_hash=member.avatar,
+                avatar_hash=getattr(member.avatar, "key", None),
+                guild_avatar_hash=getattr(member.guild_avatar, "key", None),
                 joined_at=member.joined_at,
                 created_at=member.created_at,
                 is_staff=BotConfig.staff_role_id in roles,
@@ -271,7 +276,8 @@ async def on_member_update(before: Member, member: Member) -> None:
             await User.create(
                 id=str(member.id),
                 name=member.name,
-                avatar_hash=member.avatar,
+                avatar_hash=getattr(member.avatar, "key", None),
+                guild_avatar_hash=getattr(member.guild_avatar, "key", None),
                 joined_at=member.joined_at,
                 created_at=member.created_at,
                 is_staff=BotConfig.staff_role_id in roles,
diff --git a/metricity/models.py b/metricity/models.py
index 1ee2900..239d936 100644
--- a/metricity/models.py
+++ b/metricity/models.py
@@ -40,6 +40,7 @@ class User(db.Model):
     id = db.Column(db.String, primary_key=True)
     name = db.Column(db.String, nullable=False)
     avatar_hash = db.Column(db.String, nullable=True)
+    guild_avatar_hash = db.Column(db.String, nullable=True)
     joined_at = db.Column(db.DateTime, nullable=False)
     created_at = db.Column(db.DateTime, nullable=False)
     is_staff = db.Column(db.Boolean, nullable=False)
@@ -56,6 +57,7 @@ class User(db.Model):
         update_cols = [
             "name",
             "avatar_hash",
+            "guild_avatar_hash",
             "joined_at",
             "is_staff",
             "bot",
@@ -65,8 +67,8 @@ class User(db.Model):
         ]
 
         return qs.on_conflict_do_update(
-                index_elements=[cls.id],
-                set_={k: getattr(qs.excluded, k) for k in update_cols}
+            index_elements=[cls.id],
+            set_={k: getattr(qs.excluded, k) for k in update_cols}
         ).returning(cls.__table__).gino.all()