Working on database for messages

src/commands/mod.rs

@@ -2,3 +2,4 @@ pub mod audio;
 pub mod help;
 pub mod oai;
 pub mod ping;
+pub mod schedule;

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<ConnectionManager<PgConnection>>) {
+  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<ConnectionManager<PgConnection>> {
+  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::<PgConnection>::new(database_url);
+  Pool::builder().build(manager).expect("Failed to create pool.")
+}

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<String>,
+  pub response_tags: Vec<String>,
+}
+
+#[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>,
+}

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<Text>,
+    response_tags -> Array<Text>,
+  }
+}

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<commands::oai::OAI>
+  oai: Option<commands::oai::OAI>,
+  pool: Pool<ConnectionManager<PgConnection>>
 }
 
 #[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");