kevy_cli/
migrate.rs

1//! **v2.10** — migration toolchain: `export` / `import` (RFC D1/D2).
2//!
3//! Wire format = a RESP command stream of rebuild frames (SET / HSET /
4//! RPUSH / SADD / ZADD / PEXPIREAT) — bidirectionally compatible with
5//! `redis-cli --pipe`. Export walks SCAN cursors (per-key
6//! point-in-time; SCAN-class consistency). Import pipelines 512
7//! commands per batch with a fsynced progress file for `--resume`.
8//! Every key's frames start with DEL, so replay REBUILDS the key from
9//! scratch — genuinely idempotent for every type (RPUSH would
10//! otherwise append on re-import; the round-trip test caught exactly
11//! that). No cross-batch atomicity to lose.
12
13use std::fs::{File, OpenOptions};
14use std::io::{self, BufWriter, Read, Seek, SeekFrom, Write};
15use std::path::Path;
16
17use kevy_resp::{Reply, encode_command_borrowed};
18use kevy_resp_client::RespClient;
19
20const PIPELINE: usize = 512;
21
22/// Run `export` — walk the keyspace (optionally under `prefix`) and
23/// write rebuild frames to `out_path`. Returns exported key count.
24pub fn run_export(
25    client: &mut RespClient,
26    prefix: Option<&[u8]>,
27    out_path: &Path,
28) -> io::Result<u64> {
29    let mut out = BufWriter::new(File::create(out_path)?);
30    let mut cursor: Vec<u8> = b"0".to_vec();
31    let mut pattern = prefix.unwrap_or_default().to_vec();
32    pattern.push(b'*');
33    let mut n = 0u64;
34    loop {
35        let reply = client.request_borrowed(&[b"SCAN", &cursor, b"MATCH", &pattern, b"COUNT", b"512"])?;
36        let Reply::Array(items) = reply else {
37            return Err(io::Error::new(io::ErrorKind::InvalidData, "SCAN reply shape"));
38        };
39        let (Some(Reply::Bulk(next)), Some(Reply::Array(keys))) = (items.first(), items.get(1))
40        else {
41            return Err(io::Error::new(io::ErrorKind::InvalidData, "SCAN reply shape"));
42        };
43        let next = next.clone();
44        for k in keys {
45            let Reply::Bulk(key) = k else { continue };
46            let key = key.clone();
47            if export_key(client, &key, &mut out)? {
48                n += 1;
49            }
50        }
51        cursor = next;
52        if cursor == b"0" {
53            break;
54        }
55    }
56    out.flush()?;
57    Ok(n)
58}
59
60/// Emit one key's rebuild frames. Returns false if the key vanished
61/// between SCAN and read (point-in-time per key).
62fn export_key(client: &mut RespClient, key: &[u8], out: &mut impl Write) -> io::Result<bool> {
63    match rebuild_frames(client, key, key)? {
64        Some(frame) => {
65            out.write_all(&frame)?;
66            Ok(true)
67        }
68        None => Ok(false),
69    }
70}
71
72/// Read `key` and produce DEL+rebuild frames addressed to `dst`
73/// (`dst == key` for export; a re-prefixed name for copy-prefix —
74/// the server has no COPY verb, so copying IS read+rebuild).
75pub(crate) fn rebuild_frames(
76    client: &mut RespClient,
77    key: &[u8],
78    dst: &[u8],
79) -> io::Result<Option<Vec<u8>>> {
80    let ty = match client.request_borrowed(&[b"TYPE", key])? {
81        Reply::Simple(t) => t,
82        _ => return Ok(None),
83    };
84    let mut frame = Vec::new();
85    // DEL first: replay rebuilds from scratch (idempotence for
86    // append-shaped verbs like RPUSH).
87    encode_command_borrowed(&mut frame, &[b"DEL", dst]);
88    match ty.as_slice() {
89        b"string" => {
90            let Reply::Bulk(v) = client.request_borrowed(&[b"GET", key])? else {
91                return Ok(None);
92            };
93            encode_command_borrowed(&mut frame, &[b"SET", dst, &v]);
94        }
95        b"hash" => {
96            let Reply::Array(flat) = client.request_borrowed(&[b"HGETALL", key])? else {
97                return Ok(None);
98            };
99            if flat.is_empty() {
100                return Ok(None);
101            }
102            let mut argv: Vec<&[u8]> = vec![b"HSET", dst];
103            let items: Vec<Vec<u8>> = flat
104                .into_iter()
105                .filter_map(|r| if let Reply::Bulk(b) = r { Some(b) } else { None })
106                .collect();
107            argv.extend(items.iter().map(Vec::as_slice));
108            encode_command_borrowed(&mut frame, &argv);
109        }
110        b"list" => {
111            let Reply::Array(items) = client.request_borrowed(&[b"LRANGE", key, b"0", b"-1"])?
112            else {
113                return Ok(None);
114            };
115            if items.is_empty() {
116                return Ok(None);
117            }
118            let vals: Vec<Vec<u8>> = items
119                .into_iter()
120                .filter_map(|r| if let Reply::Bulk(b) = r { Some(b) } else { None })
121                .collect();
122            let mut argv: Vec<&[u8]> = vec![b"RPUSH", dst];
123            argv.extend(vals.iter().map(Vec::as_slice));
124            encode_command_borrowed(&mut frame, &argv);
125        }
126        b"set" => {
127            let Reply::Array(items) = client.request_borrowed(&[b"SMEMBERS", key])? else {
128                return Ok(None);
129            };
130            if items.is_empty() {
131                return Ok(None);
132            }
133            let ms: Vec<Vec<u8>> = items
134                .into_iter()
135                .filter_map(|r| if let Reply::Bulk(b) = r { Some(b) } else { None })
136                .collect();
137            let mut argv: Vec<&[u8]> = vec![b"SADD", dst];
138            argv.extend(ms.iter().map(Vec::as_slice));
139            encode_command_borrowed(&mut frame, &argv);
140        }
141        b"zset" => {
142            let Reply::Array(items) =
143                client.request_borrowed(&[b"ZRANGE", key, b"0", b"-1", b"WITHSCORES"])?
144            else {
145                return Ok(None);
146            };
147            if items.is_empty() {
148                return Ok(None);
149            }
150            let flat: Vec<Vec<u8>> = items
151                .into_iter()
152                .filter_map(|r| if let Reply::Bulk(b) = r { Some(b) } else { None })
153                .collect();
154            // ZADD wants score member; ZRANGE gives member score
155            let mut argv: Vec<&[u8]> = vec![b"ZADD", dst];
156            for pair in flat.chunks(2) {
157                if pair.len() == 2 {
158                    argv.push(&pair[1]);
159                    argv.push(&pair[0]);
160                }
161            }
162            encode_command_borrowed(&mut frame, &argv);
163        }
164        _ => return Ok(None), // streams etc. — out of the rebuild set
165    }
166    // TTL rides as an absolute PEXPIREAT follow-up
167    if let Reply::Int(ms) = client.request_borrowed(&[b"PTTL", key])?
168        && ms > 0
169    {
170        let now = std::time::SystemTime::now()
171            .duration_since(std::time::UNIX_EPOCH)
172            .map_err(io::Error::other)?
173            .as_millis() as i64;
174        encode_command_borrowed(
175            &mut frame,
176            &[b"PEXPIREAT", dst, (now + ms).to_string().as_bytes()],
177        );
178    }
179    Ok(Some(frame))
180}
181
182/// Import stats.
183pub struct ImportReport {
184    /// Commands sent successfully.
185    pub sent: u64,
186    /// -ERR replies (counted, not fatal unless `strict`).
187    pub errors: u64,
188    /// Byte offset reached in the source file.
189    pub offset: u64,
190}
191
192/// Run `import` — stream `src` (a RESP command file) into the server,
193/// `PIPELINE` commands per batch. The progress file `<src>.progress`
194/// records the safely-applied byte offset after every batch (fsynced);
195/// `resume` starts there. Idempotent replay.
196pub fn run_import(
197    client: &mut RespClient,
198    src: &Path,
199    resume: bool,
200    strict: bool,
201) -> io::Result<ImportReport> {
202    let progress_path = src.with_extension("progress");
203    let mut start = 0u64;
204    if resume && let Ok(text) = std::fs::read_to_string(&progress_path) {
205        start = text.trim().parse().unwrap_or(0);
206    }
207    let mut f = File::open(src)?;
208    f.seek(SeekFrom::Start(start))?;
209    let mut pending: Vec<u8> = Vec::with_capacity(1 << 20);
210    let mut report = ImportReport { sent: 0, errors: 0, offset: start };
211    let mut chunk = vec![0u8; 1 << 20];
212    let mut batch_bytes = 0usize;
213    let mut batch_cmds = 0usize;
214    let mut progress = OpenOptions::new().create(true).truncate(false).write(true).open(&progress_path)?;
215    loop {
216        let n = f.read(&mut chunk)?;
217        if n == 0 {
218            break;
219        }
220        pending.extend_from_slice(&chunk[..n]);
221        // carve complete commands off `pending`
222        while let Some(used) = command_len(&pending[batch_bytes..]) {
223            batch_bytes += used;
224            batch_cmds += 1;
225            if batch_cmds == PIPELINE {
226                flush_batch(client, &pending[..batch_bytes], batch_cmds, strict, &mut report)?;
227                pending.drain(..batch_bytes);
228                write_progress(&mut progress, report.offset)?;
229                batch_bytes = 0;
230                batch_cmds = 0;
231            }
232        }
233    }
234    if batch_cmds > 0 {
235        flush_batch(client, &pending[..batch_bytes], batch_cmds, strict, &mut report)?;
236        write_progress(&mut progress, report.offset)?;
237    }
238    Ok(report)
239}
240
241fn flush_batch(
242    client: &mut RespClient,
243    raw: &[u8],
244    n: usize,
245    strict: bool,
246    report: &mut ImportReport,
247) -> io::Result<()> {
248    let replies = client.pipeline_raw(raw, n)?;
249    for r in replies {
250        if let Reply::Error(e) = r {
251            report.errors += 1;
252            if strict {
253                return Err(io::Error::new(
254                    io::ErrorKind::InvalidData,
255                    format!("server error (strict): {}", String::from_utf8_lossy(&e)),
256                ));
257            }
258        } else {
259            report.sent += 1;
260        }
261    }
262    report.offset += raw.len() as u64;
263    Ok(())
264}
265
266fn write_progress(f: &mut File, offset: u64) -> io::Result<()> {
267    f.set_len(0)?;
268    f.seek(SeekFrom::Start(0))?;
269    f.write_all(offset.to_string().as_bytes())?;
270    f.sync_data()
271}
272
273/// Crate-visible alias for [`command_len`] (bulk copy counts frames).
274pub(crate) fn command_len_pub(b: &[u8]) -> Option<usize> {
275    command_len(b)
276}
277
278/// Length of one complete RESP command at the head of `b`, or `None`.
279fn command_len(b: &[u8]) -> Option<usize> {
280    let mut pos = 0usize;
281    let line = take_line(b, &mut pos)?;
282    if line.first() != Some(&b'*') {
283        return None;
284    }
285    let n: usize = std::str::from_utf8(&line[1..]).ok()?.trim().parse().ok()?;
286    for _ in 0..n {
287        let hdr = take_line(b, &mut pos)?;
288        if hdr.first() != Some(&b'$') {
289            return None;
290        }
291        let len: usize = std::str::from_utf8(&hdr[1..]).ok()?.trim().parse().ok()?;
292        if b.len() < pos + len + 2 {
293            return None;
294        }
295        pos += len + 2;
296    }
297    Some(pos)
298}
299
300fn take_line<'b>(b: &'b [u8], pos: &mut usize) -> Option<&'b [u8]> {
301    let rest = &b[*pos..];
302    let idx = rest.windows(2).position(|w| w == b"\r\n")?;
303    let line = &rest[..idx];
304    *pos += idx + 2;
305    Some(line)
306}
kevy_cli/migrate.rs

kevy_cli/
migrate.rs
