zig - postgres driver (simple example)

zig-postgres-todo.zig

const std = @import("std");
const Io = std.Io;

const MessageType = enum(u8) {
    BackendKeyData = 'K',
    CommandComplete = 'C',
    DataRow = 'D',
    ErrorResponse = 'E',
    ParameterStatus = 'S',
    ReadyForQuery = 'Z',
    RowDescription = 'T',
    AuthenticationRequest = 'R',
    _,
};

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    const allocator = gpa.allocator();

    var output_buffer: [4096]u8 = undefined;
    var stdout_impl = std.fs.File.stdout().writer(&output_buffer);
    const stdout = &stdout_impl.interface;
    defer stdout.flush() catch {};

    try stdout.print("=== Section 13-04: PostgreSQL CRUD Terminal App ===\n\n", .{});

    // 1. Setup Io
    var io_instance = Io.Threaded.init(allocator);
    defer io_instance.deinit();
    const io = io_instance.io();

    // 2. Connect
    const address: Io.net.IpAddress = .{ .ip4 = try Io.net.Ip4Address.parse("127.0.0.1", 5432) };
    var stream = address.connect(io, .{ .mode = .stream }) catch |err| {
        try stdout.print("  ! Connection failed: {any}\n", .{err});
        try stdout.print("  ℹ Make sure Postgres is running (docker run -d --name pg-crud -e POSTGRES_HOST_AUTH_METHOD=trust -p 5432:5432 postgres:latest)\n", .{});
        return;
    };
    defer stream.close(io);

    var rb: [16384]u8 = undefined;
    var wb: [16384]u8 = undefined;
    var reader = stream.reader(io, &rb);
    var writer = stream.writer(io, &wb);

    // 3. Handshake (Startup)
    try startup(allocator, &reader, &writer);
    try stdout.print("  ✓ Connected to PostgreSQL.\n", .{});

    // 4. Initialize Table
    try executeQuery(allocator, &reader, &writer, "CREATE TABLE IF NOT EXISTS todos (id SERIAL PRIMARY KEY, task TEXT NOT NULL, done BOOLEAN DEFAULT FALSE);", stdout);
    try stdout.print("  ✓ Database initialized.\n\n", .{});

    // 5. Interactive Loop
    var stb: [4096]u8 = undefined;
    var stdin_wrapper = std.fs.File.stdin().reader(io, &stb);
    const stdin = &stdin_wrapper.interface;

    while (true) {
        try stdout.print("\n--- TODO MENU ---\n", .{});
        try stdout.print("1. List Todos\n", .{});
        try stdout.print("2. Add Todo\n", .{});
        try stdout.print("3. Toggle Done (by ID)\n", .{});
        try stdout.print("4. Delete Todo (by ID)\n", .{});
        try stdout.print("5. Quit\n", .{});
        try stdout.print("Selection: ", .{});
        try stdout.flush();

        const line = (try stdin.takeDelimiter('\n')) orelse break;
        const choice = std.mem.trim(u8, line, " \r\n");
        if (choice.len == 0) continue;

        if (std.mem.eql(u8, choice, "1")) {
            try listTodos(allocator, &reader, &writer, stdout);
        } else if (std.mem.eql(u8, choice, "2")) {
            try stdout.print("Enter task description: ", .{});
            try stdout.flush();
            const task_line = (try stdin.takeDelimiter('\n')) orelse break;
            const task = std.mem.trim(u8, task_line, " \r\n");
            if (task.len > 0) {
                const sql = try std.fmt.allocPrint(allocator, "INSERT INTO todos (task) VALUES ('{s}');", .{task});
                defer allocator.free(sql);
                try executeQuery(allocator, &reader, &writer, sql, stdout);
            }
        } else if (std.mem.eql(u8, choice, "3")) {
            try stdout.print("Enter Todo ID: ", .{});
            try stdout.flush();
            const id_line = (try stdin.takeDelimiter('\n')) orelse break;
            const id = std.mem.trim(u8, id_line, " \r\n");
            if (id.len > 0) {
                const sql = try std.fmt.allocPrint(allocator, "UPDATE todos SET done = NOT done WHERE id = {s};", .{id});
                defer allocator.free(sql);
                try executeQuery(allocator, &reader, &writer, sql, stdout);
            }
        } else if (std.mem.eql(u8, choice, "4")) {
            try stdout.print("Enter Todo ID to delete: ", .{});
            try stdout.flush();
            const id_line = (try stdin.takeDelimiter('\n')) orelse break;
            const id = std.mem.trim(u8, id_line, " \r\n");
            if (id.len > 0) {
                const sql = try std.fmt.allocPrint(allocator, "DELETE FROM todos WHERE id = {s};", .{id});
                defer allocator.free(sql);
                try executeQuery(allocator, &reader, &writer, sql, stdout);
            }
        } else if (std.mem.eql(u8, choice, "5")) {
            break;
        } else {
            try stdout.print("Invalid choice.\n", .{});
        }
    }

    try stdout.print("\nExiting. Goodbye!\n", .{});
}

