feat: bit more robust register

.gitignore

@@ -24,4 +24,4 @@ rust-project.json
 
 # End of https://www.toptal.com/developers/gitignore/api/rust,rust-analyzer
 
-*.db
+*.db*

Cargo.lock

@@ -36,6 +36,7 @@ dependencies = [
  "async-minecraft-ping",
  "chrono",
  "diesel",
+ "diesel_migrations",
  "poise",
  "tokio",
  "uuid",
@@ -422,6 +423,7 @@ checksum = "72eb77396836a4505da85bae0712fa324b74acfe1876d7c2f7e694ef3d0ee373"
 dependencies = [
  "diesel_derives",
  "libsqlite3-sys",
+ "r2d2",
 ]
 
 [[package]]
@@ -436,6 +438,17 @@ dependencies = [
  "syn 1.0.109",
 ]
 
+[[package]]
+name = "diesel_migrations"
+version = "2.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e9ae22beef5e9d6fab9225ddb073c1c6c1a7a6ded5019d5da11d1e5c5adc34e2"
+dependencies = [
+ "diesel",
+ "migrations_internals",
+ "migrations_macros",
+]
+
 [[package]]
 name = "digest"
 version = "0.10.6"
@@ -866,6 +879,27 @@ version = "2.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d"
 
+[[package]]
+name = "migrations_internals"
+version = "2.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c493c09323068c01e54c685f7da41a9ccf9219735c3766fbfd6099806ea08fbc"
+dependencies = [
+ "serde",
+ "toml",
+]
+
+[[package]]
+name = "migrations_macros"
+version = "2.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8a8ff27a350511de30cdabb77147501c36ef02e0451d957abea2f30caffb2b58"
+dependencies = [
+ "migrations_internals",
+ "proc-macro2",
+ "quote",
+]
+
 [[package]]
 name = "mime"
 version = "0.3.17"
@@ -1073,6 +1107,17 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "r2d2"
+version = "0.8.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93"
+dependencies = [
+ "log",
+ "parking_lot",
+ "scheduled-thread-pool",
+]
+
 [[package]]
 name = "rand"
 version = "0.8.5"
@@ -1231,6 +1276,15 @@ version = "1.0.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f91339c0467de62360649f8d3e185ca8de4224ff281f66000de5eb2a77a79041"
 
+[[package]]
+name = "scheduled-thread-pool"
+version = "0.2.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19"
+dependencies = [
+ "parking_lot",
+]
+
 [[package]]
 name = "scopeguard"
 version = "1.1.0"
@@ -1554,6 +1608,15 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "toml"
+version = "0.5.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f4f7f0dd8d50a853a531c426359045b1998f04219d88799810762cd4ad314234"
+dependencies = [
+ "serde",
+]
+
 [[package]]
 name = "tower-service"
 version = "0.3.2"

Cargo.toml

@@ -22,4 +22,5 @@ chrono = "*"
 uuid = "*"
 poise = "0.5.5"
 async-minecraft-ping = "0.8.0"
-diesel = { version = "2.0.4", features = ["sqlite", "returning_clauses_for_sqlite_3_35"] }
+diesel = { version = "2.0.4", features = ["sqlite", "returning_clauses_for_sqlite_3_35", "r2d2"] }
+diesel_migrations = { version = "2.0.0", features = [ "sqlite" ] }

migrations/2023-05-15-142901_create_users/up.sql

@@ -1,7 +1,11 @@
 -- Your SQL goes here
 CREATE TABLE users (
-    discord_id UNSIGNED BIG INT PRIMARY KEY NOT NULL,
+    id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
+    discord_id UNSIGNED BIG INT NOT NULL,
+    guild_id UNSIGNED BIG INT NOT NULL,
     email TEXT UNIQUE NOT NULL,
     first_name TEXT NOT NULL,
-    last_name TEXT NOT NULL
+    last_name TEXT NOT NULL,
+
+    UNIQUE(discord_id, guild_id)
 );

src/build.rs

@@ -0,0 +1,3 @@
+fn main() {
+   println!("cargo:rerun-if-changed=migrations");
+}

src/commands/users.rs

@@ -1,5 +1,6 @@
-use crate::db::users::{user_all, user_insert, User};
+use crate::db::users::{NewUser, User};
 use crate::{Context, Error};
+use diesel::RunQueryDsl;
 
 #[poise::command(prefix_command, slash_command)]
 pub async fn register(
@@ -8,16 +9,33 @@ pub async fn register(
     last_name: String,
     email: String,
 ) -> Result<(), Error> {
-    let user = User {
-        discord_id: ctx.author().id.0 as i64,
-        first_name,
-        last_name,
-        email,
-    };
+    if let Some(guild_id) = ctx.guild_id() {
+        let discord_id = ctx.author().id.0 as i64;
 
-    {
-        let mut conn = ctx.data().conn.lock().unwrap();
-        user_insert(&mut conn, &user);
+        {
+            let mut conn = ctx.data().pool.get()?;
+
+            if User::get(&mut conn, guild_id.into(), discord_id)?.is_some() {
+                ctx.say("You've already been registered.").await?;
+
+                return Ok(());
+            }
+
+            let new_user = NewUser {
+                discord_id,
+                guild_id: guild_id.into(),
+                first_name,
+                last_name,
+                email,
+            };
+
+            new_user.insert(&mut conn)?;
+        }
+
+        ctx.say("You have been registered.").await?;
+    } else {
+        ctx.say("You have to send this message from a guild.")
+            .await?;
     }
 
     Ok(())
@@ -25,21 +43,25 @@ pub async fn register(
 
 #[poise::command(prefix_command, slash_command)]
 pub async fn registered(ctx: Context<'_>) -> Result<(), Error> {
-    let users = {
-        let mut conn = ctx.data().conn.lock().unwrap();
-        user_all(&mut conn)
-    };
+    if let Some(guild_id) = ctx.guild_id() {
+        let users = {
+            let mut conn = ctx.data().pool.get()?;
+            User::by_guild_id(guild_id.into()).load(&mut conn)?
+        };
 
-    ctx.send(|f| {
-        f.embed(|e| {
-            e.description("Registered users").fields(
-                users
-                    .into_iter()
-                    .map(|u| (format!("{} {}", u.first_name, u.last_name), u.email, false)),
-            )
+        ctx.send(|f| {
+            f.embed(|e| {
+                e.description("Registered users").fields(
+                    users
+                        .into_iter()
+                        .map(|u| (format!("{} {}", u.first_name, u.last_name), u.email, false)),
+                )
+            })
         })
-    })
-    .await?;
+        .await?;
+    } else {
+        ctx.say("You are not in a guild.").await?;
+    }
 
     Ok(())
 }

src/db/mod.rs

@@ -1,2 +1,43 @@
 mod schema;
 pub mod users;
+
+use diesel::connection::SimpleConnection;
+use diesel::r2d2::{ConnectionManager, Pool};
+use diesel::sqlite::{Sqlite, SqliteConnection};
+use std::error::Error;
+
+use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness};
+pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations");
+
+type DbError = Box<dyn Error + Send + Sync + 'static>;
+
+fn run_migrations(connection: &mut impl MigrationHarness<Sqlite>) -> Result<(), DbError> {
+    // This will run the necessary migrations.
+    //
+    // See the documentation for `MigrationHarness` for
+    // all available methods.
+    connection.run_pending_migrations(MIGRATIONS)?;
+
+    Ok(())
+}
+
+fn initialize_db(conn: &mut SqliteConnection) -> Result<(), DbError> {
+    // Enable WAL mode and enforce foreign keys
+    conn.batch_execute(
+        "PRAGMA journal_mode = WAL; PRAGMA synchronous = NORMAL; PRAGMA foreign_keys = ON;",
+    )?;
+    run_migrations(conn)?;
+
+    Ok(())
+}
+
+pub fn initialize_pool(url: &str) -> Result<Pool<ConnectionManager<SqliteConnection>>, DbError> {
+    let manager = ConnectionManager::new(url);
+
+    let pool = Pool::builder().test_on_check_out(true).build(manager)?;
+
+    let mut conn = pool.get()?;
+    initialize_db(&mut conn)?;
+
+    Ok(pool)
+}

src/db/schema.rs

@@ -1,8 +1,10 @@
 // @generated automatically by Diesel CLI.
 
 diesel::table! {
-    users (discord_id) {
+    users (id) {
+        id -> Integer,
         discord_id -> BigInt,
+        guild_id -> BigInt,
         email -> Text,
         first_name -> Text,
         last_name -> Text,

src/db/users.rs

@@ -1,22 +1,79 @@
-use super::schema::users;
+use super::schema::users::{self, dsl::*};
+use diesel::dsl::Eq;
+use diesel::dsl::{AsSelect, Select};
+use diesel::helper_types::Filter;
 use diesel::prelude::*;
+use diesel::sqlite::Sqlite;
 use diesel::sqlite::SqliteConnection;
 
-#[derive(Queryable, Insertable)]
+#[derive(Queryable, Selectable, AsChangeset)]
+#[diesel(table_name = users)]
 pub struct User {
+    pub id: i32,
     pub discord_id: i64,
+    pub guild_id: i64,
     pub email: String,
     pub first_name: String,
     pub last_name: String,
 }
 
-pub fn user_insert(conn: &mut SqliteConnection, user: &User) -> User {
-    diesel::insert_into(users::table)
-        .values(user)
-        .get_result(conn)
-        .expect("fuck")
+#[derive(Insertable)]
+#[diesel(table_name = users)]
+pub struct NewUser {
+    pub discord_id: i64,
+    pub guild_id: i64,
+    pub email: String,
+    pub first_name: String,
+    pub last_name: String,
 }
 
-pub fn user_all(conn: &mut SqliteConnection) -> Vec<User> {
-    users::table.load::<User>(conn).expect("nou")
+type All = Select<users::table, AsSelect<User, Sqlite>>;
+type WithGuild<T> = Eq<guild_id, T>;
+type ByGuild<T> = Filter<All, WithGuild<T>>;
+
+impl User {
+    pub fn all() -> All {
+        users::table.select(User::as_select())
+    }
+
+    // pub fn by_guild<T>(guild_id_: T) -> ByGuild<T>
+    //     where T: AsExpression<BigInt>
+    //  {
+    //     Self::all().filter(guild_id.eq(guild_id_))
+    // }
+
+    pub fn by_guild_id(guild_id_: i64) -> ByGuild<i64> {
+        Self::all().filter(guild_id.eq(guild_id_))
+    }
+
+    pub fn get(
+        conn: &mut SqliteConnection,
+        guild_id_: i64,
+        discord_id_: i64,
+    ) -> Result<Option<User>, diesel::result::Error> {
+        Self::all()
+            .filter(guild_id.eq(guild_id_))
+            .filter(discord_id.eq(discord_id_))
+            .first(conn)
+            .optional()
+    }
+
+    pub fn get_by_id(
+        conn: &mut SqliteConnection,
+        id_: i32,
+    ) -> Result<Option<User>, diesel::result::Error> {
+        Self::all().find(id_).first(conn).optional()
+    }
+
+    pub fn update(&self, conn: &mut SqliteConnection) -> Result<usize, diesel::result::Error> {
+        diesel::update(users::table).set(self).execute(conn)
+    }
+}
+
+impl NewUser {
+    pub fn insert(&self, conn: &mut SqliteConnection) -> Result<User, diesel::result::Error> {
+        diesel::insert_into(users::table)
+            .values(self)
+            .get_result(conn)
+    }
 }

src/main.rs

@@ -2,9 +2,9 @@ mod commands;
 mod db;
 
 use affluences_api::AffluencesClient;
-use diesel::Connection;
+use diesel::r2d2::{ConnectionManager, Pool};
+use diesel::sqlite::SqliteConnection;
 use poise::serenity_prelude as serenity;
-use std::sync::Mutex;
 use std::{env::var, time::Duration};
 
 // Types used by all command functions
@@ -14,7 +14,7 @@ type Context<'a> = poise::Context<'a, Data, Error>;
 // Custom user data passed to all command functions
 pub struct Data {
     client: AffluencesClient,
-    conn: Mutex<diesel::sqlite::SqliteConnection>,
+    pool: Pool<ConnectionManager<SqliteConnection>>,
 }
 
 async fn on_error(error: poise::FrameworkError<'_, Data, Error>) {
@@ -84,6 +84,8 @@ async fn main() {
         ..Default::default()
     };
 
+    let pool = db::initialize_pool("affy.db").expect("Failed to initialize database.");
+
     poise::Framework::builder()
         .token(
             var("DISCORD_TOKEN")
@@ -95,9 +97,7 @@ async fn main() {
                 poise::builtins::register_globally(ctx, &framework.options().commands).await?;
                 Ok(Data {
                     client: AffluencesClient::new(),
-                    conn: Mutex::new(
-                        diesel::sqlite::SqliteConnection::establish("affy.db").unwrap(),
-                    ),
+                    pool,
                 })
             })
         })