discord-bot/src/bin/update_users.rs

use serenity::model::id::GuildId;
use serenity::{
  async_trait,
  client::{Context, EventHandler},
  model::{
    gateway::{GatewayIntents, Ready},
    prelude::RoleId,
  },
  Client,
};
use skynet_discord_bot::{db_init, get_config, Accounts, Config, DataBase, Servers};
use sqlx::{Pool, Sqlite};
use std::process;
use std::sync::Arc;
use tokio::sync::RwLock;

#[tokio::main]
async fn main() {
  let config = get_config();
  let db = match db_init(&config).await {
    Ok(x) => x,
    Err(_) => return,
  };

  // Intents are a bitflag, bitwise operations can be used to dictate which intents to use
  let intents = GatewayIntents::GUILDS | GatewayIntents::GUILD_MESSAGES | GatewayIntents::MESSAGE_CONTENT | GatewayIntents::GUILD_MEMBERS;
  // Build our client.
  let mut client = Client::builder(&config.discord_token, intents)
    .event_handler(Handler {})
    .await
    .expect("Error creating client");

  {
    let mut data = client.data.write().await;

    data.insert::<Config>(Arc::new(RwLock::new(config)));
    data.insert::<DataBase>(Arc::new(RwLock::new(db)));
  }

  if let Err(why) = client.start().await {
    println!("Client error: {:?}", why);
  }
}

struct Handler;
#[async_trait]
impl EventHandler for Handler {
  async fn ready(&self, ctx: Context, ready: Ready) {
    let ctx = Arc::new(ctx);
    println!("{} is connected!", ready.user.name);

    bulk_check(Arc::clone(&ctx)).await;

    // finish up
    process::exit(0);
  }
}

async fn bulk_check(ctx: Arc<Context>) {
  let db_lock = {
    let data_read = ctx.data.read().await;
    data_read.get::<DataBase>().expect("Expected Config in TypeMap.").clone()
  };

  let db = db_lock.read().await;

  for server_config in get_server_config_bulk(&db).await {
    let Servers {
      server,
      role_past,
      role_current,
      ..
    } = server_config;

    let mut roles_set = [0, 0, 0];
    let mut members = vec![];

    let server = GuildId::from(server as u64);
    for member in get_server_member_bulk(&db, &server).await {
      if let Some(x) = member.discord {
        members.push(x);
      }
    }
    let mut members_all = members.len();

    if let Ok(x) = server.members(&ctx, None, None).await {
      for mut member in x {
        if members.contains(&member.user.name) {
          let mut roles = vec![];

          if let Some(role) = &role_past {
            let role = RoleId::from(*role as u64);
            if !member.roles.contains(&role) {
              roles_set[0] += 1;
              roles.push(role.to_owned());
            }
          }

          if let Some(role) = &role_current {
            let role = RoleId::from(*role as u64);
            if !member.roles.contains(&role) {
              roles_set[1] += 1;
              roles.push(role.to_owned());
            }
          }

          if let Err(e) = member.add_roles(&ctx, &roles).await {
            println!("{:?}", e);
          }
        } else {
          // old and never

          if let Some(role) = &role_past {
            let role = RoleId::from(*role as u64);
            if member.roles.contains(&role) {
              members_all += 1;
            }
          }

          if let Some(role) = &role_current {
            let role = RoleId::from(*role as u64);
            if member.roles.contains(&role) {
              roles_set[2] += 1;
              // if theya re not a current member and have the role then remove it
              if let Err(e) = member.remove_role(&ctx, role).await {
                println!("{:?}", e);
              }
            }
          }
        }
      }
    }

    set_server_numbers(&db, &server, members_all as i64, members.len() as i64).await;

    // small bit of logging to note changes over time
    println!("{:?} Changes: New: +{}, Current: +{}/-{}", server.as_u64(), roles_set[0], roles_set[1], roles_set[2]);
  }
}
async fn get_server_config_bulk(db: &Pool<Sqlite>) -> Vec<Servers> {
  sqlx::query_as::<_, Servers>(
    r#"
        SELECT *
        FROM servers
        "#,
  )
  .fetch_all(db)
  .await
  .unwrap_or_default()
}

async fn get_server_member_bulk(db: &Pool<Sqlite>, server: &GuildId) -> Vec<Accounts> {
  sqlx::query_as::<_, Accounts>(
    r#"
        SELECT *
        FROM accounts
        WHERE server = ? AND discord IS NOT NULL
        "#,
  )
  .bind(*server.as_u64() as i64)
  .fetch_all(db)
  .await
  .unwrap_or_default()
}

async fn set_server_numbers(db: &Pool<Sqlite>, server: &GuildId, past: i64, current: i64) {
  match sqlx::query_as::<_, Accounts>(
    "
      UPDATE servers
      SET member_past = ?, member_current = ?
      WHERE server = ?
    ",
  )
  .bind(past)
  .bind(current)
  .bind(*server.as_u64() as i64)
  .fetch_optional(db)
  .await
  {
    Ok(_) => {}
    Err(e) => {
      println!("Failure to insert into {}", server.as_u64());
      println!("{:?}", e);
    }
  }
}