Merge pull request #20 from theseus-rs/refactor-pgpassword

refactor: update psql to manage setting the PGPASSWORD environment variable when pg_password is set

Author: brianheineman

Cargo.lock

@@ -41,6 +41,7 @@ pub struct PsqlBuilder {
     username: Option<OsString>,
     no_password: bool,
     password: bool,
+    pg_password: Option<OsString>,
 }
 
 impl PsqlBuilder {
@@ -267,6 +268,12 @@ impl PsqlBuilder {
         self.password = true;
         self
     }
+
+    /// user password
+    pub fn pg_password<S: AsRef<OsStr>>(mut self, pg_password: S) -> Self {
+        self.pg_password = Some(pg_password.as_ref().to_os_string());
+        self
+    }
 }
 
 impl CommandBuilder for PsqlBuilder {
@@ -440,6 +447,17 @@ impl CommandBuilder for PsqlBuilder {
 
         args
     }
+
+    /// Get the environment variables for the command
+    fn get_envs(&self) -> Vec<(OsString, OsString)> {
+        let mut envs: Vec<(OsString, OsString)> = Vec::new();
+
+        if let Some(password) = &self.pg_password {
+            envs.push(("PGPASSWORD".into(), password.into()));
+        }
+
+        envs
+    }
 }
 
 #[cfg(test)]
@@ -496,10 +514,11 @@ mod tests {
             .username("postgres")
             .no_password()
             .password()
+            .pg_password("password")
             .build();
 
         assert_eq!(
-            r#""psql" "--command" "SELECT * FROM test" "--dbname" "dbname" "--file" "test.sql" "--list" "--variable" "ON_ERROR_STOP=1" "--version" "--no-psqlrc" "--single-transaction" "--help" "options" "--echo-all" "--echo-errors" "--echo-queries" "--echo-hidden" "--log-file" "psql.log" "--no-readline" "--output" "output.txt" "--quiet" "--single-step" "--single-line" "--no-align" "--csv" "--field-separator" "|" "--html" "--pset" "border=1" "--record-separator" "\n" "--tuples-only" "--table-attr" "width=100" "--expanded" "--field-separator-zero" "--record-separator-zero" "--host" "localhost" "--port" "5432" "--username" "postgres" "--no-password" "--password""#,
+            r#"PGPASSWORD="password" "psql" "--command" "SELECT * FROM test" "--dbname" "dbname" "--file" "test.sql" "--list" "--variable" "ON_ERROR_STOP=1" "--version" "--no-psqlrc" "--single-transaction" "--help" "options" "--echo-all" "--echo-errors" "--echo-queries" "--echo-hidden" "--log-file" "psql.log" "--no-readline" "--output" "output.txt" "--quiet" "--single-step" "--single-line" "--no-align" "--csv" "--field-separator" "|" "--html" "--pset" "border=1" "--record-separator" "\n" "--tuples-only" "--table-attr" "width=100" "--expanded" "--field-separator-zero" "--record-separator-zero" "--host" "localhost" "--port" "5432" "--username" "postgres" "--no-password" "--password""#,
             command.to_command_string()
         );
     }

postgresql_embedded/src/command/psql.rs

@@ -23,7 +23,14 @@ pub trait CommandBuilder {
     }
 
     /// Get the arguments for the command
-    fn get_args(&self) -> Vec<OsString>;
+    fn get_args(&self) -> Vec<OsString> {
+        vec![]
+    }
+
+    /// Get the environment variables for the command
+    fn get_envs(&self) -> Vec<(OsString, OsString)> {
+        vec![]
+    }
 
     /// Build a standard Command
     fn build(self) -> std::process::Command
@@ -34,6 +41,7 @@ pub trait CommandBuilder {
         let mut command = std::process::Command::new(program_file);
 
         command.args(self.get_args());
+        command.envs(self.get_envs());
         command
     }
 
@@ -47,6 +55,7 @@ pub trait CommandBuilder {
         let mut command = tokio::process::Command::new(program_file);
 
         command.args(self.get_args());
+        command.envs(self.get_envs());
         command
     }
 }
