Something went wrong. Try again.
Native PostgreSQL driver / client for Zig
Something went wrong. Try again.
Zig
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280const std = @import("std");const lib = @import("lib.zig");const Buffer = @import("buffer").Buffer;
const proto = lib.proto;const Conn = lib.Conn;const Reader = lib.Reader;const NotificationResponse = lib.proto.NotificationResponse;
const Stream = lib.Stream;const Allocator = std.mem.Allocator;const Io = std.Io;
const ListenError = union(enum) { err: anyerror, pg: lib.proto.Error,};
pub const Listener = struct { err: ?ListenError = null, closed: bool = false,
_stream: Stream,
// A buffer used for writing to PG. This can grow dynamically as needed. _buf: Buffer,
// Used to read data from PG. Has its own buffer which can grow dynamically _reader: Reader,
// If we get a PG error, we'll return a LIstenError.pg, and we'll own its // memory. _err_data: ?[]const u8 = null,
_allocator: Allocator,
_io: Io,
pub fn open(io: Io, allocator: Allocator, opts: Conn.Opts) !Listener { var stream = try Stream.connect(io, allocator, opts, null); errdefer stream.close();
const buf = try Buffer.init(allocator, opts.write_buffer orelse 2048); errdefer buf.deinit();
const reader = try Reader.init(allocator, opts.read_buffer orelse 4096, stream); errdefer reader.deinit();
return .{ ._buf = buf, ._stream = stream, ._reader = reader, ._allocator = allocator, ._io = io, }; }
pub fn deinit(self: *Listener) void { if (self._err_data) |err_data| { self._allocator.free(err_data); } self._buf.deinit(); self._reader.deinit();
self.stop() catch {}; self._stream.close(); }
pub fn stop(self: *Listener) !void { if (@atomicRmw(bool, &self.closed, .Xchg, true, .monotonic) == true) { return; }
lib.sendTerminate(&self._stream, self._io); return self._stream.shutdown(.both); }
pub fn auth(self: *Listener, opts: Conn.AuthOpts) !void { if (try lib.auth.auth(self._io, &self._stream, &self._buf, &self._reader, opts)) |raw_pg_err| { return self.setErr(raw_pg_err); }
while (true) { const msg = try self.read(); switch (msg.type) { 'Z' => return, 'K' => {}, // TODO: BackendKeyData 'S' => {}, // TODO: ParameterStatus, else => return error.UnexpectedDBMessage, } } }
const ListenOpts = struct { timeout: u32 = 0, }; pub fn listen(self: *Listener, channel: []const u8, opts: ListenOpts) !void { // LISTEN doesn't support parameterized queries. It has to be a simple query. // We don't use proto.Query because we want to quote the identifier.
const buf = &self._buf; buf.reset();
// "LISTEN " = 7 // "IDENTIFIER" = 128 // max identifier size is 63, but if we need to quote every character, that's // 126. + 2 for the opening and closing quote // + 1 for null terminator try buf.ensureTotalCapacity(136); buf.writeByteAssumeCapacity('Q');
var len_view = try buf.skip(4);
buf.writeAssumeCapacity("LISTEN \"");
// + 4 for the length itself // + 7 for the LISTEN // + 2 for the quotes // + 1 for the null terminator var len = 11 + channel.len + 3; for (channel) |c| { if (c == '"') { len += 1; buf.writeAssumeCapacity("\"\""); } else { buf.writeByteAssumeCapacity(c); } } buf.writeByteAssumeCapacity('"'); buf.writeByteAssumeCapacity(0);
// fill in the length len_view.writeIntBig(u32, @intCast(len));
try self._stream.writeAll(buf.string());
{ // we expect a command complete ('C') const msg = try self.read(); switch (msg.type) { 'C' => {}, else => return error.UnexpectedDBMessage, } }
{ // followed by a ReadyForQuery ('Z') const msg = try self.read(); switch (msg.type) { 'Z' => {}, else => return error.UnexpectedDBMessage, } }
try self._reader.startFlow(null, opts.timeout); }
pub fn next(self: *Listener) ?NotificationResponse { if (@atomicLoad(bool, &self.closed, .acquire) == true) { return null; }
const msg = self.read() catch |err| { self.err = .{ .err = err }; return null; };
switch (msg.type) { 'A' => return NotificationResponse.parse(msg.data) catch |err| { self.err = .{ .err = err }; return null; }, else => { self.err = .{ .err = error.UnexpectedDBMessage }; return null; }, } }
fn read(self: *Listener) !lib.Message { var reader = &self._reader; while (true) { const msg = try reader.next(); switch (msg.type) { 'N' => {}, // TODO: NoticeResponse 'E' => return self.setErr(msg.data), else => return msg, } } }
fn setErr(self: *Listener, data: []const u8) error{ PG, OutOfMemory } { const allocator = self._allocator;
// The proto.Error that we're about to create is going to reference data. // But data is owned by our Reader and its lifetime doesn't necessarily match // what we want here. So we're going to dupe it and make the connection own // the data so it can tie its lifecycle to the error.
// That means clearing out any previous duped error data we had if (self._err_data) |err_data| { allocator.free(err_data); }
const owned = try allocator.dupe(u8, data); self._err_data = owned; self.err = .{ .pg = proto.Error.parse(owned) }; return error.PG; }};
const t = lib.testing;test "Listener" { var l = try Listener.open(t.io, t.allocator, t.connectOpts()); defer l.deinit(); try l.auth(t.authOpts(.{})); try testListener(&l);}
test "Listener: from Pool" { var pool = try lib.Pool.init(t.io, t.allocator, .{ .size = 1, .auth = t.authOpts(.{}), .connect = t.connectOpts(), }); defer pool.deinit();
var l = try pool.newListener(); defer l.deinit();
try testListener(&l);}
fn testListener(l: *Listener) !void { const io = t.io; var reset: std.Io.Event = .unset; var tt = try std.Thread.spawn(.{}, struct { fn shutdown(io_p: Io, ll: *Listener, r: *std.Io.Event) !void { try r.wait(io_p); try ll.stop(); } }.shutdown, .{ io, l, &reset }); tt.detach();
try l.listen("chan-1", .{}); try l.listen("chan_2", .{});
const thrd = try std.Thread.spawn(.{}, testNotifier, .{}); { const notification = l.next().?; try t.expectString("chan-1", notification.channel); try t.expectString("pl-1", notification.payload); }
{ const notification = l.next().?; try t.expectString("chan_2", notification.channel); try t.expectString("pl-2", notification.payload); }
{ const notification = l.next().?; try t.expectString("chan-1", notification.channel); try t.expectString("", notification.payload); }
reset.set(io); try t.expectEqual(null, l.next()); thrd.join();}
fn testNotifier() !void { var c = try t.connect(.{}); defer c.deinit(); _ = c.exec("select pg_notify($1, $2)", .{ "chan_x", "pl-x" }) catch unreachable; _ = c.exec("select pg_notify($1, $2)", .{ "chan-1", "pl-1" }) catch unreachable; _ = c.exec("select pg_notify($1, $2)", .{ "chan_2", "pl-2" }) catch unreachable; _ = c.exec("select pg_notify($1, null)", .{"chan-1"}) catch unreachable;}