Discord Bot - 1on1/Advisor Tickets

Python SQL Discord.py

---

About the project

In order for members of Notify to get better more focused support we introduced a service that comes with their membership that allows them to get unlimited 1on1/advisor sessions related to services we provide. When we first introduced the service we had a staff member manually assign an advisor to members based on their support needs. This quickly became inefficient so I created a bot to automate the process, removing the need of manual placement. 

 

 

1. Bot Features

The bot is very simple on the user side since we want the experience to be quick but informative. Some of the features the bot has are:

• Embeds to display advisor info.
• Buttons to interact with and move to the next embeds.
• Direct linking to advisor calendars for easy signup.
• Session limit updating for advisors via a Discord command.

Below you can find a gif of the full usage of the bot.

 

 

 

2. Code

Below you can find the code for this project. The bot has been operational for over 2 years now but recently got a huge rewrite. Below you can find some of the key takeaways of the update. The new code can also be found below.

• Modularity
  • Made the code more reusable by making new libraries that all my bots access.
• Efficiency
  • Reduced the size and extra steps some of the button callbacks used, resulting in better efficiency.
  • Reduced the number of database calls.
• Useless/redundant Code
  • Removed useless code such as conn.close(), as AsyncConnection will close automatically when exited.
• Packing more into queries
  • As some of my old code resorted to sorting data outside of a query, I made my queries more complex to extract what I wanted in one simple line of code.
• Cleanliness
  • The bot originally was ~800 lines of code and was way too long.
  • With the update the total length of the code is approximately 400 lines (~50% reduction in size), reusing many function from other bots as well.
• Environment variable configuration
  • Moved any sensitive info and environment-specific variables into a .env which will also be used for all the bots used within a Discord server.
  • There was a need for some global variables as well, so I created a global_vars.py file to contain and make accessing them throughout each file more efficient.

Note, I will not be including the global_vars.py and embeds.py file since there is nothing special about them, they contain embeds for the Discord bot to send and some dicts for the other files to access.

 

Main function containing the Discord bot.

import os
import discord
from discord.ui import Button, View
import logging
import asyncio
from discord import app_commands
import psycopg
from psycopg import sql
import logging
from dotenv import load_dotenv

from global_vars import GlobalVariables
import one_on_one_utils
import dbaccess

logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"), filename='1on1ticketbotdb.log', filemode='a', format='%(asctime)s - [Line:%(lineno)d - Function:%(funcName)s In:%(filename)s From:%(name)s] \n[%(levelname)s] - %(message)s \n', datefmt='%d-%b-%y %H:%M:%S')

load_dotenv()
TOKEN = os.environ.get('ONE_ON_ONE_TOKEN')
intents = discord.Intents().all()
intents.members = True
intents.guilds = True

NOTIFYGUILDID = int(os.environ.get('NOTIFYGUILDID'))

class Client(discord.Client):
    def __init__(self):
        super().__init__(intents = intents)
        self.synced = False
        asyncio.set_event_loop_policy(
            asyncio.WindowsSelectorEventLoopPolicy()
        )

    async def on_ready(self):
        await self.wait_until_ready()
        if not self.synced:
            await client.sync()
            self.synced = True
        print(f'{self.user} has connected to Discord!')

        GlobalVariables().advisorDict = await one_on_one_utils.get_all_values()

        Client.add_view(one_on_one_utils.AdvisorMain())
        Client.add_view(one_on_one_utils.StaffView(category=0))

Client = Client()
client = app_commands.CommandTree(Client)

# Used by staff to update their sessions
@client.command(name = "updatesessions")
async def updatesessions(interaction: discord.Interaction, number: int):
    guild = Client.get_guild(NOTIFYGUILDID)
    role = discord.utils.find(lambda r: r.name == "Team: Coaches", interaction.channel.guild.roles)

    if role in interaction.user.roles:
        staff = interaction.user.id

        query = """UPDATE oneonone set sessions = %s WHERE id = %s """

        await dbaccess.write_data(query, (number, staff))

        # Update dict
        GlobalVariables().advisorDict[staff]["remaining"] = number

        await interaction.response.send_message("Set total to " + str(number) + " sessions.")
    else:
        await interaction.response.send_message("You don't have access to this command.")

