Merge pull request #76 from theseus-rs/add-command-envars

fix: correct errors when PGDATABASE envar is set

Author: brianheineman

Cargo.lock

@@ -6,8 +6,10 @@ use std::path::PathBuf;
 
 /// `clusterdb` clusters all previously clustered tables in a database.
 #[derive(Clone, Debug, Default)]
+#[allow(clippy::struct_excessive_bools)]
 pub struct ClusterDbBuilder {
     program_dir: Option<PathBuf>,
+    envs: Vec<(OsString, OsString)>,
     all: bool,
     dbname: Option<OsString>,
     echo: bool,
@@ -26,12 +28,13 @@ pub struct ClusterDbBuilder {
 }
 
 impl ClusterDbBuilder {
-    /// Create a new [ClusterDbBuilder]
+    /// Create a new [`ClusterDbBuilder`]
+    #[must_use]
     pub fn new() -> Self {
         Self::default()
     }
 
-    /// Create a new [ClusterDbBuilder] from [Settings]
+    /// Create a new [`ClusterDbBuilder`] from [Settings]
     pub fn from(settings: &dyn Settings) -> Self {
         Self::new()
             .program_dir(settings.get_binary_dir())
@@ -42,96 +45,112 @@ impl ClusterDbBuilder {
     }
 
     /// Location of the program binary
+    #[must_use]
     fn program_dir<P: Into<PathBuf>>(mut self, path: P) -> Self {
         self.program_dir = Some(path.into());
         self
     }
 
     /// Cluster all databases
+    #[must_use]
     pub fn all(mut self) -> Self {
         self.all = true;
         self
     }
 
     /// Database to cluster
+    #[must_use]
     pub fn dbname<S: AsRef<OsStr>>(mut self, dbname: S) -> Self {
         self.dbname = Some(dbname.as_ref().to_os_string());
         self
     }
 
     /// Show the commands being sent to the server
+    #[must_use]
     pub fn echo(mut self) -> Self {
         self.echo = true;
         self
     }
 
     /// Don't write any messages
+    #[must_use]
     pub fn quiet(mut self) -> Self {
         self.quiet = true;
         self
     }
 
     /// Cluster specific table(s) only
+    #[must_use]
     pub fn table<S: AsRef<OsStr>>(mut self, table: S) -> Self {
         self.table = Some(table.as_ref().to_os_string());
         self
     }
 
     /// Write a lot of output
+    #[must_use]
     pub fn verbose(mut self) -> Self {
         self.verbose = true;
         self
     }
 
     /// Output version information, then exit
+    #[must_use]
     pub fn version(mut self) -> Self {
         self.version = true;
         self
     }
 
     /// Show help, then exit
+    #[must_use]
     pub fn help(mut self) -> Self {
         self.help = true;
         self
     }
 
     /// Database server host or socket directory
+    #[must_use]
     pub fn host<S: AsRef<OsStr>>(mut self, host: S) -> Self {
         self.host = Some(host.as_ref().to_os_string());
         self
     }
 
     /// Database server port
+    #[must_use]
     pub fn port(mut self, port: u16) -> Self {
         self.port = Some(port);
         self
     }
 
     /// User name to connect as
+    #[must_use]
     pub fn username<S: AsRef<OsStr>>(mut self, username: S) -> Self {
         self.username = Some(username.as_ref().to_os_string());
         self
     }
 
     /// Never prompt for password
+    #[must_use]
     pub fn no_password(mut self) -> Self {
         self.no_password = true;
         self
     }
 
     /// Force password prompt
+    #[must_use]
     pub fn password(mut self) -> Self {
         self.password = true;
         self
     }
 
     /// user password
+    #[must_use]
     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
     }
 
     /// Alternate maintenance database
+    #[must_use]
     pub fn maintenance_db<S: AsRef<OsStr>>(mut self, db: S) -> Self {
         self.maintenance_db = Some(db.as_ref().to_os_string());
         self
@@ -189,7 +208,7 @@ impl CommandBuilder for ClusterDbBuilder {
 
         if let Some(host) = &self.host {
             args.push("--host".into());
-            args.push(host.into())
+            args.push(host.into());
         }
 
         if let Some(port) = &self.port {
@@ -220,14 +239,21 @@ impl CommandBuilder for ClusterDbBuilder {
 
     /// Get the environment variables for the command
     fn get_envs(&self) -> Vec<(OsString, OsString)> {
-        let mut envs: Vec<(OsString, OsString)> = Vec::new();
+        let mut envs: Vec<(OsString, OsString)> = self.envs.clone();
 
         if let Some(password) = &self.pg_password {
             envs.push(("PGPASSWORD".into(), password.into()));
         }
 
         envs
     }
+
+    /// Set an environment variable for the command
+    fn env<S: AsRef<OsStr>>(mut self, key: S, value: S) -> Self {
+        self.envs
+            .push((key.as_ref().to_os_string(), value.as_ref().to_os_string()));
+        self
+    }
 }
 
 #[cfg(test)]
@@ -258,6 +284,7 @@ mod tests {
     #[test]
     fn test_builder() {
         let command = ClusterDbBuilder::new()
+            .env("PGDATABASE", "database")
             .all()
             .dbname("dbname")
             .echo()
@@ -276,7 +303,7 @@ mod tests {
             .build();
 
         assert_eq!(
-            r#"PGPASSWORD="password" "clusterdb" "--all" "--dbname" "dbname" "--echo" "--quiet" "--table" "table" "--verbose" "--version" "--help" "--host" "localhost" "--port" "5432" "--username" "postgres" "--no-password" "--password" "--maintenance-db" "postgres""#,
+            r#"PGDATABASE="database" PGPASSWORD="password" "clusterdb" "--all" "--dbname" "dbname" "--echo" "--quiet" "--table" "table" "--verbose" "--version" "--help" "--host" "localhost" "--port" "5432" "--username" "postgres" "--no-password" "--password" "--maintenance-db" "postgres""#,
             command.to_command_string()
         );
     }