Add postgres database via sqlx

Author: lucawen

Cargo.toml

@@ -17,6 +17,7 @@ anyhow = "1.0.33"
 config = "0.10.1"
 serde = "1.0.117"
 serde_derive = "1.0.117"
+sqlx = { version = "0.4.0-beta.1", default-features = false, features = [ "runtime-tokio", "macros", "postgres", "uuid" ] }
 
 [dependencies.serenity]
 version = "0.9.0-rc.2"

config/development.toml

@@ -1,3 +1,5 @@
+database_url="postgres://postgres:[email protected]/sheepbot"
+
 [discord]
 prefix = "?"
 

config/example.toml

@@ -1,3 +1,5 @@
+database_url="postgres://[email protected]/example"
+
 [discord]
 prefix = "?"
 token = "DISCORD_TOKEN"

docker-compose.prod.yml

@@ -9,3 +9,6 @@ services:
     image: fredboat/lavalink:master
     volumes:
       - ./lavalink/application.yml:/opt/Lavalink/application.yml
+  db:
+    image: postgres:12
+    restart: always

docker-compose.yml

@@ -17,4 +17,12 @@ services:
     volumes:
       - ./lavalink/application.yml:/opt/Lavalink/application.yml
     ports:
-      - 2333:2333
+      - 2333:2333
+  db:
+    image: postgres:12
+    restart: always
+    environment:
+      POSTGRES_PASSWORD: passworddevelop
+      POSTGRES_DB: sheepbot
+    ports:
+      - 5432:5432

migrations/20201016231711_add_only_link_channel.sql

@@ -0,0 +1,13 @@
+-- 
+-- Editor SQL for DB table only_link_channel
+-- Created by http://editor.datatables.net/generator
+-- 
+
+CREATE TABLE IF NOT EXISTS "only_link_channel" (
+	"id" serial,
+	"guild_id" numeric(9,2),
+	"user_id" numeric(9,2),
+	"url" text,
+	"channel_id" numeric(9,2),
+	PRIMARY KEY( id )
+);

migrations/20201016235905_add_config_and_prefixes.sql

@@ -0,0 +1,11 @@
+-- Add migration script here
+
+CREATE TABLE config (
+    guild_id bigint NOT NULL PRIMARY KEY,
+    host_role_id bigint
+);
+
+CREATE TABLE prefixes (
+    guild_id bigint NOT NULL PRIMARY KEY,
+    prefix text
+);

src/handler.rs

@@ -1,3 +1,5 @@
+use tracing::{error, info};
+
 use std::{
     sync::Arc,
     collections::HashSet,
@@ -28,6 +30,7 @@ use serenity::{
         gateway::Ready,
         id::GuildId,
         event::VoiceServerUpdateEvent,
+        prelude::{Guild},
     },
 };
 
@@ -37,6 +40,10 @@ use lavalink_rs::{
     model::*,
     gateway::*,
 };
+use sqlx::PgPool;
+use crate::settings::Settings;
+
+use crate::utils::database::{initialize_tables};
 
 
 pub(crate) struct VoiceManager;
@@ -55,14 +62,25 @@ impl TypeMapKey for VoiceGuildUpdate {
     type Value = Arc<RwLock<HashSet<GuildId>>>;
 }
 
