Add support for loading environment variables from .env file (#208)

* feat: support loading variables from .env file

* Fix. Remove debug logs

* Remove unused string

* refactor: removed duplicate call to dotenv::var

* Add CLI argument for custom .env file name and update help text

* clippy: use next back

* fmt

* dotenv test

* fix the test

* update tests

---------

Co-authored-by: Leo5878 <[email protected]>
Co-authored-by: JasonShin <[email protected]>

Author: 3 people

src/common/cli.rs

@@ -74,4 +74,8 @@ pub struct Cli {
   /// log level to be used for the CLI debug > info > warning > error
   #[clap(value_enum, long)]
   pub log_level: Option<LogLevel>,
+
+  /// Dotfile name (example: .env or .env.dev) [default: .env or environment]
+  #[clap(long)]
+  pub env: Option<String>,
 }

src/common/config.rs

@@ -87,7 +87,7 @@ impl Default for Config {
 
 impl Config {
   pub fn new() -> Config {
-    let dotenv = Dotenv::new();
+    let dotenv = Dotenv::new(CLI_ARGS.env.clone());
 
     let default_config_path = PathBuf::from_str(".sqlxrc.json").unwrap();
     let default_ignore_config_path = PathBuf::from_str(".sqlxignore").unwrap();

src/common/dotenv.rs

@@ -1,5 +1,5 @@
 use crate::common::types::DatabaseType;
-use std::env::var;
+use dotenv;
 
 #[derive(Clone, Debug)]
 pub struct Dotenv {
@@ -14,14 +14,22 @@ pub struct Dotenv {
 
 impl Default for Dotenv {
   fn default() -> Self {
-    Self::new()
+    Self::new(None)
   }
 }
 
 impl Dotenv {
-  pub fn new() -> Dotenv {
+  fn get_var(key: &str) -> Option<String> {
+    dotenv::var(key).ok()
+  }
+
+  pub fn new(path_to_dotenv: Option<String>) -> Dotenv {
+    if let Some(value) = path_to_dotenv {
+      dotenv::from_path(value).ok();
+    }
+
     Dotenv {
-      db_type: match var("DB_TYPE").ok() {
+      db_type: match Self::get_var("DB_TYPE") {
         None => None,
         Some(val) => {
           if val == "mysql" {
@@ -31,14 +39,12 @@ impl Dotenv {
           }
         }
       },
-      db_user: var("DB_USER").ok(),
-      db_host: var("DB_HOST").ok(),
-      db_port: var("DB_PORT")
-        .ok()
-        .map(|val| val.parse::<u16>().expect("DB_PORT is not a valid integer")),
-      db_pass: var("DB_PASS").ok(),
-      db_name: var("DB_NAME").ok(),
-      pg_search_path: var("PG_SEARCH_PATH").ok(),
+      db_user: Self::get_var("DB_USER"),
+      db_host: Self::get_var("DB_HOST"),
+      db_port: Self::get_var("DB_PORT").map(|val| val.parse::<u16>().expect("DB_PORT is not a valid integer")),
+      db_pass: Self::get_var("DB_PASS"),
+      db_name: Self::get_var("DB_NAME"),
+      pg_search_path: Self::get_var("PG_SEARCH_PATH"),
     }
   }
 }

src/core/postgres/pool.rs

@@ -20,8 +20,6 @@ impl bb8::ManageConnection for PostgresConnectionManager {
   type Error = Error;
 
   async fn connect(&self) -> Result<Client, Error> {
-    let conn_url = self.conn_url.clone();
-
     let (client, connection) =
       tokio_postgres::connect(&self.conn_url, NoTls)
         .await

src/ts_generator/generator.rs

@@ -33,7 +33,7 @@ pub fn get_query_name(sql: &SQL) -> Result<String> {
       .unwrap()
       .as_str()
       .split(':')
-      .last()
+      .next_back()
       .unwrap()
       .to_string();
 

tests/dotenv.rs

@@ -0,0 +1,96 @@
+#[cfg(test)]
+mod dotenv_test {
+  use assert_cmd::Command;
+  use std::fs::{self};
+  use std::io::Write;
+  use tempfile::tempdir;
+
+  #[test]
+  fn loads_env_vars_from_dotenv_file() -> Result<(), Box<dyn std::error::Error>> {
+    let temp_dir = tempdir()?;
+    let dotenv_path = temp_dir.path().join(".env");
+
+    let dotenv_content = r#"
+DB_TYPE=postgres
+DB_HOST=127.0.0.1
+DB_PORT=54321
+DB_USER=postgres
+DB_PASS=postgres
+DB_NAME=postgres
+"#;
+    fs::write(&dotenv_path, dotenv_content)?;
+
+    let sample_dir = temp_dir.path().join("sample");
+    fs::create_dir_all(&sample_dir)?;
+    fs::copy("tests/sample/sample.ts", sample_dir.join("sample.ts"))?;
+
+    let mut cmd = Command::cargo_bin("sqlx-ts")?;
+    cmd
+      .current_dir(temp_dir.path())
+      .arg("--ext=ts")
+      .arg("-g")
+      .arg(sample_dir.to_str().unwrap());
+
+    cmd
+      .assert()
+      .success()
+      .stdout(predicates::str::contains("No SQL errors detected!")); // adjust if your CLI errors differently
+
+    Ok(())
+  }
+
+  #[test]
+  fn loads_env_vars_from_dotenv_file_but_wrong_config() -> Result<(), Box<dyn std::error::Error>> {
+    let temp_dir = tempdir()?;
+    let dotenv_path = temp_dir.path().join(".env");
+
+    let dotenv_content = r#"
+DB_TYPE=postgres
+DB_HOST=127.0.0.1
+DB_PORT=54322
+DB_USER=postgres
+DB_PASS=postgres
+DB_NAME=postgres
+"#;
+    fs::write(&dotenv_path, dotenv_content)?;
+
+    let sample_dir = temp_dir.path().join("sample");
+    fs::create_dir_all(&sample_dir)?;
+    fs::copy("tests/sample/sample.ts", sample_dir.join("sample.ts"))?;
+
+    let mut cmd = Command::cargo_bin("sqlx-ts")?;
+    cmd
+      .current_dir(temp_dir.path())
+      .arg("--ext=ts")
+      .arg("-g")
+      .arg(sample_dir.to_str().unwrap());
+
+    cmd.assert().failure().stderr(predicates::str::contains(
+      "Postgres database connection error: error connecting to server: Connection refused",
+    )); // adjust if your CLI errors differently
+
+    Ok(())
+  }
+
+  #[test]
+  fn fails_if_dotenv_missing() -> Result<(), Box<dyn std::error::Error>> {
+    let temp_dir = tempdir()?;
+
+    let sample_dir = temp_dir.path().join("sample");
+    fs::create_dir_all(&sample_dir)?;
+    fs::copy("tests/sample/sample.ts", sample_dir.join("sample.ts"))?;
+
+    let mut cmd = Command::cargo_bin("sqlx-ts")?;
+    cmd
+      .current_dir(temp_dir.path())
+      .arg("--ext=ts")
+      .arg("-g")
+      .arg(sample_dir.to_str().unwrap());
+
+    cmd.assert().failure().stderr(predicates::str::contains(
+      "Unable to retrieve a database type, please check your configuration and try again",
+    ));
+
+    Ok(())
+  }
+}