Files
cobalt/src/repl.zig
T

325 lines
14 KiB
Zig

const std = @import("std");
const dbmod = @import("db.zig");
const txnmod = @import("txn.zig");
const rowmod = @import("row.zig");
const idx = @import("index.zig");
const querymod = @import("query.zig");
const txntablemod = @import("txn_table.zig");
const collectionsmod = @import("collections.zig");
const collttlmod = @import("collttl.zig");
const ttlmod = @import("ttl.zig");
const exprmod = @import("expr.zig");
const joinmod = @import("join.zig");
const aggregatemod = @import("aggregate.zig");
const windowmod = @import("window.zig");
const Allocator = std.mem.Allocator;
const Db = dbmod.Db;
const Transaction = txnmod.Transaction;
const Table = querymod.Table;
const TableTxn = txntablemod.TableTxn;
const Value = rowmod.Value;
const Field = rowmod.Field;
const Row = rowmod.Row;
const CollKind = collectionsmod.Kind;
/// Errors `Repl.dispatch` can raise on top of whatever the underlying `Db`/
/// `Transaction`/`Table` call it wraps returns (those propagate as-is, e.g.
/// `error.Conflict` from a failed commit).
pub const ReplError = error{
UnknownCommand,
WrongArgCount,
NoActiveTransaction,
TransactionAlreadyActive,
InvalidField,
UnknownCollectionKind,
UnknownCollectionOp,
InvalidArgument,
NoSuchKey,
};
/// Interactive session bound to one already-open `*Db`. The one-shot CLI in
/// main.zig opens a fresh `Db` per process per command, and its `txn
/// <script-file>` command runs a whole transaction non-interactively from a
/// static file. `Repl` instead keeps the store open across many commands and
/// lets `begin`/`commit`/`rollback` bracket an interactive `Transaction`:
/// a `get` between `begin` and `commit` can steer what the *next* `put`
/// does (e.g. "read balance, then decide how much to add"), which a script
/// file evaluated top to bottom cannot express.
///
/// `table put/get/del` follow the active transaction automatically: with no
/// transaction open they go straight through `query.Table` (each call is
/// independently durable the moment it returns, exactly like `table put` on
/// the plain CLI); with a transaction open they route through a `TableTxn`
/// bound to it instead, so multiple `table put` calls across multiple
/// tables in the same `begin`/`commit` bracket commit atomically as one
/// `WriteBatch`, or none of them do on `error.Conflict` - the interactive
/// equivalent of `txn <script-file>`.
///
/// `scan` always reads the live, already-committed `Db` state regardless of
/// any open transaction: `Transaction` has no scan/range primitive (only
/// `get`, which resolves the transaction's own staged writes first), so a
/// staged-but-uncommitted put is invisible to `scan` until `commit` lands it
/// for real. This is a deliberate, documented limitation, not an oversight.
/// `table select`/`table join` are the same story one level up: `TableTxn`
/// only ever brackets single-row put/get/delete, so both always run against
/// `query.Table.open(self.db, ...)` directly, never through the active
/// transaction, exactly like `scan`.
pub const Repl = struct {
db: *Db,
allocator: Allocator,
txn: ?Transaction = null,
tableTxns: std.StringHashMap(TableTxn),
pub fn init(allocator: Allocator, db: *Db) Repl {
return .{
.db = db,
.allocator = allocator,
.tableTxns = std.StringHashMap(TableTxn).init(allocator),
};
}
pub fn deinit(self: *Repl) void {
self.discardTxnState();
self.tableTxns.deinit();
}
pub fn inTransaction(self: *const Repl) bool {
return self.txn != null;
}
fn discardTxnState(self: *Repl) void {
var it = self.tableTxns.iterator();
while (it.next()) |entry| {
self.allocator.free(entry.key_ptr.*);
entry.value_ptr.deinit();
}
self.tableTxns.clearRetainingCapacity();
if (self.txn) |*t| {
t.deinit();
self.txn = null;
}
}
/// Runs one line of input, writing command output or `error: <name>` to
/// `writer`. Returns `false` on `exit`/`quit` (caller should stop
/// reading), `true` otherwise - including after a reported error, so one
/// bad line never kills the session.
pub fn runLine(self: *Repl, line: []const u8, writer: anytype) !bool {
const trimmed = std.mem.trim(u8, line, " \t\r\n");
if (trimmed.len == 0 or trimmed[0] == '#') return true;
var tokens = std.ArrayList([]const u8).init(self.allocator);
defer tokens.deinit();
var it = std.mem.tokenizeAny(u8, trimmed, " \t");
while (it.next()) |tok| try tokens.append(tok);
if (tokens.items.len == 0) return true;
const cmd = tokens.items[0];
if (std.mem.eql(u8, cmd, "exit") or std.mem.eql(u8, cmd, "quit")) return false;
self.dispatch(cmd, tokens.items[1..], writer) catch |err| {
try writer.print("error: {s}\n", .{@errorName(err)});
};
return true;
}
/// Reads lines from `reader` until EOF, running each through `runLine`.
/// Writes a `cobalt>`/`cobalt(txn)>` prompt before each read when
/// `prompt` is set (interactive stdin use); script files pass `false`.
pub fn runFromReader(self: *Repl, reader: anytype, writer: anytype, prompt: bool) !void {
var buf: [8192]u8 = undefined;
while (true) {
if (prompt) {
try writer.writeAll(if (self.inTransaction()) "cobalt(txn)> " else "cobalt> ");
}
const line = (try reader.readUntilDelimiterOrEof(&buf, '\n')) orelse break;
if (!try self.runLine(line, writer)) break;
}
}
fn dispatch(self: *Repl, cmd: []const u8, args: []const []const u8, writer: anytype) !void {
if (std.mem.eql(u8, cmd, "begin")) return self.cmdBegin(args, writer);
if (std.mem.eql(u8, cmd, "commit")) return self.cmdCommit(args, writer);
if (std.mem.eql(u8, cmd, "rollback")) return self.cmdRollback(args, writer);
if (std.mem.eql(u8, cmd, "put")) return self.cmdPut(args, writer);
if (std.mem.eql(u8, cmd, "get")) return self.cmdGet(args, writer);
if (std.mem.eql(u8, cmd, "del")) return self.cmdDel(args, writer);
if (std.mem.eql(u8, cmd, "scan")) return self.cmdScan(args, writer);
if (std.mem.eql(u8, cmd, "stats")) return self.cmdStats(args, writer);
if (std.mem.eql(u8, cmd, "flush")) return self.cmdFlush(args, writer);
if (std.mem.eql(u8, cmd, "compact")) return self.cmdCompact(args, writer);
if (std.mem.eql(u8, cmd, "table")) return self.cmdTable(args, writer);
if (std.mem.eql(u8, cmd, "coll")) return self.cmdColl(args, writer);
if (std.mem.eql(u8, cmd, "expire")) return self.cmdExpire(args, writer);
if (std.mem.eql(u8, cmd, "persist")) return self.cmdPersist(args, writer);
if (std.mem.eql(u8, cmd, "ttl")) return self.cmdTtl(args, writer);
if (std.mem.eql(u8, cmd, "expire-sweep")) return self.cmdExpireSweep(args, writer);
if (std.mem.eql(u8, cmd, "help")) return self.cmdHelp(writer);
return ReplError.UnknownCommand;
}
fn cmdBegin(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len != 0) return ReplError.WrongArgCount;
if (self.txn != null) return ReplError.TransactionAlreadyActive;
self.txn = Transaction.begin(self.db);
try writer.writeAll("OK\n");
}
fn cmdCommit(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len != 0) return ReplError.WrongArgCount;
if (self.txn == null) return ReplError.NoActiveTransaction;
try self.txn.?.commit();
self.discardTxnState();
try writer.writeAll("OK\n");
}
fn cmdRollback(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len != 0) return ReplError.WrongArgCount;
if (self.txn == null) return ReplError.NoActiveTransaction;
self.discardTxnState();
try writer.writeAll("OK\n");
}
fn cmdPut(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len != 2) return ReplError.WrongArgCount;
if (self.txn) |*t| {
try t.put(args[0], args[1]);
} else {
try self.db.put(args[0], args[1]);
}
try writer.writeAll("OK\n");
}
fn cmdGet(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len != 1) return ReplError.WrongArgCount;
const value = if (self.txn) |*t|
try t.get(self.allocator, args[0])
else
try self.db.get(self.allocator, args[0]);
defer if (value) |v| self.allocator.free(v);
if (value) |v| {
try writer.print("{s}\n", .{v});
} else {
try writer.writeAll("(nil)\n");
}
}
fn cmdDel(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len != 1) return ReplError.WrongArgCount;
if (self.txn) |*t| {
try t.delete(args[0]);
} else {
try self.db.delete(args[0]);
}
try writer.writeAll("OK\n");
}
fn cmdScan(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len > 1) return ReplError.WrongArgCount;
const prefix = if (args.len == 1) args[0] else "";
var iter = try self.db.scanPrefix(self.allocator, prefix);
defer iter.deinit();
while (try iter.next()) |entry| {
try writer.print("{s}\t{s}\n", .{ entry.key, entry.value });
}
}
/// `expire`/`persist`/`ttl`/`expire-sweep` mirror the one-shot CLI's
/// same-named commands (main.zig) over `ttl.zig`'s flat-key TTL layer.
/// Like `scan` above, these always act on the live, already-committed
/// `Db` regardless of any open transaction: a TTL marker lives under its
/// own keyspace tag, set/cleared directly against `self.db`, never
/// staged through `Transaction` - there is no OCC-safe way to make
/// "expire this key in N seconds" conditional on a commit that hasn't
/// happened yet, so (deliberately, like `coll`'s ops) this bypasses the
/// active transaction entirely rather than pretending to participate in
/// it.
fn cmdExpire(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len != 2) return ReplError.WrongArgCount;
const seconds = std.fmt.parseInt(u64, args[1], 10) catch return ReplError.InvalidArgument;
const existing = try self.db.get(self.allocator, args[0]) orelse return ReplError.NoSuchKey;
self.allocator.free(existing);
try ttlmod.setExpireSeconds(self.db, self.allocator, args[0], std.time.nanoTimestamp(), seconds);
try writer.writeAll("OK\n");
}
fn cmdPersist(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len != 1) return ReplError.WrongArgCount;
try ttlmod.clearExpire(self.db, self.allocator, args[0]);
try writer.writeAll("OK\n");
}
fn cmdTtl(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len != 1) return ReplError.WrongArgCount;
switch (try ttlmod.remaining(self.db, self.allocator, args[0], std.time.nanoTimestamp())) {
.noKey => try writer.writeAll("-2\n"),
.noExpiry => try writer.writeAll("-1\n"),
.seconds => |s| try writer.print("{d}\n", .{s}),
}
}
fn cmdExpireSweep(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len != 0) return ReplError.WrongArgCount;
const reaped = try ttlmod.sweepExpired(self.db, self.allocator, std.time.nanoTimestamp());
try writer.print("reaped: {d}\n", .{reaped});
}
fn cmdStats(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len != 0) return ReplError.WrongArgCount;
try self.db.stats().print(writer);
}
fn cmdFlush(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len != 0) return ReplError.WrongArgCount;
try self.db.flush();
try writer.writeAll("OK\n");
}
fn cmdCompact(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len != 0) return ReplError.WrongArgCount;
try self.db.compact();
try writer.writeAll("OK\n");
}
fn cmdTable(self: *Repl, args: []const []const u8, writer: anytype) !void {
if (args.len < 1) return ReplError.WrongArgCount;
const op = args[0];
const rest = args[1..];
if (std.mem.eql(u8, op, "select")) return self.cmdTableSelect(rest, writer);
if (std.mem.eql(u8, op, "explain")) return self.cmdTableExplain(rest, writer);
if (std.mem.eql(u8, op, "update")) return self.cmdTableUpdate(rest, writer);
if (std.mem.eql(u8, op, "delete")) return self.cmdTableDelete(rest, writer);
if (std.mem.eql(u8, op, "join")) return self.cmdTableJoin(rest, writer);
if (std.mem.eql(u8, op, "aggregate")) return self.cmdTableAggregate(rest, writer);
if (std.mem.eql(u8, op, "window")) return self.cmdTableWindow(rest, writer);
if (rest.len < 2) return ReplError.WrongArgCount;
const tableName = rest[0];
const pk = try Value.parseCli(rest[1]);
if (std.mem.eql(u8, op, "put")) {
var fields = std.ArrayList(Field).init(self.allocator);
defer fields.deinit();
for (rest[2..]) |raw| try fields.append(try parseField(raw));
try self.tablePut(tableName, pk, fields.items);
try writer.writeAll("OK\n");
} else if (std.mem.eql(u8, op, "get")) {
if (rest.len != 2) return ReplError.WrongArgCount;
var row = try self.tableGet(tableName, pk) orelse {
try writer.writeAll("(nil)\n");
return;
};
defer row.deinit();
try printRow(writer, row);
} else if (std.mem.eql(u8, op, "del")) {
if (rest.len != 2) return ReplError.WrongArgCount;
try self.tableDel(tableName, pk);
try writer.writeAll("OK\n");
} else {
return ReplError.UnknownCommand;
}
}