Skip to main content

oliphaunt_tools/
lib.rs

1#![deny(unsafe_code)]
2
3mod arguments;
4
5use std::error::Error as StdError;
6use std::ffi::OsString;
7use std::fmt;
8use std::io::Write;
9use std::path::{Path, PathBuf};
10use std::process::{Command, Stdio};
11use std::thread;
12
13use arguments::{validate_pg_dump_arguments, validate_psql_arguments};
14
15/// Product id for the native PostgreSQL client tools artifact family.
16pub const PRODUCT: &str = "oliphaunt-tools";
17
18/// Artifact kind relayed by this facade crate.
19pub const KIND: &str = "native-tools";
20
21/// Options for a plain-text `pg_dump`.
22#[derive(Debug, Clone, Default, PartialEq, Eq)]
23pub struct PgDumpOptions {
24    args: Vec<String>,
25}
26
27impl PgDumpOptions {
28    /// Create default plain-text dump options.
29    pub fn new() -> Self {
30        Self::default()
31    }
32
33    /// Add one PostgreSQL `pg_dump` argument.
34    pub fn arg(mut self, argument: impl Into<String>) -> Self {
35        self.args.push(argument.into());
36        self
37    }
38
39    /// Add PostgreSQL `pg_dump` arguments.
40    pub fn args(mut self, arguments: impl IntoIterator<Item = impl Into<String>>) -> Self {
41        self.args.extend(arguments.into_iter().map(Into::into));
42        self
43    }
44}
45
46/// Options for a non-interactive `psql` invocation.
47#[derive(Debug, Clone, Default, PartialEq, Eq)]
48pub struct PsqlOptions {
49    args: Vec<String>,
50    input: Option<PsqlInput>,
51}
52
53#[derive(Debug, Clone, PartialEq, Eq)]
54enum PsqlInput {
55    Command(String),
56    Script(String),
57}
58
59impl PsqlOptions {
60    /// Create default non-interactive psql options.
61    pub fn new() -> Self {
62        Self::default()
63    }
64
65    /// Add one PostgreSQL `psql` argument.
66    pub fn arg(mut self, argument: impl Into<String>) -> Self {
67        self.args.push(argument.into());
68        self
69    }
70
71    /// Add PostgreSQL `psql` arguments.
72    pub fn args(mut self, arguments: impl IntoIterator<Item = impl Into<String>>) -> Self {
73        self.args.extend(arguments.into_iter().map(Into::into));
74        self
75    }
76
77    /// Run one command through `psql -c`.
78    pub fn command(mut self, sql: impl Into<String>) -> Self {
79        self.input = Some(PsqlInput::Command(sql.into()));
80        self
81    }
82
83    /// Run a complete SQL script through psql standard input.
84    pub fn script(mut self, sql: impl Into<String>) -> Self {
85        self.input = Some(PsqlInput::Script(sql.into()));
86        self
87    }
88}
89
90/// Failure returned by a PostgreSQL frontend program.
91#[derive(Debug)]
92pub struct PostgresToolError {
93    /// Program name (`pg_dump` or `psql`).
94    pub tool: &'static str,
95    /// Process exit status when the program started.
96    pub exit_code: Option<i32>,
97    /// UTF-8 standard output captured before failure.
98    pub stdout: String,
99    /// UTF-8 standard error captured before failure.
100    pub stderr: String,
101    source: Option<std::io::Error>,
102}
103
104impl fmt::Display for PostgresToolError {
105    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
106        if let Some(source) = &self.source {
107            return write!(formatter, "could not run {}: {source}", self.tool);
108        }
109        write!(
110            formatter,
111            "{} exited with status {}{}",
112            self.tool,
113            self.exit_code
114                .map_or_else(|| "unknown".to_owned(), |status| status.to_string()),
115            if self.stderr.trim().is_empty() {
116                String::new()
117            } else {
118                format!(": {}", self.stderr.trim())
119            }
120        )
121    }
122}
123
124impl StdError for PostgresToolError {
125    fn source(&self) -> Option<&(dyn StdError + 'static)> {
126        self.source
127            .as_ref()
128            .map(|source| source as &(dyn StdError + 'static))
129    }
130}
131
132/// Run packaged `pg_dump` against a PostgreSQL connection string.
133pub fn pg_dump(
134    connection_string: &str,
135    options: PgDumpOptions,
136) -> Result<String, PostgresToolError> {
137    validate_connection_string("pg_dump", connection_string)?;
138    validate_pg_dump_arguments(&options.args)
139        .map_err(|message| configuration_error("pg_dump", &message))?;
140    let mut arguments = options
141        .args
142        .into_iter()
143        .map(OsString::from)
144        .collect::<Vec<_>>();
145    arguments.push(OsString::from("--encoding=UTF8"));
146    arguments.push(OsString::from("--no-password"));
147    arguments.push(OsString::from(format!("--dbname={connection_string}")));
148    run_tool("pg_dump", arguments, None)
149}
150
151/// Run packaged non-interactive `psql` against a PostgreSQL connection string.
152pub fn psql(connection_string: &str, options: PsqlOptions) -> Result<String, PostgresToolError> {
153    validate_connection_string("psql", connection_string)?;
154    validate_psql_arguments(&options.args)
155        .map_err(|message| configuration_error("psql", &message))?;
156    if options.input.is_none() && options.args.is_empty() {
157        return Err(configuration_error(
158            "psql",
159            "psql requires command(), script(), or a non-input argument",
160        ));
161    }
162    match options.input.as_ref() {
163        Some(PsqlInput::Command(command)) => validate_text("psql", "command", command)?,
164        Some(PsqlInput::Script(script)) => validate_text("psql", "script", script)?,
165        None => {}
166    }
167    let (arguments, stdin) = psql_invocation(connection_string, options);
168    run_tool("psql", arguments, stdin)
169}
170
171fn psql_invocation(
172    connection_string: &str,
173    options: PsqlOptions,
174) -> (Vec<OsString>, Option<Vec<u8>>) {
175    let mut arguments = options
176        .args
177        .into_iter()
178        .map(OsString::from)
179        .collect::<Vec<_>>();
180    arguments.extend([
181        OsString::from("--no-psqlrc"),
182        OsString::from("--no-password"),
183        OsString::from("--set=ON_ERROR_STOP=1"),
184        OsString::from(format!("--dbname={connection_string}")),
185    ]);
186    let stdin = match options.input {
187        Some(PsqlInput::Command(command)) => {
188            arguments.push(OsString::from("--command"));
189            arguments.push(OsString::from(command));
190            None
191        }
192        Some(PsqlInput::Script(script)) => {
193            arguments.push(OsString::from("--file=-"));
194            Some(script.into_bytes())
195        }
196        None => None,
197    };
198    (arguments, stdin)
199}
200
201fn run_tool(
202    tool: &'static str,
203    arguments: Vec<OsString>,
204    stdin: Option<Vec<u8>>,
205) -> Result<String, PostgresToolError> {
206    let executable = resolve_tool(tool)?;
207    let mut command = Command::new(&executable);
208    command
209        .args(arguments)
210        .env("PGCLIENTENCODING", "UTF8")
211        .stdin(if stdin.is_some() {
212            Stdio::piped()
213        } else {
214            Stdio::null()
215        })
216        .stdout(Stdio::piped())
217        .stderr(Stdio::piped());
218    configure_runtime_environment(&mut command, &executable);
219    let mut child = command.spawn().map_err(|source| PostgresToolError {
220        tool,
221        exit_code: None,
222        stdout: String::new(),
223        stderr: String::new(),
224        source: Some(source),
225    })?;
226    // Drain stdout/stderr while a potentially large psql script is written.
227    // Writing all stdin first can deadlock when the child fills an output pipe.
228    let input_writer = stdin.and_then(|input| {
229        child
230            .stdin
231            .take()
232            .map(|mut writer| thread::spawn(move || writer.write_all(&input)))
233    });
234    let output = child
235        .wait_with_output()
236        .map_err(|source| PostgresToolError {
237            tool,
238            exit_code: None,
239            stdout: String::new(),
240            stderr: String::new(),
241            source: Some(source),
242        })?;
243    let input_failure = input_writer.and_then(|writer| match writer.join() {
244        Ok(Ok(())) => None,
245        Ok(Err(error)) => Some(error),
246        Err(_) => Some(std::io::Error::other("psql stdin writer panicked")),
247    });
248    if !output.status.success() {
249        return Err(PostgresToolError {
250            tool,
251            exit_code: output.status.code(),
252            stdout: String::from_utf8_lossy(&output.stdout).into_owned(),
253            stderr: String::from_utf8_lossy(&output.stderr).into_owned(),
254            source: None,
255        });
256    }
257    if let Some(source) = input_failure {
258        return Err(PostgresToolError {
259            tool,
260            exit_code: output.status.code(),
261            stdout: String::from_utf8_lossy(&output.stdout).into_owned(),
262            stderr: String::from_utf8_lossy(&output.stderr).into_owned(),
263            source: Some(source),
264        });
265    }
266    String::from_utf8(output.stdout).map_err(|error| PostgresToolError {
267        tool,
268        exit_code: output.status.code(),
269        stdout: String::from_utf8_lossy(error.as_bytes()).into_owned(),
270        stderr: format!(
271            "{}{} produced non-UTF-8 output: {error}",
272            String::from_utf8_lossy(&output.stderr),
273            tool
274        ),
275        source: None,
276    })
277}
278
279fn resolve_tool(tool: &'static str) -> Result<PathBuf, PostgresToolError> {
280    let executable = if cfg!(windows) {
281        format!("{tool}.exe")
282    } else {
283        tool.to_owned()
284    };
285    let mut roots = Vec::new();
286    if let Some(directory) = std::env::var_os("OLIPHAUNT_TOOLS_DIR") {
287        roots.push(PathBuf::from(directory));
288    }
289    if let Some(directory) = option_env!("OLIPHAUNT_PACKAGED_TOOLS_DIR") {
290        roots.push(PathBuf::from(directory));
291    }
292    if let Some(directory) = option_env!("OLIPHAUNT_RESOURCES_DIR") {
293        roots.push(PathBuf::from(directory).join("native-tools/oliphaunt-tools/runtime"));
294    }
295    if let Ok(current) = std::env::current_exe()
296        && let Some(directory) = current.parent()
297    {
298        roots.push(directory.join("oliphaunt-tools/runtime"));
299        roots.push(directory.join("runtime"));
300    }
301    for root in roots {
302        let candidate = root.join("bin").join(&executable);
303        if candidate.is_file() {
304            return Ok(candidate);
305        }
306    }
307    Err(configuration_error(
308        tool,
309        &format!(
310            "could not locate packaged {tool}; add the oliphaunt-tools artifact facade or set OLIPHAUNT_TOOLS_DIR"
311        ),
312    ))
313}
314
315fn configure_runtime_environment(command: &mut Command, executable: &Path) {
316    let Some(runtime) = executable.parent().and_then(Path::parent) else {
317        return;
318    };
319    let library = runtime.join("lib");
320    prepend_environment_path(command, "PATH", executable.parent().unwrap_or(runtime));
321    if cfg!(target_os = "macos") {
322        prepend_environment_path(command, "DYLD_LIBRARY_PATH", &library);
323    } else if cfg!(unix) {
324        prepend_environment_path(command, "LD_LIBRARY_PATH", &library);
325    }
326    let icu = runtime.join("share/icu");
327    if icu.is_dir() {
328        command.env("ICU_DATA", icu);
329    }
330}
331
332fn prepend_environment_path(command: &mut Command, name: &str, value: &Path) {
333    let mut paths = vec![value.to_path_buf()];
334    if let Some(existing) = std::env::var_os(name) {
335        paths.extend(std::env::split_paths(&existing));
336    }
337    if let Ok(joined) = std::env::join_paths(paths) {
338        command.env(name, joined);
339    }
340}
341
342fn validate_connection_string(tool: &'static str, value: &str) -> Result<(), PostgresToolError> {
343    if value.trim().is_empty() || value.as_bytes().contains(&0) {
344        return Err(configuration_error(
345            tool,
346            "connection string must not be empty or contain NUL bytes",
347        ));
348    }
349    Ok(())
350}
351
352fn validate_text(tool: &'static str, label: &str, value: &str) -> Result<(), PostgresToolError> {
353    if value.as_bytes().contains(&0) {
354        return Err(configuration_error(
355            tool,
356            &format!("{label} must not contain NUL bytes"),
357        ));
358    }
359    Ok(())
360}
361
362fn configuration_error(tool: &'static str, message: &str) -> PostgresToolError {
363    PostgresToolError {
364        tool,
365        exit_code: None,
366        stdout: String::new(),
367        stderr: message.to_owned(),
368        source: None,
369    }
370}
371
372#[cfg(test)]
373mod tests {
374    use super::*;
375
376    #[test]
377    fn psql_scripts_explicitly_read_standard_input() {
378        let (script, stdin) = psql_invocation(
379            "postgresql://localhost/postgres",
380            PsqlOptions::new().script("SELECT 1;"),
381        );
382        assert!(script.iter().any(|argument| argument == "--file=-"));
383        assert_eq!(stdin.as_deref(), Some(b"SELECT 1;".as_slice()));
384
385        let (command, stdin) = psql_invocation(
386            "postgresql://localhost/postgres",
387            PsqlOptions::new().command("SELECT 1"),
388        );
389        assert!(!command.iter().any(|argument| argument == "--file=-"));
390        assert!(
391            command
392                .windows(2)
393                .any(|arguments| arguments == ["--command", "SELECT 1"])
394        );
395        assert!(stdin.is_none());
396    }
397}
398
399// Generated release-only native target guard.
400#[cfg(not(any(all(target_os = "linux", target_arch = "aarch64", target_env = "gnu"), all(target_os = "linux", target_arch = "x86_64", target_env = "gnu"), all(target_os = "macos", target_arch = "aarch64"), all(target_os = "windows", target_arch = "x86_64", target_env = "msvc"))))]
401compile_error!("oliphaunt-tools supports only linux-arm64-gnu, linux-x64-gnu, macos-arm64, windows-x64-msvc; use one of these declared native targets; this package has no portable fallback.");