feat: extend main with runReplicate and parseHostPort

This commit is contained in:
dualfroz committed 2026-09-07 15:22:50 +02:00
1 parent 601a264b78
commit f9b2e7df9c
1 file changed
+173
+173
View File
@@ -2214,6 +2214,179 @@ fn runTxn(allocator: std.mem.Allocator, stdout: anytype, stderr: anytype, db: *c
} }
/// Dispatches `replicate <replica-dir> [--once] [--interval-ms <n>]`: opens /// Dispatches `replicate <replica-dir> [--once] [--interval-ms <n>]`: opens
/// (or creates) a second Db at <replica-dir> and keeps it mirroring `db` by
/// physically tailing `db`'s wal.log (replicate.zig), recovering via a full
/// snapshot resync whenever the tailer falls behind a primary flush.
fn runReplicate(allocator: std.mem.Allocator, stdout: anytype, db: *cobalt.Db, replicaDir: []const u8, rest: []const []const u8) !void {
var once = false;
var intervalMs: u64 = 200;
var i: usize = 0;
while (i < rest.len) : (i += 1) {
if (std.mem.eql(u8, rest[i], "--once")) {
once = true;
} else if (std.mem.eql(u8, rest[i], "--interval-ms")) {
i += 1;
if (i >= rest.len) {
try stdout.writeAll("error: --interval-ms needs a value\n");
std.process.exit(2);
}
intervalMs = std.fmt.parseInt(u64, rest[i], 10) catch {
try stdout.writeAll("error: --interval-ms needs an integer\n");
std.process.exit(2);
};
} else {
try stdout.print("error: unknown replicate flag '{s}'\n", .{rest[i]});
std.process.exit(2);
}
}
var replica = try cobalt.Db.open(allocator, replicaDir, .{});
defer replica.close();
const resumeState = cobalt.replicate.loadOffset(allocator, replicaDir, db.dirPath) catch |err| switch (err) {
cobalt.replicate.StateError.PrimaryMismatch => {
try stdout.print("error: replica at '{s}' has persisted state for a different primary than '{s}' - refusing to resume against the wrong primary\n", .{ replicaDir, db.dirPath });
std.process.exit(1);
},
else => return err,
};
var rep = if (resumeState) |state|
try cobalt.Replicator.initFromOffset(allocator, db.dirPath, &replica, state.offset, state.generation)
else
try cobalt.Replicator.init(allocator, db.dirPath, &replica);
defer rep.deinit();
if (resumeState) |state| {
try stdout.print("resuming tail from persisted offset={d} generation={d}\n", .{ state.offset, state.generation });
} else {
const bootstrapped = try rep.fullResync(db, 500);
try stdout.print("bootstrap: {d} entries copied from snapshot\n", .{bootstrapped});
try cobalt.replicate.saveState(allocator, replicaDir, db.dirPath, rep.offset(), rep.generation());
}
while (true) {
const applied = rep.syncOnce() catch |err| switch (err) {
cobalt.ReplicateError.WalTruncatedAheadOfTailer, cobalt.ReplicateError.WalGenerationChanged => blk: {
const recovered = try rep.fullResync(db, 500);
try stdout.print("resync after truncation: {d} entries\n", .{recovered});
// fullResync always re-bases the tailer to offset 0 and the
// primary's current generation, which must be persisted
// immediately even if nothing new was applied afterwards -
// otherwise a later process would resume from the stale
// pre-truncation offset/generation instead of the fresh one.
try cobalt.replicate.saveState(allocator, replicaDir, db.dirPath, rep.offset(), rep.generation());
break :blk @as(usize, 0);
},
else => return err,
};
if (applied > 0) {
try stdout.print("applied {d} record(s), tail offset={d}\n", .{ applied, rep.offset() });
try cobalt.replicate.saveState(allocator, replicaDir, db.dirPath, rep.offset(), rep.generation());
}
if (once) break;
std.time.sleep(intervalMs * std.time.ns_per_ms);
}
}
/// Parses a `host:port` string as used by both `serve` (bind address) and
/// `remote-select` (address to connect to). Exits with usage code 2 on a
/// malformed address rather than propagating a parse error, matching every
/// other CLI argument validation in this file.
fn parseHostPort(stdout: anytype, spec: []const u8) !std.net.Address {
const colon = std.mem.lastIndexOfScalar(u8, spec, ':') orelse {
try stdout.print("error: address must be host:port, got '{s}'\n", .{spec});
std.process.exit(2);
};
const host = spec[0..colon];
const port = std.fmt.parseInt(u16, spec[colon + 1 ..], 10) catch {
try stdout.print("error: invalid port in '{s}'\n", .{spec});
std.process.exit(2);
};
return std.net.Address.parseIp(host, port);
}
const AuthFlags = struct { username: ?[]const u8 = null, password: ?[]const u8 = null, token: ?[]const u8 = null };
/// Pulls `--user NAME` / `--pass PASS` / `--token TOKEN` out of a
/// `remote-select`/`remote-user` argument list, same shape and same "strip
/// recognized flags, keep the rest" idea as `extractPageFlags`. `--token`
/// is the CLI counterpart of `client.authenticateWithToken` - a previously
/// issued session token (see `serve --session-ttl`) used in place of a
/// username/password pair. Caller owns and frees `rest`.
fn extractAuthFlags(allocator: std.mem.Allocator, args: []const []const u8) !struct { auth: AuthFlags, rest: [][]const u8 } {
var authFlags = AuthFlags{};
var rest = std.ArrayList([]const u8).init(allocator);
errdefer rest.deinit();
var i: usize = 0;
while (i < args.len) : (i += 1) {
const isUser = std.mem.eql(u8, args[i], "--user");
const isPass = std.mem.eql(u8, args[i], "--pass");
const isToken = std.mem.eql(u8, args[i], "--token");
if (isUser or isPass or isToken) {
i += 1;
if (i >= args.len) return error.InvalidAuthFlag;
if (isUser) authFlags.username = args[i] else if (isPass) authFlags.password = args[i] else authFlags.token = args[i];
} else {
try rest.append(args[i]);
}
}
return .{ .auth = authFlags, .rest = try rest.toOwnedSlice() };
}
/// Shared by `runRemoteSelect` and `runRemoteUser`: authenticates the
/// just-connected `client` using whichever credential `authFlags` carried
/// (a session token takes priority over a username/password pair - a
/// caller passing both almost certainly means "use the token, the
/// user/pass was left over from a script template"), or does nothing if
/// neither was given (server has no --auth enabled, or the op needs none).
/// A fresh password login against a server with sessions enabled gets a
/// brand new token back (`client.sessionToken()`) - printed to stderr (kept
/// out of stdout so scripts piping row/user output are not polluted) so the
/// caller can pass it to a later invocation via `--token` instead of
/// re-sending the password. Logging in with `--token` itself never mints a
/// new one (see `server.zig`'s `handleAuthToken`), so nothing is printed
/// in that branch.
fn authenticateIfRequested(client: *cobalt.Client, stderr: anytype, authFlags: AuthFlags) !void {
if (authFlags.token) |token| {
client.authenticateWithToken(token) catch |err| {
try stderr.print("error: token authentication failed: {s}\n", .{@errorName(err)});
std.process.exit(1);
};
} else if (authFlags.username) |username| {
client.authenticate(username, authFlags.password orelse "") catch |err| {
try stderr.print("error: authentication failed: {s}\n", .{@errorName(err)});
std.process.exit(1);
};
if (client.sessionToken()) |token| {
try stderr.print("session token (use with --token): {s}\n", .{token});
}
}
}
/// Connects to a remote `cobalt serve` instance and runs `table select`
/// there instead of against a local Db - the client-side counterpart of
/// `server.zig`'s `table_select` opcode, so a query engine that was
/// previously CLI-local-only is now reachable over the wire with the exact
/// same expression grammar as `table select`. `--user`/`--pass` authenticate
/// first, needed only against a server started with `serve --auth`.
fn runRemoteSelect(allocator: std.mem.Allocator, stdout: anytype, stderr: anytype, addressSpec: []const u8, rest: []const []const u8) !void {
if (rest.len < 1) return fail(stderr, "remote-select needs <table> [expr...] [--limit N] [--offset N] [--order-by field[:desc][,field[:desc]...]]");
const address = try parseHostPort(stdout, addressSpec);
var client = cobalt.Client.connect(allocator, address) catch |err| {
try stderr.print("error: could not connect to {s}: {s}\n", .{ addressSpec, @errorName(err) });
std.process.exit(1);
};
defer client.close();
const authFlags = extractAuthFlags(allocator, rest[1..]) catch return fail(stderr, "invalid --user/--pass/--token value");
defer allocator.free(authFlags.rest);
try authenticateIfRequested(&client, stderr, authFlags.auth);
const sel = extractSelectFlags(allocator, authFlags.rest) catch return fail(stderr, "invalid --limit/--offset/--order-by value");
} }
fn runVerify(allocator: std.mem.Allocator, stdout: anytype, dir: []const u8, count: usize) !void { fn runVerify(allocator: std.mem.Allocator, stdout: anytype, dir: []const u8, count: usize) !void {