1use std::collections::BTreeSet;
24use std::io::{self, Write};
25use std::process::ExitCode;
26
27use kevy_resp_client::RespClient;
28
29pub enum Source {
31 Index(String),
33 Prefix {
35 prefix: String,
37 keep: bool,
44 },
45 File(String),
47}
48
49impl Source {
50 pub fn label(&self) -> String {
52 match self {
53 Source::Index(k) => format!("index {k}"),
54 Source::Prefix { prefix, keep } => {
55 format!("prefix {prefix}{}", if *keep { " (whole keys)" } else { "" })
56 }
57 Source::File(p) => format!("file {p}"),
58 }
59 }
60}
61
62pub struct SourceReport {
64 pub label: String,
66 pub total: usize,
68 pub unique: usize,
71}
72
73pub struct Union {
75 pub names: Vec<Vec<u8>>,
77 pub sources: Vec<SourceReport>,
79}
80
81pub fn collect(client: &mut RespClient, sources: &[Source]) -> io::Result<Union> {
83 let mut per_source: Vec<BTreeSet<Vec<u8>>> = Vec::with_capacity(sources.len());
84 let mut names: Vec<Vec<u8>> = Vec::new();
85 let mut seen: BTreeSet<Vec<u8>> = BTreeSet::new();
86 for s in sources {
87 let got = read_source(client, s)?;
88 for n in &got {
89 if seen.insert(n.clone()) {
90 names.push(n.clone());
91 }
92 }
93 per_source.push(got.into_iter().collect());
94 }
95 let labels: Vec<String> = sources.iter().map(Source::label).collect();
96 Ok(Union { names, sources: account(&labels, &per_source) })
97}
98
99fn account(labels: &[String], per_source: &[BTreeSet<Vec<u8>>]) -> Vec<SourceReport> {
104 labels
105 .iter()
106 .zip(per_source)
107 .enumerate()
108 .map(|(i, (label, mine))| SourceReport {
109 label: label.clone(),
110 total: mine.len(),
111 unique: mine
112 .iter()
113 .filter(|n| !per_source.iter().enumerate().any(|(j, o)| j != i && o.contains(*n)))
114 .count(),
115 })
116 .collect()
117}
118
119fn read_source(client: &mut RespClient, s: &Source) -> io::Result<Vec<Vec<u8>>> {
120 match s {
121 Source::Index(key) => crate::collections::members(client, key),
122 Source::Prefix { prefix, keep } => read_prefix(client, prefix, *keep),
123 Source::File(path) => Ok(std::fs::read_to_string(path)?
124 .lines()
125 .map(str::trim)
126 .filter(|l| !l.is_empty())
127 .map(|l| l.as_bytes().to_vec())
128 .collect()),
129 }
130}
131
132fn read_prefix(client: &mut RespClient, prefix: &str, keep: bool) -> io::Result<Vec<Vec<u8>>> {
136 Ok(crate::collections::scan_prefix(client, prefix)?
137 .into_iter()
138 .map(|k| if keep { k } else { k[prefix.len().min(k.len())..].to_vec() })
139 .collect())
140}
141
142pub fn print_report(u: &Union) {
144 let e = io::stderr();
145 let mut e = e.lock();
146 let _ = writeln!(e, "{} name(s) in the union", u.names.len());
147 for s in &u.sources {
148 let _ = writeln!(e, " {:<32} {} name(s), {} only here", s.label, s.total, s.unique);
149 }
150 let missed: usize = u.sources.iter().map(|s| s.unique).sum();
151 let _ = if missed == 0 && u.sources.len() > 1 {
152 writeln!(e, "every source named the same items — no source alone would have missed a row")
153 } else if u.sources.len() > 1 {
154 writeln!(
155 e,
156 "{missed} name(s) appear in only one source — backfilling from any single one \
157 would have missed them"
158 )
159 } else {
160 writeln!(e, "one source given; there is nothing to union it against")
161 };
162}
163
164fn parse_args(args: &[String]) -> (String, u16, Vec<Source>) {
166 let (mut host, mut port) = (crate::DEFAULT_HOST.to_string(), crate::DEFAULT_PORT);
167 let (mut sources, mut keep) = (Vec::new(), false);
168 let mut i = 0;
169 while i < args.len() {
170 let has_val = i + 1 < args.len();
171 match args[i].as_str() {
172 "-h" if has_val => {
173 host = args[i + 1].clone();
174 i += 2;
175 }
176 "-p" if has_val => {
177 port = args[i + 1].parse().unwrap_or(crate::DEFAULT_PORT);
178 i += 2;
179 }
180 "--keep-prefix" => {
181 keep = true;
182 i += 1;
183 }
184 "--from-index" if has_val => {
185 sources.push(Source::Index(args[i + 1].clone()));
186 i += 2;
187 }
188 "--from-prefix" if has_val => {
189 sources.push(Source::Prefix { prefix: args[i + 1].clone(), keep: false });
190 i += 2;
191 }
192 "--from-file" if has_val => {
193 sources.push(Source::File(args[i + 1].clone()));
194 i += 2;
195 }
196 _ => i += 1,
197 }
198 }
199 if keep {
200 for s in &mut sources {
201 if let Source::Prefix { keep: k, .. } = s {
202 *k = true;
203 }
204 }
205 }
206 (host, port, sources)
207}
208
209pub fn run_backfill_keys_cli(args: &[String]) -> ExitCode {
212 let (host, port, sources) = parse_args(args);
213 if sources.is_empty() {
214 eprintln!("kevy-cli backfill-keys: give at least one source");
215 eprintln!(
216 "usage: kevy-cli backfill-keys [-h host] [-p port] \
217 [--from-index <key>] [--from-prefix <p> [--keep-prefix]] [--from-file <path>] …"
218 );
219 return ExitCode::FAILURE;
220 }
221 emit(&host, port, &sources)
222}
223
224fn emit(host: &str, port: u16, sources: &[Source]) -> ExitCode {
228 let mut client = match RespClient::connect(host, port) {
229 Ok(c) => c,
230 Err(e) => {
231 eprintln!("kevy-cli backfill-keys: could not connect to {host}:{port}: {e}");
232 return ExitCode::FAILURE;
233 }
234 };
235 let u = match collect(&mut client, sources) {
236 Ok(u) => u,
237 Err(e) => {
238 eprintln!("kevy-cli backfill-keys: {e}");
239 return ExitCode::FAILURE;
240 }
241 };
242 let out = io::stdout();
243 let mut out = out.lock();
244 for n in &u.names {
245 let _ = out.write_all(n);
246 let _ = out.write_all(b"\n");
247 }
248 let _ = out.flush();
249 print_report(&u);
250 ExitCode::SUCCESS
251}
252
253#[cfg(test)]
254mod tests {
255 use super::*;
256
257 fn set(names: &[&str]) -> BTreeSet<Vec<u8>> {
258 names.iter().map(|n| n.as_bytes().to_vec()).collect()
259 }
260
261 fn labels(n: usize) -> Vec<String> {
262 (0..n).map(|i| format!("s{i}")).collect()
263 }
264
265 #[test]
269 fn unique_means_no_other_source_named_it() {
270 let sources =
271 [set(&["1", "2", "3"]), set(&["3", "4", "7"]), set(&["1", "2", "3", "4", "5"])];
272 let r = account(&labels(3), &sources);
273 assert_eq!((r[0].total, r[0].unique), (3, 0), "all of s0 is covered elsewhere");
274 assert_eq!((r[1].total, r[1].unique), (3, 1), "only s1 names 7");
275 assert_eq!((r[2].total, r[2].unique), (5, 1), "only s2 names 5");
276 }
277
278 #[test]
281 fn sources_that_agree_have_nothing_unique() {
282 let sources = [set(&["a", "b"]), set(&["b", "a"])];
283 for r in account(&labels(2), &sources) {
284 assert_eq!(r.unique, 0);
285 }
286 }
287
288 #[test]
292 fn a_lone_source_owns_everything_it_names() {
293 let r = account(&labels(1), &[set(&["a", "b"])]);
294 assert_eq!((r[0].total, r[0].unique), (2, 2));
295 }
296}