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
15pub const PRODUCT: &str = "oliphaunt-tools";
17
18pub const KIND: &str = "native-tools";
20
21#[derive(Debug, Clone, Default, PartialEq, Eq)]
23pub struct PgDumpOptions {
24 args: Vec<String>,
25}
26
27impl PgDumpOptions {
28 pub fn new() -> Self {
30 Self::default()
31 }
32
33 pub fn arg(mut self, argument: impl Into<String>) -> Self {
35 self.args.push(argument.into());
36 self
37 }
38
39 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#[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 pub fn new() -> Self {
62 Self::default()
63 }
64
65 pub fn arg(mut self, argument: impl Into<String>) -> Self {
67 self.args.push(argument.into());
68 self
69 }
70
71 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 pub fn command(mut self, sql: impl Into<String>) -> Self {
79 self.input = Some(PsqlInput::Command(sql.into()));
80 self
81 }
82
83 pub fn script(mut self, sql: impl Into<String>) -> Self {
85 self.input = Some(PsqlInput::Script(sql.into()));
86 self
87 }
88}
89
90#[derive(Debug)]
92pub struct PostgresToolError {
93 pub tool: &'static str,
95 pub exit_code: Option<i32>,
97 pub stdout: String,
99 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
132pub 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
151pub 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 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#[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.");