pq-zig

main.zig (5666B)

---

 const std = @import("std");
 const pq = @cImport(
     @cInclude("libpq-fe.h"),
 );
 
 pub const PqError = error{PqError};
 
 // libpq wrapper
 // later, this could be a pure-zig client implementation
 pub const Db = struct {
     c_conn: *pq.PGconn,
     pub fn init(connect_url: [:0]const u8) !Db {
         if (pq.PQisthreadsafe() == 0) {
             std.log.err("PQisthreadsafe returned 0, can't use libpq in this program", .{});
             return PqError.PqError;
         }
         const maybe_conn: ?*pq.PGconn = pq.PQconnectdb(connect_url);
         if (maybe_conn == null) {
             std.log.err("PQconnectdb returned null", .{});
             return PqError.PqError;
         }
         if (pq.PQstatus(maybe_conn) == pq.CONNECTION_BAD) {
             std.log.err("PQstatus returned CONNECTION_BAD: {s}", .{pq.PQerrorMessage(maybe_conn)});
 
             return PqError.PqError;
         } else if (pq.PQstatus(maybe_conn) != pq.CONNECTION_OK) {
             std.log.err("PQstatus returned unknown status {}: {s}", .{ pq.PQstatus(maybe_conn), pq.PQerrorMessage(maybe_conn) });
             return PqError.PqError;
         }
         return Db{
             .c_conn = maybe_conn.?,
         };
     }
 
     pub fn deinit(self: Db) void {
         pq.PQfinish(self.c_conn);
     }
 
     pub fn exec(self: Db, query: [:0]const u8) !void {
         const res: ?*pq.PGresult = pq.PQexec(self.c_conn, query);
         defer pq.PQclear(res);
         const est: pq.ExecStatusType = pq.PQresultStatus(res);
         if (est != pq.PGRES_COMMAND_OK) {
             std.log.err("PQexec error code {} message {s}", .{ est, pq.PQerrorMessage(self.c_conn) });
             return PqError.PqError;
         }
     }
 
     pub fn prepare_statement(self: Db, allocator: std.mem.Allocator, query: [:0]const u8) PqError!Stmt {
         return Stmt{
             .db = self,
             .aa = std.heap.ArenaAllocator.init(allocator),
             .query = query,
         };
     }
 };
 
 pub const Stmt = struct {
     const MAX_PARAMS = 128;
     db: Db,
     query: [:0]const u8,
     aa: std.heap.ArenaAllocator,
 
     n_params: usize = 0,
     param_values: [MAX_PARAMS][*c]const u8 = undefined,
     did_exec: bool = false,
     c_res: ?*pq.PGresult = null,
     res_index: c_int = -1,
     n_tuples: c_int = -1,
     n_fields: c_int = -1,
 
     pub fn deinit(self: *Stmt) void {
         self.aa.deinit();
         if (self.c_res != null) {
             pq.PQclear(self.c_res);
         }
     }
 
     pub fn step(self: *Stmt) !bool {
         if (!self.did_exec) {
             self.did_exec = true;
             self.c_res = pq.PQexecParams(self.db.c_conn, self.query, @as(c_int, @intCast(self.n_params)), null, &self.param_values, null, null, 0);
             const rs = pq.PQresultStatus(self.c_res);
             if (rs != pq.PGRES_TUPLES_OK and rs != pq.PGRES_SINGLE_TUPLE and rs != pq.PGRES_COMMAND_OK) {
                 std.log.err("PQresultStatus {} error: {s}", .{ rs, pq.PQerrorMessage(self.db.c_conn) });
                 return PqError.PqError;
             }
             self.n_tuples = pq.PQntuples(self.c_res);
             self.n_fields = pq.PQnfields(self.c_res);
         }
         self.res_index = self.res_index + 1;
         return self.res_index < self.n_tuples;
     }
 
     pub fn read_column(self: Stmt, idx: c_int, comptime T: type) !T {
         const value_c: [*c]u8 = pq.PQgetvalue(self.c_res, self.res_index, idx);
         const value: []const u8 = std.mem.sliceTo(value_c, 0);
         return switch (@typeInfo(T)) {
             .Int => std.fmt.parseInt(T, value, 10),
             .Pointer => |ptrinfo| blk: {
                 if (ptrinfo.child != u8) {
                     @compileError("pointer type []const u8 only is supported in read_column");
                 }
                 if (ptrinfo.size != .Slice) {
                     @compileError("pointer type []const u8 only is supported in read_column");
                 }
                 break :blk value;
             },
             else => @compileError("unhandled type " ++ @tagName(@typeInfo(T)) ++ " in read_column"),
         };
     }
 
     pub fn read_columnN(self: Stmt, name: [:0]const u8, comptime T: type) !T {
         const idx = pq.PQfnumber(self.c_res, name.ptr);
         if (idx == -1) {
             std.log.err("read_columnN ColumnNotFound [{s}]", .{name});
             return error.ColumnNotFound;
         }
         return read_column(self, idx, T);
     }
 
     pub fn bind(self: *Stmt, idx: usize, t: anytype) !void {
         const T = @TypeOf(t);
         const value: [:0]const u8 = switch (@typeInfo(T)) {
             .Pointer => try std.fmt.allocPrintZ(self.aa.allocator(), "{s}", .{t}),
             .Int => try std.fmt.allocPrintZ(self.aa.allocator(), "{d}", .{t}),
             else => @compileError("unhandled type " ++ @tagName(@typeInfo(T) ++ " in bind")),
         };
         self.param_values[idx] = value.ptr;
         self.n_params = @max(self.n_params, idx + 1);
     }
 
     pub fn read_struct(self: Stmt, comptime T: type) !T {
         const ti = @typeInfo(T);
         var t: T = undefined;
         inline for (ti.Struct.fields) |field| {
             const name: [:0]const u8 = &addZ(field.name.len, field.name[0..].*);
             const val = try self.read_columnN(name, field.type);
             @field(t, field.name) = val;
         }
         return t;
     }
 };
 
 // https://github.com/ziglang/zig/issues/16116
 pub fn addZ(comptime length: usize, value: [length]u8) [length:0]u8 {
     var terminated_value: [length:0]u8 = undefined;
     terminated_value[length] = 0;
     @memcpy(&terminated_value, &value);
     return terminated_value;
 }
 
 test "connect" {
     var db = try Db.init("postgresql://localhost/comments");
     db.deinit();
 }