feat: databases now automatically close at the end of the script

Author: ElhamAryanpur

Cargo.toml

@@ -95,6 +95,7 @@ sha3 = "0.10.8"
 base64 = "0.22.1"
 bytes = { version = "1.11.1", features = ["serde"] }
 chrono = { version = "0.4.43", features = ["serde"] }
+markdown = { version = "1.0.0", features = ["serde"] }
 
 # runtime
 tokio = { version = "1.49.0", features = [
@@ -123,7 +124,6 @@ sqlx = { version = "0.8.6", features = [
 ] }
 include_dir = "0.7.4"
 sqlite-vec = "0.1.6"
-markdown = { version = "1.0.0", features = ["serde"] }
 
 [profile.release]
 opt-level = 3

src/commands/run.rs

@@ -49,27 +49,31 @@ pub async fn run_command(
     spawn_termination_task();
 
     // Remove the Shebang lines
-    let user_file = user_file
+    let mut user_file = user_file
         .lines()
         .filter(|line| !line.starts_with("#!"))
-        .collect::<Vec<_>>()
-        .join("\n");
+        .collect::<Vec<_>>();
+    user_file.push("astra_internal__close_all_databases()");
+    let user_file = user_file.join("\n");
 
-    if actual_path_str == "<commandline>" {
-        if let Err(e) = lua.load(user_file).set_name("<commandline>").exec_async().await {
-            error!("{}", e);
-        }
-    } else if let Some(is_teal) = PathBuf::from(&actual_path_str).extension()
+    if let Some(is_teal) = PathBuf::from(&actual_path_str).extension()
         && is_teal == "tl"
     {
         if let Err(e) = crate::components::load_teal(lua).await {
             error!("{}", e);
         }
 
-        if let Err(e) = crate::components::execute_teal_code(lua, &actual_path_str, &user_file).await {
+        if let Err(e) =
+            crate::components::execute_teal_code(lua, &actual_path_str, &user_file).await
+        {
             error!("{}", e);
         }
-    } else if let Err(e) = lua.load(user_file).set_name(actual_path_str).exec_async().await {
+    } else if let Err(e) = lua
+        .load(user_file)
+        .set_name(actual_path_str)
+        .exec_async()
+        .await
+    {
         error!("{}", e);
     }
 

src/components/global.rs

@@ -1,3 +1,4 @@
+use crate::components::database::DATABASE_POOLS;
 use mlua::{LuaSerdeExt, UserData};
 
 pub fn register_to_lua(lua: &mlua::Lua) -> mlua::Result<()> {
@@ -7,6 +8,7 @@ pub fn register_to_lua(lua: &mlua::Lua) -> mlua::Result<()> {
     pprintln(lua)?;
     AstraRegex::register_to_lua(lua)?;
     uuid_v4(lua)?;
+    close_dbs(lua)?;
     // env
     getenv(lua)?;
     setenv(lua)?;
@@ -18,6 +20,23 @@ pub fn register_to_lua(lua: &mlua::Lua) -> mlua::Result<()> {
     Ok(())
 }
 
+// At the end of the script, the database files should be closed automatically
+pub fn close_dbs(lua: &mlua::Lua) -> mlua::Result<()> {
+    lua.globals().set(
+        "astra_internal__close_all_databases",
+        lua.create_async_function(|_, _: ()| async {
+            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,
+                }
+            }
+            Ok(())
+        })?,
+    )
+}
+
 pub fn dotenv_function(lua: &mlua::Lua) -> mlua::Result<()> {
     lua.globals().set(
         "astra_internal__dotenv_load",
@@ -29,15 +48,19 @@ pub fn dotenv_function(lua: &mlua::Lua) -> mlua::Result<()> {
 }
 
 fn lua_print(args: mlua::MultiValue) {
+    let mut padding = " ";
+    if args.len() == 1 {
+        padding = ""
+    }
     for input in args.iter() {
         if input.is_string()
             && let Ok(s) = input.to_string()
         {
-            print!("{} ", s);
+            print!("{}{padding}", s);
         } else if input.is_userdata() {
-            print!("{input:?} ")
+            print!("{input:?}{padding}")
         } else {
-            print!("{input:#?} ")
+            print!("{input:#?}{padding}")
         }
     }
 }