@Client.event
async def on_message(message):

    # Initializes the bot flow
    if ('Initiating Jarvis 1on1...' in message.content) and (message.channel.category_id == 915303264889749544):

        await one_on_one_utils.initial_func(message.channel)

Client.run(TOKEN)

 

The next part contains the utility functions/classes that the main bot uses to handle the ticket flow. The original callback flow for each interaction was the bulk of the code and made up about 400-500 lines, the new code is only ~110 lines.

import discord
from discord.ui import Button, View
import asyncio
import functools

from global_vars import GlobalVariables
import dbaccess
import embeds

async def get_all_values():
    users_data = {}

    query = "select * from oneonone"

    records = await dbaccess.get_data(query, None)

    for row in records:
        user_id = row[0]
        user_data = {
            "remaining": row[2],
            "calendar": row[3],
            "description_sneakers": row[4] if len(row) > 4 else None
        }
        users_data[user_id] = user_data

    return users_data

async def initial_func(ch):
    embedVar = await embeds.advisor_initial()
    view = AdvisorMain()
    await ch.send(embed=embedVar, view=view)

class AdvisorMain(View):
    def __init__(self):
        super().__init__(timeout=None)

    # Define Buttons
    @discord.ui.button(label = "Sneakers", style = discord.ButtonStyle.grey, emoji = "👟", custom_id="sneakers")
    async def sneakers_callback(self, interaction, button):
        await interaction.response.edit_message(view = None)
        embedVar, embedVar1 = await embeds.advisor_sneaker()
        view = StaffView(0)
        await interaction.channel.send(embed=embedVar)
        await interaction.channel.send(embed=embedVar1, view=view)
    @discord.ui.button(label = "Flips", style = discord.ButtonStyle.grey, emoji = "⚡", custom_id="flips")
    async def flips_callback(self, interaction, button):
        await interaction.response.edit_message(view = None)
        embedVar = await embeds.advisor_flips()
        view = StaffView(1)
        await interaction.channel.send(embed=embedVar, view=view)
    @discord.ui.button(label = "Amazon FBA/Freebies", style = discord.ButtonStyle.grey, emoji = "📦", custom_id="amazon")
    async def amazon_callback(self, interaction, button):
        await interaction.response.edit_message(view = None)
        embedVar = await embeds.advisor_amazon()
        view = StaffView(2)
        await interaction.channel.send(embed=embedVar, view=view)
    @discord.ui.button(label = "Sports Betting", style = discord.ButtonStyle.grey, emoji = "🥇", custom_id="sports")
    async def sports_betting_callback(self, interaction, button):
        await interaction.response.edit_message(view = None)
        embedVar = await embeds.advisor_sports()
        view = StaffView(3)
        await interaction.channel.send(embed=embedVar, view=view)
    @discord.ui.button(label = "General/Other", style = discord.ButtonStyle.grey, emoji = "💸", custom_id="general")
    async def general_callback(self, interaction, button):
        await interaction.response.edit_message(view = None)
        embedVar = await embeds.advisor_general()
        view = StaffView(4)
        await interaction.channel.send(embed=embedVar, view=view)

class StaffView(View):
    def __init__(self, category):
        super().__init__(timeout=None)
        self.category = category

        self.matching_ids = []
        for user_id, session_info in GlobalVariables().staff_session_types.items():
            if session_info[self.category] == 1 and GlobalVariables().advisorDict[user_id]["remaining"] != 0:
                self.matching_ids.append(user_id)

        for user_id in self.matching_ids:
            member_name = GlobalVariables().id_to_member[user_id]
            button = discord.ui.Button(
                label=member_name,
                style=discord.ButtonStyle.grey,
                emoji=GlobalVariables().staff_emotes[member_name],
                custom_id=member_name
            )
            button.callback = functools.partial(self.handle_callback, user_id=user_id)
            self.add_item(button)

        button = discord.ui.Button(
            style = discord.ButtonStyle.red,
            label = "Restart",
            custom_id="restart"
        )
        button.callback = self.restart_callback
        self.add_item(button)

    async def handle_callback(self, interaction, user_id):
        member_name = GlobalVariables().id_to_member[user_id]
        remaining = GlobalVariables().advisorDict[user_id]["remaining"]
        await interaction.response.send_message(f"{GlobalVariables().staff_pings[member_name]} Will be with you shortly.\n\nPlease fill this out now: {GlobalVariables().advisorDict[user_id]['calendar']} and let your coach know when you have finished.\n\nMake sure to schedule it at least 48 hours from now.\n\nIf the advisor has a full schedule, feel free to choose someone else.")
        GlobalVariables().advisorDict[user_id]["remaining"] = remaining - 1
        query = "UPDATE oneonone set sessions = %s WHERE id = %s"
        await dbaccess.write_data(query, (remaining - 1, user_id))

    async def restart_callback(self, interaction):
        await interaction.response.edit_message(view = None)
        embedVar = await embeds.advisor_initial()
        view = AdvisorMain()
        await interaction.channel.send(embed=embedVar, view=view)

 

