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
     }
+
+    /// database user name
+    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,16 @@ impl CommandBuilder for PsqlBuilder {
 
         args
     }
+
+    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 +513,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
     }
 }

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
     }
 }