From 17e98f9160970d13d7ec5000d4f5cafdb1a5d999 Mon Sep 17 00:00:00 2001 From: natto1784 Date: Sat, 12 Feb 2022 23:24:32 +0530 Subject: [PATCH] added initial sql script and some minor change --- Cargo.lock | 4 +- Cargo.toml | 10 +-- init.sql | 5 ++ src/commands/count.rs | 155 ++++++++++++++++++------------------------ src/handler/count.rs | 74 ++++++++------------ src/handler/mod.rs | 17 +++-- src/main.rs | 32 +++++++-- 7 files changed, 146 insertions(+), 151 deletions(-) create mode 100644 init.sql diff --git a/Cargo.lock b/Cargo.lock index 84564a5..29a2260 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -662,9 +662,9 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.8.0" +version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "692fcb63b64b1758029e0a96ee63e049ce8c5948587f2f7208df04625e5f6b56" +checksum = "da32515d9f6e6e489d7bc9d84c71b060db7247dc035bbe44eac88cf87486d8d5" [[package]] name = "opaque-debug" diff --git a/Cargo.toml b/Cargo.toml index 1fa64ab..3717f54 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,13 +5,13 @@ authors = [ "Amneesh Singh " ] edition = "2018" [dependencies] -tracing = "0.1.22" -regex = "1.5.4" -tokio-postgres = "0.7.2" -rand = "0.8.4" +tracing = "*" +regex = "*" +tokio-postgres = "*" +rand = "*" [dependencies.serenity] -version = "0.10.8" +version = "0.10.*" features = ["cache", "framework", "standard_framework", "rustls_backend", "unstable_discord_api", "collector"] [dependencies.tokio] diff --git a/init.sql b/init.sql new file mode 100644 index 0000000..9111751 --- /dev/null +++ b/init.sql @@ -0,0 +1,5 @@ +CREATE TABLE IF NOT EXISTS words ( + id serial primary key, + name varchar not null, + reg varchar not null, + owner varchar ); diff --git a/src/commands/count.rs b/src/commands/count.rs index d07c663..c4936e8 100644 --- a/src/commands/count.rs +++ b/src/commands/count.rs @@ -4,8 +4,6 @@ use serenity::{ prelude::*, utils::Colour, }; -use std::env; -use tokio_postgres::NoTls; #[command] pub async fn kitna(ctx: &Context, msg: &Message, args: Args) -> CommandResult { @@ -14,17 +12,14 @@ pub async fn kitna(ctx: &Context, msg: &Message, args: Args) -> CommandResult { msg.reply(ctx, "bruh kitna kya?").await?; return Ok(()); } - let db: String = env::var("DB_URL").expect("bhay DB_URL daal na"); - let (client, conn) = tokio_postgres::connect(&db, NoTls) - .await - .expect("cant connect bha"); - tokio::spawn(async move { - if let Err(e) = conn.await { - eprintln!("connection error: {}", e); - } - }); + let data_read = ctx.data.read().await; + let db = data_read + .get::() + .expect("Expected Database in TypeMap.") + .clone(); + let id = msg.author.id.to_string(); - let mut query_helper = client + let mut query_helper = db .query( format!("select name from words where '{}' ~ reg", query).as_str(), &[], @@ -32,7 +27,7 @@ pub async fn kitna(ctx: &Context, msg: &Message, args: Args) -> CommandResult { .await .expect("helper query to select count failed"); if query_helper.is_empty() { - query_helper = client + query_helper = db .query( format!("select name from words where name='{}'", query).as_str(), &[], @@ -43,7 +38,7 @@ pub async fn kitna(ctx: &Context, msg: &Message, args: Args) -> CommandResult { msg.reply( ctx, format!( - "No entry for '{}' found. If you want to add it, run 'xxadd {}&'", + "No entry for '{}' found. If you want to add it, run ',count add {}&'", query, query ), ) @@ -58,7 +53,7 @@ pub async fn kitna(ctx: &Context, msg: &Message, args: Args) -> CommandResult { }; for row in query_helper { let name: &str = row.get(0); - let query_result: i32 = client + let query_result: i32 = db .query_one( format!("select count from user{} where name='{}'", id, name).as_str(), &[], @@ -75,9 +70,9 @@ pub async fn kitna(ctx: &Context, msg: &Message, args: Args) -> CommandResult { #[command] pub async fn add(ctx: &Context, msg: &Message, args: Args) -> CommandResult { let query: String = args.raw().collect::>().join(" "); - let queries = query.split("&").collect::>(); + let queries = query.splitn(2, "&").collect::>(); if queries.len() != 2 { - msg.reply(ctx, "Please use the proper syntax: xxadd &\nIf you don't know what regex is, just do: xxadd &") + msg.reply(ctx, "Please use the proper syntax: `,count add &`\nIf you don't know what regex is, just do: `,count add &`") .await?; return Ok(()); } @@ -85,16 +80,12 @@ pub async fn add(ctx: &Context, msg: &Message, args: Args) -> CommandResult { msg.reply(ctx, "Not a valid regex").await?; return Ok(()); } - let db: String = env::var("DB_URL").expect("bhay DB_URL daal na"); - let (client, conn) = tokio_postgres::connect(&db, NoTls) - .await - .expect("cant connect bha"); - tokio::spawn(async move { - if let Err(e) = conn.await { - eprintln!("connection error: {}", e); - } - }); - let check_existense = client + let data_read = ctx.data.read().await; + let db = data_read + .get::() + .expect("Expected Database in TypeMap.") + .clone(); + let check_existense = db .query( format!("select name, reg from words where name='{}'", queries[0]).as_str(), &[], @@ -109,18 +100,17 @@ pub async fn add(ctx: &Context, msg: &Message, args: Args) -> CommandResult { .await?; return Ok(()); } - client - .execute( - format!( - "insert into words(name, reg, owner) values('{}','(?i){}', '{}')", - queries[0], - queries[1], - msg.author.id.to_string() - ) - .as_str(), - &[], + db.execute( + format!( + "insert into words(name, reg, owner) values('{}','(?i){}', '{}')", + queries[0], + queries[1], + msg.author.id.to_string() ) - .await?; + .as_str(), + &[], + ) + .await?; msg.reply(ctx, "Added").await?; Ok(()) } @@ -128,16 +118,12 @@ pub async fn add(ctx: &Context, msg: &Message, args: Args) -> CommandResult { #[command] pub async fn rm(ctx: &Context, msg: &Message, args: Args) -> CommandResult { let query: String = args.raw().collect::>().join(" "); - let db: String = env::var("DB_URL").expect("bhay DB_URL daal na"); - let (client, conn) = tokio_postgres::connect(&db, NoTls) - .await - .expect("cant connect bha"); - tokio::spawn(async move { - if let Err(e) = conn.await { - eprintln!("connection error: {}", e); - } - }); - let owner = client + let data_read = ctx.data.read().await; + let db = data_read + .get::() + .expect("Expected Database in TypeMap.") + .clone(); + let owner = db .query( format!("select owner from words where name = '{}'", query).as_str(), &[], @@ -150,12 +136,11 @@ pub async fn rm(ctx: &Context, msg: &Message, args: Args) -> CommandResult { return Ok(()); } } - client - .execute( - format!("delete from words where name='{}'", query,).as_str(), - &[], - ) - .await?; + db.execute( + format!("delete from words where name='{}'", query,).as_str(), + &[], + ) + .await?; msg.reply(ctx, "Deleted if it existed").await?; Ok(()) } @@ -163,26 +148,25 @@ pub async fn rm(ctx: &Context, msg: &Message, args: Args) -> CommandResult { #[command] pub async fn change(ctx: &Context, msg: &Message, args: Args) -> CommandResult { let query: String = args.raw().collect::>().join(" "); - let queries = query.split("&").collect::>(); + let queries = query.splitn(2, "&").collect::>(); if queries.len() != 2 { - msg.reply(ctx, "Please use the proper syntax\nxxchange &") - .await?; + msg.reply( + ctx, + "Please use the proper syntax\n,count change &", + ) + .await?; return Ok(()); } if queries[1].contains(" ") { msg.reply(ctx, "Not a valid regex").await?; return Ok(()); } - let db: String = env::var("DB_URL").expect("bhay DB_URL daal na"); - let (client, conn) = tokio_postgres::connect(&db, NoTls) - .await - .expect("cant connect bha"); - tokio::spawn(async move { - if let Err(e) = conn.await { - eprintln!("connection error: {}", e); - } - }); - let owner = client + let data_read = ctx.data.read().await; + let db = data_read + .get::() + .expect("Expected Database in TypeMap.") + .clone(); + let owner = db .query( format!("select owner from words where name = '{}'", queries[0]).as_str(), &[], @@ -195,32 +179,27 @@ pub async fn change(ctx: &Context, msg: &Message, args: Args) -> CommandResult { return Ok(()); } } - client - .execute( - format!( - "update words set reg='(?i){}' where name='{}'", - queries[1], queries[0] - ) - .as_str(), - &[], + db.execute( + format!( + "update words set reg='(?i){}' where name='{}'", + queries[1], queries[0] ) - .await?; + .as_str(), + &[], + ) + .await?; msg.reply(ctx, "Changed the value if it existed").await?; Ok(()) } #[command] -pub async fn list(ctx: &Context, msg: &Message, _: Args) -> CommandResult { - let db: String = env::var("DB_URL").expect("bhay DB_URL daal na"); - let (client, conn) = tokio_postgres::connect(&db, NoTls) - .await - .expect("cant connect bha"); - tokio::spawn(async move { - if let Err(e) = conn.await { - eprintln!("connection error: {}", e); - } - }); - let rows = client.query("select * from words", &[]).await?; +pub async fn ls(ctx: &Context, msg: &Message, _: Args) -> CommandResult { + let data_read = ctx.data.read().await; + let db = data_read + .get::() + .expect("Expected Database in TypeMap.") + .clone(); + let rows = db.query("select * from words", &[]).await?; msg.channel_id .send_message(ctx, |mut m| { let mut a: u32 = 1; @@ -231,7 +210,7 @@ pub async fn list(ctx: &Context, msg: &Message, _: Args) -> CommandResult { .color(Colour::TEAL); a += 1; for row in group { - let idx: u32 = row.get(0); + let idx: i32 = row.get(0); let name: String = row.get(1); let _reg: String = row.get(2); let owner_id: String = row.get(3); diff --git a/src/handler/count.rs b/src/handler/count.rs index 2082e83..879b658 100644 --- a/src/handler/count.rs +++ b/src/handler/count.rs @@ -1,38 +1,26 @@ use regex::Regex; use serenity::model::channel::Message; -use std::env; -use tokio_postgres::NoTls; +use tokio_postgres::Client; -pub async fn count(msg: Message) { - let db: String = env::var("DB_URL").expect("bhay DB_URL daal na"); - let (client, conn) = tokio_postgres::connect(&db, NoTls) - .await - .expect("cant connect bha"); - tokio::spawn(async move { - if let Err(e) = conn.await { - eprintln!("connection error: {}", e); - } - }); +pub async fn count(msg: Message, db: std::sync::Arc) { let id = msg.author.id.as_u64().to_owned().to_string(); - client - .execute( - format!( - " - CREATE TABLE IF NOT EXISTS user{} ( + db.execute( + format!( + r#" + CREATE TABLE IF NOT EXISTS user{} ( id SERIAL PRIMARY KEY, name VARCHAR NOT NULL, count INTEGER NOT NULL - ) - ", - id - ) - .as_str(), - &[], + )"#, + id ) - .await - .expect("cant create table"); + .as_str(), + &[], + ) + .await + .expect("cant create a user table"); - for row in client + for row in db .query("SELECT name, reg FROM words", &[]) .await .expect("can't get the words to count") @@ -41,7 +29,7 @@ pub async fn count(msg: Message) { let regex: Regex = Regex::new(row.get(1)).unwrap(); let count = regex.captures_iter(&msg.content).count(); if count > 0 { - let query_result = client + let query_result = db .query( format!("SELECT count FROM user{} where name='{}'", id, name).as_str(), &[], @@ -49,29 +37,27 @@ pub async fn count(msg: Message) { .await .expect("cant select the count"); if query_result.is_empty() { - client - .execute( - format!( - "insert into user{} (name, count) values ('{}', 0)", - id, name - ) - .as_str(), - &[], - ) - .await - .expect("cant insert shit"); - } - client - .execute( + db.execute( format!( - "UPDATE user{} SET count = count + {} where name='{}'", - id, count, name + "insert into user{} (name, count) values ('{}', 0)", + id, name ) .as_str(), &[], ) .await - .expect("cant update"); + .expect("cant insert shit"); + } + db.execute( + format!( + "UPDATE user{} SET count = count + {} where name='{}'", + id, count, name + ) + .as_str(), + &[], + ) + .await + .expect("cant update"); } } } diff --git a/src/handler/mod.rs b/src/handler/mod.rs index c32403c..844efd9 100644 --- a/src/handler/mod.rs +++ b/src/handler/mod.rs @@ -1,11 +1,11 @@ -mod interactions; mod count; +mod interactions; use serenity::{ async_trait, model::{ + channel::Message, event::ResumedEvent, gateway::Ready, - channel::Message, interactions::{ ApplicationCommand, Interaction, InteractionData, InteractionResponseType, InteractionType, @@ -29,8 +29,13 @@ impl EventHandler for Handler { async fn resume(&self, _: Context, _: ResumedEvent) { info!("how th when the"); } - async fn message(&self, _: Context, msg: Message) { - count::count(msg).await; + async fn message(&self, ctx: Context, msg: Message) { + let data_read = ctx.data.read().await; + let db_client = data_read + .get::() + .expect("Expected Database in TypeMap.") + .clone(); + count::count(msg, db_client).await; } async fn interaction_create(&self, ctx: Context, interaction: Interaction) { @@ -40,7 +45,9 @@ impl EventHandler for Handler { .create_interaction_response(&ctx.http, |response| { response .kind(InteractionResponseType::ChannelMessageWithSource) - .interaction_response_data(|message| interactions::responses(data.name.to_string(), message)) + .interaction_response_data(|message| { + interactions::responses(data.name.to_string(), message) + }) }) .await { diff --git a/src/main.rs b/src/main.rs index be460d9..0a2d2f2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,9 @@ mod commands; mod handler; +use commands::count::*; +use commands::general::*; +use commands::minigames::*; +use handler::Handler; use serenity::{ client::bridge::gateway::ShardManager, framework::{standard::macros::group, StandardFramework}, @@ -8,22 +12,24 @@ use serenity::{ }; use std::{collections::HashSet, env, sync::Arc}; use tracing::error; -pub struct ShardManagerContainer; -use commands::general::*; -use commands::count::*; -use commands::minigames::*; -use handler::Handler; +struct ShardManagerContainer; impl TypeMapKey for ShardManagerContainer { type Value = Arc>; } +struct Database; +impl TypeMapKey for Database { + type Value = Arc; +} + #[group] #[commands(ping)] struct General; #[group] -#[commands(kitna,add,rm,change)] +#[prefix = "count"] +#[commands(kitna, add, rm, change, ls)] struct Count; #[group] #[commands(challenge)] @@ -46,7 +52,7 @@ async fn main() { }; let framework = StandardFramework::new() - .configure(|c| c.owners(owners).prefix("xx")) + .configure(|c| c.owners(owners).prefix(",")) .group(&GENERAL_GROUP) .group(&COUNT_GROUP) .group(&MINIGAMES_GROUP); @@ -59,7 +65,19 @@ async fn main() { .expect("Client no wokey"); { + let db_url: String = env::var("DB_URL").expect("DB_URL not found"); + let (db_client, conn) = tokio_postgres::connect(&db_url, tokio_postgres::NoTls) + .await + .expect("cant connect bha"); + tokio::spawn(async move { + if let Err(e) = conn.await { + eprintln!("connection error: {}", e); + } + }); + let init_script = std::include_str!("../init.sql"); + db_client.batch_execute(init_script).await.expect("Couldn't run the init script"); let mut data = client.data.write().await; + data.insert::(Arc::new(db_client)); data.insert::(client.shard_manager.clone()); }