efficient DB connection & query cache (#72)

* cleanup

* chore: cache fetch table result in information_schema

* chore: fmt

* chore: refactoring inefficient connection logic

* wip

* wip: postgres need to be using an async client

* works

* wip: saving changes for now - going to try connection pooling

* chore: moving connection init to static lifetime scope

* chore: lazy connection fixed

* chore: clean up

* chore: fix test

* chore: clippy

* chore: cleanup

* chore: clean up

* chore: add information on DB_CONN_CACHE lazy static

* chore: remove SOME_INT

* chore: add docs on DB_CONNECTIONS

* chore: clean up

* chore: cleanup unused imports

* chore: clean up

src/common/config.rs

@@ -230,7 +230,7 @@ impl Config {
     ///     -- @db: postgres
     ///     SELECT * FROM some_table;
     ///
-    /// The method figures out the correct database to connect in order to validate the SQL query
+    /// The method figures out the connection name to connect in order to validate the SQL query
     ///
     /// If you pass down a query with a annotation to specify a DB
     /// e.g.
@@ -243,31 +243,18 @@ impl Config {
     /// e.g.
     ///     SELECT * FROM some_table;
     ///
-    /// It should return the default connection configured by your configuration settings
-    pub fn get_correct_db_connection(&self, raw_sql: &str) -> DbConnectionConfig {
+    /// It should return the connection name that is available based on your connection configurations
+    pub fn get_correct_db_connection(&self, raw_sql: &str) -> String {
         let re = Regex::new(r"(/*|//|--) @db: (?P<conn>[\w]+)( */){0,}").unwrap();
         let found_matches = re.captures(raw_sql);
 
         if let Some(found_match) = &found_matches {
             let detected_conn_name = &found_match[2];
-            return self
-                .connections
-                .get(detected_conn_name)
-                .unwrap_or_else(|| {
-                    panic!("Failed to find a matching connection type - connection name: {detected_conn_name}")
-                })
-                .clone();
+
+            return detected_conn_name.to_string();
         }
 
-        self.connections
-            .get("default")
-            .expect(
-                r"Failed to find the default connection configuration - check your configuration
-              CLI options: https://jasonshin.github.io/sqlx-ts/user-guide/2.1.cli-options.html
-              File based config: https://jasonshin.github.io/sqlx-ts/reference-guide/2.configs-file-based.html
-            ",
-            )
-            .clone()
+        return "default".to_string();
     }
 
     pub fn get_postgres_cred(&self, conn: &DbConnectionConfig) -> String {
@@ -295,6 +282,7 @@ impl Config {
             .db_name(db_name.clone())
     }
 
+    // TODO: update this to also factor in env variable
     pub fn get_log_level(file_config_path: &PathBuf) -> LogLevel {
         let file_based_config = fs::read_to_string(file_config_path);
         let file_based_config = &file_based_config.map(|f| serde_json::from_str::<SqlxConfig>(f.as_str()).unwrap());

src/common/lazy.rs

@@ -1,18 +1,52 @@
 use crate::common::cli::Cli;
 use crate::common::config::Config;
+use crate::common::types::DatabaseType;
+use crate::core::connection::{DBConn, DBConnections};
 use crate::ts_generator::information_schema::DBSchema;
 use clap::Parser;
 use lazy_static::lazy_static;
+use mysql::Conn as MySQLConn;
+use postgres::{Client as PGClient, NoTls as PGNoTls};
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
 
 // The file contains all implicitly dependent variables or state that files need for the logic
 // We have a lot of states that we need to drill down into each methods
 lazy_static! {
-    pub static ref SOME_INT: i32 = 5;
-
     pub static ref CLI_ARGS: Cli = Cli::parse();
     pub static ref CONFIG: Config =  Config::new();
 
     // This is a holder for shared DBSChema used to fetch information for information_schema table
     // By having a singleton, we can think about caching the result if we are fetching a query too many times
-    pub static ref DB_SCHEMA: DBSchema = DBSchema::new();
+    pub static ref DB_SCHEMA: Mutex<DBSchema> = Mutex::new(DBSchema::new());
+
+    // This variable holds database connections for each connection name that is defined in the config
+    // We are using lazy_static to initialize the connections once and use them throughout the application
+    static ref DB_CONN_CACHE: HashMap<String, Arc<Mutex<DBConn>>> = {
+        let mut cache = HashMap::new();
+        for connection in CONFIG.connections.keys() {
+            let connection_config = CONFIG.connections.get(connection).unwrap();
+            let db_type = connection_config.db_type.to_owned();
+            let conn = match db_type {
+                DatabaseType::Mysql => {
+                    let opts = CONFIG.get_mysql_cred(&connection_config);
+                    let mut conn = MySQLConn::new(opts).unwrap();
+                    DBConn::MySQLPooledConn(Mutex::new(conn))
+                }
+                DatabaseType::Postgres => {
+                    let postgres_cred = &CONFIG.get_postgres_cred(&connection_config);
+                    DBConn::PostgresConn(Mutex::new(PGClient::connect(postgres_cred, PGNoTls).unwrap()))
+                }
+            };
+            cache.insert(connection.to_owned(), Arc::new(Mutex::new(conn)));
+        };
+        cache
+    };
+
+    // This variable holds a singleton of DBConnections that is used to get a DBConn from the cache
+    // DBConn is used to access the raw connection to the database or run `prepare` statement against each connection
+    pub static ref DB_CONNECTIONS: Mutex<DBConnections<'static>> = {
+        let db_connections = DBConnections::new(&DB_CONN_CACHE);
+        Mutex::new(db_connections)
+    };
 }

src/core/connection.rs

@@ -0,0 +1,55 @@
+use crate::common::lazy::CONFIG;
+use crate::common::SQL;
+use crate::core::mysql::prepare as mysql_explain;
+use crate::core::postgres::prepare as postgres_explain;
+use crate::ts_generator::types::ts_query::TsQuery;
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::Mutex;
+
+use color_eyre::Result;
+use mysql::Conn as MySQLConn;
+use postgres::Client as PostgresConn;
+use swc_common::errors::Handler;
+
+/// Enum to hold a specific database connection instance
+pub enum DBConn {
+    MySQLPooledConn(Mutex<MySQLConn>),
+    PostgresConn(Mutex<PostgresConn>),
+}
+
+impl DBConn {
+    pub fn prepare(
+        &self,
+        sql: &SQL,
+        should_generate_types: &bool,
+        handler: &Handler,
+    ) -> Result<(bool, Option<TsQuery>)> {
+        let (explain_failed, ts_query) = match &self {
+            DBConn::MySQLPooledConn(_conn) => mysql_explain::prepare(&self, sql, should_generate_types, handler)?,
+            DBConn::PostgresConn(_conn) => postgres_explain::prepare(&self, sql, should_generate_types, handler)?,
+        };
+
+        Ok((explain_failed, ts_query))
+    }
+}
+
+pub struct DBConnections<'a> {
+    pub cache: &'a HashMap<String, Arc<Mutex<DBConn>>>,
+}
+
+impl<'a> DBConnections<'a> {
+    pub fn new(cache: &'a HashMap<String, Arc<Mutex<DBConn>>>) -> Self {
+        Self { cache }
+    }
+
+    pub fn get_connection(&mut self, raw_sql: &str) -> Arc<Mutex<DBConn>> {
+        let db_conn_name = &CONFIG.get_correct_db_connection(raw_sql);
+
+        let conn = self
+            .cache
+            .get(db_conn_name)
+            .expect("Failed to get the connection from cache");
+        conn.to_owned()
+    }
+}

src/core/execute.rs

@@ -1,12 +1,15 @@
-use crate::common::lazy::{CLI_ARGS, CONFIG};
+use super::connection::DBConn;
+use crate::common::lazy::{CLI_ARGS, CONFIG, DB_CONNECTIONS};
 use crate::common::types::DatabaseType;
 use crate::common::SQL;
 use crate::core::mysql::prepare as mysql_explain;
 use crate::core::postgres::prepare as postgres_explain;
 use crate::ts_generator::generator::{write_colocated_ts_file, write_single_ts_file};
 
 use color_eyre::eyre::Result;
+use std::borrow::BorrowMut;
 use std::collections::HashMap;
+use std::sync::Arc;
 
 use std::path::PathBuf;
 use swc_common::errors::Handler;
@@ -22,19 +25,19 @@ pub fn execute(queries: &HashMap<PathBuf, Vec<SQL>>, handler: &Handler) -> Resul
     for (file_path, sqls) in queries {
         let mut sqls_to_write: Vec<String> = vec![];
         for sql in sqls {
-            let connection = &CONFIG.get_correct_db_connection(&sql.query);
+            let mut connection = DB_CONNECTIONS.lock().unwrap();
+            let connection = &connection.get_connection(&sql.query).clone();
+            let connection = &connection.lock().unwrap();
 
-            let (explain_failed, ts_query) = match connection.db_type {
-                DatabaseType::Postgres => postgres_explain::prepare(sql, should_generate_types, handler)?,
-                DatabaseType::Mysql => mysql_explain::prepare(sql, should_generate_types, handler)?,
-            };
+            let (explain_failed, ts_query) = &connection.prepare(&sql, &should_generate_types, &handler)?;
 
             // If any prepare statement fails, we should set the failed flag as true
-            failed = explain_failed;
+            failed = explain_failed.clone();
 
             if *should_generate_types {
-                let ts_query = ts_query.expect("Failed to generate types from query").to_string();
-                sqls_to_write.push(ts_query);
+                let ts_query = &ts_query.clone().expect("Failed to generate types from query");
+                let ts_query = &ts_query.to_string();
+                sqls_to_write.push(ts_query.to_owned());
             }
         }
 

src/core/mod.rs

@@ -1,3 +1,4 @@
+pub mod connection;
 pub mod execute;
 pub mod mysql;
 pub mod postgres;

src/core/mysql/prepare.rs

@@ -1,29 +1,34 @@
-use crate::common::lazy::CONFIG;
+use crate::common::lazy::{CONFIG, DB_CONNECTIONS};
 use crate::common::SQL;
+use crate::core::connection::DBConn;
 use crate::ts_generator::generator::generate_ts_interface;
-use crate::ts_generator::types::db_conn::DBConn;
 use crate::ts_generator::types::ts_query::TsQuery;
 use color_eyre::eyre::Result;
 use mysql::prelude::*;
 use mysql::*;
 
-use std::cell::RefCell;
+use std::borrow::BorrowMut;
 use swc_common::errors::Handler;
 
 /// Runs the prepare statement on the input SQL.
 /// Validates the query is right by directly connecting to the configured database.
 /// It also processes ts interfaces if the configuration is set to generate_types = true
-pub fn prepare(sql: &SQL, should_generate_types: &bool, handler: &Handler) -> Result<(bool, Option<TsQuery>)> {
-    let connection_config = CONFIG.get_correct_db_connection(&sql.query);
-    let opts = CONFIG.get_mysql_cred(&connection_config);
-    let mut conn = Conn::new(opts)?;
-
+pub fn prepare(
+    db_conn: &DBConn,
+    sql: &SQL,
+    should_generate_types: &bool,
+    handler: &Handler,
+) -> Result<(bool, Option<TsQuery>)> {
     let mut failed = false;
 
     let span = sql.span.to_owned();
     let explain_query = format!("PREPARE stmt FROM \"{}\"", sql.query);
 
-    let result: Result<Vec<Row>, _> = conn.query(explain_query);
+    let mut conn = match &db_conn {
+        DBConn::MySQLPooledConn(conn) => conn,
+        _ => panic!("Invalid connection type"),
+    };
+    let result: Result<Vec<Row>, _> = conn.lock().unwrap().borrow_mut().query(explain_query);
 
     if let Err(err) = result {
         handler.span_bug_no_panic(span, err.to_string().as_str());
@@ -33,10 +38,7 @@ pub fn prepare(sql: &SQL, should_generate_types: &bool, handler: &Handler) -> Re
     let mut ts_query = None;
 
     if should_generate_types == &true {
-        ts_query = Some(generate_ts_interface(
-            sql,
-            &DBConn::MySQLPooledConn(&mut RefCell::new(&mut conn)),
-        )?);
+        ts_query = Some(generate_ts_interface(sql, &db_conn)?);
     }
 
     Ok((failed, ts_query))

src/core/postgres/prepare.rs

@@ -1,42 +1,47 @@
-use crate::common::lazy::CONFIG;
 use crate::common::SQL;
+use crate::core::connection::DBConn;
 use crate::ts_generator::generator::generate_ts_interface;
-use crate::ts_generator::types::db_conn::DBConn;
 use crate::ts_generator::types::ts_query::TsQuery;
 use color_eyre::eyre::Result;
-use postgres::{Client, NoTls};
-use std::cell::RefCell;
+use std::borrow::BorrowMut;
 
 use swc_common::errors::Handler;
 
 /// Runs the prepare statement on the input SQL. Validates the query is right by directly connecting to the configured database.
 /// It also processes ts interfaces if the configuration is set to `generate_types = true`
-pub fn prepare<'a>(sql: &SQL, should_generate_types: &bool, handler: &Handler) -> Result<(bool, Option<TsQuery>)> {
-    let connection = &CONFIG.get_correct_db_connection(&sql.query);
-    let postgres_cred = &CONFIG.get_postgres_cred(connection);
-    let mut conn = Client::connect(postgres_cred, NoTls).unwrap();
-
+pub fn prepare(
+    db_conn: &DBConn,
+    sql: &SQL,
+    should_generate_types: &bool,
+    handler: &Handler,
+) -> Result<(bool, Option<TsQuery>)> {
     let mut failed = false;
 
+    let mut conn = match &db_conn {
+        DBConn::PostgresConn(conn) => conn,
+        _ => panic!("Invalid connection type"),
+    };
     let span = sql.span.to_owned();
     let prepare_query = format!("PREPARE sqlx_stmt AS {}", sql.query);
-    let result = conn.query(prepare_query.as_str(), &[]);
+    let result = conn.lock().unwrap().borrow_mut().query(prepare_query.as_str(), &[]);
 
     if let Err(e) = result {
         handler.span_bug_no_panic(span, e.as_db_error().unwrap().message());
         failed = true;
     } else {
         // We should only deallocate if the prepare statement was executed successfully
-        conn.query("DEALLOCATE sqlx_stmt", &[]).unwrap();
+        let _ = &conn
+            .lock()
+            .unwrap()
+            .borrow_mut()
+            .query("DEALLOCATE sqlx_stmt", &[])
+            .unwrap();
     }
 
     let mut ts_query = None;
 
     if should_generate_types == &true {
-        ts_query = Some(generate_ts_interface(
-            sql,
-            &DBConn::PostgresConn(&mut RefCell::new(&mut conn)),
-        )?);
+        ts_query = Some(generate_ts_interface(sql, &db_conn)?);
     }
 
     Ok((failed, ts_query))

src/main.rs

@@ -9,15 +9,13 @@ extern crate dotenv;
 
 use crate::core::execute::execute;
 
-use dotenv::dotenv;
 use sqlx_ts::ts_generator::generator::clear_single_ts_file_if_exists;
 use std::env;
-use std::io::{stderr, stdout, Write};
 
 use crate::common::lazy::CLI_ARGS;
 use crate::common::logger::*;
 use crate::{parser::parse_source, scan_folder::scan_folder};
-use color_eyre::{eyre::eyre, eyre::Result};
+use color_eyre::eyre::Result;
 
 fn set_default_env_var() {
     if env::var("SQLX_TS_LOG").is_err() {

src/ts_generator/generator.rs

@@ -6,10 +6,10 @@ use std::{
 };
 
 use super::annotations::extract_param_annotations;
-use super::types::db_conn::DBConn;
 
 use crate::common::lazy::CONFIG;
 use crate::common::SQL;
+use crate::core::connection::DBConn;
 use crate::ts_generator::annotations::extract_result_annotations;
 use crate::ts_generator::sql_parser::translate_stmt::translate_stmt;
 use crate::ts_generator::types::ts_query::TsQuery;