From 9a8587e4b875c0607820d14fffd20a6325ae952e Mon Sep 17 00:00:00 2001 From: Benjamin Sherriff Date: Thu, 6 Jul 2023 21:07:54 -0400 Subject: [PATCH] Working on database for messages --- Cargo.toml | 7 +++++- Dockerfile | 4 +-- README.md | 14 ++++++++--- migrations/create_messages/down.sql | 1 + migrations/create_messages/up.sql | 12 +++++++++ src/commands/mod.rs | 1 + src/commands/schedule.rs | 0 src/database/mod.rs | 38 +++++++++++++++++++++++++++++ src/database/models.rs | 33 +++++++++++++++++++++++++ src/database/schema.rs | 14 +++++++++++ src/main.rs | 24 +++++++++++------- 11 files changed, 132 insertions(+), 16 deletions(-) create mode 100644 migrations/create_messages/down.sql create mode 100644 migrations/create_messages/up.sql create mode 100644 src/commands/schedule.rs create mode 100644 src/database/mod.rs create mode 100644 src/database/models.rs create mode 100644 src/database/schema.rs diff --git a/Cargo.toml b/Cargo.toml index 990393f..58dd4bb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,4 +33,9 @@ features = ["derive"] [dependencies.reqwest] version = "0.11.18" default-features = false -features = ["json", "rustls-tls"] \ No newline at end of file +features = ["json", "rustls-tls"] + +[dependencies.diesel] +version = "2.1.0" +default-features = false +features = ["postgres", "32-column-tables", "serde_json", "r2d2", "with-deprecated"] \ No newline at end of file diff --git a/Dockerfile b/Dockerfile index ed1b0e4..3a5db82 100644 --- a/Dockerfile +++ b/Dockerfile @@ -7,7 +7,7 @@ RUN cargo build --release --bin siren FROM debian:bullseye-slim as packages WORKDIR /packages -RUN apt-get update && apt-get install -y libopus-dev curl tar xz-utils +RUN apt-get update && apt-get install -y libopus-dev libpq5 libpq-dev curl tar xz-utils RUN curl -L https://github.com/yt-dlp/yt-dlp/releases/latest/download/yt-dlp_linux > yt-dlp && \ chmod +x yt-dlp RUN curl -L https://github.com/yt-dlp/FFmpeg-Builds/releases/download/latest/ffmpeg-master-latest-linux64-gpl.tar.xz > ffmpeg.tar.xz && \ @@ -15,7 +15,7 @@ RUN curl -L https://github.com/yt-dlp/FFmpeg-Builds/releases/download/latest/ffm FROM debian:bullseye-slim as runtime WORKDIR /siren -# RUN apt-get update && apt-get install -y libopus-dev ffmpeg youtube-dl COPY --from=builder /siren/target/release/siren /usr/local/bin/siren COPY --from=packages /packages /usr/bin +ADD migrations ./migrations/ CMD ["siren"] diff --git a/README.md b/README.md index c51bc87..1c4f331 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,7 @@ Example Invite: ``` https://discord.com/api/oauth2/authorize?client_id=&permissions=40671259392832&scope=bot%20applications.commands ``` - - The CLIENT_ID can be found in the General Information tab on the Discord Developer Portal for your application, under `Application ID` +The CLIENT_ID can be found in the General Information tab on the Discord Developer Portal for your application, under `Application ID` 1. Copy `.env.TEMPLATE` to `.env` and fill out the fields 2. Build the [Docker](https://www.docker.com/) application with `make build` @@ -27,16 +27,22 @@ https://discord.com/api/oauth2/authorize?client_id=&permissions=40671 ## Contributing [Rust](https://www.rust-lang.org/) must be installed to run locally. See [serenity-rs/serenity](https://github.com/serenity-rs/serenity) for more information about Rust Discord API Library. -Furthermore, the following packages must be installed for [serenity-rs/songbird](https://github.com/serenity-rs/songbird). View the repository for additional installation and setup information on other operating systems. +The following packages must be installed for [serenity-rs/songbird](https://github.com/serenity-rs/songbird). View the repository for additional installation and setup information on other operating systems. ``` sudo apt install libopus-dev sudo apt install ffmpeg sudo apt apt install youtube-dl ``` + - Potentially requires [yt-dlp](https://github.com/yt-dlp/yt-dlp#installation) and [yt-dlp FFmpeg Static Auto-Builds](https://github.com/yt-dlp/FFmpeg-Builds). -Potentially requires [yt-dlp](https://github.com/yt-dlp/yt-dlp#installation) and [yt-dlp FFmpeg Static Auto-Builds](https://github.com/yt-dlp/FFmpeg-Builds). +The following packages must be installed for [PostgreSQL](). +``` +sudo apt install libpq5 +sudo apt install libpq-dev +``` -Begin the application with `cargo run` +Begin the application with `cargo run` (note the database must still be running) + - `docker compose up -d db` The application can also be tested from within a Docker container: ``` diff --git a/migrations/create_messages/down.sql b/migrations/create_messages/down.sql new file mode 100644 index 0000000..90b6925 --- /dev/null +++ b/migrations/create_messages/down.sql @@ -0,0 +1 @@ +DROP TABLE messages \ No newline at end of file diff --git a/migrations/create_messages/up.sql b/migrations/create_messages/up.sql new file mode 100644 index 0000000..f5965cf --- /dev/null +++ b/migrations/create_messages/up.sql @@ -0,0 +1,12 @@ +CREATE TABLE IF NOT EXISTS messages ( + id TEXT PRIMARY KEY NOT NULL, + guild_id BIGINT NOT NULL, + channel_id BIGINT NOT NULL, + user_id BIGINT NOT NULL, + created BIGINT NOT NULL, + model TEXT NOT NULL, + request TEXT NOT NULL, + response TEXT NOT NULL, + request_tags TEXT[] NOT NULL, + response_tags TEXT[] NOT NULL +) \ No newline at end of file diff --git a/src/commands/mod.rs b/src/commands/mod.rs index fd3b9e5..6f92333 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -2,3 +2,4 @@ pub mod audio; pub mod help; pub mod oai; pub mod ping; +pub mod schedule; diff --git a/src/commands/schedule.rs b/src/commands/schedule.rs new file mode 100644 index 0000000..e69de29 diff --git a/src/database/mod.rs b/src/database/mod.rs new file mode 100644 index 0000000..1223358 --- /dev/null +++ b/src/database/mod.rs @@ -0,0 +1,38 @@ +use std::env; +use std::path::Path; + +use diesel::r2d2::{Pool, ConnectionManager}; +use diesel::pg::PgConnection; + +pub mod models; +pub mod schema; + +pub fn run_migrations(pool: &Pool>) { + let mut connection = pool.get().unwrap(); + let migrations_dir = Path::new("./migrations"); + let migrations = std::fs::read_dir(&migrations_dir).unwrap(); + + for migration in migrations { + if migration.as_ref().unwrap().file_type().unwrap().is_dir() { + let migration_paths = std::fs::read_dir(&migration.unwrap().path()).unwrap(); + + for migration_path in migration_paths { + if migration_path.as_ref().unwrap().file_name().eq_ignore_ascii_case("up.sql") { + let path = &migration_path.unwrap().path(); + let contents = std::fs::read_to_string(path).expect("Unable to read from file"); + // connection.build_transaction() + } + } + } + } +} + +pub fn establish_connection() -> Pool> { + let database_user = env::var("POSTGRES_USER").expect("Expected a user in the environment"); + let database_password = env::var("POSTGRES_PASSWORD").expect("Expected a password in the environment"); + let database_name = env::var("POSTGRES_DB").expect("Expected a database name in the environment"); + + let database_url = format!("postgres://{}:{}@localhost/{}", database_user, database_password, database_name); + let manager = ConnectionManager::::new(database_url); + Pool::builder().build(manager).expect("Failed to create pool.") +} \ No newline at end of file diff --git a/src/database/models.rs b/src/database/models.rs new file mode 100644 index 0000000..56947bb --- /dev/null +++ b/src/database/models.rs @@ -0,0 +1,33 @@ +use diesel::prelude::*; + +use super::schema::messages; + +#[derive(Queryable, Selectable)] +#[diesel(table_name = messages)] +pub struct MessageDB { + pub id: i64, + pub guild_id: i64, + pub channel_id: i64, + pub user_id: i64, + pub created: i64, + pub model: String, + pub request: String, + pub response: String, + pub request_tags: Vec, + pub response_tags: Vec, +} + +#[derive(Insertable)] +#[diesel(table_name = messages)] +pub struct NewMessageDB<'a> { + pub id: i64, + pub guild_id: i64, + pub channel_id: i64, + pub user_id: i64, + pub created: i64, + pub model: &'a str, + pub request: &'a str, + pub response: &'a str, + pub request_tags: Vec<&'a str>, + pub response_tags: Vec<&'a str>, +} \ No newline at end of file diff --git a/src/database/schema.rs b/src/database/schema.rs new file mode 100644 index 0000000..47676b5 --- /dev/null +++ b/src/database/schema.rs @@ -0,0 +1,14 @@ +diesel::table! { + messages (id) { + id -> BigInt, + guild_id -> BigInt, + channel_id -> BigInt, + user_id -> BigInt, + created -> BigInt, + model -> Text, + request -> Text, + response -> Text, + request_tags -> Array, + response_tags -> Array, + } +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 8fdc3a2..e3a23d6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,8 @@ use std::collections::HashSet; use std::env; use commands::audio::create_response; +use diesel::r2d2::{Pool, ConnectionManager}; +use diesel::pg::PgConnection; use dotenv::dotenv; use log::{error, warn, info}; use serenity::async_trait; @@ -14,9 +16,12 @@ use serenity::prelude::*; use songbird::SerenityInit; mod commands; +mod database; + struct Handler { // Open AI Config - oai: Option + oai: Option, + pool: Pool> } #[async_trait] @@ -110,28 +115,29 @@ async fn main() { Err(why) => panic!("Could not access application info: {:?}", why) }; - let framework = StandardFramework::new() - .configure(|c| c - .owners(owners) - .prefix("!") - ); + let pool = database::establish_connection(); + database::run_migrations(&pool); let handler = match env::var("OPENAI_API_KEY") { Ok(token) => { info!("Loaded OpenAI token"); Handler { - oai: Some(commands::oai::OAI { client: reqwest::Client::new(), base_url: "https://api.openai.com/v1".to_string(), max_attempts: 5, token }) + oai: Some(commands::oai::OAI { client: reqwest::Client::new(), base_url: "https://api.openai.com/v1".to_string(), max_attempts: 5, token }), + pool } } Err(err) => { warn!("Could not load OpenAI token: {}", err); - Handler { oai: None } + Handler { oai: None, pool } } }; let mut client = Client::builder(token, intents) .event_handler(handler) - .framework(framework) + .framework(StandardFramework::new() + .configure(|c| c + .owners(owners) + )) .register_songbird() .await .expect("Error creating client");