@@ -147,9 +156,32 @@ mod test {
     use super::*;
     use test_log::test;
 
+    #[test]
+    fn test_command_builder_defaults() {
+        struct DefaultCommandBuilder {
+            program_dir: Option<PathBuf>,
+        }
+
+        impl CommandBuilder for DefaultCommandBuilder {
+            fn get_program(&self) -> &'static OsStr {
+                "test".as_ref()
+            }
+
+            fn get_program_dir(&self) -> &Option<PathBuf> {
+                &self.program_dir
+            }
+        }
+
+        let builder = DefaultCommandBuilder { program_dir: None };
+        let command = builder.build();
+
+        assert_eq!(r#""test""#, command.to_command_string());
+    }
+
     struct TestCommandBuilder {
         program_dir: Option<PathBuf>,
         args: Vec<OsString>,
+        envs: Vec<(OsString, OsString)>,
     }
 
     impl CommandBuilder for TestCommandBuilder {
@@ -164,18 +196,26 @@ mod test {
         fn get_args(&self) -> Vec<OsString> {
             self.args.clone()
         }
+
+        fn get_envs(&self) -> Vec<(OsString, OsString)> {
+            self.envs.clone()
+        }
     }
 
     #[test]
     fn test_standard_command_builder() {
         let builder = TestCommandBuilder {
             program_dir: None,
             args: vec!["--help".to_string().into()],
+            envs: vec![(OsString::from("PASSWORD"), OsString::from("foo"))],
         };
         let command = builder.build();
 
         assert_eq!(
-            format!(r#""{}" "--help""#, PathBuf::from("test").to_string_lossy()),
+            format!(
+                r#"PASSWORD="foo" "{}" "--help""#,
+                PathBuf::from("test").to_string_lossy()
+            ),
             command.to_command_string()
         );
     }
@@ -186,11 +226,15 @@ mod test {
         let builder = TestCommandBuilder {
             program_dir: None,
             args: vec!["--help".to_string().into()],
+            envs: vec![(OsString::from("PASSWORD"), OsString::from("foo"))],
         };
         let command = builder.build_tokio();
 
         assert_eq!(
-            format!(r#""{}" "--help""#, PathBuf::from("test").to_string_lossy()),
+            format!(
+                r#"PASSWORD="foo" "{}" "--help""#,
+                PathBuf::from("test").to_string_lossy()
+            ),
             command.to_command_string()
         );
     }

postgresql_embedded/src/command/traits.rs

@@ -338,6 +338,7 @@ impl PostgreSQL {
             .host(&self.settings.host)
             .port(self.settings.port)
             .username(&self.settings.username)
+            .pg_password(&self.settings.password)
             .no_psqlrc()
             .no_align()
             .tuples_only();
@@ -373,6 +374,7 @@ impl PostgreSQL {
             .host(&self.settings.host)
             .port(self.settings.port)
             .username(&self.settings.username)
+            .pg_password(&self.settings.password)
             .no_psqlrc()
             .no_align()
             .tuples_only();
@@ -403,6 +405,7 @@ impl PostgreSQL {
             .host(&self.settings.host)
             .port(self.settings.port)
             .username(&self.settings.username)
+            .pg_password(&self.settings.password)
             .no_psqlrc()
             .no_align()
             .tuples_only();
@@ -428,8 +431,6 @@ impl PostgreSQL {
         command_builder: B,
     ) -> Result<(String, String)> {
         let mut command = command_builder.build();
-        // TODO: move this into the command builder
-        command.env("PGPASSWORD", &self.settings.password);
         command.execute(self.settings.timeout).await
     }
 
@@ -440,8 +441,6 @@ impl PostgreSQL {
         command_builder: B,
     ) -> Result<(String, String)> {
         let mut command = command_builder.build_tokio();
-        // TODO: move this into the command builder
-        command.env("PGPASSWORD", &self.settings.password);
         command.execute(self.settings.timeout).await
     }
 }