update commands

This commit is contained in:
DataHearth 2023-05-20 20:25:32 +02:00
parent 27ac44182d
commit e89b0fe30e
No known key found for this signature in database
GPG Key ID: E88FD356ACC5F3C4
4 changed files with 134 additions and 79 deletions

4
.gitignore vendored
View File

@ -1,2 +1,4 @@
/target
.env
.env
*.db
*.wal

View File

@ -1,9 +1,6 @@
use std::sync::Mutex;
use once_cell::sync::OnceCell;
use poise::command;
use crate::database::Database;
use crate::database::*;
const HELP_MESSAGE: &str = "
Hello fellow human! I am a bot that can help you adding new technologies to a git repository.
@ -34,15 +31,12 @@ To get help, just type:
/help
```
";
const AUTHORIZED_USERS: [u64; 1] = [252497456447750144];
type Error = Box<dyn std::error::Error + Send + Sync>;
type Context<'a> = poise::Context<'a, MsgData, Error>;
pub struct MsgData {}
pub static DB: OnceCell<Mutex<Database>> = OnceCell::new();
/// Show help for all commands
#[command(slash_command, prefix_command)]
pub async fn help(ctx: Context<'_>) -> Result<(), Error> {
@ -58,7 +52,9 @@ pub async fn add(
#[description = "Technology name"] technology: String,
#[description = "Git repository link"] link: String,
) -> Result<(), Error> {
ctx.say(format!("Added {technology} with link {link}"))
add_tech(link.clone(), technology.clone())?;
ctx.say(format!("Added {technology} with link {link}",))
.await?;
Ok(())
@ -67,7 +63,21 @@ pub async fn add(
/// List all available technologies
#[command(slash_command, prefix_command)]
pub async fn list(ctx: Context<'_>) -> Result<(), Error> {
ctx.say("Listed all technologies").await?;
let techs = list_tech()?;
if techs.len() == 0 {
ctx.say("No technologies saved").await?;
return Ok(());
}
ctx.say(format!(
"Saved technologies: {}",
techs
.iter()
.map(|tech| format!("[{}]({})", tech.name, tech.link))
.collect::<Vec<String>>()
.join(", ")
))
.await?;
Ok(())
}
@ -78,7 +88,12 @@ pub async fn search(
ctx: Context<'_>,
#[description = "Technology name"] technology: String,
) -> Result<(), Error> {
ctx.say(format!("Found {technology}")).await?;
if let Some(tech) = search_tech(technology.clone())? {
ctx.say(format!("Name: {}\nLink: {}", tech.name, tech.link))
.await?;
} else {
ctx.say("No technology found").await?;
}
Ok(())
}
@ -89,13 +104,37 @@ pub async fn remove(
ctx: Context<'_>,
#[description = "Technology name"] technology: String,
) -> Result<(), Error> {
if !AUTHORIZED_USERS.contains(&ctx.author().id.0) {
if is_auth_user(ctx.author().id.to_string())? {
ctx.say("You don't have permission to remove a technology")
.await?;
return Ok(());
}
ctx.say(format!("Removed {technology}")).await?;
remove_tech(technology)?;
ctx.say("Technology removed").await?;
Ok(())
}
/// Remove a technology from the technologies list
#[command(slash_command, prefix_command)]
pub async fn add_auth_user(
ctx: Context<'_>,
#[description = "Discord user ID"] id: String,
) -> Result<(), Error> {
if !std::env::var("ADMIN_USERS")
.expect("missing ADMIN_USERS")
.contains(&ctx.author().id.to_string())
{
ctx.say("You don't have permission to add a new user")
.await?;
return Ok(());
}
set_auth_user(id)?;
ctx.say("User added!").await?;
Ok(())
}

View File

@ -1,78 +1,88 @@
use polodb_core::{
bson::{doc, oid::ObjectId},
Database as DB,
};
use once_cell::sync::OnceCell;
use polodb_core::{bson::doc, Database};
use serde::{Deserialize, Serialize};
use std::{
io::{Error, ErrorKind},
str::FromStr,
};
use std::io::{Error, ErrorKind};
pub static DB: OnceCell<Database> = OnceCell::new();
#[derive(Default, Debug, Serialize, Deserialize)]
struct Technology {
link: String,
name: String,
pub struct Technology {
pub link: String,
pub name: String,
}
#[derive(Debug, Serialize, Deserialize)]
struct AuthorizedUser {
name: String,
discord_id: String,
}
pub struct Database {
pub db: DB,
/// Add a new technology to the database.
pub fn add_tech(link: String, name: String) -> Result<(), Error> {
DB.get()
.unwrap()
.collection::<Technology>("technologies")
.insert_one(Technology { link, name })
.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?;
Ok(())
}
impl Database {
/// Create a new database instance.
pub fn new() -> Self {
Self {
db: DB::open_file("test-polo.db").expect("failed to initialize database"),
}
}
pub fn remove_tech(name: String) -> Result<(), Error> {
DB.get()
.unwrap()
.collection::<Technology>("technologies")
.delete_one(doc! { "name": name })
.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?;
/// Add a new technology to the database.
pub fn add_tech(&self, link: String, name: String) -> Result<String, Error> {
Ok(self
.db
.collection::<Technology>("technologies")
.insert_one(Technology { link, name })
.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?
.inserted_id
.to_string())
}
Ok(())
}
pub fn remove_tech(&self, id: &str) -> Result<(), Error> {
self.db
.collection::<Technology>("technologies")
.delete_one(doc! { "_id": Some(ObjectId::from_str(id).map_err(|err| Error::new(ErrorKind::InvalidInput, err))?) })
.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?;
pub fn list_tech() -> Result<Vec<Technology>, Error> {
Ok(DB
.get()
.unwrap()
.collection("technologies")
.find(None)
.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?
.filter(|doc| doc.is_ok())
.map(|doc| doc.unwrap())
.collect())
}
Ok(())
}
pub fn list_tech(&self) -> Result<Vec<Technology>, Error> {
Ok(self
.db
.collection("technologies")
.find(None)
.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?
.filter(|doc| doc.is_ok())
.map(|doc| doc.unwrap())
.collect())
}
pub fn search_tech(&self, name: String) -> Result<Technology, Error> {
if let Some(tech) = self
.db
.collection::<Technology>("technologies")
.find(doc! { "$eq": [{"name": name}] })
.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?
.next()
{
Ok(tech.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?)
} else {
Err(Error::new(ErrorKind::NotFound, "Technology not found"))
}
pub fn search_tech(name: String) -> Result<Option<Technology>, Error> {
if let Some(tech) = DB
.get()
.unwrap()
.collection::<Technology>("technologies")
.find(doc! { "name": {"$eq": name} })
.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?
.next()
{
Ok(Some(
tech.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?,
))
} else {
Ok(None)
}
}
pub fn set_auth_user(discord_id: String) -> Result<(), Error> {
DB.get()
.unwrap()
.collection("authorized_users")
.insert_one(AuthorizedUser { discord_id })
.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?;
Ok(())
}
pub fn is_auth_user(discord_id: String) -> Result<bool, Error> {
Ok(DB
.get()
.unwrap()
.collection::<AuthorizedUser>("authorized_users")
.find(doc! { "$eq": [{"discord_id": discord_id}] })
.map_err(|err| Error::new(ErrorKind::InvalidInput, err))?
.next()
.is_some())
}

View File

@ -3,15 +3,15 @@ mod database;
use poise::serenity_prelude as serenity;
use commands::{add, help, list, remove, search, MsgData, DB};
use commands::{add, add_auth_user, help, list, remove, search, MsgData};
use database::DB;
use polodb_core::Database;
#[tokio::main]
async fn main() {
DB.set(Mutex)
let framework = poise::Framework::builder()
.options(poise::FrameworkOptions {
commands: vec![help(), add(), list(), search(), remove()],
commands: vec![help(), add(), list(), search(), remove(), add_auth_user()],
..Default::default()
})
.token(std::env::var("DISCORD_TOKEN").expect("missing DISCORD_TOKEN"))
@ -25,5 +25,9 @@ async fn main() {
})
});
DB.get_or_init(|| {
Database::open_file(std::env::var("DB_PATH").expect("missing DB_PATH"))
.expect("failed to initialize database")
});
framework.run().await.unwrap();
}