|
| 1 | +// Neon branch-per-attempt disposable DB fork for APR execution (#7858, implements #7649's ratified decision). |
| 2 | +// Mirrors worktree-pool.ts's per-attempt code-checkout isolation one level deeper: where that module gives |
| 3 | +// each attempt its own git worktree (a filesystem-level fork), this module gives an attempt its own Neon |
| 4 | +// branch (a storage-level fork) off the OPERATOR's already-provisioned tenant branch -- never off Neon's |
| 5 | +// project-level default branch, and never shared between concurrent attempts. |
| 6 | +// |
| 7 | +// Self-host scope only (see this issue's own scope note): a bare containerized Node.js process connects to a |
| 8 | +// Neon branch over the plain Postgres wire protocol via `connectionString` below -- no Cloudflare Hyperdrive |
| 9 | +// binding is used or required. Hyperdrive is an optional, Workers-runtime-specific connection-pooling layer; |
| 10 | +// Neon branches are independently connectable with any standard Postgres client regardless of whether one |
| 11 | +// exists. The hosted path (control-plane's AmsTenantContainer) has its own separate, still-open blocker before |
| 12 | +// ANY database credential can reach a running hosted container at all (#8202) -- unrelated to this module. |
| 13 | +// |
| 14 | +// Endpoint paths/response shapes below mirror control-plane/src/neon-database-driver.ts's already-reviewed |
| 15 | +// pattern (Neon's public v2 API, https://api-docs.neon.tech/reference, as documented at the time this was |
| 16 | +// written) -- same caveat that file states applies here too: verify against a live account before the first |
| 17 | +// real deploy; the test suite mocks every call, no live Neon credentials are used anywhere in this repo. |
| 18 | +import { createHash } from "node:crypto"; |
| 19 | + |
| 20 | +const DEFAULT_API_BASE_URL = "https://console.neon.tech/api/v2"; |
| 21 | +const DEFAULT_TIMEOUT_MS = 10_000; |
| 22 | +const DEFAULT_OPERATION_POLL_INTERVAL_MS = 500; |
| 23 | +const DEFAULT_OPERATION_POLL_TIMEOUT_MS = 30_000; |
| 24 | + |
| 25 | +export type AttemptDbForkConfig = { |
| 26 | + apiKey: string; |
| 27 | + projectId: string; |
| 28 | + /** The operator's own already-provisioned Neon branch ID to fork attempt branches FROM -- never Neon's |
| 29 | + * project-level default branch, so an attempt fork always starts from the operator's real, current tenant |
| 30 | + * data, not an unrelated/empty baseline. */ |
| 31 | + parentBranchId: string; |
| 32 | + /** Override for tests only -- production always uses Neon's real API. */ |
| 33 | + apiBaseUrl?: string; |
| 34 | + /** Override for tests only -- keeps operation-polling tests fast. */ |
| 35 | + operationPollIntervalMs?: number; |
| 36 | + operationPollTimeoutMs?: number; |
| 37 | +}; |
| 38 | + |
| 39 | +export type AttemptDbFork = { |
| 40 | + branchId: string; |
| 41 | + connectionString: string; |
| 42 | +}; |
| 43 | + |
| 44 | +type NeonOperation = { id: string; status: string }; |
| 45 | +type NeonBranch = { id: string; name: string }; |
| 46 | +type NeonEndpoint = { host: string }; |
| 47 | +type NeonRole = { name: string; password?: string }; |
| 48 | + |
| 49 | +// Neon branch names are case-sensitive and length-limited (63 chars) -- mirrors neon-database-driver.ts's own |
| 50 | +// #8026 collision-guard reasoning: only truncate names that actually need it, and suffix a truncated one with |
| 51 | +// a short hash of the untruncated name so two long, prefix-similar attempt ids can never collide on the same |
| 52 | +// branch name. |
| 53 | +const NEON_BRANCH_NAME_MAX_LENGTH = 63; |
| 54 | +const NEON_BRANCH_NAME_COLLISION_SUFFIX_LENGTH = 8; |
| 55 | + |
| 56 | +function sanitizeForBranchName(raw: string): string { |
| 57 | + return raw |
| 58 | + .toLowerCase() |
| 59 | + .replaceAll(/[^a-z0-9_-]+/g, "-") |
| 60 | + .replaceAll(/-{2,}/g, "-") |
| 61 | + .replace(/^-+|-+$/g, ""); |
| 62 | +} |
| 63 | + |
| 64 | +function branchNameFor(attemptId: string): string { |
| 65 | + const sanitized = sanitizeForBranchName(`attempt-${attemptId}`); |
| 66 | + if (sanitized.length <= NEON_BRANCH_NAME_MAX_LENGTH) return sanitized; |
| 67 | + const suffix = createHash("sha256").update(sanitized).digest("hex").slice(0, NEON_BRANCH_NAME_COLLISION_SUFFIX_LENGTH); |
| 68 | + const prefixLength = NEON_BRANCH_NAME_MAX_LENGTH - 1 - suffix.length; |
| 69 | + return `${sanitized.slice(0, prefixLength)}-${suffix}`; |
| 70 | +} |
| 71 | + |
| 72 | +/** Every attempt branch's role is named identically to its branch -- one branch, one role, no separate naming |
| 73 | + * scheme to keep in sync (same convention as neon-database-driver.ts's tenant-level roleNameFor). */ |
| 74 | +function roleNameForBranch(branchName: string): string { |
| 75 | + return branchName; |
| 76 | +} |
| 77 | + |
| 78 | +class NeonApiError extends Error { |
| 79 | + constructor(method: string, path: string, status: number, body: string) { |
| 80 | + super(`Neon API ${method} ${path} failed (${status}): ${body.slice(0, 500)}`); |
| 81 | + this.name = "NeonApiError"; |
| 82 | + } |
| 83 | +} |
| 84 | + |
| 85 | +async function neonFetch<T>(config: AttemptDbForkConfig, method: string, path: string, body?: unknown): Promise<T> { |
| 86 | + const baseUrl = config.apiBaseUrl ?? DEFAULT_API_BASE_URL; |
| 87 | + const response = await fetch(`${baseUrl}${path}`, { |
| 88 | + method, |
| 89 | + headers: { |
| 90 | + authorization: `Bearer ${config.apiKey}`, |
| 91 | + "content-type": "application/json", |
| 92 | + accept: "application/json", |
| 93 | + }, |
| 94 | + ...(body !== undefined ? { body: JSON.stringify(body) } : {}), |
| 95 | + signal: AbortSignal.timeout(DEFAULT_TIMEOUT_MS), |
| 96 | + }); |
| 97 | + const text = await response.text(); |
| 98 | + if (!response.ok) throw new NeonApiError(method, path, response.status, text); |
| 99 | + return (text ? JSON.parse(text) : undefined) as T; |
| 100 | +} |
| 101 | + |
| 102 | +/** Same async-operation-polling contract as neon-database-driver.ts's identical helper: branch/role mutations |
| 103 | + * return pending `operations[]` that must reach `"finished"` before the resource is actually usable. */ |
| 104 | +async function waitForOperations(config: AttemptDbForkConfig, operations: readonly NeonOperation[]): Promise<void> { |
| 105 | + const intervalMs = config.operationPollIntervalMs ?? DEFAULT_OPERATION_POLL_INTERVAL_MS; |
| 106 | + const timeoutMs = config.operationPollTimeoutMs ?? DEFAULT_OPERATION_POLL_TIMEOUT_MS; |
| 107 | + const deadline = Date.now() + timeoutMs; |
| 108 | + let pending = operations.filter((operation) => operation.status !== "finished"); |
| 109 | + while (pending.length > 0) { |
| 110 | + if (Date.now() >= deadline) { |
| 111 | + throw new Error(`Neon operation(s) did not finish within ${timeoutMs}ms: ${pending.map((operation) => operation.id).join(", ")}`); |
| 112 | + } |
| 113 | + await new Promise((resolve) => setTimeout(resolve, intervalMs)); |
| 114 | + const refreshed = await Promise.all( |
| 115 | + pending.map((operation) => neonFetch<{ operation: NeonOperation }>(config, "GET", `/projects/${config.projectId}/operations/${operation.id}`)), |
| 116 | + ); |
| 117 | + for (const { operation } of refreshed) { |
| 118 | + if (operation.status === "failed") throw new Error(`Neon operation ${operation.id} failed`); |
| 119 | + } |
| 120 | + pending = refreshed.map(({ operation }) => operation).filter((operation) => operation.status !== "finished"); |
| 121 | + } |
| 122 | +} |
| 123 | + |
| 124 | +async function findBranchByName(config: AttemptDbForkConfig, name: string): Promise<NeonBranch | undefined> { |
| 125 | + const { branches } = await neonFetch<{ branches: NeonBranch[] }>(config, "GET", `/projects/${config.projectId}/branches`); |
| 126 | + return branches.find((branch) => branch.name === name); |
| 127 | +} |
| 128 | + |
| 129 | +/** Create a disposable Neon branch forked off `config.parentBranchId` for exactly one attempt, with its own |
| 130 | + * freshly-created role (explicit, not assumed-inherited from the parent -- this repo has no live Neon account |
| 131 | + * to verify role-inheritance semantics against, so this mirrors neon-database-driver.ts's own always-explicit |
| 132 | + * role-creation pattern rather than relying on unverified copy-on-write assumptions). The database itself |
| 133 | + * (and its data) IS inherited from the parent branch via Neon's storage-level branching -- unlike a role, |
| 134 | + * Postgres databases are catalog objects that live inside the branched storage itself, so no separate |
| 135 | + * database-creation call is made here. |
| 136 | + * |
| 137 | + * Idempotent on `attemptId`: a retried call for the same attempt finds its already-created branch by name |
| 138 | + * (via {@link findBranchByName}) instead of creating a duplicate, mirroring provisionNeonDatabase's own |
| 139 | + * existing-branch-reuse path. The attempt's own database name is resolved from the PARENT branch's own |
| 140 | + * database list (the branch inherits the same database(s) the parent already has), not re-derived from |
| 141 | + * attemptId, since it must match whatever database the coding agent's own connection actually expects. */ |
| 142 | +export async function createAttemptDbFork(config: AttemptDbForkConfig, attemptId: string): Promise<AttemptDbFork> { |
| 143 | + const branchName = branchNameFor(attemptId); |
| 144 | + const roleName = roleNameForBranch(branchName); |
| 145 | + |
| 146 | + const existing = await findBranchByName(config, branchName); |
| 147 | + if (existing) { |
| 148 | + const { endpoints } = await neonFetch<{ endpoints: NeonEndpoint[] }>(config, "GET", `/projects/${config.projectId}/branches/${existing.id}/endpoints`); |
| 149 | + const host = endpoints[0]?.host; |
| 150 | + if (!host) throw new Error(`Neon attempt branch ${existing.id} has no compute endpoint`); |
| 151 | + const { role } = await neonFetch<{ role: NeonRole }>(config, "GET", `/projects/${config.projectId}/branches/${existing.id}/roles/${roleName}/reveal_password`); |
| 152 | + if (!role.password) throw new Error(`Neon role ${roleName} on attempt branch ${existing.id} has no revealable password`); |
| 153 | + const databaseName = await parentDatabaseName(config); |
| 154 | + return { branchId: existing.id, connectionString: connectionStringFor(host, databaseName, roleName, role.password) }; |
| 155 | + } |
| 156 | + |
| 157 | + const databaseName = await parentDatabaseName(config); |
| 158 | + const created = await neonFetch<{ branch: NeonBranch; endpoints: NeonEndpoint[]; operations: NeonOperation[] }>( |
| 159 | + config, |
| 160 | + "POST", |
| 161 | + `/projects/${config.projectId}/branches`, |
| 162 | + { branch: { name: branchName, parent_id: config.parentBranchId }, endpoints: [{ type: "read_write" }] }, |
| 163 | + ); |
| 164 | + await waitForOperations(config, created.operations); |
| 165 | + const host = created.endpoints[0]?.host; |
| 166 | + if (!host) throw new Error(`Neon attempt branch ${created.branch.id} was created without a compute endpoint`); |
| 167 | + |
| 168 | + const roleCreated = await neonFetch<{ role: NeonRole; operations: NeonOperation[] }>( |
| 169 | + config, |
| 170 | + "POST", |
| 171 | + `/projects/${config.projectId}/branches/${created.branch.id}/roles`, |
| 172 | + { role: { name: roleName } }, |
| 173 | + ); |
| 174 | + await waitForOperations(config, roleCreated.operations); |
| 175 | + if (!roleCreated.role.password) throw new Error(`Neon role ${roleName} was created without a password`); |
| 176 | + |
| 177 | + return { branchId: created.branch.id, connectionString: connectionStringFor(host, databaseName, roleName, roleCreated.role.password) }; |
| 178 | +} |
| 179 | + |
| 180 | +async function parentDatabaseName(config: AttemptDbForkConfig): Promise<string> { |
| 181 | + const { databases } = await neonFetch<{ databases: { name: string }[] }>( |
| 182 | + config, |
| 183 | + "GET", |
| 184 | + `/projects/${config.projectId}/branches/${config.parentBranchId}/databases`, |
| 185 | + ); |
| 186 | + const database = databases[0]; |
| 187 | + if (!database) throw new Error(`Neon parent branch ${config.parentBranchId} has no database to fork`); |
| 188 | + return database.name; |
| 189 | +} |
| 190 | + |
| 191 | +function connectionStringFor(host: string, database: string, user: string, password: string): string { |
| 192 | + return `postgres://${user}:${password}@${host}:5432/${database}`; |
| 193 | +} |
| 194 | + |
| 195 | +/** Discard an attempt's disposable branch. Idempotent: an attempt whose branch was never created (blocked |
| 196 | + * before {@link createAttemptDbFork} ran) or already discarded is a safe no-op. Deleting a Neon branch |
| 197 | + * cascades to its role/database/endpoint together -- there is nothing else to clean up separately. Never |
| 198 | + * merges the branch's data back into the parent; ratified explicitly by #7649 as a hard requirement, not a |
| 199 | + * default that happens to be convenient here. */ |
| 200 | +export async function discardAttemptDbFork(config: AttemptDbForkConfig, attemptId: string): Promise<void> { |
| 201 | + const branchName = branchNameFor(attemptId); |
| 202 | + const existing = await findBranchByName(config, branchName); |
| 203 | + if (!existing) return; |
| 204 | + |
| 205 | + // Tolerates a body-less success response (e.g. 204 No Content) -- some APIs return nothing for a DELETE that |
| 206 | + // completed synchronously, with no operation left to poll. |
| 207 | + const result = await neonFetch<{ operations?: NeonOperation[] } | undefined>(config, "DELETE", `/projects/${config.projectId}/branches/${existing.id}`); |
| 208 | + await waitForOperations(config, result?.operations ?? []); |
| 209 | +} |
0 commit comments