karlseguin.tngl.sh
Native PostgreSQL driver / client for Zig
1const std = @import("std");
2const lib = @import("lib.zig");
3const Buffer = @import("buffer").Buffer;
4
5const proto = lib.proto;
6const Conn = lib.Conn;
7const Reader = lib.Reader;
8const NotificationResponse = lib.proto.NotificationResponse;
9
10const Stream = lib.Stream;
11const Allocator = std.mem.Allocator;
12const Io = std.Io;
13
14const ListenError = union(enum) {
15 err: anyerror,
16 pg: lib.proto.Error,
17};
18
19pub const Listener = struct {
20 err: ?ListenError = null,
21 closed: bool = false,
22
23 _stream: Stream,
24
25 // A buffer used for writing to PG. This can grow dynamically as needed.
26 _buf: Buffer,
27
28 // Used to read data from PG. Has its own buffer which can grow dynamically
29 _reader: Reader,
30
31 // If we get a PG error, we'll return a LIstenError.pg, and we'll own its
32 // memory.
33 _err_data: ?[]const u8 = null,
34
35 _allocator: Allocator,
36
37 _io: Io,
38
39 pub fn open(io: Io, allocator: Allocator, opts: Conn.Opts) !Listener {
40 var stream = try Stream.connect(io, allocator, opts, null);
41 errdefer stream.close();
42
43 const buf = try Buffer.init(allocator, opts.write_buffer orelse 2048);
44 errdefer buf.deinit();
45
46 const reader = try Reader.init(allocator, opts.read_buffer orelse 4096, stream);
47 errdefer reader.deinit();
48
49 return .{
50 ._buf = buf,
51 ._stream = stream,
52 ._reader = reader,
53 ._allocator = allocator,
54 ._io = io,
55 };
56 }
57
58 pub fn deinit(self: *Listener) void {
59 if (self._err_data) |err_data| {
60 self._allocator.free(err_data);
61 }
62 self._buf.deinit();
63 self._reader.deinit();
64
65 self.stop() catch {};
66 self._stream.close();
67 }
68
69 pub fn stop(self: *Listener) !void {
70 if (@atomicRmw(bool, &self.closed, .Xchg, true, .monotonic) == true) {
71 return;
72 }
73
74 lib.sendTerminate(&self._stream, self._io);
75 return self._stream.shutdown(.both);
76 }
77
78 pub fn auth(self: *Listener, opts: Conn.AuthOpts) !void {
79 if (try lib.auth.auth(self._io, &self._stream, &self._buf, &self._reader, opts)) |raw_pg_err| {
80 return self.setErr(raw_pg_err);
81 }
82
83 while (true) {
84 const msg = try self.read();
85 switch (msg.type) {
86 'Z' => return,
87 'K' => {}, // TODO: BackendKeyData
88 'S' => {}, // TODO: ParameterStatus,
89 else => return error.UnexpectedDBMessage,
90 }
91 }
92 }
93
94 const ListenOpts = struct {
95 timeout: u32 = 0,
96 };
97 pub fn listen(self: *Listener, channel: []const u8, opts: ListenOpts) !void {
98 // LISTEN doesn't support parameterized queries. It has to be a simple query.
99 // We don't use proto.Query because we want to quote the identifier.
100
101 const buf = &self._buf;
102 buf.reset();
103
104 // "LISTEN " = 7
105 // "IDENTIFIER" = 128
106 // max identifier size is 63, but if we need to quote every character, that's
107 // 126. + 2 for the opening and closing quote
108 // + 1 for null terminator
109 try buf.ensureTotalCapacity(136);
110 buf.writeByteAssumeCapacity('Q');
111
112 var len_view = try buf.skip(4);
113
114 buf.writeAssumeCapacity("LISTEN \"");
115
116 // + 4 for the length itself
117 // + 7 for the LISTEN
118 // + 2 for the quotes
119 // + 1 for the null terminator
120 var len = 11 + channel.len + 3;
121 for (channel) |c| {
122 if (c == '"') {
123 len += 1;
124 buf.writeAssumeCapacity("\"\"");
125 } else {
126 buf.writeByteAssumeCapacity(c);
127 }
128 }
129 buf.writeByteAssumeCapacity('"');
130 buf.writeByteAssumeCapacity(0);
131
132 // fill in the length
133 len_view.writeIntBig(u32, @intCast(len));
134
135 try self._stream.writeAll(buf.string());
136
137 {
138 // we expect a command complete ('C')
139 const msg = try self.read();
140 switch (msg.type) {
141 'C' => {},
142 else => return error.UnexpectedDBMessage,
143 }
144 }
145
146 {
147 // followed by a ReadyForQuery ('Z')
148 const msg = try self.read();
149 switch (msg.type) {
150 'Z' => {},
151 else => return error.UnexpectedDBMessage,
152 }
153 }
154
155 try self._reader.startFlow(null, opts.timeout);
156 }
157
158 pub fn next(self: *Listener) ?NotificationResponse {
159 if (@atomicLoad(bool, &self.closed, .acquire) == true) {
160 return null;
161 }
162
163 const msg = self.read() catch |err| {
164 self.err = .{ .err = err };
165 return null;
166 };
167
168 switch (msg.type) {
169 'A' => return NotificationResponse.parse(msg.data) catch |err| {
170 self.err = .{ .err = err };
171 return null;
172 },
173 else => {
174 self.err = .{ .err = error.UnexpectedDBMessage };
175 return null;
176 },
177 }
178 }
179
180 fn read(self: *Listener) !lib.Message {
181 var reader = &self._reader;
182 while (true) {
183 const msg = try reader.next();
184 switch (msg.type) {
185 'N' => {}, // TODO: NoticeResponse
186 'E' => return self.setErr(msg.data),
187 else => return msg,
188 }
189 }
190 }
191
192 fn setErr(self: *Listener, data: []const u8) error{ PG, OutOfMemory } {
193 const allocator = self._allocator;
194
195 // The proto.Error that we're about to create is going to reference data.
196 // But data is owned by our Reader and its lifetime doesn't necessarily match
197 // what we want here. So we're going to dupe it and make the connection own
198 // the data so it can tie its lifecycle to the error.
199
200 // That means clearing out any previous duped error data we had
201 if (self._err_data) |err_data| {
202 allocator.free(err_data);
203 }
204
205 const owned = try allocator.dupe(u8, data);
206 self._err_data = owned;
207 self.err = .{ .pg = proto.Error.parse(owned) };
208 return error.PG;
209 }
210};
211
212const t = lib.testing;
213test "Listener" {
214 var l = try Listener.open(t.io, t.allocator, .{ .host = "127.0.0.1" });
215 defer l.deinit();
216 try l.auth(t.authOpts(.{}));
217 try testListener(&l);
218}
219
220test "Listener: from Pool" {
221 var pool = try lib.Pool.init(t.io, t.allocator, .{
222 .size = 1,
223 .auth = t.authOpts(.{}),
224 });
225 defer pool.deinit();
226
227 var l = try pool.newListener();
228 defer l.deinit();
229
230 try testListener(&l);
231}
232
233fn testListener(l: *Listener) !void {
234 const io = t.io;
235 var reset: std.Io.Event = .unset;
236 var tt = try std.Thread.spawn(.{}, struct {
237 fn shutdown(io_p: Io, ll: *Listener, r: *std.Io.Event) !void {
238 try r.wait(io_p);
239 try ll.stop();
240 }
241 }.shutdown, .{ io, l, &reset });
242 tt.detach();
243
244 try l.listen("chan-1", .{});
245 try l.listen("chan_2", .{});
246
247 const thrd = try std.Thread.spawn(.{}, testNotifier, .{});
248 {
249 const notification = l.next().?;
250 try t.expectString("chan-1", notification.channel);
251 try t.expectString("pl-1", notification.payload);
252 }
253
254 {
255 const notification = l.next().?;
256 try t.expectString("chan_2", notification.channel);
257 try t.expectString("pl-2", notification.payload);
258 }
259
260 {
261 const notification = l.next().?;
262 try t.expectString("chan-1", notification.channel);
263 try t.expectString("", notification.payload);
264 }
265
266 reset.set(io);
267 try t.expectEqual(null, l.next());
268 thrd.join();
269}
270
271fn testNotifier() !void {
272 var c = try t.connect(.{});
273 defer c.deinit();
274 _ = c.exec("select pg_notify($1, $2)", .{ "chan_x", "pl-x" }) catch unreachable;
275 _ = c.exec("select pg_notify($1, $2)", .{ "chan-1", "pl-1" }) catch unreachable;
276 _ = c.exec("select pg_notify($1, $2)", .{ "chan_2", "pl-2" }) catch unreachable;
277 _ = c.exec("select pg_notify($1, null)", .{"chan-1"}) catch unreachable;
278}