feat: Database connections now automatically closed on program closing

Author: ElhamAryanpur

src/commands/run.rs

@@ -1,4 +1,4 @@
-use crate::{LUA, RUNTIME_FLAGS};
+use crate::{LUA, RUNTIME_FLAGS, components::database::DATABASE_POOLS};
 use std::path::PathBuf;
 use tracing::error;
 
@@ -94,6 +94,14 @@ pub async fn run_command(
             error!("{e}");
         }
 
+        let database_pools = DATABASE_POOLS.lock().await.clone();
+        for i in database_pools {
+            match i {
+                crate::components::database::DatabaseType::Postgres(pool) => pool.close().await,
+                crate::components::database::DatabaseType::Sqlite(pool) => pool.close().await,
+            }
+        }
+
         std::process::exit(
             #[cfg(unix)]
             0,

src/components/database.rs

@@ -1,5 +1,10 @@
 use mlua::{LuaSerdeExt, UserData};
 use sqlx::{Pool, Postgres, Row, Sqlite, migrate::MigrateDatabase};
+use std::sync::LazyLock;
+use tokio::sync::Mutex;
+
+pub static DATABASE_POOLS: LazyLock<Mutex<Vec<DatabaseType>>> =
+    LazyLock::new(|| Mutex::new(Vec::new()));
 
 #[derive(Debug, Clone)]
 pub enum DatabaseType {
@@ -39,9 +44,14 @@ impl Database {
                             .connect(format!("sqlite:{url}").as_str())
                             .await
                         {
-                            Ok(pool) => Ok(Database {
-                                db: Some(DatabaseType::Sqlite(pool)),
-                            }),
+                            Ok(pool) => {
+                                let pool = DatabaseType::Sqlite(pool);
+
+                                let mut database_pools = DATABASE_POOLS.lock().await;
+                                database_pools.push(pool.clone());
+
+                                Ok(Database { db: Some(pool) })
+                            }
                             Err(e) => Err(mlua::Error::runtime(format!(
                                 "Error connecting to Sqlite: {e:#?}"
                             ))),
@@ -52,9 +62,14 @@ impl Database {
                         .connect(url.as_str())
                         .await
                     {
-                        Ok(pool) => Ok(Database {
-                            db: Some(DatabaseType::Postgres(pool)),
-                        }),
+                        Ok(pool) => {
+                            let pool = DatabaseType::Postgres(pool);
+
+                            let mut database_pools = DATABASE_POOLS.lock().await;
+                            database_pools.push(pool.clone());
+
+                            Ok(Database { db: Some(pool) })
+                        }
                         Err(e) => Err(mlua::Error::runtime(format!(
                             "Error connecting to Postgres: {e:#?}"
                         ))),

src/components/mod.rs

@@ -1,7 +1,7 @@
 use mlua::{ExternalError, LuaSerdeExt};
 
 mod crypto;
-mod database;
+pub mod database;
 mod datetime;
 pub mod global;
 pub mod http;