A Postgres Client

const std = @import("std");

// ------------------------------------------------- framing, as the recipe
// Postgres frames everything big-endian: a tag byte, an Int32 length that
// counts itself but not the tag, then the payload.

fn putInt(comptime T: type, w: *std.Io.Writer, v: T) !void {
    var b: [@divExact(@typeInfo(T).int.bits, 8)]u8 = undefined;
    std.mem.writeInt(T, &b, v, .big);
    try w.writeAll(&b);
}

fn sendMsg(w: *std.Io.Writer, tag: u8, payload: []const u8) !void {
    try w.writeAll(&.{tag});
    try putInt(i32, w, @intCast(payload.len + 4));
    try w.writeAll(payload);
}

const Cursor = struct {
    b: []const u8,
    i: usize = 0,
    fn int(self: *Cursor, comptime T: type) T {
        const n = @divExact(@typeInfo(T).int.bits, 8);
        const v = std.mem.readInt(T, self.b[self.i..][0..n], .big);
        self.i += n;
        return v;
    }
    fn cstr(self: *Cursor) []const u8 {
        const start = self.i;
        while (self.b[self.i] != 0) self.i += 1;
        defer self.i += 1;
        return self.b[start..self.i];
    }
};

// ------------------------------------------------------------- the client

const max_columns = 8;

/// A Postgres connection: a reader, a writer, and enough state to know
/// what the server last said. It never touches a socket. Whoever owns
/// the connection hands in the two streams, which is why the demo below
/// can hand it a script instead.
const PgClient = struct {
    r: *std.Io.Reader,
    w: *std.Io.Writer,
    msg_buf: [512]u8 = undefined,
    asked_for_password: bool = false,
    /// From the last ReadyForQuery: 'I' idle, 'T' in transaction, 'E' failed.
    ready_status: u8 = 0,
    err_code: [5]u8 = undefined,
    err_msg: [256]u8 = undefined,
    err_msg_len: usize = 0,

    const Msg = struct { tag: u8, payload: []const u8 };

    fn readMsg(self: *PgClient) !Msg {
        var tag: [1]u8 = undefined;
        try self.r.readSliceAll(&tag);
        var len_bytes: [4]u8 = undefined;
        try self.r.readSliceAll(&len_bytes);
        const len = std.mem.readInt(i32, &len_bytes, .big);
        const payload = self.msg_buf[0..@intCast(len - 4)];
        try self.r.readSliceAll(payload);
        return .{ .tag = tag[0], .payload = payload };
    }

    /// Startup and authentication, up to the first ReadyForQuery. The
    /// server chooses the auth method; this client speaks "none" and
    /// "cleartext password", which is what loopback setups use. SCRAM,
    /// the production method, slots in as one more case here.
    fn connect(r: *std.Io.Reader, w: *std.Io.Writer, user: []const u8, database: []const u8, password: []const u8) !PgClient {
        var body_buf: [128]u8 = undefined;
        var body: std.Io.Writer = .fixed(&body_buf);
        try putInt(i32, &body, 196608); // protocol 3.0
        for ([_][]const u8{ "user", user, "database", database }) |s| {
            try body.writeAll(s);
            try body.writeAll("\x00");
        }
        try body.writeAll("\x00");
        // The startup message alone has no tag byte.
        try putInt(i32, w, @intCast(4 + body.buffered().len));
        try w.writeAll(body.buffered());
        try w.flush();

        var client = PgClient{ .r = r, .w = w };
        while (true) {
            const m = try client.readMsg();
            var c = Cursor{ .b = m.payload };
            switch (m.tag) {
                'R' => switch (c.int(i32)) {
                    0 => {}, // AuthenticationOk
                    3 => { // AuthenticationCleartextPassword
                        client.asked_for_password = true;
                        var pw_buf: [64]u8 = undefined;
                        var pw: std.Io.Writer = .fixed(&pw_buf);
                        try pw.writeAll(password);
                        try pw.writeAll("\x00");
                        try sendMsg(w, 'p', pw.buffered());
                        try w.flush();
                    },
                    else => return error.UnsupportedAuth,
                },
                'S', 'K' => {}, // parameters and cancel key: noted, unused
                'E' => {
                    client.storeError(m.payload);
                    return error.ConnectFailed;
                },
                'Z' => {
                    client.ready_status = m.payload[0];
                    return client;
                },
                else => {},
            }
        }
    }

    fn storeError(self: *PgClient, payload: []const u8) void {
        // ErrorResponse is a list of fields: a type byte, a C string,
        // ending at a zero byte. Code and message are the useful two.
        var c = Cursor{ .b = payload };
        while (c.b[c.i] != 0) {
            const kind = c.b[c.i];
            c.i += 1;
            const text = c.cstr();
            switch (kind) {
                'C' => @memcpy(&self.err_code, text[0..5]),
                'M' => {
                    const n = @min(text.len, self.err_msg.len);
                    @memcpy(self.err_msg[0..n], text[0..n]);
                    self.err_msg_len = n;
                },
                else => {},
            }
        }
    }

    fn errorMessage(self: *const PgClient) []const u8 {
        return self.err_msg[0..self.err_msg_len];
    }

    fn query(self: *PgClient, sql: []const u8) !Rows {
        var sql_buf: [256]u8 = undefined;
        var q: std.Io.Writer = .fixed(&sql_buf);
        try q.writeAll(sql);
        try q.writeAll("\x00");
        try sendMsg(self.w, 'Q', q.buffered());
        try self.w.flush();
        return .{ .client = self };
    }
};