fn startup(allocator: std.mem.Allocator, reader: anytype, writer: anytype) !void {
    var payload = std.ArrayListUnmanaged(u8){};
    defer payload.deinit(allocator);

    try payload.appendSlice(allocator, &[_]u8{ 0, 3, 0, 0 }); // Protocol 3.0
    try payload.appendSlice(allocator, "user\x00postgres\x00database\x00postgres\x00\x00");

    const total_len = @as(i32, @intCast(payload.items.len + 4));
    var len_buf: [4]u8 = undefined;
    std.mem.writeInt(i32, &len_buf, total_len, .big);
    
    try writer.interface.writeAll(&len_buf);
    try writer.interface.writeAll(payload.items);
    try writer.interface.flush();

    while (true) {
        var msg_type_buf: [1]u8 = undefined;
        try reader.interface.readSliceAll(&msg_type_buf);
        const msg_type: MessageType = @enumFromInt(msg_type_buf[0]);

        var msg_len_buf: [4]u8 = undefined;
        try reader.interface.readSliceAll(&msg_len_buf);
        const msg_len = std.mem.readInt(i32, &msg_len_buf, .big);
        const payload_len = @as(usize, @intCast(msg_len - 4));

        switch (msg_type) {
            .AuthenticationRequest => {
                var auth_type_buf: [4]u8 = undefined;
                try reader.interface.readSliceAll(&auth_type_buf);
                const auth_type = std.mem.readInt(i32, &auth_type_buf, .big);
                if (auth_type != 0) return error.AuthRequired;
            },
            .ReadyForQuery => {
                try reader.interface.discardAll(1);
                return;
            },
            .ErrorResponse => {
                const err = try reader.interface.readAlloc(allocator, payload_len);
                defer allocator.free(err);
                return error.BackendError;
            },
            else => try reader.interface.discardAll(payload_len),
        }
    }
}

fn executeQuery(allocator: std.mem.Allocator, reader: anytype, writer: anytype, sql: []const u8, stdout: anytype) !void {
    const q_len = @as(i32, @intCast(sql.len + 5)); // + type byte, length field, null terminator
    var q_len_buf: [4]u8 = undefined;
    std.mem.writeInt(i32, &q_len_buf, q_len, .big);

    try writer.interface.writeByte('Q');
    try writer.interface.writeAll(&q_len_buf);
    try writer.interface.writeAll(sql);
    try writer.interface.writeByte(0); // Null terminator
    try writer.interface.flush();

    while (true) {
        var msg_type_buf: [1]u8 = undefined;
        try reader.interface.readSliceAll(&msg_type_buf);
        const msg_type: MessageType = @enumFromInt(msg_type_buf[0]);

        var msg_len_buf: [4]u8 = undefined;
        try reader.interface.readSliceAll(&msg_len_buf);
        const msg_len = std.mem.readInt(i32, &msg_len_buf, .big);
        const payload_len = @as(usize, @intCast(msg_len - 4));

        switch (msg_type) {
            .CommandComplete => {
                const res = try reader.interface.readAlloc(allocator, payload_len);
                defer allocator.free(res);
                try stdout.print("  ✓ {s}\n", .{res});
            },
            .ReadyForQuery => {
                try reader.interface.discardAll(1);
                return;
            },
            .ErrorResponse => {
                const err = try reader.interface.readAlloc(allocator, payload_len);
                defer allocator.free(err);
                try stdout.print("  ! Error: {s}\n", .{err});
                return;
            },
            else => try reader.interface.discardAll(payload_len),
        }
    }
}

fn listTodos(allocator: std.mem.Allocator, reader: anytype, writer: anytype, stdout: anytype) !void {
    const sql = "SELECT id, task, done FROM todos ORDER BY id ASC;";
    const q_len = @as(i32, @intCast(sql.len + 5));
    var q_len_buf: [4]u8 = undefined;
    std.mem.writeInt(i32, &q_len_buf, q_len, .big);

    try writer.interface.writeByte('Q');
    try writer.interface.writeAll(&q_len_buf);
    try writer.interface.writeAll(sql);
    try writer.interface.writeByte(0);
    try writer.interface.flush();

    try stdout.print("\n{s: <5} | {s: <30} | {s}\n", .{ "ID", "TASK", "DONE" });
    try stdout.print("{s:-<5}-+-{s:-<30}-+-{s:-<4}\n", .{ "", "", "" });

    while (true) {
        var msg_type_buf: [1]u8 = undefined;
        try reader.interface.readSliceAll(&msg_type_buf);
        const msg_type: MessageType = @enumFromInt(msg_type_buf[0]);

        var msg_len_buf: [4]u8 = undefined;
        try reader.interface.readSliceAll(&msg_len_buf);
        const msg_len = std.mem.readInt(i32, &msg_len_buf, .big);
        const payload_len = @as(usize, @intCast(msg_len - 4));

        switch (msg_type) {
            .RowDescription => try reader.interface.discardAll(payload_len),
            .DataRow => {
                var cols_buf: [2]u8 = undefined;
                try reader.interface.readSliceAll(&cols_buf);
                const num_cols = std.mem.readInt(u16, &cols_buf, .big);
                
                var col: u16 = 0;
                var row_data: [3][]u8 = undefined;
                while (col < num_cols) : (col += 1) {
                    var col_len_buf: [4]u8 = undefined;
                    try reader.interface.readSliceAll(&col_len_buf);
                    const col_len = std.mem.readInt(i32, &col_len_buf, .big);
                    if (col_len == -1) {
                        row_data[col] = try allocator.dupe(u8, "NULL");
                    } else {
                        row_data[col] = try reader.interface.readAlloc(allocator, @as(usize, @intCast(col_len)));
                    }
                }
                
                try stdout.print("{s: <5} | {s: <30} | {s}\n", .{ row_data[0], row_data[1], if (std.mem.eql(u8, row_data[2], "t")) "✓" else " " });
                
                for (row_data) |d| allocator.free(d);
            },
            .ReadyForQuery => {
                try reader.interface.discardAll(1);
                return;
            },
            .CommandComplete => try reader.interface.discardAll(payload_len),
            .ErrorResponse => {
                const err = try reader.interface.readAlloc(allocator, payload_len);
                defer allocator.free(err);
                try stdout.print("  ! Error: {s}\n", .{err});
                return;
            },
            else => try reader.interface.discardAll(payload_len),
        }
    }
}