Thanks to visit codestin.com
Credit goes to docs.rs

postgresql_commands/
traits.rs

1use crate::error::{Error, Result};
2use std::env::consts::OS;
3use std::ffi::{OsStr, OsString};
4use std::fmt::Debug;
5#[cfg(target_os = "windows")]
6use std::os::windows::process::CommandExt;
7use std::path::PathBuf;
8use std::process::ExitStatus;
9use std::time::Duration;
10use tracing::debug;
11
12/// Constant for the `CREATE_NO_WINDOW` flag on Windows to prevent the creation of a console window
13/// when executing commands. This is useful for background processes or services that do not require
14/// user interaction.
15///
16/// # References
17///
18/// - [Windows API: Process Creation Flags](https://learn.microsoft.com/en-us/windows/win32/procthread/process-creation-flags#flags)
19#[cfg(target_os = "windows")]
20const CREATE_NO_WINDOW: u32 = 0x0800_0000;
21
22/// Interface for `PostgreSQL` settings
23pub trait Settings {
24    fn get_binary_dir(&self) -> PathBuf;
25    fn get_host(&self) -> OsString;
26    fn get_port(&self) -> u16;
27    fn get_username(&self) -> OsString;
28    fn get_password(&self) -> OsString;
29}
30
31#[cfg(test)]
32pub struct TestSettings;
33
34#[cfg(test)]
35impl Settings for TestSettings {
36    fn get_binary_dir(&self) -> PathBuf {
37        PathBuf::from(".")
38    }
39
40    fn get_host(&self) -> OsString {
41        "localhost".into()
42    }
43
44    fn get_port(&self) -> u16 {
45        5432
46    }
47
48    fn get_username(&self) -> OsString {
49        "postgres".into()
50    }
51
52    fn get_password(&self) -> OsString {
53        "password".into()
54    }
55}
56
57/// Trait to build a command
58pub trait CommandBuilder: Debug {
59    /// Get the program name
60    fn get_program(&self) -> &'static OsStr;
61
62    /// Location of the program binary
63    fn get_program_dir(&self) -> &Option<PathBuf>;
64
65    /// Fully qualified path to the program binary
66    fn get_program_file(&self) -> PathBuf {
67        let program_name = &self.get_program();
68        match self.get_program_dir() {
69            Some(program_dir) => program_dir.join(program_name),
70            None => PathBuf::from(program_name),
71        }
72    }
73
74    /// Get the arguments for the command
75    fn get_args(&self) -> Vec<OsString> {
76        vec![]
77    }
78
79    /// Get the environment variables for the command
80    fn get_envs(&self) -> Vec<(OsString, OsString)>;
81
82    /// Set an environment variable for the command
83    #[must_use]
84    fn env<S: AsRef<OsStr>>(self, key: S, value: S) -> Self;
85
86    /// Build a standard Command
87    fn build(self) -> std::process::Command
88    where
89        Self: Sized,
90    {
91        let program_file = self.get_program_file();
92        let mut command = std::process::Command::new(program_file);
93
94        #[cfg(target_os = "windows")]
95        {
96            command.creation_flags(CREATE_NO_WINDOW);
97        }
98
99        command.args(self.get_args());
100        command.envs(self.get_envs());
101        command
102    }
103
104    #[cfg(feature = "tokio")]
105    /// Build a tokio Command
106    fn build_tokio(self) -> tokio::process::Command
107    where
108        Self: Sized,
109    {
110        let program_file = self.get_program_file();
111        let mut command = tokio::process::Command::new(program_file);
112
113        #[cfg(target_os = "windows")]
114        {
115            command.creation_flags(CREATE_NO_WINDOW);
116        }
117
118        command.args(self.get_args());
119        command.envs(self.get_envs());
120        command
121    }
122}
123
124/// Trait to convert a command to a string representation
125pub trait CommandToString {
126    fn to_command_string(&self) -> String;
127}
128
129/// Implement the [`CommandToString`] trait for [`Command`](std::process::Command)
130impl CommandToString for std::process::Command {
131    fn to_command_string(&self) -> String {
132        format!("{self:?}")
133    }
134}
135
136#[cfg(feature = "tokio")]
137/// Implement the [`CommandToString`] trait for [`Command`](tokio::process::Command)
138impl CommandToString for tokio::process::Command {
139    fn to_command_string(&self) -> String {
140        format!("{self:?}")
141            .replace("Command { std: ", "")
142            .replace(", kill_on_drop: false }", "")
143    }
144}
145
146/// Interface for executing a command
147pub trait CommandExecutor {
148    /// Execute the command and return the stdout and stderr
149    ///
150    /// # Errors
151    ///
152    /// Returns an error if the command fails
153    fn execute(&mut self) -> Result<(String, String)>;
154}
155
156/// Interface for executing a command
157pub trait AsyncCommandExecutor {
158    /// Execute the command and return the stdout and stderr
159    async fn execute(&mut self, timeout: Option<Duration>) -> Result<(String, String)>;
160}
161
162/// Implement the [`CommandExecutor`] trait for [`Command`](std::process::Command)
163impl CommandExecutor for std::process::Command {
164    /// Execute the command and return the stdout and stderr
165    fn execute(&mut self) -> Result<(String, String)> {
166        debug!("Executing command: {}", self.to_command_string());
167        let program = self.get_program().to_string_lossy().to_string();
168        let stdout: String;
169        let stderr: String;
170        let status: ExitStatus;
171
172        if OS == "windows" && program.as_str().ends_with("pg_ctl") {
173            // The pg_ctl process can hang on Windows when attempting to get stdout/stderr.
174            let mut process = self
175                .stdout(std::process::Stdio::piped())
176                .stderr(std::process::Stdio::piped())
177                .spawn()?;
178            stdout = String::new();
179            stderr = String::new();
180            status = process.wait()?;
181        } else {
182            let output = self.output()?;
183            stdout = String::from_utf8_lossy(&output.stdout).into_owned();
184            stderr = String::from_utf8_lossy(&output.stderr).into_owned();
185            status = output.status;
186        }
187        debug!(
188            "Result: {}\nstdout: {}\nstderr: {}",
189            status.code().map_or("None".to_string(), |c| c.to_string()),
190            stdout,
191            stderr
192        );
193
194        if status.success() {
195            Ok((stdout, stderr))
196        } else {
197            Err(Error::CommandError { stdout, stderr })
198        }
199    }
200}
201
202#[cfg(feature = "tokio")]
203/// Implement the [`CommandExecutor`] trait for [`Command`](tokio::process::Command)
204impl AsyncCommandExecutor for tokio::process::Command {
205    /// Execute the command and return the stdout and stderr
206    async fn execute(&mut self, timeout: Option<Duration>) -> Result<(String, String)> {
207        debug!("Executing command: {}", self.to_command_string());
208        let program = self.as_std().get_program().to_string_lossy().to_string();
209        let stdout: String;
210        let stderr: String;
211        let status: ExitStatus;
212
213        if OS == "windows" && program.as_str().ends_with("pg_ctl") {
214            // The pg_ctl process can hang on Windows when attempting to get stdout/stderr.
215            let mut process = self
216                .stdout(std::process::Stdio::piped())
217                .stderr(std::process::Stdio::piped())
218                .spawn()?;
219            stdout = String::new();
220            stderr = String::new();
221            status = process.wait().await?;
222        } else {
223            let output = match timeout {
224                Some(duration) => tokio::time::timeout(duration, self.output()).await?,
225                None => self.output().await,
226            }?;
227            stdout = String::from_utf8_lossy(&output.stdout).into_owned();
228            stderr = String::from_utf8_lossy(&output.stderr).into_owned();
229            status = output.status;
230        }
231
232        debug!(
233            "Result: {}\nstdout: {}\nstderr: {}",
234            status.code().map_or("None".to_string(), |c| c.to_string()),
235            stdout,
236            stderr
237        );
238
239        if status.success() {
240            Ok((stdout, stderr))
241        } else {
242            Err(Error::CommandError { stdout, stderr })
243        }
244    }
245}
246#[cfg(test)]
247mod test {
248    use super::*;
249    use test_log::test;
250
251    #[test]
252    fn test_command_builder_defaults() {
253        #[derive(Debug, Default)]
254        struct DefaultCommandBuilder {
255            program_dir: Option<PathBuf>,
256            envs: Vec<(OsString, OsString)>,
257        }
258
259        impl CommandBuilder for DefaultCommandBuilder {
260            fn get_program(&self) -> &'static OsStr {
261                "test".as_ref()
262            }
263
264            fn get_program_dir(&self) -> &Option<PathBuf> {
265                &self.program_dir
266            }
267
268            fn get_envs(&self) -> Vec<(OsString, OsString)> {
269                self.envs.clone()
270            }
271
272            fn env<S: AsRef<OsStr>>(mut self, key: S, value: S) -> Self {
273                self.envs
274                    .push((key.as_ref().to_os_string(), value.as_ref().to_os_string()));
275                self
276            }
277        }
278
279        let builder = DefaultCommandBuilder::default();
280        let command = builder.env("ENV", "foo").build();
281        #[cfg(not(target_os = "windows"))]
282        let command_prefix = r#"ENV="foo" "#;
283        #[cfg(target_os = "windows")]
284        let command_prefix = String::new();
285
286        assert_eq!(
287            format!(r#"{command_prefix}"test""#),
288            command.to_command_string()
289        );
290    }
291
292    #[derive(Debug)]
293    struct TestCommandBuilder {
294        program_dir: Option<PathBuf>,
295        args: Vec<OsString>,
296        envs: Vec<(OsString, OsString)>,
297    }
298
299    impl CommandBuilder for TestCommandBuilder {
300        fn get_program(&self) -> &'static OsStr {
301            "test".as_ref()
302        }
303
304        fn get_program_dir(&self) -> &Option<PathBuf> {
305            &self.program_dir
306        }
307
308        fn get_args(&self) -> Vec<OsString> {
309            self.args.clone()
310        }
311
312        fn get_envs(&self) -> Vec<(OsString, OsString)> {
313            self.envs.clone()
314        }
315
316        fn env<S: AsRef<OsStr>>(mut self, key: S, value: S) -> Self {
317            self.envs
318                .push((key.as_ref().to_os_string(), value.as_ref().to_os_string()));
319            self
320        }
321    }
322
323    #[test]
324    fn test_standard_command_builder() {
325        let builder = TestCommandBuilder {
326            program_dir: None,
327            args: vec!["--help".to_string().into()],
328            envs: vec![],
329        };
330        let command = builder.env("PASSWORD", "foo").build();
331        #[cfg(not(target_os = "windows"))]
332        let command_prefix = r#"PASSWORD="foo" "#;
333        #[cfg(target_os = "windows")]
334        let command_prefix = String::new();
335
336        assert_eq!(
337            format!(
338                r#"{command_prefix}"{}" "--help""#,
339                PathBuf::from("test").to_string_lossy()
340            ),
341            command.to_command_string()
342        );
343    }
344
345    #[cfg(feature = "tokio")]
346    #[test]
347    fn test_tokio_command_builder() {
348        let builder = TestCommandBuilder {
349            program_dir: None,
350            args: vec!["--help".to_string().into()],
351            envs: vec![],
352        };
353        let command = builder.env("PASSWORD", "foo").build_tokio();
354
355        assert_eq!(
356            format!(
357                r#"PASSWORD="foo" "{}" "--help""#,
358                PathBuf::from("test").to_string_lossy()
359            ),
360            command.to_command_string()
361        );
362    }
363
364    #[test]
365    fn test_standard_to_command_string() {
366        let mut command = std::process::Command::new("test");
367        command.arg("-l");
368        assert_eq!(r#""test" "-l""#, command.to_command_string(),);
369    }
370
371    #[cfg(feature = "tokio")]
372    #[test]
373    fn test_tokio_to_command_string() {
374        let mut command = tokio::process::Command::new("test");
375        command.arg("-l");
376        assert_eq!(r#""test" "-l""#, command.to_command_string(),);
377    }
378
379    #[test(tokio::test)]
380    async fn test_standard_command_execute() -> Result<()> {
381        #[cfg(not(target_os = "windows"))]
382        let mut command = std::process::Command::new("sh");
383        #[cfg(not(target_os = "windows"))]
384        command.args(["-c", "echo foo"]);
385
386        #[cfg(target_os = "windows")]
387        let mut command = std::process::Command::new("cmd");
388        #[cfg(target_os = "windows")]
389        command.args(["/C", "echo foo"]);
390
391        let (stdout, stderr) = command.execute()?;
392        assert!(stdout.starts_with("foo"));
393        assert!(stderr.is_empty());
394        Ok(())
395    }
396
397    #[test(tokio::test)]
398    async fn test_standard_command_execute_error() {
399        let mut command = std::process::Command::new("bogus_command");
400        assert!(command.execute().is_err());
401    }
402
403    #[cfg(feature = "tokio")]
404    #[test(tokio::test)]
405    async fn test_tokio_command_execute() -> Result<()> {
406        #[cfg(not(target_os = "windows"))]
407        let mut command = tokio::process::Command::new("sh");
408        #[cfg(not(target_os = "windows"))]
409        command.args(["-c", "echo foo"]);
410
411        #[cfg(target_os = "windows")]
412        let mut command = tokio::process::Command::new("cmd");
413        #[cfg(target_os = "windows")]
414        command.args(["/C", "echo foo"]);
415
416        let (stdout, stderr) = command.execute(None).await?;
417        assert!(stdout.starts_with("foo"));
418        assert!(stderr.is_empty());
419        Ok(())
420    }
421
422    #[cfg(feature = "tokio")]
423    #[test(tokio::test)]
424    async fn test_tokio_command_execute_error() -> Result<()> {
425        let mut command = tokio::process::Command::new("bogus_command");
426        assert!(command.execute(None).await.is_err());
427        Ok(())
428    }
429}