postgres.bend

bend-kit-postgres@0.3.0.0 · 677 lines · 82 definitions

1# Postgres client: protocol 3.0 over TCP or TLS, SCRAM-SHA-256, prepared queries, and a pool. Source: https://github.com/paymog/bend-kit/tree/main/postgres
2import Base
3import bend-kit-bytes@0.3.3.0/bytes.bend as Bytes
4import bend-kit-wire@0.4.6.2/wire.bend as Wire
5import bend-kit-dns@0.7.0.0/dns.bend as Dns
6import bend-kit-postgres@0.3.0.0/codec.bend as Codec
7import bend-kit-postgres@0.3.0.0/scram.bend as Scram
8
9# host is a name or a numeric address; TLS checks the certificate against it. With tls
10# set, a server that refuses SSLRequest fails with NoTls rather than falling back to TCP.
11# ms bounds each connect and read (0: none). user, pass, and db are byte strings (one Char per octet).
12type Config is Data:
13  Config{host: String, port: U32, tls: Bool, user: String, pass: String, db: String, ms: U32}
14
15# TLS on, 10 s deadlines.
16def config(+host: String, +port: U32, +user: String, +pass: String, +db: String) -> Config:
17  Config{host, port, True{}, user, pass, db, 10000}
18
19# A server error: the SQLSTATE code (for example "42P01") and the message.
20type Err is Data:
21  Err{code: String, message: String}
22
23# A failed call closes its connection. Rejected carries the server's error during startup;
24# Auth is an authentication method this client does not speak, or a SCRAM check that failed.
25type Failure is Type:
26  IoFail{code: U32, why: String}
27  Closed{}
28  Malformed{}
29  NoTls{}
30  Auth{why: String}
31  Rejected{err: Err}
32
33# A query's result. Rows holds the columns, the rows in order (None is NULL, values in text
34# format), and the command tag ("SELECT 2", "INSERT 0 1"). Failed is an error the server
35# reported for this query; the connection stays in step and can run the next one.
36type Reply is Type:
37  Rows{fields: List<&1, Codec.Field>, rows: List<&1, List<&1, Maybe<&1, Bytes.Bytes>>>, tag: Bytes.Bytes}
38  Failed{err: Err}
39
40# key names the pool bucket: TLS, user, host, port, and db (never the password).
41# status is the last ReadyForQuery byte: 73 'I' idle, 84 'T' in a transaction, 69 'E' failed one.
42type Conn is Type:
43  Conn{tls: Bool, sock: Socket, rd: Codec.Reader, key: String, ms: U32, status: U32}
44
45# A parameter in text format.
46def text(+s: String) -> Maybe<&1, Bytes.Bytes>:
47  Some{Bytes.from_string(s)}
48
49def bytes(+s: String) -> Bytes.Bytes:
50  Bytes.from_string(s)
51
52# ---- errors
53
54def err.put(+k: U32, +s: String, e: Err) -> Err:
55  Err{+code, +message} = e
56  Bool.pick(Err, U32.is_eq(k, 67), Err{s, message}, Bool.pick(Err, U32.is_eq(k, 77), Err{code, s}, Err{code, message}))
57
58def err.go(fs: List<&1, Codec.Notice>, e: Err) -> Err:
59  match fs:
60    case Nil{}:
61      e
62    case Con{Codec.Notice{+k, v}, t}:
63      err.go(t, err.put(k, Bytes.to_string(v), e))
64
65# The C (SQLSTATE) and M (message) fields of an ErrorResponse.
66def err.of(fs: List<&1, Codec.Notice>) -> Err:
67  err.go(fs, Err{"", ""})
68
69# ---- socket IO
70
71def io.close(tls: Bool, s: Socket) -> IO(Unit):
72  match tls:
73    case True{}:
74      Wire.tls.close(s)
75    case False{}:
76      Socket.close(s)
77
78def io.send(tls: Bool, s: Socket, +len: U32, buf: Array<U32>) -> IO(Socket & Result<&1, &1, U32 & String, Unit>):
79  match tls:
80    case True{}:
81      Wire.tls.send.words(s, len, buf)
82    case False{}:
83      Wire.send.words(s, len, buf)
84
85def io.write(tls: Bool, s: Socket, data: Bytes.Bytes) -> IO(Socket & Result<&1, &1, U32 & String, Unit>):
86  Bytes.Bytes{+len, buf} = data
87  io.send(tls, s, len, buf)
88
89def io.recv(tls: Bool, s: Socket, +ms: U32) -> IO(Socket & Result<&1, &1, U32 & String, U32 & Array<U32>>):
90  match tls:
91    case True{}:
92      Wire.tls.recv.words(s, 65536, ms)
93    case False{}:
94      Wire.recv.words(s, 65536, ms)
95
96def io.close.after(tls: Bool, m: Socket & Result<&1, &1, U32 & String, Unit>) -> IO(Unit):
97  (s, r) = m
98  io.close(tls, s)
99
100# Sends Terminate, then closes the socket.
101def close(c: Conn) -> IO(Unit):
102  Conn{+tls, sock, rd, key, ms, status} = c
103  do IO<Unit>:
104    m : Socket & Result<&1, &1, U32 & String, Unit> <- io.write(tls, sock, Codec.encode(Codec.Terminate{}))
105    io.close.after(tls, m)
106
107def abort(-A: Type, c: Conn, why: Failure) -> IO(Result<&1, &1, Failure, A>):
108  do IO<Result<&1, &1, Failure, A>>:
109    close(c)
110    return Fail{why}
111
112def reject(-A: Type, c: Conn, fs: List<&1, Codec.Notice>) -> IO(Result<&1, &1, Failure, A>):
113  abort(A, c, Rejected{err.of(fs)})
114
115def send.after(+tls: Bool, rd: Codec.Reader, +key: String, +ms: U32, +status: U32, m: Socket & Result<&1, &1, U32 & String, Unit>) -> IO(Result<&1, &1, Failure, Conn>):
116  (s, r) = m
117  match r:
118    case Fail{(code, why)}:
119      do IO<Result<&1, &1, Failure, Conn>>:
120        io.close(tls, s)
121        return Fail{IoFail{code, why}}
122    case Done{u}:
123      IO.pure(Result<&1, &1, Failure, Conn>, Done{Conn{tls, s, rd, key, ms, status}})
124
125def send(c: Conn, data: Bytes.Bytes) -> IO(Result<&1, &1, Failure, Conn>):
126  Conn{+tls, sock, rd, +key, +ms, +status} = c
127  do IO<Result<&1, &1, Failure, Conn>>:
128    m : Socket & Result<&1, &1, U32 & String, Unit> <- io.write(tls, sock, data)
129    send.after(tls, rd, key, ms, status, m)
130
131# RNext holds the reader after its last decode attempt; RFail a socket to close.
132type Rd is Type:
133  RNext{tls: Bool, sock: Socket, key: String, ms: U32, status: U32, r: Codec.Reader & Codec.Next}
134  RFail{tls: Bool, sock: Socket, why: Failure}
135
136def got.len(zero: Bool, +tls: Bool, +key: String, +ms: U32, +status: U32, rd: Codec.Reader, s: Socket, +len: U32, words: Array<U32>) -> Rd:
137  match zero:
138    case True{}:
139      RFail{tls, s, Closed{}}
140    case False{}:
141      RNext{tls, s, key, ms, status, Codec.next(Codec.feed(rd, Bytes.Bytes{len, words}))}
142
143def got(+tls: Bool, +key: String, +ms: U32, +status: U32, rd: Codec.Reader, m: Socket & Result<&1, &1, U32 & String, U32 & Array<U32>>) -> Rd:
144  (s, r) = m
145  match r:
146    case Fail{(code, why)}:
147      RFail{tls, s, IoFail{code, why}}
148    case Done{(+len, words)}:
149      got.len(U32.is_zero(len), tls, key, ms, status, rd, s, len, words)
150
151def read.fail(tls: Bool, sock: Socket, why: Failure) -> IO(Result<&1, &1, Failure, Conn & Codec.Back>):
152  do IO<Result<&1, &1, Failure, Conn & Codec.Back>>:
153    io.close(tls, sock)
154    return Fail{why}
155
156# One recv per step. NoticeResponse and ParameterStatus may arrive at any time; they are skipped.
157def read.loop(fuel: Nat, st: Rd) -> IO(Result<&1, &1, Failure, Conn & Codec.Back>):
158  match fuel:
159    case 0n:
160      match st:
161        case RNext{tls, sock, key, ms, status, r}:
162          read.fail(tls, sock, Malformed{})
163        case RFail{tls, sock, why}:
164          read.fail(tls, sock, why)
165    case 1n+f:
166      match st:
167        case RFail{tls, sock, why}:
168          read.fail(tls, sock, why)
169        case RNext{+tls, sock, +key, +ms, +status, r}:
170          (rd, nx) = r
171          match nx:
172            case Codec.Got{Codec.NoticeResponse{fs}}:
173              read.loop(f, RNext{tls, sock, key, ms, status, Codec.next(rd)})
174            case Codec.Got{Codec.ParameterStatus{name, value}}:
175              read.loop(f, RNext{tls, sock, key, ms, status, Codec.next(rd)})
176            case Codec.Got{m}:
177              IO.pure(Result<&1, &1, Failure, Conn & Codec.Back>, Done{(Conn{tls, sock, rd, key, ms, status}, m)})
178            case Codec.Bad{}:
179              read.fail(tls, sock, Malformed{})
180            case Codec.Want{}:
181              do IO<Result<&1, &1, Failure, Conn & Codec.Back>>:
182                m : Socket & Result<&1, &1, U32 & String, U32 & Array<U32>> <- io.recv(tls, sock, ms)
183                read.loop(f, got(tls, key, ms, status, rd, m))
184
185# The next message. Bytes already read are decoded before the socket is asked for more.
186# 2^24 reads of up to 64 KiB pass Postgres's 1 GiB field limit.
187def read(c: Conn) -> IO(Result<&1, &1, Failure, Conn & Codec.Back>):
188  Conn{tls, sock, rd, key, ms, status} = c
189  read.loop(Nat.mul(65536n, 256n), RNext{tls, sock, key, ms, status, Codec.next(rd)})
190
191# The next message after a send.
192def reply(r: Result<&1, &1, Failure, Conn>) -> IO(Result<&1, &1, Failure, Conn & Codec.Back>):
193  match r:
194    case Fail{e}:
195      IO.pure(Result<&1, &1, Failure, Conn & Codec.Back>, Fail{e})
196    case Done{c}:
197      read(c)
198
199def exchange(c: Conn, m: Codec.Front) -> IO(Result<&1, &1, Failure, Conn & Codec.Back>):
200  do IO<Result<&1, &1, Failure, Conn & Codec.Back>>:
201    s : Result<&1, &1, Failure, Conn> <- send(c, Codec.encode(m))
202    reply(s)
203
204def set.status(c: Conn, +status: U32) -> Conn:
205  Conn{tls, sock, rd, key, ms, old} = c
206  Conn{tls, sock, rd, key, ms, status}
207
208# ---- startup
209
210# BackendKeyData, then ReadyForQuery.
211def ready.loop(fuel: Nat, r: Result<&1, &1, Failure, Conn & Codec.Back>) -> IO(Result<&1, &1, Failure, Conn>):
212  match fuel:
213    case 0n:
214      match r:
215        case Fail{e}:
216          IO.pure(Result<&1, &1, Failure, Conn>, Fail{e})
217        case Done{(c, m)}:
218          abort(Conn, c, Malformed{})
219    case 1n+f:
220      match r:
221        case Fail{e}:
222          IO.pure(Result<&1, &1, Failure, Conn>, Fail{e})
223        case Done{(c, m)}:
224          match m:
225            case Codec.ReadyForQuery{+status}:
226              IO.pure(Result<&1, &1, Failure, Conn>, Done{set.status(c, status)})
227            case Codec.BackendKeyData{pid, key}:
228              do IO<Result<&1, &1, Failure, Conn>>:
229                m2 : Result<&1, &1, Failure, Conn & Codec.Back> <- read(c)
230                ready.loop(f, m2)
231            case Codec.ErrorResponse{fs}:
232              reject(Conn, c, fs)
233            case _:
234              abort(Conn, c, Malformed{})
235
236def ready(c: Conn) -> IO(Result<&1, &1, Failure, Conn>):
237  do IO<Result<&1, &1, Failure, Conn>>:
238    m : Result<&1, &1, Failure, Conn & Codec.Back> <- read(c)
239    ready.loop(16n, m)
240
241# AuthenticationOk, after a password or the SCRAM exchange.
242def auth.done(r: Result<&1, &1, Failure, Conn & Codec.Back>) -> IO(Result<&1, &1, Failure, Conn>):
243  match r:
244    case Fail{e}:
245      IO.pure(Result<&1, &1, Failure, Conn>, Fail{e})
246    case Done{(c, m)}:
247      match m:
248        case Codec.AuthOk{}:
249          ready(c)
250        case Codec.ErrorResponse{fs}:
251          reject(Conn, c, fs)
252        case _:
253          abort(Conn, c, Malformed{})
254
255def sasl.verified(c: Conn, v: Result<&1, &1, U32 & String, Unit>) -> IO(Result<&1, &1, Failure, Conn>):
256  match v:
257    case Fail{(code, why)}:
258      abort(Conn, c, Auth{why})
259    case Done{u}:
260      do IO<Result<&1, &1, Failure, Conn>>:
261        m : Result<&1, &1, Failure, Conn & Codec.Back> <- read(c)
262        auth.done(m)
263
264# AuthenticationSASLFinal carries the server signature, which proves the server knew the password.
265def sasl.final(fin: Scram.Final, r: Result<&1, &1, Failure, Conn & Codec.Back>) -> IO(Result<&1, &1, Failure, Conn>):
266  match r:
267    case Fail{e}:
268      IO.pure(Result<&1, &1, Failure, Conn>, Fail{e})
269    case Done{(c, m)}:
270      match m:
271        case Codec.AuthSASLFinal{data}:
272          do IO<Result<&1, &1, Failure, Conn>>:
273            v : Result<&1, &1, U32 & String, Unit> <- Scram.verify(fin, Bytes.to_string(data))
274            sasl.verified(c, v)
275        case Codec.ErrorResponse{fs}:
276          reject(Conn, c, fs)
277        case _:
278          abort(Conn, c, Malformed{})
279
280def sasl.respond(c: Conn, f: Result<&1, &1, U32 & String, Scram.Final & String>) -> IO(Result<&1, &1, Failure, Conn>):
281  match f:
282    case Fail{(code, why)}:
283      abort(Conn, c, Auth{why})
284    case Done{(fin, msg)}:
285      do IO<Result<&1, &1, Failure, Conn>>:
286        m : Result<&1, &1, Failure, Conn & Codec.Back> <- exchange(c, Codec.SASLResponse{bytes(msg)})
287        sasl.final(fin, m)
288
289def sasl.cont(st: Scram.First, +pass: String, r: Result<&1, &1, Failure, Conn & Codec.Back>) -> IO(Result<&1, &1, Failure, Conn>):
290  match r:
291    case Fail{e}:
292      IO.pure(Result<&1, &1, Failure, Conn>, Fail{e})
293    case Done{(c, m)}:
294      match m:
295        case Codec.AuthSASLContinue{data}:
296          do IO<Result<&1, &1, Failure, Conn>>:
297            f : Result<&1, &1, U32 & String, Scram.Final & String> <- Scram.respond(st, Bytes.to_string(data), pass)
298            sasl.respond(c, f)
299        case Codec.ErrorResponse{fs}:
300          reject(Conn, c, fs)
301        case _:
302          abort(Conn, c, Malformed{})
303
304def sasl.first(c: Conn, +pass: String, f: Result<&1, &1, U32 & String, Scram.First & String>) -> IO(Result<&1, &1, Failure, Conn>):
305  match f:
306    case Fail{(code, why)}:
307      abort(Conn, c, Auth{why})
308    case Done{(st, msg)}:
309      do IO<Result<&1, &1, Failure, Conn>>:
310        m : Result<&1, &1, Failure, Conn & Codec.Back> <- exchange(c, Codec.SASLInitial{bytes(Scram.mechanism()), bytes(msg)})
311        sasl.cont(st, pass, m)
312
313def sasl(ok: Bool, c: Conn, +user: String, +pass: String) -> IO(Result<&1, &1, Failure, Conn>):
314  match ok:
315    case True{}:
316      do IO<Result<&1, &1, Failure, Conn>>:
317        f : Result<&1, &1, U32 & String, Scram.First & String> <- Scram.first(user)
318        sasl.first(c, pass, f)
319    case False{}:
320      abort(Conn, c, Auth{"server offers no " ++ Scram.mechanism()})
321
322def has.mech(xs: List<&1, Bytes.Bytes>) -> Bool:
323  match xs:
324    case Nil{}:
325      False{}
326    case Con{h, t}:
327      Bool.or(String.eq(Bytes.to_string(h), Scram.mechanism()), has.mech(t))
328
329# The first message after StartupMessage picks the authentication method.
330def auth(+user: String, +pass: String, r: Result<&1, &1, Failure, Conn & Codec.Back>) -> IO(Result<&1, &1, Failure, Conn>):
331  match r:
332    case Fail{e}:
333      IO.pure(Result<&1, &1, Failure, Conn>, Fail{e})
334    case Done{(c, m)}:
335      match m:
336        case Codec.AuthOk{}:
337          ready(c)
338        case Codec.AuthCleartext{}:
339          do IO<Result<&1, &1, Failure, Conn>>:
340            a : Result<&1, &1, Failure, Conn & Codec.Back> <- exchange(c, Codec.Password{bytes(pass)})
341            auth.done(a)
342        case Codec.AuthSASL{mechs}:
343          sasl(has.mech(mechs), c, user, pass)
344        case Codec.ErrorResponse{fs}:
345          reject(Conn, c, fs)
346        case Codec.AuthOther{+code, data}:
347          abort(Conn, c, Auth{"unsupported authentication request " ++ U32.show(code)})
348        case _:
349          abort(Conn, c, Malformed{})
350
351def startup(c: Conn, +user: String, +pass: String, +db: String) -> IO(Result<&1, &1, Failure, Conn>):
352  do IO<Result<&1, &1, Failure, Conn>>:
353    m : Result<&1, &1, Failure, Conn & Codec.Back> <- exchange(c, Codec.Startup{[Codec.Param{bytes("user"), bytes(user)}, Codec.Param{bytes("database"), bytes(db)}]})
354    auth(user, pass, m)
355
356def tls.done(+key: String, +ms: U32, +user: String, +pass: String, +db: String, m: Socket & Result<&1, &1, U32 & String, Unit>) -> IO(Result<&1, &1, Failure, Conn>):
357  (s, r) = m
358  match r:
359    case Fail{(code, why)}:
360      do IO<Result<&1, &1, Failure, Conn>>:
361        Wire.tls.close(s)
362        return Fail{IoFail{code, why}}
363    case Done{u}:
364      startup(Conn{True{}, s, Codec.reader(), key, ms, 0}, user, pass, db)
365
366def plain.fail(s: Socket, why: Failure) -> IO(Result<&1, &1, Failure, Conn>):
367  do IO<Result<&1, &1, Failure, Conn>>:
368    Socket.close(s)
369    return Fail{why}
370
371# 'S' starts TLS; 'N' refuses it. Anything else (an old server's error) is not an answer.
372def ssl.byte(s: Socket, +key: String, +host: String, +user: String, +pass: String, +db: String, +ms: U32, r: Bytes.Bytes & Maybe<&2, U32>) -> IO(Result<&1, &1, Failure, Conn>):
373  (b, got) = r
374  match got:
375    case None{}:
376      plain.fail(s, Closed{})
377    case Some{c}:
378      match c:
379        case 83:
380          do IO<Result<&1, &1, Failure, Conn>>:
381            m : Socket & Result<&1, &1, U32 & String, Unit> <- Wire.tls.connect(s, host, ms)
382            tls.done(key, ms, user, pass, db, m)
383        case 78:
384          plain.fail(s, NoTls{})
385        case _:
386          plain.fail(s, Malformed{})
387
388# The answer is one byte, read alone so no byte sent before the TLS handshake is taken as data.
389def ssl.answer(+key: String, +host: String, +user: String, +pass: String, +db: String, +ms: U32, a: Socket & Result<&1, &1, U32 & String, U32 & Array<U32>>) -> IO(Result<&1, &1, Failure, Conn>):
390  (s, r) = a
391  match r:
392    case Fail{(code, why)}:
393      plain.fail(s, IoFail{code, why})
394    case Done{(+len, words)}:
395      ssl.byte(s, key, host, user, pass, db, ms, Bytes.get(Bytes.Bytes{len, words}, 0))
396
397def ssl.sent(+key: String, +host: String, +user: String, +pass: String, +db: String, +ms: U32, m: Socket & Result<&1, &1, U32 & String, Unit>) -> IO(Result<&1, &1, Failure, Conn>):
398  (s, r) = m
399  match r:
400    case Fail{(code, why)}:
401      plain.fail(s, IoFail{code, why})
402    case Done{u}:
403      do IO<Result<&1, &1, Failure, Conn>>:
404        a : Socket & Result<&1, &1, U32 & String, U32 & Array<U32>> <- Wire.recv.words(s, 1, ms)
405        ssl.answer(key, host, user, pass, db, ms, a)
406
407def ssl.ask(s: Socket, +key: String, +host: String, +user: String, +pass: String, +db: String, +ms: U32) -> IO(Result<&1, &1, Failure, Conn>):
408  do IO<Result<&1, &1, Failure, Conn>>:
409    m : Socket & Result<&1, &1, U32 & String, Unit> <- io.write(False{}, s, Codec.encode(Codec.SSLRequest{}))
410    ssl.sent(key, host, user, pass, db, ms, m)
411
412def dialed(r: Result<&1, &1, U32 & String, Socket>, tls: Bool, +key: String, +host: String, +user: String, +pass: String, +db: String, +ms: U32) -> IO(Result<&1, &1, Failure, Conn>):
413  match r:
414    case Fail{(code, why)}:
415      IO.pure(Result<&1, &1, Failure, Conn>, Fail{IoFail{code, why}})
416    case Done{s}:
417      match tls:
418        case True{}:
419          ssl.ask(s, key, host, user, pass, db, ms)
420        case False{}:
421          startup(Conn{False{}, s, Codec.reader(), key, ms, 0}, user, pass, db)
422
423def dial.after(r: Result<&1, &1, U32 & String, Socket>, +ip: String, +port: U32, +ms: U32) -> IO(Result<&1, &1, U32 & String, Socket>):
424  match r:
425    case Done{s}:
426      IO.pure(Result<&1, &1, U32 & String, Socket>, Done{s})
427    case Fail{e}:
428      Wire.connect(ip, port, ms)
429
430def dial.more(prev: IO(Result<&1, &1, U32 & String, Socket>), +ip: String, +port: U32, +ms: U32) -> IO(Result<&1, &1, U32 & String, Socket>):
431  do IO<Result<&1, &1, U32 & String, Socket>>:
432    r : Result<&1, &1, U32 & String, Socket> <- prev
433    dial.after(r, ip, port, ms)
434
435# Each address in turn, until one connects.
436def dial(ips: List<&2, String>, +port: U32, +ms: U32, acc: IO(Result<&1, &1, U32 & String, Socket>)) -> IO(Result<&1, &1, U32 & String, Socket>):
437  match ips:
438    case Nil{}:
439      acc
440    case Con{ip, t}:
441      dial(t, port, ms, dial.more(acc, ip, port, ms))
442
443def key(cfg: Config) -> String:
444  Config{+host, +port, +tls, +user, pass, +db, ms} = cfg
445  +sep = {SCon{Chr{0}, SNil{}} : String}
446  Bool.pick(String, tls, "tls", "tcp") ++ sep ++ user ++ sep ++ host ++ sep ++ U32.show(port) ++ sep ++ db
447
448def connect.to(+k: String, cfg: Config) -> IO(Result<&1, &1, Failure, Conn>):
449  Config{+host, +port, +tls, +user, +pass, +db, +ms} = cfg
450  do IO<Result<&1, &1, Failure, Conn>>:
451    ips : List<&2, String> <- Dns.resolve.all(host)
452    r : Result<&1, &1, U32 & String, Socket> <- dial(ips, port, ms, IO.pure(Result<&1, &1, U32 & String, Socket>, Fail{(0, "no address for " ++ host)}))
453    dialed(r, tls, k, host, user, pass, db, ms)
454
455# TCP, SSLRequest and TLS when cfg.tls (checking the certificate and host name), StartupMessage,
456# cleartext or SCRAM-SHA-256 authentication, then ReadyForQuery. MD5 is refused as Auth.
457def connect(+cfg: Config) -> IO(Result<&1, &1, Failure, Conn>):
458  connect.to(key(cfg), cfg)
459
460# ---- queries
461
462type Acc is Type:
463  Acc{fields: List<&1, Codec.Field>, rows: List<&1, List<&1, Maybe<&1, Bytes.Bytes>>>, tag: Bytes.Bytes, err: Maybe<&1, Err>}
464
465type Step is Type:
466  Go{acc: Acc}
467  Stop{acc: Acc, status: U32}
468  Odd{}
469
470def acc.fields(a: Acc, fs: List<&1, Codec.Field>) -> Acc:
471  Acc{old, rows, tag, err} = a
472  Acc{fs, rows, tag, err}
473
474# Rows are kept newest first until ReadyForQuery.
475def acc.row(a: Acc, cols: List<&1, Maybe<&1, Bytes.Bytes>>) -> Acc:
476  Acc{fs, rows, tag, err} = a
477  Acc{fs, Con{cols, rows}, tag, err}
478
479def acc.tag(a: Acc, t: Bytes.Bytes) -> Acc:
480  Acc{fs, rows, old, err} = a
481  Acc{fs, rows, t, err}
482
483def err.first(old: Maybe<&1, Err>, e: Err) -> Maybe<&1, Err>:
484  match old:
485    case None{}:
486      Some{e}
487    case Some{x}:
488      Some{x}
489
490def acc.err(a: Acc, e: Err) -> Acc:
491  Acc{fs, rows, tag, old} = a
492  Acc{fs, rows, tag, err.first(old, e)}
493
494def absorb(a: Acc, m: Codec.Back) -> Step:
495  match m:
496    case Codec.DataRow{cols}:
497      Go{acc.row(a, cols)}
498    case Codec.RowDescription{fs}:
499      Go{acc.fields(a, fs)}
500    case Codec.CommandComplete{t}:
501      Go{acc.tag(a, t)}
502    case Codec.ErrorResponse{fs}:
503      Go{acc.err(a, err.of(fs))}
504    case Codec.ReadyForQuery{+s}:
505      Stop{a, s}
506    case Codec.ParseComplete{}:
507      Go{a}
508    case Codec.BindComplete{}:
509      Go{a}
510    case Codec.NoData{}:
511      Go{a}
512    case Codec.EmptyQuery{}:
513      Go{a}
514    case Codec.CloseComplete{}:
515      Go{a}
516    case Codec.PortalSuspended{}:
517      Go{a}
518    case _:
519      Odd{}
520
521def absorbed(a: Acc, r: Result<&1, &1, Failure, Conn & Codec.Back>) -> Result<&1, &1, Failure, Conn & Step>:
522  match r:
523    case Fail{e}:
524      Fail{e}
525    case Done{(c, m)}:
526      Done{(c, absorb(a, m))}
527
528def finish(c: Conn, +status: U32, a: Acc) -> IO(Result<&1, &1, Failure, Conn & Reply>):
529  Acc{fs, rows, tag, err} = a
530  match err:
531    case None{}:
532      IO.pure(Result<&1, &1, Failure, Conn & Reply>, Done{(set.status(c, status), Rows{fs, List.reverse(&1, List<&1, Maybe<&1, Bytes.Bytes>>, rows), tag})})
533    case Some{e}:
534      IO.pure(Result<&1, &1, Failure, Conn & Reply>, Done{(set.status(c, status), Failed{e})})
535
536# After an ErrorResponse the server skips to Sync, so reading on to ReadyForQuery keeps
537# the connection in step. A message that has no place in the reply closes it.
538def collect(fuel: Nat, r: Result<&1, &1, Failure, Conn & Step>) -> IO(Result<&1, &1, Failure, Conn & Reply>):
539  match fuel:
540    case 0n:
541      match r:
542        case Fail{e}:
543          IO.pure(Result<&1, &1, Failure, Conn & Reply>, Fail{e})
544        case Done{(c, s)}:
545          abort(Conn & Reply, c, Malformed{})
546    case 1n+f:
547      match r:
548        case Fail{e}:
549          IO.pure(Result<&1, &1, Failure, Conn & Reply>, Fail{e})
550        case Done{(c, s)}:
551          match s:
552            case Go{a}:
553              do IO<Result<&1, &1, Failure, Conn & Reply>>:
554                m : Result<&1, &1, Failure, Conn & Codec.Back> <- read(c)
555                collect(f, absorbed(a, m))
556            case Stop{a, +status}:
557              finish(c, status, a)
558            case Odd{}:
559              abort(Conn & Reply, c, Malformed{})
560
561# The whole extended-query exchange for one statement, in one write.
562def query.bytes(+sql: String, params: List<&1, Maybe<&1, Bytes.Bytes>>) -> Bytes.Bytes:
563  Bytes.concat([
564    Codec.encode(Codec.Parse{Bytes.new(0), bytes(sql), Nil{}}),
565    Codec.encode(Codec.Bind{Bytes.new(0), Bytes.new(0), Nil{}, params, Nil{}}),
566    Codec.encode(Codec.Describe{80, Bytes.new(0)}),
567    Codec.encode(Codec.Execute{Bytes.new(0), 0}),
568    Codec.encode(Codec.Sync{})])
569
570# Runs sql (a byte string, parameters as $1, $2, ...) as an unnamed prepared statement:
571# Parse, Bind, Describe, Execute, Sync. Parameters and results are in text format; None is NULL.
572# A server error is a Failed reply on a live connection; a Failure has closed it.
573def query(c: Conn, +sql: String, params: List<&1, Maybe<&1, Bytes.Bytes>>) -> IO(Result<&1, &1, Failure, Conn & Reply>):
574  do IO<Result<&1, &1, Failure, Conn & Reply>>:
575    s : Result<&1, &1, Failure, Conn> <- send(c, query.bytes(sql, params))
576    m : Result<&1, &1, Failure, Conn & Codec.Back> <- reply(s)
577    collect(Nat.mul(65536n, 65536n), absorbed(Acc{Nil{}, Nil{}, Bytes.new(0), None{}}, m))
578
579# ---- pool
580
581# Idle connections per key, up to cap per key.
582type Pool is Type:
583  Pool{cap: U32, idle: Map<&1, List<&1, Conn>>}
584
585def pool.new.with(+cap: U32) -> Pool:
586  Pool{cap, Map.new(&1, List<&1, Conn>)}
587
588def pool.new() -> Pool:
589  pool.new.with(8)
590
591def pool.take(+cap: U32, +k: String, r: Map<&1, List<&1, Conn>> & Maybe<&1, List<&1, Conn>>) -> Pool & Maybe<&1, Conn>:
592  (m, got) = r
593  match got:
594    case None{}:
595      (Pool{cap, m}, None{})
596    case Some{xs}:
597      match xs:
598        case Nil{}:
599          (Pool{cap, m}, None{})
600        case Con{c, t}:
601          (Pool{cap, Map.set(&1, List<&1, Conn>, m, k, t)}, Some{c})
602
603def pool.opened(p: Pool, r: Result<&1, &1, Failure, Conn>) -> IO(Pool & Result<&1, &1, Failure, Conn>):
604  IO.pure(Pool & Result<&1, &1, Failure, Conn>, (p, r))
605
606def pool.got(cfg: Config, r: Pool & Maybe<&1, Conn>) -> IO(Pool & Result<&1, &1, Failure, Conn>):
607  (p, m) = r
608  match m:
609    case Some{c}:
610      IO.pure(Pool & Result<&1, &1, Failure, Conn>, (p, Done{c}))
611    case None{}:
612      do IO<Pool & Result<&1, &1, Failure, Conn>>:
613        c : Result<&1, &1, Failure, Conn> <- connect(cfg)
614        pool.opened(p, c)
615
616# The idle connection given back last, or a new one.
617# ponytail: an idle connection the server has since closed fails on its first query; ping before reuse if that bites.
618def pool.get(p: Pool, +cfg: Config) -> IO(Pool & Result<&1, &1, Failure, Conn>):
619  Pool{+cap, idle} = p
620  +k = key(cfg)
621  pool.got(cfg, pool.take(cap, k, Map.pop(&1, List<&1, Conn>, idle, k)))
622
623# The first n connections; the rest are closed.
624def conns.cut(xs: List<&1, Conn>, n: Nat) -> IO(List<&1, Conn>):
625  match xs:
626    case Nil{}:
627      IO.pure(List<&1, Conn>, Nil{})
628    case Con{c, t}:
629      match n:
630        case 0n:
631          do IO<List<&1, Conn>>:
632            close(c)
633            conns.cut(t, 0n)
634        case 1n+k:
635          do IO<List<&1, Conn>>:
636            r : List<&1, Conn> <- conns.cut(t, k)
637            return c <> r
638
639def pool.keep(+cap: U32, +k: String, c: Conn, r: Map<&1, List<&1, Conn>> & Maybe<&1, List<&1, Conn>>) -> IO(Pool):
640  (m, old) = r
641  do IO<Pool>:
642    xs : List<&1, Conn> <- conns.cut(c <> Maybe.default(&1, List<&1, Conn>, old, Nil{}), U32.to_nat(cap))
643    return Pool{cap, Map.set(&1, List<&1, Conn>, m, k, xs)}
644
645def pool.clean(clean: Bool, p: Pool, c: Conn) -> IO(Pool):
646  match clean:
647    case True{}:
648      Pool{+cap, idle} = p
649      Conn{tls, sock, rd, +k, ms, status} = c
650      pool.keep(cap, k, Conn{tls, sock, rd, k, ms, status}, Map.pop(&1, List<&1, Conn>, idle, k))
651    case False{}:
652      do IO<Pool>:
653        close(c)
654        return p
655
656def pool.left(p: Pool, +tls: Bool, sock: Socket, +key: String, +ms: U32, +status: U32, r: Codec.Reader & U32) -> IO(Pool):
657  (rd, +n) = r
658  pool.clean(Bool.and(U32.is_zero(n), U32.is_eq(status, 73)), p, Conn{tls, sock, rd, key, ms, status})
659
660# Gives c back for reuse. A connection with unread bytes is out of step, and one inside a
661# transaction would leak it to the next user, so either is closed instead.
662def pool.put(p: Pool, c: Conn) -> IO(Pool):
663  Conn{+tls, sock, rd, +key, +ms, +status} = c
664  pool.left(p, tls, sock, key, ms, status, Codec.pending(rd))
665
666def pool.close.go(xs: List<&1, List<&1, Conn>>) -> IO(Unit):
667  match xs:
668    case Nil{}:
669      IO.pure(Unit, Unit{})
670    case Con{cs, t}:
671      do IO<Unit>:
672        none : List<&1, Conn> <- conns.cut(cs, 0n)
673        pool.close.go(t)
674
675def pool.close(p: Pool) -> IO(Unit):
676  Pool{cap, idle} = p
677  pool.close.go(Map.values(&1, List<&1, Conn>, idle))