discord-bot/src/common/database.rs

use crate::Config;
use serde::{Deserialize, Serialize};
use serenity::model::guild;
use serenity::model::id::{ChannelId, GuildId, RoleId, UserId};
use serenity::prelude::TypeMapKey;
use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions, SqliteRow};
use sqlx::{Error, FromRow, Pool, Row, Sqlite};
use std::str::FromStr;
use std::sync::Arc;
use tokio::sync::RwLock;

pub struct DataBase;
impl TypeMapKey for DataBase {
  type Value = Arc<RwLock<Pool<Sqlite>>>;
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ServerMembers {
  pub server: GuildId,
  pub id_wolves: i64,
  pub expiry: String,
}

impl<'r> FromRow<'r, SqliteRow> for ServerMembers {
  fn from_row(row: &'r SqliteRow) -> Result<Self, Error> {
    let server_tmp: i64 = row.try_get("server")?;
    let server = GuildId::from(server_tmp as u64);

    Ok(Self {
      server,
      id_wolves: row.try_get("id_wolves")?,
      expiry: row.try_get("expiry")?,
    })
  }
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ServerMembersWolves {
  pub server: GuildId,
  pub id_wolves: i64,
  pub expiry: String,
  pub email: String,
  pub discord: Option<UserId>,
  pub minecraft: Option<String>,
  pub minecraft_uid: Option<String>,
}

impl<'r> FromRow<'r, SqliteRow> for ServerMembersWolves {
  fn from_row(row: &'r SqliteRow) -> Result<Self, Error> {
    let server_tmp: i64 = row.try_get("server")?;
    let server = GuildId::from(server_tmp as u64);

    Ok(Self {
      server,
      id_wolves: row.try_get("id_wolves")?,
      expiry: row.try_get("expiry")?,
      email: row.try_get("email")?,
      discord: get_discord_from_row(row),
      minecraft: row.try_get("minecraft")?,
      minecraft_uid: row.try_get("minecraft_uid")?,
    })
  }
}

fn get_discord_from_row(row: &SqliteRow) -> Option<UserId> {
  match row.try_get("discord") {
    Ok(x) => {
      let tmp: i64 = x;
      if tmp == 0 {
        None
      } else {
        Some(UserId::from(tmp as u64))
      }
    }
    _ => None,
  }
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Wolves {
  pub id_wolves: i64,
  pub email: String,
  pub discord: Option<UserId>,
  pub minecraft: Option<String>,
}

impl<'r> FromRow<'r, SqliteRow> for Wolves {
  fn from_row(row: &'r SqliteRow) -> Result<Self, Error> {
    Ok(Self {
      id_wolves: row.try_get("id_wolves")?,
      email: row.try_get("email")?,
      discord: get_discord_from_row(row),
      minecraft: row.try_get("minecraft")?,
    })
  }
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct WolvesVerify {
  pub email: String,
  pub discord: UserId,
  pub auth_code: String,
  pub date_expiry: String,
}

impl<'r> FromRow<'r, SqliteRow> for WolvesVerify {
  fn from_row(row: &'r SqliteRow) -> Result<Self, Error> {
    let user_tmp: i64 = row.try_get("discord")?;
    let discord = UserId::from(user_tmp as u64);

    Ok(Self {
      email: row.try_get("email")?,
      discord,
      auth_code: row.try_get("auth_code")?,
      date_expiry: row.try_get("date_expiry")?,
    })
  }
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Servers {
  pub server: GuildId,
  pub wolves_api: String,
  pub wolves_id: i64,
  pub role_past: Option<RoleId>,
  pub role_current: RoleId,
  pub member_past: i64,
  pub member_current: i64,
  pub bot_channel_id: ChannelId,
}

impl<'r> FromRow<'r, SqliteRow> for Servers {
  fn from_row(row: &'r SqliteRow) -> Result<Self, Error> {
    let server_tmp: i64 = row.try_get("server")?;
    let server = GuildId::from(server_tmp as u64);
    let role_past = match row.try_get("role_past") {
      Ok(x) => {
        let tmp: i64 = x;
        if tmp == 0 {
          None
        } else {
          Some(RoleId::from(tmp as u64))
        }
      }
      _ => None,
    };
    let role_current = match row.try_get("role_current") {
      Ok(x) => {
        let tmp: i64 = x;
        RoleId::from(tmp as u64)
      }
      _ => RoleId::from(0u64),
    };

    let bot_channel_tmp: i64 = row.try_get("bot_channel_id")?;
    let bot_channel_id = ChannelId::from(bot_channel_tmp as u64);

    Ok(Self {
      server,
      wolves_api: row.try_get("wolves_api")?,
      wolves_id: row.try_get("wolves_id").unwrap_or(0),
      role_past,
      role_current,
      member_past: row.try_get("member_past")?,
      member_current: row.try_get("member_current")?,
      bot_channel_id,
    })
  }
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct RoleAdder {
  pub server: GuildId,
  pub role_a: RoleId,
  pub role_b: RoleId,
  pub role_c: RoleId,
}

impl<'r> FromRow<'r, SqliteRow> for RoleAdder {
  fn from_row(row: &'r SqliteRow) -> Result<Self, Error> {
    let server_tmp: i64 = row.try_get("server")?;
    let server = GuildId::from(server_tmp as u64);

    Ok(Self {
      server,
      role_a: get_role_from_row(row, "role_a"),
      role_b: get_role_from_row(row, "role_b"),
      role_c: get_role_from_row(row, "role_c"),
    })
  }
}

fn get_role_from_row(row: &SqliteRow, col: &str) -> RoleId {
  match row.try_get(col) {
    Ok(x) => {
      let tmp: i64 = x;
      RoleId::new(tmp as u64)
    }
    _ => RoleId::from(0u64),
  }
}

pub async fn db_init(config: &Config) -> Result<Pool<Sqlite>, Error> {
  let database = format!("{}/{}", &config.home, &config.database);

  let pool = SqlitePoolOptions::new()
    .max_connections(5)
    .connect_with(
      SqliteConnectOptions::from_str(&format!("sqlite://{}", database))?
        .foreign_keys(true)
        .create_if_missing(true),
    )
    .await?;

  // migrations are amazing!
  sqlx::migrate!("./db/migrations").run(&pool).await?;

  Ok(pool)
}

pub async fn get_server_config(db: &Pool<Sqlite>, server: &GuildId) -> Option<Servers> {
  sqlx::query_as::<_, Servers>(
    r#"
    SELECT *
    FROM servers
    WHERE server = ?
    "#,
  )
  .bind(server.get() as i64)
  .fetch_one(db)
  .await
  .ok()
}

pub async fn get_server_member(db: &Pool<Sqlite>, server: &GuildId, member: &guild::Member) -> Result<ServerMembersWolves, Error> {
  sqlx::query_as::<_, ServerMembersWolves>(
    r#"
    SELECT *
    FROM server_members
    JOIN wolves USING (id_wolves)
    WHERE server = ? AND discord = ?
    "#,
  )
  .bind(server.get() as i64)
  .bind(member.user.id.get() as i64)
  .fetch_one(db)
  .await
}

pub 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()
}