1use crate::cli::{PsqlAdminAction, PsqlAdminRequest};
2use serde::Serialize;
3use serde_json::{Value, json};
4use std::path::{Path, PathBuf};
5
6const MARKER: &str = "afpsql-managed: true";
7const GENERATED_BY: &str = "Generated by afpsql psql install";
8
9pub fn run(req: PsqlAdminRequest) -> i32 {
10 let result = handle_action(req.action);
11 let mut emitter =
12 agent_first_data::CliEmitter::from_output_to(crate::emit::output_to(), req.output)
13 .with_strict_protocol();
14 match result {
15 Ok(value) => match emitter.emit_result(value) {
16 Ok(()) => 0,
17 Err(_) => 4,
18 },
19 Err(err) => match agent_first_data::json_error("invalid_request", &err.message)
20 .hint_if_some(err.hint.as_deref())
21 .build()
22 .map_err(agent_first_data::CliEmitterError::Build)
23 {
24 Ok(event) => match emitter.emit(event) {
25 Ok(()) => 1,
26 Err(_) => 4,
27 },
28 Err(_) => 4,
29 },
30 }
31}
32
33pub(crate) fn handle_action(action: PsqlAdminAction) -> Result<Value, AdminError> {
34 match action {
35 PsqlAdminAction::Status { bin_dir } => status(bin_dir),
36 PsqlAdminAction::Install { bin_dir } => install(bin_dir),
37 PsqlAdminAction::Uninstall { bin_dir } => uninstall(bin_dir),
38 }
39}
40
41fn status(bin_dir: Option<String>) -> Result<Value, AdminError> {
42 let target = resolve_target(bin_dir)?;
43 let managed = is_managed_wrapper(&target.wrapper_path).unwrap_or(false);
44 let installed = target.wrapper_path.is_file();
45 let current_psql = find_on_path(wrapper_file_name())?;
46 let active_in_path = current_psql
47 .as_ref()
48 .map(|p| same_path(p, &target.wrapper_path))
49 .unwrap_or(false);
50
51 Ok(admin_result(
52 "psql_status",
53 json!({
54 "wrapper_path": target.wrapper_path,
55 "bin_dir": target.bin_dir,
56 "installed": installed,
57 "managed": managed,
58 "active_in_path": active_in_path,
59 "current_psql": current_psql,
60 "afpsql": current_afpsql_path()?,
61 }),
62 ))
63}
64
65fn install(bin_dir: Option<String>) -> Result<Value, AdminError> {
66 let target = resolve_target(bin_dir)?;
67 std::fs::create_dir_all(&target.bin_dir).map_err(|e| AdminError::io("create bin dir", e))?;
68
69 if target.wrapper_path.exists() && !is_managed_wrapper(&target.wrapper_path)? {
70 return Err(AdminError::invalid_request(
71 format!(
72 "refusing to overwrite unmanaged psql at {}",
73 target.wrapper_path.display()
74 ),
75 Some("choose another --bin-dir or remove the existing psql yourself".to_string()),
76 ));
77 }
78
79 let afpsql = current_afpsql_path()?;
80 std::fs::write(&target.wrapper_path, wrapper_contents(&afpsql))
81 .map_err(|e| AdminError::io("write psql wrapper", e))?;
82 set_executable(&target.wrapper_path)?;
83
84 let current_psql = find_on_path(wrapper_file_name())?;
85 let active_in_path = current_psql
86 .as_ref()
87 .map(|p| same_path(p, &target.wrapper_path))
88 .unwrap_or(false);
89
90 Ok(admin_result(
91 "psql_install",
92 json!({
93 "wrapper_path": target.wrapper_path,
94 "bin_dir": target.bin_dir,
95 "installed": true,
96 "managed": true,
97 "active_in_path": active_in_path,
98 "current_psql": current_psql,
99 "afpsql": afpsql,
100 "hint": if active_in_path { Value::Null } else { json!(format!("add {} to the front of PATH for this psql to take effect", target.bin_dir.display())) },
101 }),
102 ))
103}
104
105fn uninstall(bin_dir: Option<String>) -> Result<Value, AdminError> {
106 let target = resolve_target(bin_dir)?;
107 if !target.wrapper_path.exists() {
108 return Ok(admin_result(
109 "psql_uninstall",
110 json!({
111 "wrapper_path": target.wrapper_path,
112 "bin_dir": target.bin_dir,
113 "removed": false,
114 "message": "psql wrapper is not installed"
115 }),
116 ));
117 }
118 if !is_managed_wrapper(&target.wrapper_path)? {
119 return Err(AdminError::invalid_request(
120 format!(
121 "refusing to remove unmanaged psql at {}",
122 target.wrapper_path.display()
123 ),
124 Some("only wrappers generated by afpsql psql install can be removed".to_string()),
125 ));
126 }
127
128 std::fs::remove_file(&target.wrapper_path)
129 .map_err(|e| AdminError::io("remove psql wrapper", e))?;
130 Ok(admin_result(
131 "psql_uninstall",
132 json!({
133 "wrapper_path": target.wrapper_path,
134 "bin_dir": target.bin_dir,
135 "removed": true
136 }),
137 ))
138}
139
140fn admin_result(code: &str, payload: Value) -> Value {
141 let mut payload = payload;
142 if let Value::Object(fields) = &mut payload {
143 fields.insert("code".to_string(), Value::String(code.to_string()));
144 }
145 payload
146}
147
148struct TargetPath {
149 bin_dir: PathBuf,
150 wrapper_path: PathBuf,
151}
152
153fn resolve_target(bin_dir: Option<String>) -> Result<TargetPath, AdminError> {
154 let bin_dir = match bin_dir {
155 Some(dir) => expand_tilde(&dir)?,
156 None => default_bin_dir()?,
157 };
158 let wrapper_path = bin_dir.join(wrapper_file_name());
159 Ok(TargetPath {
160 bin_dir,
161 wrapper_path,
162 })
163}
164
165fn default_bin_dir() -> Result<PathBuf, AdminError> {
166 let afpsql = current_afpsql_path()?;
167 afpsql.parent().map(Path::to_path_buf).ok_or_else(|| {
168 AdminError::invalid_request(
169 "cannot determine afpsql executable directory".to_string(),
170 Some("pass --bin-dir explicitly".to_string()),
171 )
172 })
173}
174
175fn home_dir() -> Result<PathBuf, AdminError> {
176 std::env::var_os("HOME")
177 .or_else(|| std::env::var_os("USERPROFILE"))
178 .map(PathBuf::from)
179 .ok_or_else(|| {
180 AdminError::invalid_request(
181 "cannot determine home directory".to_string(),
182 Some("pass --bin-dir explicitly".to_string()),
183 )
184 })
185}
186
187fn expand_tilde(input: &str) -> Result<PathBuf, AdminError> {
188 if input == "~" {
189 return home_dir();
190 }
191 if let Some(rest) = input.strip_prefix("~/") {
192 return Ok(home_dir()?.join(rest));
193 }
194 Ok(PathBuf::from(input))
195}
196
197fn wrapper_file_name() -> &'static str {
198 if cfg!(windows) { "psql.cmd" } else { "psql" }
199}
200
201fn current_afpsql_path() -> Result<PathBuf, AdminError> {
202 std::env::current_exe().map_err(|e| AdminError::io("resolve current afpsql path", e))
203}
204
205fn wrapper_contents(afpsql: &Path) -> String {
206 let afpsql = afpsql.to_string_lossy();
207 if cfg!(windows) {
208 format!(
209 "@echo off\r\nREM {GENERATED_BY}\r\nREM {MARKER}\r\n\"{}\" --mode psql %*\r\n",
210 afpsql.replace('"', "\"\"")
211 )
212 } else {
213 format!(
214 "#!/usr/bin/env bash\n# {GENERATED_BY}\n# {MARKER}\nexec \"{}\" --mode psql \"$@\"\n",
215 shell_double_quote_escape(&afpsql)
216 )
217 }
218}
219
220fn shell_double_quote_escape(value: &str) -> String {
221 value
222 .replace('\\', "\\\\")
223 .replace('"', "\\\"")
224 .replace('$', "\\$")
225 .replace('`', "\\`")
226}
227
228fn is_managed_wrapper(path: &Path) -> Result<bool, AdminError> {
229 if !path.exists() {
230 return Ok(false);
231 }
232 let text = std::fs::read_to_string(path).map_err(|e| AdminError::io("read psql wrapper", e))?;
233 Ok(text.contains(MARKER) && text.contains(GENERATED_BY))
234}
235
236#[cfg(unix)]
237fn set_executable(path: &Path) -> Result<(), AdminError> {
238 use std::os::unix::fs::PermissionsExt;
239 let mut permissions = std::fs::metadata(path)
240 .map_err(|e| AdminError::io("read wrapper metadata", e))?
241 .permissions();
242 permissions.set_mode(0o755);
243 std::fs::set_permissions(path, permissions)
244 .map_err(|e| AdminError::io("set wrapper executable bit", e))
245}
246
247#[cfg(not(unix))]
248fn set_executable(_path: &Path) -> Result<(), AdminError> {
249 Ok(())
250}
251
252fn find_on_path(name: &str) -> Result<Option<PathBuf>, AdminError> {
253 let Some(paths) = std::env::var_os("PATH") else {
254 return Ok(None);
255 };
256 for dir in std::env::split_paths(&paths) {
257 for candidate_name in path_candidate_names(name) {
258 let candidate = dir.join(candidate_name);
259 if candidate.is_file() {
260 return Ok(Some(candidate));
261 }
262 }
263 }
264 Ok(None)
265}
266
267fn path_candidate_names(name: &str) -> Vec<String> {
268 if !cfg!(windows) {
269 return vec![name.to_string()];
270 }
271 let lower = name.to_ascii_lowercase();
272 if lower.ends_with(".exe") || lower.ends_with(".cmd") || lower.ends_with(".bat") {
273 return vec![name.to_string()];
274 }
275 vec![
276 name.to_string(),
277 format!("{name}.exe"),
278 format!("{name}.cmd"),
279 format!("{name}.bat"),
280 ]
281}
282
283fn same_path(a: &Path, b: &Path) -> bool {
284 match (a.canonicalize(), b.canonicalize()) {
285 (Ok(a), Ok(b)) => a == b,
286 _ => a == b,
287 }
288}
289
290#[derive(Debug, Serialize)]
291pub(crate) struct AdminError {
292 message: String,
293 #[serde(skip_serializing_if = "Option::is_none")]
294 hint: Option<String>,
295}
296
297impl AdminError {
298 fn invalid_request(message: String, hint: Option<String>) -> Self {
299 Self { message, hint }
300 }
301
302 fn io(action: &str, err: std::io::Error) -> Self {
303 Self {
304 message: format!("{action} failed: {err}"),
305 hint: None,
306 }
307 }
308}
309
310#[cfg(test)]
311mod tests {
312 use super::*;
313 use std::time::{SystemTime, UNIX_EPOCH};
314
315 fn temp_bin_dir(name: &str) -> PathBuf {
316 let suffix = SystemTime::now()
317 .duration_since(UNIX_EPOCH)
318 .map(|d| d.as_nanos())
319 .unwrap_or(0);
320 std::env::temp_dir().join(format!("afpsql_{name}_{}_{}", std::process::id(), suffix))
321 }
322
323 #[test]
324 fn default_bin_dir_uses_current_exe_parent() {
325 let expected = current_afpsql_path()
326 .ok()
327 .and_then(|p| p.parent().map(Path::to_path_buf));
328 assert_eq!(default_bin_dir().ok(), expected);
329 }
330
331 #[test]
332 fn install_status_uninstall_managed_wrapper() {
333 let dir = temp_bin_dir("shim");
334 let dir_arg = dir.to_string_lossy().to_string();
335
336 let installed = handle_action(PsqlAdminAction::Install {
337 bin_dir: Some(dir_arg.clone()),
338 });
339 assert!(installed.is_ok());
340 let wrapper = dir.join(wrapper_file_name());
341 assert!(wrapper.is_file());
342 assert!(is_managed_wrapper(&wrapper).unwrap_or(false));
343
344 let status = handle_action(PsqlAdminAction::Status {
345 bin_dir: Some(dir_arg.clone()),
346 });
347 assert!(status.is_ok());
348 if let Ok(value) = status {
349 assert_eq!(value["installed"], true);
350 assert_eq!(value["managed"], true);
351 }
352
353 let removed = handle_action(PsqlAdminAction::Uninstall {
354 bin_dir: Some(dir_arg),
355 });
356 assert!(removed.is_ok());
357 assert!(!wrapper.exists());
358 let _ = std::fs::remove_dir_all(dir);
359 }
360
361 #[test]
362 fn install_and_uninstall_refuse_unmanaged_psql() {
363 let dir = temp_bin_dir("unmanaged");
364 assert!(std::fs::create_dir_all(&dir).is_ok());
365 let wrapper = dir.join(wrapper_file_name());
366 assert!(std::fs::write(&wrapper, "#!/bin/sh\nexit 0\n").is_ok());
367 let dir_arg = dir.to_string_lossy().to_string();
368
369 let install = handle_action(PsqlAdminAction::Install {
370 bin_dir: Some(dir_arg.clone()),
371 });
372 assert!(install.is_err());
373
374 let uninstall = handle_action(PsqlAdminAction::Uninstall {
375 bin_dir: Some(dir_arg),
376 });
377 assert!(uninstall.is_err());
378 assert!(wrapper.exists());
379 let _ = std::fs::remove_dir_all(dir);
380 }
381}