Updated dependencies, fixed migrations usage

src/commands/oai.rs

@@ -2,8 +2,7 @@
 use std::error::Error;
 use std::fmt;
 
-use diesel::{prelude::*, PgConnection, insert_into};
-use diesel::r2d2::{Pool, ConnectionManager};
+use diesel::{prelude::*, insert_into};
 use log::{error, debug, trace, warn};
 
 use serde::{Serialize, Deserialize};
@@ -13,7 +12,8 @@ use serenity::model::channel::Message;
 use serenity::model::prelude::{ChannelType, PermissionOverwrite, PermissionOverwriteType};
 use serenity::prelude::*;
 
-use crate::database::models::{NewMessageDB, MessageDB};
+use crate::db;
+use crate::messages::{NewMessageDB, MessageDB};
 
 pub struct OAI {
   pub client: reqwest::Client,
@@ -187,26 +187,26 @@ impl OAI {
   }
 }
 
-pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI, pool: &Pool<ConnectionManager<PgConnection>>) {
+pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) {
   debug!("Generating response for message: {}", msg.content);
+  let mut connection = db::connection().unwrap();
 
   let guild_id = msg.guild_id.unwrap();
   let channel_id = msg.channel_id;
   let author_id = msg.author.id;
-  let mut connection = pool.get().unwrap();
 
   // Parse out the bot mention from the message
   let bot_mention: String = format!("<@{}>", ctx.cache.current_user_id().0);
   let parsed_content = msg.content.replace(bot_mention.as_str(), "");
 
   // Setup the request messages
-  let result: Result<Vec<MessageDB>, diesel::result::Error> = crate::database::schema::messages::table
+  let result: Result<Vec<MessageDB>, diesel::result::Error> = crate::schema::messages::table
     .select(MessageDB::as_select())
-    .filter((crate::database::schema::messages::guild_id.eq(guild_id.0 as i64))
-      .and(crate::database::schema::messages::channel_id.eq(channel_id.0  as i64))
-      .and(crate::database::schema::messages::user_id.eq(author_id.0  as i64))
+    .filter((crate::schema::messages::guild_id.eq(guild_id.0 as i64))
+      .and(crate::schema::messages::channel_id.eq(channel_id.0  as i64))
+      .and(crate::schema::messages::user_id.eq(author_id.0  as i64))
     )
-    .order(crate::database::schema::messages::created.asc())
+    .order(crate::schema::messages::created.asc())
     .limit(oai.max_context_questions)
     .load(&mut connection);
   
@@ -284,7 +284,7 @@ pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI, pool: &P
       if !r.choices.is_empty() {
         let res = r.choices[0].message.content.clone();
         // Insert the message into the messages database table
-        if let Err(err) = insert_into(crate::database::schema::messages::table).values(NewMessageDB {
+        if let Err(err) = insert_into(crate::schema::messages::table).values(NewMessageDB {
           id: &r.id,
           guild_id: guild_id.0 as i64,
           channel_id: response_channel.0 as i64,

src/database/mod.rs

@@ -1,61 +0,0 @@
-use std::env;
-// use std::path::Path;
-
-use diesel::RunQueryDsl;
-use diesel::r2d2::{Pool, ConnectionManager};
-use diesel::pg::PgConnection;
-use log::{error, info};
-
-pub mod models;
-pub mod schema;
-
-pub fn run_migrations(pool: &Pool<ConnectionManager<PgConnection>>) {
-  let mut connection = pool.get().unwrap();
-  if let Err(err) = diesel::sql_query("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
-  )").execute(&mut connection) {
-    error!("Could not create messages table: {}", err);
-  } else {
-    info!("Successfully created messages table");
-  }
-  // 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");
-  //         if let Err(err) = diesel::sql_query(&contents).execute(&mut connection) {
-  //           error!("Could not run migration: {}", err);
-  //         } else {
-  //           info!("Successfully ran migration: {}", path.display());
-  //         }
-  //       }
-  //     }
-  //   }
-  // }
-}
-
-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_host = env::var("POSTGRES_HOST").unwrap_or("localhost".to_string());
-
-  let database_url = format!("postgres://{}:{}@{}/{}", database_user, database_password, database_host, database_name);
-  let manager = ConnectionManager::<PgConnection>::new(database_url);
-  Pool::builder().build(manager).expect("Failed to create pool.")
-}

src/db.rs

@@ -0,0 +1,39 @@
+use crate::error_handler::ServiceError;
+use diesel::{r2d2::ConnectionManager, PgConnection};
+use crate::diesel_migrations::MigrationHarness;
+use lazy_static::lazy_static;
+use log::{error, info};
+use r2d2;
+use std::env;
+
+type Pool = r2d2::Pool<ConnectionManager<PgConnection>>;
+pub type DbConnection = r2d2::PooledConnection<ConnectionManager<PgConnection>>;
+
+pub const MIGRATIONS: diesel_migrations::EmbeddedMigrations = embed_migrations!();
+
+lazy_static! {
+  static ref POOL: Pool = {
+    let username = env::var("DATABASE_USER").expect("DATABASE_USERNAME is not set");
+    let password = env::var("DATABASE_PASSWORD").expect("DATABASE_PASSWORD is not set");
+    let host = env::var("DATABASE_HOST").unwrap_or("localhost".to_string());
+    let name = env::var("DATABASE_NAME").expect("DATABASE_NAME is not set");
+    let port = env::var("DATABASE_PORT").unwrap_or("5432".to_string());
+    let url = format!("postgres://{}:{}@{}:{}/{}", username, password, host, port, name);
+    let manager = ConnectionManager::<PgConnection>::new(url);
+    Pool::builder().test_on_check_out(true).build(manager).expect("Failed to create db pool")
+  };
+}
+
+pub fn init() {
+  lazy_static::initialize(&POOL);
+  let mut pool: DbConnection = connection().expect("Failed to get db connection");
+  match pool.run_pending_migrations(MIGRATIONS) {
+    Ok(_) => info!("Database initialized"),
+    Err(err) => error!("Failed to initialize database; {}", err)
+  };
+}
+
+pub fn connection() -> Result<DbConnection, ServiceError> {
+  POOL.get()
+    .map_err(|e| ServiceError::new(500, format!("Failed getting db connection: {}", e)))
+}

src/error_handler.rs

@@ -0,0 +1,36 @@
+use diesel::result::Error as DieselError;
+use serde::{Deserialize, Serialize};
+use std::fmt;
+
+#[derive(Debug, Deserialize, Serialize)]
+pub struct ServiceError {
+    pub error_status_code: u16,
+    pub error_message: String,
+}
+
+impl ServiceError {
+  pub fn new(error_status_code: u16, error_message: String) -> ServiceError {
+    ServiceError {
+      error_status_code,
+      error_message,
+    }
+  }
+}
+
+impl fmt::Display for ServiceError {
+  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+    f.write_str(self.error_message.as_str())
+  }
+}
+
+impl From<DieselError> for ServiceError {
+  fn from(error: DieselError) -> ServiceError {
+    match error {
+      DieselError::DatabaseError(_, err) => ServiceError::new(409, err.message().to_string()),
+      DieselError::NotFound => {
+        ServiceError::new(404, "The record was not found".to_string())
+      }
+      err => ServiceError::new(500, format!("Unknown Diesel error: {}", err)),
+    }
+  }
+}

src/main.rs

@@ -1,10 +1,12 @@
+extern crate diesel;
+#[macro_use]
+extern crate diesel_migrations;
+
 use std::collections::{HashSet, HashMap};
 use std::env;
 use std::sync::Arc;
 
 use commands::audio::{create_response, AudioConfig, AudioConfigs};
-use diesel::r2d2::{Pool, ConnectionManager};
-use diesel::pg::PgConnection;
 
 use dotenv::dotenv;
 use log::{error, warn, info};
@@ -18,11 +20,13 @@ use serenity::prelude::*;
 use songbird::SerenityInit;
 
 mod commands;
-mod database;
+mod error_handler;
+mod db;
+mod messages;
+mod schema;
 struct Handler {
   // Open AI Config
-  oai: Option<commands::oai::OAI>,
-  pool: Pool<ConnectionManager<PgConnection>>
+  oai: Option<commands::oai::OAI>
 }
 
 #[async_trait]
@@ -46,7 +50,7 @@ impl EventHandler for Handler {
               Err(_) => false
             };
             if mentioned || bot_in_thread {
-              commands::oai::generate_response(&ctx, &msg, oai, &self.pool).await;
+              commands::oai::generate_response(&ctx, &msg, oai).await;
             }
           }
           Err(why) => warn!("Could not check mentions: {:?}", why)
@@ -134,20 +138,18 @@ async fn main() {
     Err(why) => panic!("Could not access application info: {:?}", why)
   };
 
-  let pool = database::establish_connection();
-  database::run_migrations(&pool);
+  db::init();
   
   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 , max_context_questions: 15 }),
-        pool
+        oai: Some(commands::oai::OAI { client: reqwest::Client::new(), base_url: "https://api.openai.com/v1".to_string(), max_attempts: 5, token , max_context_questions: 15 })
       }
     }
     Err(err) => {
       warn!("Could not load OpenAI token: {}", err);
-      Handler { oai: None, pool }
+      Handler { oai: None }
     }
   };
 

src/messages/mod.rs

@@ -0,0 +1,3 @@
+mod model;
+
+pub use model::*;

src/messages/model.rs

@@ -1,6 +1,6 @@
 use diesel::prelude::*;
 
-use super::schema::messages;
+use crate::schema::messages;
 
 #[derive(Queryable, Selectable)]
 #[diesel(table_name = messages)]