fetch_ml/cli/src/commands/prune.zig
Jeremie Fraeys 94441fdc76
fix(cli): update imports after logging.zig removal
Update files still referencing deleted logging.zig:
- prune.zig: import io.zig, replace logging. with io.
- deps.zig: re-export log from io module

All zig tests now pass.
2026-03-04 21:18:18 -05:00

124 lines
4.4 KiB
Zig

const std = @import("std");
const Config = @import("../config.zig").Config;
const ws = @import("../net/ws/client.zig");
const crypto = @import("../utils/crypto.zig");
const io = @import("../utils/io.zig");
const core = @import("../core.zig");
pub fn run(allocator: std.mem.Allocator, args: []const []const u8) !void {
var flags = core.flags.CommonFlags{};
var keep_count: ?u32 = null;
var older_than_days: ?u32 = null;
// Parse flags
var i: usize = 0;
while (i < args.len) : (i += 1) {
if (std.mem.eql(u8, args[i], "--help") or std.mem.eql(u8, args[i], "-h")) {
return printUsage();
} else if (std.mem.eql(u8, args[i], "--json")) {
flags.json = true;
} else if (std.mem.eql(u8, args[i], "--keep") and i + 1 < args.len) {
keep_count = try std.fmt.parseInt(u32, args[i + 1], 10);
i += 1;
} else if (std.mem.eql(u8, args[i], "--older-than") and i + 1 < args.len) {
older_than_days = try std.fmt.parseInt(u32, args[i + 1], 10);
i += 1;
}
}
core.output.setMode(if (flags.flags.json) .flags.json else .text);
if (keep_count == null and older_than_days == null) {
core.output.usage("prune", "ml prune --keep <n> | --older-than <days>");
return error.InvalidArgs;
}
const config = try Config.load(allocator);
defer {
var mut_config = config;
mut_config.deinit(allocator);
}
// Add confirmation prompt
if (!flags.flags.json) {
if (keep_count) |count| {
if (!io.confirm("This will permanently delete all but the {d} most recent experiments. Continue?", .{count})) {
io.info("Prune cancelled.\n", .{});
return;
}
} else if (older_than_days) |days| {
if (!io.confirm("This will permanently delete all experiments older than {d} days. Continue?", .{days})) {
io.info("Prune cancelled.\n", .{});
return;
}
}
}
io.info("Pruning experiments...\n", .{});
// Use plain password for WebSocket authentication, hash for binary protocol
const api_key_plain = config.api_key; // Plain password from config
const api_key_hash = try crypto.hashApiKey(allocator, api_key_plain);
defer allocator.free(api_key_hash);
// Connect to WebSocket and send prune message
const ws_url = try config.getWebSocketUrl(allocator);
defer allocator.free(ws_url);
var client = try ws.Client.connect(allocator, ws_url, api_key_plain);
defer client.close();
// Determine prune type and send message
var prune_type: u8 = undefined;
var value: u32 = undefined;
if (keep_count) |count| {
prune_type = 0; // keep N
value = count;
io.info("Keeping {d} most recent experiments\n", .{count});
}
if (older_than_days) |days| {
prune_type = 1; // older than days
value = days;
io.info("Removing experiments older than {d} days\n", .{days});
}
try client.sendPrune(api_key_hash, prune_type, value);
// Receive response
const response = try client.receiveMessage(allocator);
defer allocator.free(response);
// Parse prune response (simplified - assumes success/failure byte)
if (response.len > 0) {
if (response[0] == 0x00) {
if (flags.json) {
std.debug.print("{\"ok\":true}\n", .{});
} else {
io.success("Prune operation completed successfully\n", .{});
}
} else {
if (flags.json) {
std.debug.print("{\"ok\":false,\"error_code\":{d}}\n", .{response[0]});
} else {
io.err("[FAIL] Prune operation failed: error code {d}\n", .{response[0]});
}
return error.PruneFailed;
}
} else {
if (flags.json) {
std.debug.print("{\"ok\":true,\"note\":\"no_response\"}\n", .{});
} else {
io.success("Prune request sent (no response received)\n", .{});
}
}
}
fn printUsage() void {
io.info("Usage: ml prune [options]\n\n", .{});
io.info("Options:\n", .{});
io.info("\t--keep <N>\t\tKeep N most recent experiments\n", .{});
io.info("\t--older-than <days>\tRemove experiments older than N days\n", .{});
io.info("\t--json\t\t\tOutput machine-readable JSON\n", .{});
io.info("\t--help, -h\t\tShow this help message\n", .{});
}