kevy_cli/
migrate.rs

1//! Migration toolchain: `export` / `import`.
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 Some(items) = fetch_bulks(client, &[b"HGETALL", key])? else {
97                return Ok(None);
98            };
99            encode_multi(&mut frame, b"HSET", dst, &items);
100        }
101        b"list" => {
102            let Some(vals) = fetch_bulks(client, &[b"LRANGE", key, b"0", b"-1"])? else {
103                return Ok(None);
104            };
105            encode_multi(&mut frame, b"RPUSH", dst, &vals);
106        }
107        b"set" => {
108            let Some(ms) = fetch_bulks(client, &[b"SMEMBERS", key])? else {
109                return Ok(None);
110            };
111            encode_multi(&mut frame, b"SADD", dst, &ms);
112        }
113        b"zset" => {
114            let zrange: &[&[u8]] = &[b"ZRANGE", key, b"0", b"-1", b"WITHSCORES"];
115            let Some(flat) = fetch_bulks(client, zrange)? else {
116                return Ok(None);
117            };
118            encode_zadd(&mut frame, dst, &flat);
119        }
120        _ => return Ok(None), // streams etc. — out of the rebuild set
121    }
122    append_ttl_frame(client, key, dst, &mut frame)?;
123    Ok(Some(frame))
124}
125
126/// Issue `cmd` and unwrap its Array reply into bulk payloads.
127/// `None` when the reply isn't an array or the array is empty (the key
128/// vanished / changed type between TYPE and read).
129fn fetch_bulks(client: &mut RespClient, cmd: &[&[u8]]) -> io::Result<Option<Vec<Vec<u8>>>> {
130    let Reply::Array(items) = client.request_borrowed(cmd)? else {
131        return Ok(None);
132    };
133    if items.is_empty() {
134        return Ok(None);
135    }
136    Ok(Some(
137        items
138            .into_iter()
139            .filter_map(|r| if let Reply::Bulk(b) = r { Some(b) } else { None })
140            .collect(),
141    ))
142}
143
144/// Encode `<verb> <dst> <vals…>` onto `frame` (HSET / RPUSH / SADD).
145fn encode_multi(frame: &mut Vec<u8>, verb: &[u8], dst: &[u8], vals: &[Vec<u8>]) {
146    let mut argv: Vec<&[u8]> = vec![verb, dst];
147    argv.extend(vals.iter().map(Vec::as_slice));
148    encode_command_borrowed(frame, &argv);
149}
150
151/// Encode `ZADD <dst> score member …` onto `frame`.
152/// ZADD wants score member; ZRANGE gives member score.
153fn encode_zadd(frame: &mut Vec<u8>, dst: &[u8], flat: &[Vec<u8>]) {
154    let mut argv: Vec<&[u8]> = vec![b"ZADD", dst];
155    for pair in flat.chunks(2) {
156        if pair.len() == 2 {
157            argv.push(&pair[1]);
158            argv.push(&pair[0]);
159        }
160    }
161    encode_command_borrowed(frame, &argv);
162}
163
164/// TTL rides as an absolute PEXPIREAT follow-up.
165fn append_ttl_frame(
166    client: &mut RespClient,
167    key: &[u8],
168    dst: &[u8],
169    frame: &mut Vec<u8>,
170) -> io::Result<()> {
171    if let Reply::Int(ms) = client.request_borrowed(&[b"PTTL", key])?
172        && ms > 0
173    {
174        let now = std::time::SystemTime::now()
175            .duration_since(std::time::UNIX_EPOCH)
176            .map_err(io::Error::other)?
177            .as_millis() as i64;
178        encode_command_borrowed(
179            frame,
180            &[b"PEXPIREAT", dst, (now + ms).to_string().as_bytes()],
181        );
182    }
183    Ok(())
184}
185
186/// Import stats.
187pub struct ImportReport {
188    /// Commands sent successfully.
189    pub sent: u64,
190    /// -ERR replies (counted, not fatal unless `strict`).
191    pub errors: u64,
192    /// Byte offset reached in the source file.
193    pub offset: u64,
194}
195
196/// Run `import` — stream `src` (a RESP command file) into the server,
197/// `PIPELINE` commands per batch. The progress file `<src>.progress`
198/// records the safely-applied byte offset after every batch (fsynced);
199/// `resume` starts there. Idempotent replay.
200pub fn run_import(
201    client: &mut RespClient,
202    src: &Path,
203    resume: bool,
204    strict: bool,
205) -> io::Result<ImportReport> {
206    let progress_path = src.with_extension("progress");
207    let mut start = 0u64;
208    if resume && let Ok(text) = std::fs::read_to_string(&progress_path) {
209        start = text.trim().parse().unwrap_or(0);
210    }
211    let mut f = File::open(src)?;
212    f.seek(SeekFrom::Start(start))?;
213    let mut pending: Vec<u8> = Vec::with_capacity(1 << 20);
214    let mut report = ImportReport { sent: 0, errors: 0, offset: start };
215    let mut chunk = vec![0u8; 1 << 20];
216    let mut batch_bytes = 0usize;
217    let mut batch_cmds = 0usize;
218    let mut progress = OpenOptions::new().create(true).truncate(false).write(true).open(&progress_path)?;
219    loop {
220        let n = f.read(&mut chunk)?;
221        if n == 0 {
222            break;
223        }
224        pending.extend_from_slice(&chunk[..n]);
225        // carve complete commands off `pending`
226        while let Some(used) = command_len(&pending[batch_bytes..]) {
227            batch_bytes += used;
228            batch_cmds += 1;
229            if batch_cmds == PIPELINE {
230                flush_batch(client, &pending[..batch_bytes], batch_cmds, strict, &mut report)?;
231                pending.drain(..batch_bytes);
232                write_progress(&mut progress, report.offset)?;
233                batch_bytes = 0;
234                batch_cmds = 0;
235            }
236        }
237    }
238    if batch_cmds > 0 {
239        flush_batch(client, &pending[..batch_bytes], batch_cmds, strict, &mut report)?;
240        write_progress(&mut progress, report.offset)?;
241    }
242    Ok(report)
243}
244
245fn flush_batch(
246    client: &mut RespClient,
247    raw: &[u8],
248    n: usize,
249    strict: bool,
250    report: &mut ImportReport,
251) -> io::Result<()> {
252    let replies = client.pipeline_raw(raw, n)?;
253    for r in replies {
254        if let Reply::Error(e) = r {
255            report.errors += 1;
256            if strict {
257                return Err(io::Error::new(
258                    io::ErrorKind::InvalidData,
259                    format!("server error (strict): {}", String::from_utf8_lossy(&e)),
260                ));
261            }
262        } else {
263            report.sent += 1;
264        }
265    }
266    report.offset += raw.len() as u64;
267    Ok(())
268}
269
270fn write_progress(f: &mut File, offset: u64) -> io::Result<()> {
271    f.set_len(0)?;
272    f.seek(SeekFrom::Start(0))?;
273    f.write_all(offset.to_string().as_bytes())?;
274    f.sync_data()
275}
276
277/// Crate-visible alias for [`command_len`] (bulk copy counts frames).
278pub(crate) fn command_len_pub(b: &[u8]) -> Option<usize> {
279    command_len(b)
280}
281
282/// Length of one complete RESP command at the head of `b`, or `None`.
283fn command_len(b: &[u8]) -> Option<usize> {
284    let mut pos = 0usize;
285    let line = take_line(b, &mut pos)?;
286    if line.first() != Some(&b'*') {
287        return None;
288    }
289    let n: usize = std::str::from_utf8(&line[1..]).ok()?.trim().parse().ok()?;
290    for _ in 0..n {
291        let hdr = take_line(b, &mut pos)?;
292        if hdr.first() != Some(&b'$') {
293            return None;
294        }
295        let len: usize = std::str::from_utf8(&hdr[1..]).ok()?.trim().parse().ok()?;
296        if b.len() < pos + len + 2 {
297            return None;
298        }
299        pos += len + 2;
300    }
301    Some(pos)
302}
303
304fn take_line<'b>(b: &'b [u8], pos: &mut usize) -> Option<&'b [u8]> {
305    let rest = &b[*pos..];
306    let idx = rest.windows(2).position(|w| w == b"\r\n")?;
307    let line = &rest[..idx];
308    *pos += idx + 2;
309    Some(line)
310}
kevy_cli/migrate.rs

kevy_cli/
migrate.rs
