1import * as bigqueryreservation from "@distilled.cloud/gcp/bigqueryreservation_v1";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
4import * as Stream from "effect/Stream";
5import { Unowned } from "../../AdoptPolicy.ts";
6import { isResolved } from "../../Diff.ts";
7import * as Provider from "../../Provider.ts";
8import { Resource } from "../../Resource.ts";
9import { GcpEnvironment } from "../Environment.ts";
10import type { Providers } from "../Providers.ts";
11import {
12 compact,
13 hasOwnershipMarker,
14 lastSegment,
15 listLocations,
16 normalizeLocation,
17 ownedByAlchemy,
18 parentOf,
19 toResourceId,
20} from "./internal.ts";
21
22export type AssignmentJobType =
23 | bigqueryreservation.AssignmentJobTypeEnum
24 | (string & {});
25
26export type AssignmentSchedulingPolicy = {
27 /**
28 * Soft per-project slot cap for jobs in this assignment. Preview.
29 */
30 maxSlots?: string;
31 /**
32 * Soft per-project job concurrency cap. Preview.
33 */
34 concurrency?: string;
35};
36
37export type AssignmentProps = {
38 /**
39 * Parent reservation resource name
40 * `projects/{project}/locations/{location}/reservations/{reservation}`
41 * or a bare reservation id (resolved in the current project and
42 * `location`). Immutable — changing it replaces the assignment.
43 */
44 reservation: string;
45 /**
46 * Resource that consumes the reservation: `projects/{project}`,
47 * `folders/{folder}`, or `organizations/{org}`. Defaults to the
48 * current GCP project. Immutable — changing it replaces the
49 * assignment. One assignment per `(assignee, jobType, location)`.
50 */
51 assignee?: string;
52 /**
53 * Job type that uses the reservation (`QUERY`, `PIPELINE`,
54 * `ML_EXTERNAL`, `BACKGROUND`, `CONTINUOUS`, …). Immutable —
55 * changing it replaces the assignment.
56 * @default "QUERY"
57 */
58 jobType?: AssignmentJobType;
59 /**
60 * Assignment id (the `{assignment}` segment). If omitted, a unique
61 * name is generated. Assignments have no labels field, so Alchemy
62 * stamps ownership into this id (`alch---…`) for `list` / nuke.
63 * Immutable — changing it replaces the assignment. Max 64
64 * characters, lowercase letters, digits, or dashes.
65 */
66 assignmentId?: string;
67 /**
68 * Location of the parent reservation. Used when `reservation` is a
69 * bare id. Immutable — changing it replaces the assignment.
70 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
71 */
72 location?: string;
73 /**
74 * Principal that must run jobs to use this assignment. When empty,
75 * any job of `jobType` on `assignee` uses the reservation.
76 */
77 principal?: string;
78 /**
79 * Per-assignment scheduling policy. Preview.
80 */
81 schedulingPolicy?: AssignmentSchedulingPolicy;
82};
83
84export type Assignment = Resource<
85 "GCP.BigQueryReservation.Assignment",
86 AssignmentProps,
87 {
88 /** Full resource name `…/reservations/{reservation}/assignments/{assignment}`. */
89 name: string;
90 /** Assignment id (last path segment). */
91 assignmentId: string;
92 /** Parent reservation resource name. */
93 reservation: string;
94 /** Reservation id. */
95 reservationId: string;
96 /** Project that owns the reservation. */
97 project: string;
98 /** Location id. */
99 location: string;
100 /** Assignee resource (`projects/…`, `folders/…`, `organizations/…`). */
101 assignee: string | undefined;
102 /** Job type. */
103 jobType: string | undefined;
104 /** Assignment state (`PENDING`, `ACTIVE`). */
105 state: string | undefined;
106 /** Principal, if set. */
107 principal: string | undefined;
108 /** Scheduling policy, if set. */
109 schedulingPolicy: bigqueryreservation.SchedulingPolicy | undefined;
110 },
111 never,
112 Providers
113>;
114
115/**
116 * A BigQuery reservation assignment — grants a project, folder, or
117 * organization slots from a {@link Reservation} for one job type.
118 *
119 * Assignments have no labels. Alchemy stamps ownership into
120 * `assignmentId` so `list` / `pnpm nuke:gcp` can find them. Changing
121 * `reservation`, `assignee`, `jobType`, `location`, or `assignmentId`
122 * replaces the assignment. `principal` and `schedulingPolicy` update
123 * in place.
124 *
125 * ### Creating an Assignment
126 * **Example:** Assign QUERY jobs in this project
127 * ```typescript
128 * const slots = yield* GCP.BigQueryReservation.Reservation("Slots", {});
129 * const assignment = yield* GCP.BigQueryReservation.Assignment("Query", {
130 * reservation: slots.name,
131 * jobType: "QUERY",
132 * });
133 * ```
134 *
135 * **Example:** None assignment (on-demand)
136 * ```typescript
137 * const none = yield* GCP.BigQueryReservation.Assignment("OnDemand", {
138 * reservation: "none",
139 * location: "US",
140 * jobType: "QUERY",
141 * });
142 * ```
143 *
144 * @resource
145 * @category BigQueryReservation
146 */
147export const Assignment = Resource<Assignment>(
148 "GCP.BigQueryReservation.Assignment",
149);
150
151export class AssignmentNotResolved extends Data.TaggedError(
152 "GCP.BigQueryReservation.AssignmentNotResolved",
153)<{
154 name: string;
155}> {}
156
157const DEFAULT_JOB_TYPE = "QUERY";
158
159const parseName = (name: string) => {
160 const parts = name.split("/").filter((part) => part.length > 0);
161 const after = (segment: string) => {
162 const at = parts.lastIndexOf(segment);
163 return at >= 0 && parts[at + 1] ? parts[at + 1]! : "";
164 };
165 const reservationId = after("reservations");
166 const location = after("locations");
167 const project = after("projects");
168 return {
169 project,
170 location,
171 reservationId,
172 assignmentId: after("assignments") || lastSegment(name),
173 reservation:
174 project && location && reservationId
175 ? `projects/${project}/locations/${location}/reservations/${reservationId}`
176 : "",
177 };
178};
179
180const reservationParent = (
181 project: string,
182 location: string,
183 reservation: string,
184) => {
185 if (reservation.includes("/reservations/")) {
186 const parsed = parseName(
187 reservation.endsWith("/assignments")
188 ? reservation
189 : `${reservation}/assignments/_`,
190 );
191 return `projects/${parsed.project || project}/locations/${parsed.location || location}/reservations/${parsed.reservationId || lastSegment(reservation)}`;
192 }
193 return `projects/${project}/locations/${location}/reservations/${lastSegment(reservation)}`;
194};
195
196const assigneeOf = (project: string, assignee: string | undefined) =>
197 assignee && assignee.length > 0 ? assignee : `projects/${project}`;
198
199const jobTypeOf = (jobType: string | undefined) =>
200 (jobType ?? DEFAULT_JOB_TYPE).toUpperCase();
201
202const jsonOf = (value: unknown) => JSON.stringify(value ?? null);
203
204const desiredSchedulingPolicy = (
205 news: AssignmentProps,
206): bigqueryreservation.SchedulingPolicy | undefined =>
207 news.schedulingPolicy === undefined
208 ? undefined
209 : compact({
210 maxSlots: news.schedulingPolicy.maxSlots,
211 concurrency: news.schedulingPolicy.concurrency,
212 });
213
214const toAttrs = (
215 current: bigqueryreservation.Assignment,
216 fallbackProject: string,
217) => {
218 const name = current.name ?? "";
219 const parsed = parseName(name);
220 return {
221 name,
222 assignmentId: parsed.assignmentId,
223 reservation: parsed.reservation,
224 reservationId: parsed.reservationId,
225 project: parsed.project || fallbackProject,
226 location: parsed.location,
227 assignee: current.assignee,
228 jobType: current.jobType,
229 state: current.state,
230 principal: current.principal,
231 schedulingPolicy: current.schedulingPolicy,
232 };
233};
234
235const listAt = (parent: string) =>
236 parent.length === 0
237 ? Effect.succeed([] as bigqueryreservation.Assignment[])
238 : bigqueryreservation.listProjectsLocationsReservationsAssignments
239 .pages({ parent, pageSize: 1000 })
240 .pipe(
241 Stream.flatMap((page) => Stream.fromIterable(page.assignments ?? [])),
242 Stream.runCollect,
243 Effect.map((chunk) => Array.from(chunk)),
244 Effect.catchTag("NotFound", () =>
245 Effect.succeed([] as bigqueryreservation.Assignment[]),
246 ),
247 );
248
249const getByName = (name: string) => {
250 if (name.length === 0) return Effect.succeed(undefined);
251 const parent = parseName(name).reservation;
252 return listAt(parent).pipe(
253 Effect.map((assignments) => assignments.find((item) => item.name === name)),
254 );
255};
256
257const findExisting = (
258 parent: string,
259 assignee: string,
260 jobType: string,
261 assignmentId?: string,
262) =>
263 listAt(parent).pipe(
264 Effect.map((assignments) =>
265 assignments.find((item) => {
266 if (assignmentId && lastSegment(item.name ?? "") === assignmentId) {
267 return true;
268 }
269 return (
270 (item.assignee ?? "") === assignee &&
271 jobTypeOf(item.jobType) === jobType
272 );
273 }),
274 ),
275 );
276
277export const AssignmentProvider = () =>
278 Provider.succeed(Assignment, {
279 stables: [
280 "name",
281 "assignmentId",
282 "reservation",
283 "reservationId",
284 "project",
285 "location",
286 ],
287
288 diff: Effect.fn(function* ({ news, olds, output }) {
289 if (!isResolved(news)) return undefined;
290 const env = yield* GcpEnvironment.current;
291 const previousReservation = olds?.reservation ?? output?.reservation;
292 const previousAssignee = olds?.assignee ?? output?.assignee;
293 const previousJobType = jobTypeOf(olds?.jobType ?? output?.jobType);
294 const previousLocation = normalizeLocation(
295 olds?.location ?? output?.location ?? env.region,
296 );
297 const previousId = olds?.assignmentId ?? output?.assignmentId;
298 const nextJobType = jobTypeOf(news.jobType ?? previousJobType);
299 const nextLocation = normalizeLocation(
300 news.location ?? output?.location ?? env.region,
301 );
302 if (
303 (previousReservation !== undefined &&
304 news.reservation !== previousReservation) ||
305 (previousAssignee !== undefined &&
306 news.assignee !== undefined &&
307 news.assignee !== previousAssignee) ||
308 previousJobType !== nextJobType ||
309 previousLocation !== nextLocation ||
310 (previousId !== undefined &&
311 news.assignmentId !== undefined &&
312 news.assignmentId !== previousId)
313 ) {
314 return { action: "replace" as const, deleteFirst: true };
315 }
316 return undefined;
317 }),
318
319 read: Effect.fn(function* ({ id, olds, output }) {
320 const env = yield* GcpEnvironment.current;
321 const location = normalizeLocation(
322 olds?.location ?? output?.location ?? env.region,
323 );
324 const parent = reservationParent(
325 env.project,
326 location,
327 olds?.reservation ?? output?.reservation ?? "",
328 );
329 const assignmentId = yield* toResourceId(
330 id,
331 olds?.assignmentId,
332 output?.assignmentId,
333 );
334 const name =
335 output?.name ?? (parent ? `${parent}/assignments/${assignmentId}` : "");
336 const existing =
337 (yield* getByName(name)) ??
338 (yield* findExisting(
339 parent,
340 assigneeOf(env.project, olds?.assignee ?? output?.assignee),
341 jobTypeOf(olds?.jobType ?? output?.jobType),
342 assignmentId,
343 ));
344 if (existing === undefined) return undefined;
345 const attrs = toAttrs(existing, env.project);
346 return (yield* ownedByAlchemy(id, attrs.assignmentId))
347 ? attrs
348 : Unowned(attrs);
349 }),
350
351 list: () =>
352 Effect.gen(function* () {
353 const env = yield* GcpEnvironment.current;
354 const pages = yield* Effect.forEach(
355 listLocations(env.region),
356 (location) =>
357 listAt(`${parentOf(env.project, location)}/reservations/-`).pipe(
358 Effect.map((assignments) =>
359 assignments
360 .filter((item) =>
361 hasOwnershipMarker(lastSegment(item.name ?? "")),
362 )
363 .map((item) => toAttrs(item, env.project)),
364 ),
365 ),
366 { concurrency: 4 },
367 );
368 const byName = new Map<string, ReturnType<typeof toAttrs>>();
369 for (const item of pages.flat()) {
370 byName.set(item.name, item);
371 }
372 return Array.from(byName.values());
373 }),
374
375 reconcile: Effect.fn(function* ({ id, news, output }) {
376 const env = yield* GcpEnvironment.current;
377 const location = normalizeLocation(
378 news.location ?? output?.location ?? env.region,
379 );
380 const parent = reservationParent(env.project, location, news.reservation);
381 const assignmentId = yield* toResourceId(
382 id,
383 news.assignmentId,
384 output?.assignmentId,
385 );
386 const name = `${parent}/assignments/${assignmentId}`;
387 const assignee = assigneeOf(env.project, news.assignee);
388 const jobType = jobTypeOf(news.jobType);
389 const body = compact({
390 assignee,
391 jobType,
392 principal: news.principal,
393 schedulingPolicy: desiredSchedulingPolicy(news),
394 });
395
396 let current =
397 (yield* getByName(output?.name ?? name)) ??
398 (yield* findExisting(parent, assignee, jobType, assignmentId));
399
400 if (current === undefined) {
401 const created = yield* bigqueryreservation
402 .createProjectsLocationsReservationsAssignments({
403 parent,
404 assignmentId,
405 body,
406 })
407 .pipe(
408 Effect.catchTag("Conflict", () =>
409 findExisting(parent, assignee, jobType, assignmentId),
410 ),
411 );
412 current = created ?? undefined;
413 }
414
415 if (current === undefined) {
416 return yield* new AssignmentNotResolved({ name });
417 }
418
419 const currentName = current.name ?? name;
420 const principalChanged =
421 news.principal !== undefined &&
422 (current.principal ?? "") !== news.principal;
423 const schedulingChanged =
424 news.schedulingPolicy !== undefined &&
425 jsonOf(current.schedulingPolicy) !==
426 jsonOf(desiredSchedulingPolicy(news));
427 const updateMask = [
428 principalChanged ? "principal" : undefined,
429 schedulingChanged ? "schedulingPolicy" : undefined,
430 ].filter((field): field is string => field !== undefined);
431
432 if (updateMask.length > 0) {
433 current =
434 yield* bigqueryreservation.patchProjectsLocationsReservationsAssignments(
435 {
436 name: currentName,
437 updateMask: updateMask.join(","),
438 body,
439 },
440 );
441 }
442
443 return toAttrs(current, env.project);
444 }),
445
446 delete: Effect.fn(function* ({ output }) {
447 if (!output.name) return;
448 yield* bigqueryreservation
449 .deleteProjectsLocationsReservationsAssignments({
450 name: output.name,
451 })
452 .pipe(Effect.catchTag("NotFound", () => Effect.void));
453 }),
454 });