1import * as ps from "@distilled.cloud/planetscale";
2import * as Effect from "effect/Effect";
3import * as Schedule from "effect/Schedule";
4import { PlanetscaleConflict, pollUntil } from "../Util.ts";
5
6/**
7 * Available PlanetScale MySQL (Vitess) cluster sizes.
8 *
9 * `PS_*` sizes are backed by network-attached storage (NAS) and can be
10 * specified either as the short size (`"PS_10"`) or the API SKU
11 * (`"PS_10_AWS_X86"`)
12 *
13 * @see https://planetscale.com/docs/concepts/planetscale-skus
14 */
15export type MySQLClusterSize =
16 | "PS_DEV"
17 | "PS_5"
18 | "PS_10"
19 | "PS_20"
20 | "PS_40"
21 | "PS_80"
22 | "PS_160"
23 | "PS_320"
24 | "PS_400"
25 | "PS_640"
26 | "PS_700"
27 | "PS_900"
28 | "PS_1280"
29 | "PS_1400"
30 | "PS_1800"
31 | "PS_2100"
32 | "PS_2560"
33 | "PS_2700"
34 | "PS_2800"
35 | (string & {});
36
37/**
38 * Polls keyspaces in a branch until the named keyspace reports
39 * `resizing === false` (i.e. any in-flight resize has completed).
40 */
41export const waitForKeyspaceReady = Effect.fn(function* (
42 organization: string,
43 database: string,
44 branch: string,
45 keyspace: string,
46) {
47 yield* pollUntil(
48 `keyspace "${keyspace}" not resizing`,
49 ps.listKeyspaces({ organization, database, branch }),
50 (page) => {
51 const ks = page.data.find((x) => x.name === keyspace);
52 // If keyspace is missing, treat as ready (caller will re-check)
53 return ks ? !ks.resizing : true;
54 },
55 );
56});
57
58/**
59 * Observes the default keyspace of a branch. The default keyspace usually
60 * shares the database's name, but a renamed database keeps its original
61 * keyspace name, so prefer the `default` flag and fall back to name
62 * matching.
63 */
64const observeDefaultKeyspace = Effect.fn(function* (
65 organization: string,
66 database: string,
67 branch: string,
68) {
69 const keyspaces = yield* ps.listKeyspaces({
70 organization,
71 database,
72 branch,
73 page: 1,
74 per_page: 100,
75 });
76 return (
77 keyspaces.data.find((x) => x.default) ??
78 keyspaces.data.find((x) => x.name === database)
79 );
80});
81
82/**
83 * Observes the total replica count of a branch's default keyspace, or
84 * `undefined` when the branch or keyspace does not exist (yet).
85 */
86export const observeDefaultKeyspaceReplicas = Effect.fn(function* (
87 organization: string,
88 database: string,
89 branch: string,
90) {
91 return yield* observeDefaultKeyspace(organization, database, branch).pipe(
92 Effect.map((keyspace) => keyspace?.replicas),
93 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
94 );
95});
96
97/**
98 * Ensures the default keyspace of a MySQL production branch has the
99 * expected cluster size and (optionally) total replica count, driving
100 * PlanetScale's in-place keyspace resize lifecycle when it doesn't.
101 * Cluster sizes can only be configured on production branches. Returns
102 * the final observed keyspace.
103 */
104export const ensureMySQLProductionBranchClusterSize = Effect.fn(function* (
105 organization: string,
106 database: string,
107 branch: string,
108 expectedClusterSize: MySQLClusterSize,
109 expectedReplicas?: number,
110) {
111 let keyspace = yield* observeDefaultKeyspace(organization, database, branch);
112 if (!keyspace) {
113 return yield* Effect.die(`No default keyspace found for branch ${branch}`);
114 }
115
116 yield* waitForKeyspaceReady(organization, database, branch, keyspace.name);
117
118 if (keyspace.cluster_name !== expectedClusterSize) {
119 yield* ps.updateBranchClusterConfig({
120 organization,
121 database,
122 branch,
123 cluster_size: expectedClusterSize,
124 });
125 yield* waitForKeyspaceReady(organization, database, branch, keyspace.name);
126 // Re-observe so the replica sync below diffs against the post-resize
127 // keyspace state.
128 keyspace =
129 (yield* observeDefaultKeyspace(organization, database, branch)) ??
130 keyspace;
131 }
132
133 // Sync replicas — MySQL databases cannot configure replicas at
134 // creation time (the API rejects the `replicas` param for mysql), so
135 // the desired total replica count is converged in place via a keyspace
136 // resize request.
137 if (
138 expectedReplicas !== undefined &&
139 keyspace.replicas !== expectedReplicas
140 ) {
141 // Each cluster size includes a fixed number of replicas (2 for
142 // production PS_* sizes); the resize API only accepts the count of
143 // additional replicas beyond that.
144 const includedReplicas = keyspace.replicas - keyspace.extra_replicas;
145 const extraReplicas = expectedReplicas - includedReplicas;
146 if (extraReplicas < 0) {
147 return yield* Effect.fail(
148 new PlanetscaleConflict({
149 message:
150 `Cannot set replicas to ${expectedReplicas} on keyspace "${keyspace.name}": ` +
151 `cluster size ${keyspace.cluster_name} always includes ${includedReplicas} replicas. ` +
152 `Set replicas to ${includedReplicas} or more (or omit it).`,
153 }),
154 );
155 }
156
157 const resize = yield* ps
158 .updateKeyspaceResizeRequest({
159 organization,
160 database,
161 branch,
162 keyspace: keyspace.name,
163 extra_replicas: extraReplicas,
164 })
165 .pipe(
166 // A prior resize (e.g. the cluster-size change above) can still
167 // be finalizing after the keyspace reports `resizing: false`;
168 // PlanetScale rejects new resize requests during that window.
169 Effect.retry({
170 while: (e): boolean =>
171 e._tag === "UnprocessableEntity" &&
172 e.message.includes("resize in progress"),
173 schedule: Schedule.max([
174 Schedule.spaced("5 seconds"),
175 Schedule.recurs(120),
176 ]),
177 }),
178 );
179
180 yield* pollUntil(
181 `keyspace "${keyspace.name}" resize completed`,
182 ps.listKeyspaceResizeRequests({
183 organization,
184 database,
185 branch,
186 keyspace: keyspace.name,
187 }),
188 (page) => {
189 const request = page.data.find((r) => r.id === resize.id);
190 // If the request is missing, treat as settled (caller re-observes).
191 return request ? request.state === "completed" : true;
192 },
193 );
194
195 keyspace =
196 (yield* observeDefaultKeyspace(organization, database, branch)) ??
197 keyspace;
198 }
199
200 return keyspace;
201});