+pub(crate) struct ConnectionPool;
+
+impl TypeMapKey for ConnectionPool {
+    type Value = PgPool;
+}
+
+pub(crate) struct SettingsConf;
+
+impl TypeMapKey for SettingsConf {
+    type Value = Settings;
+}
 
 pub(crate) struct Handler;
 
 #[async_trait]
 impl EventHandler for Handler {
     async fn ready(&self, _: Context, ready: Ready) {
         if let Some(shard) = ready.shard {
-            println!(
+            info!(
                 "{} is connected on shard {}/{}!",
                 ready.user.name,
                 shard[0],
@@ -80,25 +98,41 @@ impl EventHandler for Handler {
         }
     }
 
+    async fn guild_create(&self, ctx: Context, guild: Guild, is_new: bool) {
+        // We'll initialize the database tables for a guild if it's new.
+        if !is_new {
+            return;
+        }
+
+        initialize_tables(&ctx, &guild).await;
+    }
+
 }
 
 pub(crate) struct LavalinkHandler;
 
 #[async_trait]
 impl LavalinkEventHandler for LavalinkHandler {
     async fn track_start(&self, _client: Arc<Mutex<LavalinkClient>>, event: TrackStart) {
-        println!("Track started!\nGuild: {}", event.guild_id);
+        info!("Track started!\nGuild: {}", event.guild_id);
     }
     async fn track_finish(&self, _client: Arc<Mutex<LavalinkClient>>, event: TrackFinish) {
-        println!("Track finished!\nGuild: {}", event.guild_id);
+        info!("Track finished!\nGuild: {}", event.guild_id);
     }
 }
 
 #[hook]
-async fn after(_ctx: &Context, _msg: &Message, command_name: &str, command_result: CommandResult) {
-    match command_result {
-        Err(why) => println!("Command '{}' returned error {:?} => {}", command_name, why, why),
-        _ => (),
+pub(crate) async fn after(ctx: &Context, msg: &Message, cmd_name: &str, error: CommandResult) {
+    if let Err(why) = &error {
+        error!("Error while running command {}", &cmd_name);
+        error!("{:?}", &error);
+
+        let err = why.to_string();
+        if msg.channel_id.say(ctx, &err).await.is_err() {
+            error!(
+                "Unable to send messages on channel id {}",
+                &msg.channel_id.0
+            );
+        };
     }
 }
-

src/main.rs

@@ -10,6 +10,8 @@ mod utils;
 mod handler;
 mod settings;
 
+use tracing::{error, info};
+
 use std::{
     sync::Arc,
     collections::HashSet,
@@ -24,13 +26,12 @@ use serenity::{
     framework::standard::{
         Args, StandardFramework, CommandGroup,
         HelpOptions, help_commands, CommandResult,
-        macros::{help},
+        macros::{help, hook},
     },
     model::prelude::*,
     prelude::*,
 };
 
-use serenity::prelude::*;
 use lavalink_rs::{
     LavalinkClient
 };
@@ -43,11 +44,16 @@ use crate::handler::{
     Lavalink,
     VoiceManager,
     VoiceGuildUpdate,
-    LavalinkHandler
+    LavalinkHandler,
+    ConnectionPool,
+    after,
+    SettingsConf
 };
 
 use settings::Settings;
 
+use crate::utils::database::{obtain_pool};
+
 
 #[help]
 async fn my_help(
@@ -62,37 +68,85 @@ async fn my_help(
     Ok(())
 }
 
+#[hook]
+// Sets a custom prefix for a guild.
+pub async fn dynamic_prefix(ctx: &Context, msg: &Message) -> Option<String> {
+    let guild_id = &msg.guild_id;
+
+    let data = ctx.data.read().await;
+    let settings = data.get::<SettingsConf>().unwrap();
+    let default_prefix = settings.discord.prefix.as_str();
+
+    let prefix: String;
+    if let Some(id) = guild_id {
+        let pool = data.get::<ConnectionPool>().unwrap();
+
+        let res = sqlx::query!(
+            "SELECT prefix FROM prefixes WHERE guild_id = $1",
+            id.0 as i64
+        )
+        .fetch_one(pool)
+        .await;
+
+        prefix = if let Ok(data) = res {
+            if let Some(p) = data.prefix {
+                p
+            } else {
+                default_prefix.to_string()
+            }
+        } else {
+            error!("I couldn't query the database for getting guild prefix.");
+            default_prefix.to_string()
+        }
+    } else {
+        prefix = default_prefix.to_string();
+    };
+
+    Some(prefix)
+}
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let settings =  match Settings::new() {
         Ok(conf) => conf,
         Err(why) => panic!("Could not read config: {:?}", why),
     };
 
-    let token = settings.discord.token;
-    let lavalink_url = settings.lavalink.url;
-    let lavalink_password = settings.lavalink.password;
+    let token = &settings.discord.token;
+    let lavalink_url = &settings.lavalink.url;
+    let lavalink_password = &settings.lavalink.password;
 
-    let prefix = settings.discord.prefix.as_str();
+    let db_url = &settings.database_url;
 
-    let http = Http::new_with_token(&token);
+    let http = Http::new_with_token(token);
 
-    let bot_id = match http.get_current_application_info().await {
-        Ok(info) => info.id,
+    // We will fetch your bot's owners and id
+    let (owners, bot_id) = match http.get_current_application_info().await {
+        Ok(info) => {
+            let mut owners = HashSet::new();
+            owners.insert(info.owner.id);
+
+            (owners, info.id)
+        }
         Err(why) => panic!("Could not access application info: {:?}", why),
     };
 
 
     let framework = StandardFramework::new()
-        .configure(|c| c
-                   .prefix(prefix))
+        .configure(|c| {
+            c.owners(owners)
+                .dynamic_prefix(dynamic_prefix)
+                .with_whitespace(false)
+                .on_mention(Some(bot_id))
+        })
         .help(&MY_HELP)
+        .after(after)
         .group(&MUSIC_GROUP)
         .group(&FUN_GROUP);
 
-    let mut client = Client::new(&token)
-        .event_handler(Handler)
+    let mut client = Client::new(token)
         .framework(framework)
+        .event_handler(Handler)
         .await
         .expect("Err creating client");
 
@@ -109,8 +163,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
         let lava = lava_client.initialize(LavalinkHandler).await?;
         data.insert::<Lavalink>(lava);
+
+        let pool = obtain_pool(db_url).await?;
+        data.insert::<ConnectionPool>(pool);
+        data.insert::<SettingsConf>(settings);
     }
-    
 
     // Here we clone a lock to the Shard Manager, and then move it into a new
     // thread. The thread will unlock the manager and print shards' status on a
@@ -134,6 +191,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             }
         }
     });
+    
 
     // Start two shards. Note that there is an ~5 second ratelimit period
     // between when one shard can start after another.

src/settings.rs

@@ -16,6 +16,7 @@ pub struct Lavalink {
 
 #[derive(Debug, Deserialize)]
 pub struct Settings {
+    pub database_url: String,
     pub discord: Discord,
     pub lavalink: Lavalink,
 }