Refactor database connection and configuration handling

src-tauri/src/establish_connection.rs

@@ -1,6 +1,6 @@
 use diesel::prelude::*;
+use std::result::Result;
 
-pub fn establish_connection(database_path: &String) -> SqliteConnection {
-    SqliteConnection::establish(&database_path)
-        .unwrap_or_else(|_| panic!("Error connecting to {}", database_path))
+pub fn establish_connection(database_path: &str) -> Result<SqliteConnection, diesel::ConnectionError> {
+    SqliteConnection::establish(database_path)
 }

src-tauri/src/get_chat_history.rs

@@ -11,7 +11,7 @@ pub fn get_chat_history() -> Result<Vec<RawDatabaseChatEntry>, String> {
     // Use block_on to run the async code synchronously
     block_on(async {
         let database_path = get_database_path().map_err(|e| e.to_string())?;
-        let mut conn = establish_connection(&database_path);
+        let mut conn = establish_connection(&database_path).map_err(|e| e.to_string())?;
 
         // Pass the mutable reference to Diesel operations
         let results = chat_histories

src-tauri/src/get_chat_history_by_session.rs

@@ -14,7 +14,7 @@ pub fn get_chat_history_by_session(
     block_on(async {
         // Use the helper function to get a mutable connection
         let database_path = get_database_path().map_err(|e| e.to_string())?;
-        let mut conn = establish_connection(&database_path);
+        let mut conn = establish_connection(&database_path).map_err(|e| e.to_string())?;
 
         // Filter the chat history by the specific session_id
         let results = chat_histories

src-tauri/src/get_chatgpt_response.rs

@@ -16,7 +16,7 @@ pub async fn get_chatgpt_response(
     api_key: String,
 ) -> Result<ChatResponse, String> {
     let database_path = get_database_path().map_err(|e| e.to_string())?;
-    let mut conn = establish_connection(&database_path);
+    let mut conn = establish_connection(&database_path).map_err(|e| e.to_string())?;
 
     let session_history = chat_histories
         .filter(session_id.eq(&input_session_id))

src-tauri/src/get_config.rs

@@ -1,15 +1,23 @@
 use crate::config::Config;
+use std::path::PathBuf;
 
 pub async fn get_config() -> Result<Config, String> {
-    let mut config_path = dirs::home_dir().expect("Failed to get home directory");
+    // Attempt to get the home directory
+    let mut config_path: PathBuf = dirs::home_dir().ok_or("Failed to get home directory".to_string())?;
+
     config_path.push(".cuuri/config.toml");
 
+    // Check if the config file exists
     if !config_path.exists() {
         return Err("Configuration file does not exist".to_string());
     }
 
-    let config = Config::from_file(config_path.to_str().unwrap())
+    // Convert the path to a string and handle possible error
+    let config_path_str = config_path.to_str().ok_or("Failed to convert config path to string".to_string())?;
+
+    // Load the configuration from file and handle potential errors
+    let config = Config::from_file(config_path_str)
         .map_err(|e| format!("Failed to load configuration: {}", e))?;
 
     Ok(config)
-}
+}

src-tauri/src/get_session_id_list.rs

@@ -9,7 +9,7 @@ use tauri::async_runtime::block_on;
 pub fn get_session_id_list() -> Result<Vec<SessionId>, String> {
     block_on(async {
         let database_path = get_database_path().map_err(|e| e.to_string())?;
-        let mut conn = establish_connection(&database_path);
+        let mut conn = establish_connection(&database_path).map_err(|e| e.to_string())?;
 
         let results = chat_histories
             .select(session_id)

src-tauri/src/init_config_file.rs

@@ -1,20 +1,21 @@
 use std::fs;
-use std::io::Write;
+use std::io::{self, Write};
 
-pub fn init_config_file() {
-    let mut config_path = dirs::home_dir().expect("Failed to get home directory");
+pub fn init_config_file() -> Result<(), io::Error> {
+    let mut config_path = dirs::home_dir().ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "Failed to get home directory"))?;
     config_path.push(".cuuri/config.toml");
 
     if !config_path.exists() {
         if let Some(parent) = config_path.parent() {
-            fs::create_dir_all(parent).expect("Failed to create directories for config file");
+            fs::create_dir_all(parent)?;
         }
-        let mut file = fs::File::create(&config_path).expect("Failed to create config file");
-
-        writeln!(file, "openai_api_key = \"\"").expect("Failed to write default API key");
-        writeln!(file, "default_model = \"gpt-3.5-turbo\"").expect("Failed to write default model");
+
+        let mut file = fs::File::create(&config_path)?;
+        writeln!(file, "openai_api_key = \"\"")?;
+        writeln!(file, "default_model = \"gpt-3.5-turbo\"")?;
 
         println!("Configuration file created at {:?}", config_path);
     }
 
-}
+    Ok(())
+}

src-tauri/src/lib.rs

@@ -33,11 +33,31 @@ use set_openai_api_key::set_openai_api_key;
 use std::env;
 
 pub fn run() {
-    init_config_file();
+    if let Err(e) = init_config_file() {
+        eprintln!("Failed to initialize config file: {}", e);
+        return;
+    }
 
-    let database_path: String = get_database_path().unwrap();
-    let mut connection: diesel::SqliteConnection = establish_connection(&database_path);
-    run_migrations(&mut connection);
+    let database_path = match get_database_path() {
+        Ok(path) => path,
+        Err(e) => {
+            eprintln!("Failed to get database path: {}", e);
+            return;
+        }
+    };
+
+    let mut connection = match establish_connection(&database_path) {
+        Ok(conn) => conn,
+        Err(e) => {
+            eprintln!("Failed to establish connection: {}", e);
+            return;
+        }
+    };
+
+    if let Err(e) = run_migrations(&mut connection) {
+        eprintln!("Failed to run migrations: {}", e);
+        return;
+    }
 
     tauri::Builder::default()
         .invoke_handler(tauri::generate_handler![
@@ -53,4 +73,4 @@ pub fn run() {
         ])
         .run(tauri::generate_context!())
         .expect("error while running tauri application");
-}
+}

src-tauri/src/run_migrations.rs

@@ -1,7 +1,8 @@
 use diesel::prelude::SqliteConnection;
 use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness};
+
 pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!();
 
-pub fn run_migrations(connection: &mut SqliteConnection) {
-    connection.run_pending_migrations(MIGRATIONS).unwrap();
-}
+pub fn run_migrations(connection: &mut SqliteConnection) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
+    connection.run_pending_migrations(MIGRATIONS).map(|_| ())
+}

src-tauri/src/set_openai_api_key.rs

@@ -1,5 +1,6 @@
 use crate::config::Config;
 use std::fs;
+use std::path::PathBuf;
 
 #[tauri::command]
 pub async fn set_openai_api_key(api_key: String) -> Result<(), String> {
@@ -8,17 +9,23 @@ pub async fn set_openai_api_key(api_key: String) -> Result<(), String> {
         return Err("API key cannot be empty".to_string());
     }
 
-    let mut config_path = dirs::home_dir().expect("Failed to get home directory");
+    // Attempt to get the home directory, returning an error if it fails
+    let home_dir = dirs::home_dir()
+        .ok_or_else(|| "Failed to get home directory".to_string())?;
+
+    let mut config_path = PathBuf::from(home_dir);
     config_path.push(".cuuri/config.toml");
 
-    let mut config = Config::from_file(config_path.to_str().unwrap())
+    // Attempt to load the configuration from the file
+    let mut config = Config::from_file(config_path.to_str().ok_or_else(|| "Invalid config path".to_string())?)
         .map_err(|e| format!("Failed to load configuration: {}", e))?;
 
     config.openai_api_key = api_key;
 
-    let content =
-        toml::to_string(&config).map_err(|e| format!("Failed to serialize config: {}", e))?;
+    // Serialize the config and check for errors
+    let content = toml::to_string(&config).map_err(|e| format!("Failed to serialize config: {}", e))?;
 
+    // Write the serialized config to a file
     fs::write(&config_path, content)
         .map_err(|e| format!("Failed to write to config file: {}", e))?;