The next part contains code for the database access, any query from the main or the survey files will go through this. With the code rewrite I was able to reuse the same functions many times whereas in my old code I would remake new functions when I did not need to. 

import psycopg
from psycopg import sql
import pandas as pd
import logging
import os
from dotenv import load_dotenv

from global_vars import GlobalVariables

load_dotenv()

logging.basicConfig(level=os.environ.get("LOGLEVEL", "DEBUG"), filename='jarvis.log', filemode='a', format='%(asctime)s - [Line:%(lineno)d - Function:%(funcName)s In:%(filename)s From:%(name)s] \n[%(levelname)s] - %(message)s \n', datefmt='%d-%b-%y %H:%M:%S')

DBTable = os.environ.get('DBTable')
DBHost = os.environ.get('DBHost')
DBUsr = os.environ.get('DBUsr')
DBPass = os.environ.get('DBPass')
DBPort = os.environ.get('DBPort')

# Returns wanted information as a df from the database given a query.
async def get_data_as_dataframe(query):
    async with await psycopg.AsyncConnection.connect(
        dbname=DBTable,
        user=DBUsr,
        password=DBPass,
        host=DBHost,
        port=DBPort
    ) as conn:
        try:
            logging.info("Opened database successfully")
            async with conn.cursor() as curr:
                # Execute the query to fetch the specific data
                await curr.execute(query)
                records = await curr.fetchall()
                logging.info("Pulled values")

                # Execute a separate query to get the total number of rows in the table
                total_query = "SELECT COUNT(*) FROM ticketsurvey"
                await curr.execute(total_query)
                total_rows = await curr.fetchone()
                total_rows = total_rows[0] if total_rows else 0
                logging.info("Fetched total rows")

        except Exception as e:
            logging.error(e)
            return pd.DataFrame(), 0  # Return an empty DataFrame and total rows 0 in case of error

    # Convert fetched records into a DataFrame
    df = pd.DataFrame(records)  # Specify column names
    return df, total_rows

# Returns wanted information from the database given a query.
async def get_data(query, val):
    async with await psycopg.AsyncConnection.connect(
        dbname=DBTable,
        user=DBUsr,
        password=DBPass,
        host=DBHost,
        port=DBPort
    ) as conn:
        try:
            logging.info("Opened database successfully")
            async with conn.cursor() as curr:
                if val != None:
                    await curr.execute(query, val)
                else:
                    await curr.execute(query)
                records = await curr.fetchall()
                logging.info("Pulled values")

                logging.info("Fetched total rows")

        except Exception as e:
            logging.error(e)
            
    return records

async def write_data(query, data):
    async with await psycopg.AsyncConnection.connect(dbname = DBTable, user = DBUsr, password = DBPass, host = DBHost, port = DBPort) as conn:
        try:
            logging.info("Opened database successfully")
            async with conn.cursor() as curr:
                await curr.execute(query, data)
                await conn.commit()
                logging.info("Written values and closed")

        except Exception as e:
            logging.error(e)

async def write_to_db(query1, query2, payload, user):
    async with await psycopg.AsyncConnection.connect(dbname = DBTable, user = DBUsr, password = DBPass, host = DBHost, port = DBPort) as conn:
        try:
            logging.info("Opened database successfully")
            async with conn.cursor() as curr:

                await curr.execute(query1, user)
                ticketsurvey = await curr.fetchall()

                for row in ticketsurvey:
                    chid = row[0]

                #data for the %s value
                data = (chid, user, payload)
                await curr.execute(query2, data)
                await conn.commit()
                logging.info("Written values and closed")

        except Exception as e:
            logging.error(e)