/// Row iteration. Column values are slices into the client's message
/// buffer: valid until the next call, exactly like a real driver's
/// zero-copy row. Copy anything you keep.
const Rows = struct {
    client: *PgClient,
    column_count: usize = 0,
    names: [max_columns][24]u8 = undefined,
    name_lens: [max_columns]usize = @splat(0),

    const Row = struct { values: [max_columns][]const u8, count: usize };

    fn name(self: *const Rows, i: usize) []const u8 {
        return self.names[i][0..self.name_lens[i]];
    }

    fn next(self: *Rows) !?Row {
        while (true) {
            const m = try self.client.readMsg();
            var c = Cursor{ .b = m.payload };
            switch (m.tag) {
                'T' => { // RowDescription: names now, values in 'D' messages
                    self.column_count = @intCast(c.int(i16));
                    for (0..self.column_count) |i| {
                        const n = c.cstr();
                        const len = @min(n.len, self.names[i].len);
                        @memcpy(self.names[i][0..len], n[0..len]);
                        self.name_lens[i] = len;
                        // table oid, attnum, type oid, typlen, typmod,
                        // format: 18 bytes this client has no use for.
                        c.i += 18;
                    }
                },
                'D' => {
                    var row = Row{ .values = undefined, .count = @intCast(c.int(i16)) };
                    for (0..row.count) |i| {
                        const len = c.int(i32);
                        if (len < 0) {
                            row.values[i] = "<null>";
                            continue;
                        }
                        row.values[i] = c.b[c.i..][0..@intCast(len)];
                        c.i += @intCast(len);
                    }
                    return row;
                },
                'C' => {}, // CommandComplete; the tag repeats what we sent
                'E' => {
                    self.client.storeError(m.payload);
                    // Drain to ReadyForQuery so the connection stays usable.
                    while ((try self.client.readMsg()).tag != 'Z') {}
                    self.client.ready_status = 'I';
                    return error.QueryFailed;
                },
                'Z' => {
                    self.client.ready_status = m.payload[0];
                    return null;
                },
                else => {},
            }
        }
    }
};

// ---------------------------------------------------- a scripted backend

/// The server half of two conversations, written into a buffer in
/// advance. The client cannot tell: it parses a `Reader`, and bytes from
/// a script are bytes. This is the site's usual trick, and it is also
/// how you get regression tests for a driver without a database in CI.
fn scriptBackend(w: *std.Io.Writer) !void {
    var pbuf: [256]u8 = undefined;
    var p: std.Io.Writer = .fixed(&pbuf);

    // Connection: demand a password, accept it, ready.
    try putInt(i32, &p, 3); // AuthenticationCleartextPassword
    try sendMsg(w, 'R', p.buffered());
    p = .fixed(&pbuf);
    try putInt(i32, &p, 0); // AuthenticationOk
    try sendMsg(w, 'R', p.buffered());
    try sendMsg(w, 'S', "server_version\x0016.4\x00");
    try sendMsg(w, 'Z', "I");

    // The SELECT: three columns, one row, done.
    p = .fixed(&pbuf);
    try putInt(i16, &p, 3);
    for ([_][]const u8{ "id", "target", "hits" }) |col| {
        try p.writeAll(col);
        try p.writeAll("\x00");
        try putInt(i32, &p, 0); // table oid
        try putInt(i16, &p, 0); // column number
        try putInt(i32, &p, 25); // type oid: text
        try putInt(i16, &p, -1); // variable length
        try putInt(i32, &p, -1); // no modifier
        try putInt(i16, &p, 0); // text format
    }
    try sendMsg(w, 'T', p.buffered());
    p = .fixed(&pbuf);
    try putInt(i16, &p, 3);
    for ([_][]const u8{ "7", "https://ziglang.org", "41" }) |val| {
        try putInt(i32, &p, @intCast(val.len));
        try p.writeAll(val);
    }
    try sendMsg(w, 'D', p.buffered());
    try sendMsg(w, 'C', "SELECT 1\x00");
    try sendMsg(w, 'Z', "I");

    // The INSERT: rejected, then ready again anyway.
    try sendMsg(w, 'E', "SERROR\x00C23505\x00" ++
        "Mduplicate key value violates unique constraint \"links_slug_key\"\x00\x00");
    try sendMsg(w, 'Z', "I");
}

pub fn main(init: std.process.Init) !void {
    var buf: [2048]u8 = undefined;
    var file_writer = std.Io.File.stdout().writerStreaming(init.io, &buf);
    const out = &file_writer.interface;

    var backend_bytes: [1024]u8 = undefined;
    var backend: std.Io.Writer = .fixed(&backend_bytes);
    try scriptBackend(&backend);

    var sent: [512]u8 = undefined; // what the client says, unexamined here
    var to_server: std.Io.Writer = .fixed(&sent);
    var from_server: std.Io.Reader = .fixed(backend.buffered());

    var client = try PgClient.connect(&from_server, &to_server, "shortener", "links", "hunter2");
    try out.print("connect: password asked: {}, ready: {c}\n", .{
        client.asked_for_password, client.ready_status,
    });

    try out.writeAll("\nquery: SELECT id, target, hits FROM links WHERE slug = 'zig1'\n");
    var rows = try client.query("SELECT id, target, hits FROM links WHERE slug = 'zig1'");
    while (try rows.next()) |row| {
        try out.writeAll("  row:");
        for (0..row.count) |i| {
            try out.print(" {s}={s}", .{ rows.name(i), row.values[i] });
        }
        try out.writeAll("\n");
    }

    try out.writeAll("\nquery: INSERT INTO links (slug, target) VALUES ('zig1', ...)\n");
    var failing = try client.query("INSERT INTO links (slug, target) VALUES ('zig1', 'https://example.com')");
    _ = failing.next() catch {
        try out.print("  error {s}: {s}\n", .{ client.err_code, client.errorMessage() });
    };
    try out.print("  connection after the failure: ready: {c}\n", .{client.ready_status});

    try out.flush();
}