Skip to content

Commit

Permalink
simplify codebase structure and add additional documentation
Browse files Browse the repository at this point in the history
  • Loading branch information
ivinjabraham committed Dec 20, 2024
1 parent 4b1211e commit c657d1c
Show file tree
Hide file tree
Showing 11 changed files with 159 additions and 47 deletions.
2 changes: 1 addition & 1 deletion docs/CONTRIBUTING.md
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ The event handler takes care of the rest:
```rust
// On the event of a reaction being added
FullEvent::ReactionAdd { add_reaction } => {
let message_id = MessageId::new(ARCHIVE_MESSAGE_ID);
let message_id = MessageId::new(ROLES_MESSAGE_ID);
// Check if the reaction was added to the message we want and if it is reacted with the
// emoji we want
if add_reaction.message_id == message_id && data.reaction_roles.contains_key(&add_reaction.emoji) {
Expand Down
3 changes: 3 additions & 0 deletions src/commands.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,9 @@ async fn amdctl(ctx: Context<'_>) -> Result<(), Error> {
Ok(())
}

/// Every function that is defined *should* be added to the
/// returned vector in get_commands to ensure it is registered (available for the user)
/// when the bot goes online.
pub fn get_commands() -> Vec<poise::Command<Data, Error>> {
vec![amdctl()]
}
6 changes: 2 additions & 4 deletions src/scheduler/mod.rs → src/graphql/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,5 @@ GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
pub mod scheduler;
pub mod tasks;

pub use self::scheduler::run_scheduler;
pub mod models;
pub mod queries;
28 changes: 25 additions & 3 deletions src/scheduler/tasks/mod.rs → src/graphql/models.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,29 @@ GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
pub mod status_update;
pub mod tasks;
use serde::Deserialize;
use std::borrow::Cow;

pub use self::tasks::{get_tasks, Task};
#[derive(Deserialize)]
pub struct Member<'a> {
id: Option<i32>,
roll_num: Option<Cow<'a, str>>,
name: Option<Cow<'a, str>>,
hostel: &'a str,
email: &'a str,
sex: &'a str,
year: i32,
mac_addr: &'a str,
discord_id: &'a str,
group_id: i32,
}

#[derive(Deserialize)]
struct Data<'a> {
getMember: Vec<Member<'a>>,
}

#[derive(Deserialize)]
struct Root<'a> {
data: Data<'a>,
}
8 changes: 6 additions & 2 deletions src/utils/graphql.rs → src/graphql/queries.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,18 @@ along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
use serde_json::Value;

use super::models::Member;

const REQUEST_URL: &str = "https://root.shuttleapp.rs/";

pub async fn fetch_members() -> Result<Vec<String>, reqwest::Error> {
let client = reqwest::Client::new();
let query = r#"
query {
getMember {
name
name,
groupId,
discordId
}
}"#;

Expand All @@ -40,7 +44,7 @@ pub async fn fetch_members() -> Result<Vec<String>, reqwest::Error> {
.as_array()
.unwrap()
.iter()
.map(|member| member["name"].as_str().unwrap().to_string())
.map(Member)
.collect();

Ok(member_names)
Expand Down
6 changes: 5 additions & 1 deletion src/ids.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,10 @@ GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
pub const ARCHIVE_MESSAGE_ID: u64 = 1298636092886749294;
/// Points to the Embed in the #roles channel.
pub const ROLES_MESSAGE_ID: u64 = 1298636092886749294;

// Role IDs
pub const ARCHIVE_ROLE_ID: u64 = 1208457364274028574;
pub const MOBILE_ROLE_ID: u64 = 1298553701094395936;
pub const SYSTEMS_ROLE_ID: u64 = 1298553801191718944;
Expand All @@ -24,6 +27,7 @@ pub const RESEARCH_ROLE_ID: u64 = 1298553855474270219;
pub const DEVOPS_ROLE_ID: u64 = 1298553883169132554;
pub const WEB_ROLE_ID: u64 = 1298553910167994428;

// Channel IDs
pub const GROUP_ONE_CHANNEL_ID: u64 = 1225098248293716008;
pub const GROUP_TWO_CHANNEL_ID: u64 = 1225098298935738489;
pub const GROUP_THREE_CHANNEL_ID: u64 = 1225098353378070710;
Expand Down
116 changes: 92 additions & 24 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,25 @@ GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
/// Stores all the commands for the bot.
mod commands;
/// Responsible for queries, models and mutation requests sent to and from
/// [root's](https://www.github.com/amfoss/root) graphql interace.
mod graphql;
/// Stores Discord IDs that are needed across the bot.
mod ids;
/// This module is a simple cron equivalent. It spawns threads for the regular [`Task`]s that need to be completed.
mod scheduler;
/// An interface to define a job that needs to be executed regularly, for example checking for status updates daily.
mod tasks;
/// Misc. helper functions that don't really have a place anywhere else.
mod utils;

use crate::ids::{
AI_ROLE_ID, ARCHIVE_MESSAGE_ID, ARCHIVE_ROLE_ID, DEVOPS_ROLE_ID, MOBILE_ROLE_ID,
RESEARCH_ROLE_ID, SYSTEMS_ROLE_ID, WEB_ROLE_ID,
use ids::{
AI_ROLE_ID, ARCHIVE_ROLE_ID, DEVOPS_ROLE_ID, MOBILE_ROLE_ID, RESEARCH_ROLE_ID,
ROLES_MESSAGE_ID, SYSTEMS_ROLE_ID, WEB_ROLE_ID,
};

use anyhow::Context as _;
use poise::{Context as PoiseContext, Framework, FrameworkOptions, PrefixFrameworkOptions};
use serenity::{
Expand All @@ -36,50 +46,90 @@ use std::collections::HashMap;
pub type Error = Box<dyn std::error::Error + Send + Sync>;
pub type Context<'a> = PoiseContext<'a, Data, Error>;

/// Runtime allocated storage for the bot.
pub struct Data {
pub reaction_roles: HashMap<ReactionType, RoleId>,
}

/// This function is responsible for allocating the necessary fields
/// in [`Data`], before it is passed to the bot.
///
/// Currently, it only needs to store the (emoji, [`RoleId`]) pair used
/// for assigning roles to users who react to a particular message.
pub fn initialize_data() -> Data {
let mut data = Data {
reaction_roles: HashMap::new(),
};

let roles = [
(ReactionType::Unicode("📁".to_string()), RoleId::new(ARCHIVE_ROLE_ID)),
(ReactionType::Unicode("📱".to_string()), RoleId::new(MOBILE_ROLE_ID)),
(ReactionType::Unicode("⚙️".to_string()), RoleId::new(SYSTEMS_ROLE_ID)),
(ReactionType::Unicode("🤖".to_string()), RoleId::new(AI_ROLE_ID)),
(ReactionType::Unicode("📜".to_string()), RoleId::new(RESEARCH_ROLE_ID)),
(ReactionType::Unicode("🚀".to_string()), RoleId::new(DEVOPS_ROLE_ID)),
(ReactionType::Unicode("🌐".to_string()), RoleId::new(WEB_ROLE_ID)),
// Define the emoji-role pairs
let roles = [
(
ReactionType::Unicode("📁".to_string()),
RoleId::new(ARCHIVE_ROLE_ID),
),
(
ReactionType::Unicode("📱".to_string()),
RoleId::new(MOBILE_ROLE_ID),
),
(
ReactionType::Unicode("⚙️".to_string()),
RoleId::new(SYSTEMS_ROLE_ID),
),
(
ReactionType::Unicode("🤖".to_string()),
RoleId::new(AI_ROLE_ID),
),
(
ReactionType::Unicode("📜".to_string()),
RoleId::new(RESEARCH_ROLE_ID),
),
(
ReactionType::Unicode("🚀".to_string()),
RoleId::new(DEVOPS_ROLE_ID),
),
(
ReactionType::Unicode("🌐".to_string()),
RoleId::new(WEB_ROLE_ID),
),
];

// Populate reaction_roles map.
data.reaction_roles
.extend::<HashMap<ReactionType, RoleId>>(roles.into());

data
}

/// Sets up the bot using a [`poise::Framework`], which handles most of the
/// configuration including the command prefix, the event handler, the available commands,
/// managing [`Data`] and running the [`scheduler`].
#[shuttle_runtime::main]
async fn main(
#[shuttle_runtime::Secrets] secret_store: shuttle_runtime::SecretStore,
) -> shuttle_serenity::ShuttleSerenity {
// Uses Shuttle's environment variable storage solution SecretStore
// to access the token
let discord_token = secret_store
.get("DISCORD_TOKEN")
.context("'DISCORD_TOKEN' was not found")?;

let framework = Framework::builder()
.options(FrameworkOptions {
// Load bot commands
commands: commands::get_commands(),
// Pass the event handler function
event_handler: |ctx, event, framework, data| {
Box::pin(event_handler(ctx, event, framework, data))
},
// General bot settings, set to default except for prefix
prefix_options: PrefixFrameworkOptions {
prefix: Some(String::from("$")),
..Default::default()
},
..Default::default()
})
// This function that's passed to setup() is called just as
// the bot is ready to start.
.setup(|ctx, _ready, framework| {
Box::pin(async move {
poise::builtins::register_globally(ctx, &framework.options().commands).await?;
Expand All @@ -102,21 +152,28 @@ async fn main(

Ok(client.into())
}

/// Handles various events from Discord, such as reactions.
///
/// Current functionality includes:
/// - Adding roles to users based on reactions.
/// - Removing roles from users when their reactions are removed.
///
/// TODO: Refactor for better readability and modularity.
async fn event_handler(
ctx: &SerenityContext,
event: &FullEvent,
_framework: poise::FrameworkContext<'_, Data, Error>,
data: &Data,
) -> Result<(), Error> {
match event {
// Handle reactions being added.
FullEvent::ReactionAdd { add_reaction } => {
let message_id = MessageId::new(ARCHIVE_MESSAGE_ID);
if add_reaction.message_id == message_id
&& data.reaction_roles.contains_key(&add_reaction.emoji)
{
// Check if a role needs to be added i.e check if the reaction was added to [`ROLES_MESSAGE_ID`]
if is_relevant_reaction(add_reaction.message_id, &add_reaction.emoji, data) {
// This check for a guild_id isn't strictly necessary, since we're already checking
// if the reaction was added to the [`ROLES_MESSAGE_ID`] which *should* point to a
// message in the server.
if let Some(guild_id) = add_reaction.guild_id {
// TODO: Use try_join to await concurrently?
if let Ok(member) = guild_id.member(ctx, add_reaction.user_id.unwrap()).await {
if let Err(e) = member
.add_role(
Expand All @@ -127,18 +184,21 @@ async fn event_handler(
)
.await
{
eprintln!("Error: {:?}", e);
// TODO: Replace with tracing
eprintln!("Error adding role: {:?}", e);
}
}
}
}
}

// Handle reactions being removed.
FullEvent::ReactionRemove { removed_reaction } => {
let message_id = MessageId::new(ARCHIVE_MESSAGE_ID);
if message_id == removed_reaction.message_id
&& data.reaction_roles.contains_key(&removed_reaction.emoji)
{
// Check if a role needs to be added i.e check if the reaction was added to [`ROLES_MESSAGE_ID`]
if is_relevant_reaction(removed_reaction.message_id, &removed_reaction.emoji, data) {
// This check for a guild_id isn't strictly necessary, since we're already checking
// if the reaction was added to the [`ROLES_MESSAGE_ID`] which *should* point to a
// message in the server.
if let Some(guild_id) = removed_reaction.guild_id {
if let Ok(member) = guild_id
.member(ctx, removed_reaction.user_id.unwrap())
Expand All @@ -150,18 +210,26 @@ async fn event_handler(
*data
.reaction_roles
.get(&removed_reaction.emoji)
.expect("Hard coded value verified earlier"),
.expect("Hard coded value verified earlier."),
)
.await
{
eprintln!("Error: {:?}", e);
eprintln!("Error removing role: {:?}", e);
}
}
}
}
}

// Ignore all other events for now.
_ => {}
}

Ok(())
}

/// Helper function to check if a reaction was made to [`ROLES_MESSAGE_ID`] and if
/// [`Data::reaction_roles`] contains a relevant (emoji, role) pair.
fn is_relevant_reaction(message_id: MessageId, emoji: &ReactionType, data: &Data) -> bool {
message_id == MessageId::new(ROLES_MESSAGE_ID) && data.reaction_roles.contains_key(emoji)
}
8 changes: 7 additions & 1 deletion src/scheduler/scheduler.rs → src/scheduler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,16 @@ GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
*/
use crate::scheduler::tasks::{get_tasks, Task};
use crate::tasks::{get_tasks, Task};
use serenity::client::Context as SerenityContext;

use tokio::spawn;

/// Spawns a thread for each [`Task`].
///
/// [`SerenityContext`] is passed along with it so that they can
/// call any required Serenity functions without creating a new [`serenity::http`]
/// interface with a Discord token.
pub async fn run_scheduler(ctx: SerenityContext) {
let tasks = get_tasks();

Expand All @@ -28,6 +33,7 @@ pub async fn run_scheduler(ctx: SerenityContext) {
}
}

/// Runs the function [`Task::run`] and goes back to sleep until it's time to run again.
async fn schedule_task(ctx: SerenityContext, task: Box<dyn Task>) {
loop {
let next_run_in = task.run_in();
Expand Down
Loading

0 comments on commit c657d1c

Please sign in to comment.