| 1 | /** |
| 2 | * Reading a service's D1 database near the caller with D1's Sessions API, |
| 3 | * and saying how long an RPC took: the TypeScript side of |
| 4 | * crates/kit/src/d1.rs, with the same rules. |
| 5 | * |
| 6 | * The caller chooses, per request, with the `x-d1-bookmark` header: |
| 7 | * |
| 8 | * | Header | Reads go to | |
| 9 | * | --- | --- | |
| 10 | * | absent | the primary, with no session (as always) | |
| 11 | * | `first-primary` | the primary first, then any copy at least as new | |
| 12 | * | `first-unconstrained` | the nearest copy | |
| 13 | * | a bookmark | any copy at least as new as the bookmark | |
| 14 | * |
| 15 | * Writes always go to the primary. The session's latest bookmark comes |
| 16 | * back in the response's `x-d1-bookmark`. The site's side is |
| 17 | * apps/web/app/lib/perf.ts. |
| 18 | */ |
| 19 | |
| 20 | /** The header that carries a session's constraint or bookmark, both ways. */ |
| 21 | export const BOOKMARK_HEADER = "x-d1-bookmark"; |
| 22 | |
| 23 | const MAX_BOOKMARK = 256; |
| 24 | |
| 25 | /** |
| 26 | * What an `x-d1-bookmark` header asks for: null for no session, otherwise |
| 27 | * what `withSession` is given. Anything malformed starts on the primary. |
| 28 | */ |
| 29 | export function sessionConstraint(header: string | null | undefined): string | null { |
| 30 | const value = header?.trim(); |
| 31 | if (!value) return null; |
| 32 | if (value === "first-primary" || value === "first-unconstrained") return value; |
| 33 | return value.length <= MAX_BOOKMARK && /^[0-9A-Za-z-]+$/.test(value) ? value : "first-primary"; |
| 34 | } |
| 35 | |
| 36 | /** A D1 binding, as far as sessions need it. */ |
| 37 | type SessionCapable = { withSession(constraintOrBookmark?: string): { getBookmark(): string | null } }; |
| 38 | |
| 39 | /** How long a request waited on D1, and in how many round trips. */ |
| 40 | export type D1Timing = { trips: number; ms: number }; |
| 41 | |
| 42 | /** The methods of a prepared statement that go to the database. */ |
| 43 | const TRIPS = new Set(["all", "first", "run", "raw"]); |
| 44 | |
| 45 | /** |
| 46 | * `db` with every round trip counted into `timing`: each `all`, `first`, |
| 47 | * `run` and `raw` of a statement, and each `batch`. Statements given to |
| 48 | * `batch` are unwrapped, so the binding sees its own. The `db;dur` metric |
| 49 | * crates/kit/src/d1.rs reports for the Rust services, for the TypeScript |
| 50 | * ones; the site reads both (apps/web/app/lib/perf.ts `databaseTime`). |
| 51 | */ |
| 52 | export function timedD1<D extends object>(db: D, timing: D1Timing): D { |
| 53 | const raw = new WeakMap<object, object>(); |
| 54 | const trip = async <T>(work: () => Promise<T>): Promise<T> => { |
| 55 | const from = Date.now(); |
| 56 | try { |
| 57 | return await work(); |
| 58 | } finally { |
| 59 | timing.trips += 1; |
| 60 | timing.ms += Date.now() - from; |
| 61 | } |
| 62 | }; |
| 63 | const statement = (s: object): object => { |
| 64 | const wrapped = new Proxy(s, { |
| 65 | get(target, prop) { |
| 66 | const value = Reflect.get(target, prop, target); |
| 67 | if (typeof value !== "function") return value; |
| 68 | if (prop === "bind") return (...args: unknown[]) => statement(value.apply(target, args)); |
| 69 | if (typeof prop === "string" && TRIPS.has(prop)) return (...args: unknown[]) => trip(() => value.apply(target, args)); |
| 70 | return value.bind(target); |
| 71 | }, |
| 72 | }); |
| 73 | raw.set(wrapped, s); |
| 74 | return wrapped; |
| 75 | }; |
| 76 | return new Proxy(db, { |
| 77 | get(target, prop) { |
| 78 | const value = Reflect.get(target, prop, target); |
| 79 | if (typeof value !== "function") return value; |
| 80 | if (prop === "prepare") return (query: string) => statement(value.call(target, query)); |
| 81 | if (prop === "batch") return (statements: object[]) => trip(() => value.call(target, statements.map((s) => raw.get(s) ?? s))); |
| 82 | return value.bind(target); |
| 83 | }, |
| 84 | }); |
| 85 | } |
| 86 | |
| 87 | /** |
| 88 | * The database for an RPC `request`: a session, seen as the binding, when |
| 89 | * the caller asked for one (it answers `prepare` and `batch`, all a request |
| 90 | * path uses), the binding itself otherwise, with its round trips counted. |
| 91 | * `finish` adds the bookmark, the time taken and the database time to the |
| 92 | * answer. |
| 93 | */ |
| 94 | export function openD1<D extends SessionCapable>(db: D, request: Request): { db: D; timing: D1Timing; finish(response: Response): Response } { |
| 95 | const started = Date.now(); |
| 96 | const asked = sessionConstraint(request.headers.get(BOOKMARK_HEADER)); |
| 97 | const session = asked ? db.withSession(asked) : null; |
| 98 | const timing: D1Timing = { trips: 0, ms: 0 }; |
| 99 | return { |
| 100 | db: timedD1((session ?? db) as D, timing), |
| 101 | timing, |
| 102 | finish(response) { |
| 103 | const answered = new Response(response.body, response); |
| 104 | const db = timing.trips ? `, db;dur=${timing.ms};desc="${timing.trips} round trip${timing.trips === 1 ? "" : "s"}"` : ""; |
| 105 | answered.headers.append("server-timing", `svc;dur=${Date.now() - started};desc="${session ? "session" : "primary"}"${db}`); |
| 106 | const bookmark = session?.getBookmark(); |
| 107 | if (bookmark) answered.headers.set(BOOKMARK_HEADER, bookmark); |
| 108 | return answered; |
| 109 | }, |
| 110 | }; |
| 111 | } |