1use crate::core::system::System;
2use crate::db::config::Config;
3use crate::error::MyResult;
4use crate::regex;
5use clap::Command;
6use clap_complete::Shell;
7use std::io::{Read, Write};
8
9pub fn generate_completion<F: Read + Write, S: System<F>, W: Write>(
10 system: S,
11 writer: &mut W,
12 command: &mut Command,
13 shell: Shell,
14) -> MyResult<()> {
15 if shell == Shell::Bash {
16 let mut buffer = Vec::new();
17 clap_complete::generate(shell, command, String::from("zql"), &mut buffer);
18 #[cfg(windows)]
19 clap_complete::generate(shell, command, String::from("zql.exe"), &mut buffer);
20 let aliases = Config::new(system)?.with_config()?.get_aliases();
21 let buffer = transform_completion(buffer, aliases);
22 writer.write_all(&buffer)?;
23 return Ok(());
24 }
25 clap_complete::generate(shell, command, String::from("zql"), writer);
26 Ok(())
27}
28
29fn transform_completion(buffer: Vec<u8>, aliases: Vec<String>) -> Vec<u8> {
30 let mut lines = String::from_utf8(buffer)
31 .unwrap_or_default()
32 .lines()
33 .map(str::to_string)
34 .collect::<Vec<_>>();
35 let alias_regex = regex!(r"<ALIAS>");
36 let alias_replace = aliases.join(" ");
37 let compgen_regex = regex!(r"\bcompgen -f\b");
38 let compgen_replace = format!("compgen -W \"{}\" --", alias_replace);
39 let mut found = false;
40 for index in 0..lines.len() {
41 lines[index] = alias_regex.replace(&lines[index], &alias_replace).to_string();
42 if lines[index].trim() == "--config)" {
43 found = true;
44 } else if found {
45 lines[index] = compgen_regex.replace(&lines[index], &compgen_replace).to_string();
46 found = false;
47 }
48 }
49 let mut buffer = lines.join("\n");
50 buffer.push_str("\n");
51 buffer.into_bytes()
52}
53
54#[cfg(test)]
55mod tests {
56 use crate::cli::Cli;
57 use crate::core::system::tests::MockSystem;
58 use crate::error::MyResult;
59 use crate::util::complete::generate_completion;
60 use clap::CommandFactory;
61 use clap_complete::Shell;
62 use itertools::Itertools;
63 use pretty_assertions::assert_eq;
64
65 #[test]
66 fn test_bash_completion_is_generated_for_zero_drivers() -> MyResult<()> {
67 let lines = get_lines(Shell::Bash, "")?;
68 let add_command = find_line(&lines, "zql__config__add)");
69 let remove_command = find_line(&lines, "zql__config__remove)");
70 let clear_command = find_line(&lines, "zql__config__clear)");
71 let config_option = find_line(&lines, "--config)");
72 assert_eq!(add_command.contains("--version [ODBC]\""), true, "{}", add_command);
73 assert_eq!(remove_command.contains("--version \""), true, "{}", remove_command);
74 assert_eq!(clear_command.contains("--version\""), true, "{}", clear_command);
75 assert_eq!(config_option, "COMPREPLY=($(compgen -W \"\" -- \"${cur}\"))");
76 Ok(())
77 }
78
79 #[test]
80 fn test_bash_completion_is_generated_for_three_drivers() -> MyResult<()> {
81 let source = "\
82{
83 \"dsn\": { \"default\": false, \"odbc\": { } },
84 \"mysql\": { \"default\": false, \"odbc\": { } },
85 \"sqlite\": { \"default\": false, \"odbc\": { } }
86}
87";
88 let lines = get_lines(Shell::Bash, source)?;
89 let add_command = find_line(&lines, "zql__config__add)");
90 let remove_command = find_line(&lines, "zql__config__remove)");
91 let clear_command = find_line(&lines, "zql__config__clear)");
92 let config_option = find_line(&lines, "--config)");
93 assert_eq!(add_command.contains("--version dsn mysql sqlite [ODBC]\""), true, "{}", add_command);
94 assert_eq!(remove_command.contains("--version dsn mysql sqlite\""), true, "{}", remove_command);
95 assert_eq!(clear_command.contains("--version\""), true, "{}", clear_command);
96 assert_eq!(config_option, "COMPREPLY=($(compgen -W \"dsn mysql sqlite\" -- \"${cur}\"))");
97 Ok(())
98 }
99
100 fn get_lines(shell: Shell, source: &str) -> MyResult<Vec<String>> {
101 let system = MockSystem::new(source);
102 let mut buffer = Vec::new();
103 let mut command = Cli::command();
104 generate_completion(system, &mut buffer, &mut command, shell)?;
105 let lines = String::from_utf8(buffer)
106 .unwrap_or_default()
107 .lines()
108 .map(str::trim)
109 .map(str::to_string)
110 .collect();
111 Ok(lines)
112 }
113
114 fn find_line(lines: &[String], text: &str) -> String {
115 for (first, second) in lines.iter().tuple_windows() {
116 if first.trim() == text {
117 return second.trim().to_string()
118 }
119 }
120 String::new()
121 }
122}