add user import as well

Author: zefanjajobse

cogs/admin.py

@@ -1,13 +1,15 @@
 """User management"""
 
+import datetime
+import io
 import logging
 from typing import Optional
 import discord
 from discord import Role, app_commands
 from discord.ext import commands
 from sqlalchemy import select, func, and_
 from sqlalchemy.exc import IntegrityError
-
+import csv
 from bot import KDRBot
 from database.dto.kd_roles import KDRole
 from database.dto.users import User
@@ -16,7 +18,7 @@
 from utils.register import register
 from utils.role_management import RoleManagement
 from utils.server_settings import get_guild, update_guild
-from utils.users import get_users_csv
+from utils.users import get_users_csv, import_csv
 from utils.voice_channel import create_voice_channel
 from utils.voice_channels import add_voice_channel
 
@@ -101,7 +103,29 @@ async def export_registered_users(self, interaction: discord.Interaction) -> Non
                 await interaction.followup.send("There are currently no registered users", ephemeral=True)
                 return
             await interaction.followup.send("Registered users:", ephemeral=True, file=registered_users)
-
+    
+    @group.command(name="import", description="Import registered users")
+    @app_commands.guild_only()
+    @app_commands.default_permissions(administrator=True)
+    @app_commands.checks.has_permissions(administrator=True)
+    async def import_registered_users(self, interaction: discord.Interaction, file: discord.Attachment) -> None:
+        """Import registered users"""
+        await interaction.response.defer()
+        if interaction.guild is None:
+            return  # is already set to guild_only
+        async with self.bot.db.create_session() as session:
+            try:
+                await import_csv(session, interaction.guild_id, file)
+                await interaction.followup.send(
+                    "Users have been imported!",
+                    ephemeral=True,
+                )
+            except Exception as e:
+                await interaction.followup.send(
+                    "Something went wrong during the import, did you deliver the data in the same order as the export function?",
+                    ephemeral=True,
+                )
+                print(e)
 
     @group.command(name="unregister", description="Unregister a user")
     @app_commands.describe(username="EA username")

database/dto/users.py

@@ -24,6 +24,15 @@ class User(Base):
         UniqueConstraint(server_id, discord_id, player_id, user_id),
     )  # must be a tuple!
 
+    def import_user(self, row, header):
+        import_data = dict(zip(header, row))
+        self.server_id = int(import_data.get("server_id"))
+        self.discord_id = int(import_data.get("discord_id"))
+        self.username = import_data.get("username")
+        self.user_id = int(import_data.get("user_id"))
+        return self
+
+
     async def update_kdr(self, session: AsyncSession, kdr_role_id: int | None):
         stmt = (
             update(User)

utils/users.py

@@ -1,11 +1,14 @@
 import csv
+import datetime
 import io
 
 import discord
 from sqlalchemy import select
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from database.dto.users import User
+from database.error_handling import is_unique_violation
 
 async def get_users_csv(session: AsyncSession, server_id: int) -> tuple[int, discord.File]:
     stmt = (
@@ -20,4 +23,24 @@ async def get_users_csv(session: AsyncSession, server_id: int) -> tuple[int, dis
         for row in res:
             outcsv.writerow([row[0].server_id, row[0].discord_id, row[0].username, row[0].player_id, row[0].kdr_role_id, row[0].user_id, row[0].created_at, row[0].updated_at])
         data_stream.seek(0)
-        return (total, discord.File(data_stream, filename="channel_names.csv"))
+        return (total, discord.File(data_stream, filename="channel_names.csv"))
+    
+async def import_csv(session: AsyncSession, server_id: int, file: discord.Attachment):
+    if file.filename.endswith(".csv"):
+        res = await file.read()
+        test = io.StringIO(res.decode("utf-8"))
+        data = list(csv.reader(test, delimiter=','))
+        for i in data[1:]:
+            try:
+                user = User().import_user(i, data[0])
+                user.server_id = server_id
+                user.created_at = datetime.datetime.now()
+                user.updated_at = datetime.datetime.now()
+                session.add(user)
+                await session.commit()
+            except IntegrityError as ex:
+                await session.rollback()
+                if is_unique_violation(ex):
+                    continue # user already exists
+                else:
+                    print(ex)