1import * as dataproc from "@distilled.cloud/gcp/dataproc_v1";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
4import * as Schedule from "effect/Schedule";
5import * as Stream from "effect/Stream";
6import { waitForOperation as waitForGcpOperation } from "../Operation.ts";
7import { createPhysicalName } from "../../PhysicalName.ts";
8import { tagRecord } from "../../Tags.ts";
9import { stripInternalLabels } from "../Labels.ts";
10
11export const DEFAULT_ZONE = "us-central1-a";
12export const MAX_POLICY_ID_LENGTH = 50;
13// Cluster, batch, and session operations run for minutes; poll for up to ~10
14// minutes (at 5-8s spacing) before giving up.
15export const MAX_POLLS = 120;
16export const MAX_WORKLOAD_ID_LENGTH = 63;
17export const MIN_WORKLOAD_ID_LENGTH = 4;
18// Session template ids match `[a-zA-Z0-9][-_a-zA-Z0-9]{1,48}[a-zA-Z0-9]`.
19export const MAX_SESSION_TEMPLATE_ID_LENGTH = 50;
20// Managed cluster name prefixes match `[a-z]([-a-z0-9]{0,34}[a-z0-9])?`.
21export const MAX_MANAGED_CLUSTER_PREFIX_LENGTH = 36;
22
23export const LIST_LOCATIONS = [
24 "us-central1",
25 "us-east1",
26 "us-east4",
27 "us-west1",
28 "europe-west1",
29 "asia-east1",
30] as const;
31
32export class DataprocNotResolved extends Data.TaggedError(
33 "GCP.Dataproc.ResourceNotResolved",
34)<{
35 name: string;
36}> {}
37
38export class DataprocStillExists extends Data.TaggedError(
39 "GCP.Dataproc.ResourceStillExists",
40)<{
41 name: string;
42}> {}
43
44export const lastSegment = (value: string) => {
45 const trimmed = value.replace(/\/+$/, "");
46 const parts = trimmed.split("/");
47 return parts[parts.length - 1] || trimmed;
48};
49
50export const normalizeLocation = (
51 location: string | undefined,
52 fallback: string,
53) => lastSegment(location ?? fallback).toLowerCase();
54
55export const locationParent = (project: string, location: string) =>
56 `projects/${project}/locations/${location}`;
57
58export const regionParent = (project: string, region: string) =>
59 `projects/${project}/regions/${region}`;
60
61export const parseResourceName = (
62 name: string,
63 collection: string,
64 fallbackLocation: string,
65) => {
66 const parts = name.split("/").filter((part) => part.length > 0);
67 const collectionAt = parts.lastIndexOf(collection);
68 const locationsAt = parts.lastIndexOf("locations");
69 const regionsAt = parts.lastIndexOf("regions");
70 const projectsAt = parts.lastIndexOf("projects");
71 const placeAt = locationsAt >= 0 ? locationsAt : regionsAt;
72 return {
73 project:
74 projectsAt >= 0 && parts[projectsAt + 1] ? parts[projectsAt + 1]! : "",
75 location:
76 placeAt >= 0 && parts[placeAt + 1]
77 ? parts[placeAt + 1]!
78 : fallbackLocation,
79 id:
80 collectionAt >= 0 && parts[collectionAt + 1]
81 ? parts[collectionAt + 1]!
82 : lastSegment(name),
83 parent:
84 collectionAt > 0
85 ? parts.slice(0, collectionAt).join("/")
86 : parts.slice(0, Math.max(0, parts.length - 1)).join("/"),
87 };
88};
89
90export const rfc1035 = (
91 name: string,
92 maxLength: number,
93 fallback: string,
94): string => {
95 let next = name
96 .toLowerCase()
97 .replace(/[^a-z0-9-]/g, "-")
98 .replace(/-+/g, "-")
99 .replace(/^-+|-+$/g, "");
100 if (!/^[a-z]/.test(next)) next = `${fallback[0] ?? "d"}${next}`;
101 next = next.slice(0, maxLength).replace(/-+$/g, "");
102 if (next.length === 0) next = fallback;
103 if (next.length < MIN_WORKLOAD_ID_LENGTH) {
104 next = `${next}${fallback}xxxx`.slice(0, MIN_WORKLOAD_ID_LENGTH);
105 }
106 if (!/[a-z0-9]$/.test(next)) next = `${next.slice(0, maxLength - 1)}0`;
107 return next.slice(0, maxLength);
108};
109
110export const toPhysicalId = (
111 id: string,
112 explicit: string | undefined,
113 existing: string | undefined,
114 maxLength: number,
115 fallback: string,
116) =>
117 Effect.gen(function* () {
118 if (explicit !== undefined) return rfc1035(explicit, maxLength, fallback);
119 if (existing !== undefined) return existing;
120 return rfc1035(
121 yield* createPhysicalName({
122 id,
123 maxLength,
124 lowercase: true,
125 }),
126 maxLength,
127 fallback,
128 );
129 });
130
131export const userLabels = (
132 labels: Record<string, string | undefined> | null | undefined,
133): Record<string, string> => stripInternalLabels(tagRecord(labels));
134
135export const hasAlchemyLabelMap = (
136 labels: Record<string, string | undefined> | null | undefined,
137) => Object.keys(labels ?? {}).some((key) => key.startsWith("alchemy-"));
138
139export const canonical = (value: unknown): unknown => {
140 if (value === undefined || value === null) return undefined;
141 if (typeof value === "boolean" || typeof value === "number") return value;
142 if (typeof value === "string") {
143 return value.length === 0 ? undefined : value;
144 }
145 if (Array.isArray(value)) {
146 const items = value.map(canonical).filter((item) => item !== undefined);
147 return items.length === 0 ? undefined : items;
148 }
149 if (typeof value === "object") {
150 const entries = Object.entries(value as Record<string, unknown>)
151 .map(([key, item]) => [key, canonical(item)] as const)
152 .filter(([, item]) => item !== undefined)
153 .sort(([left], [right]) => left.localeCompare(right));
154 if (entries.length === 0) return undefined;
155 return Object.fromEntries(entries);
156 }
157 return undefined;
158};
159
160export const fingerprint = (value: unknown): string =>
161 JSON.stringify(canonical(value) ?? null);
162
163export const sameJson = (left: unknown, right: unknown) =>
164 fingerprint(left) === fingerprint(right);
165
166export const collectPages = <Page, Item, E, R>(
167 stream: Stream.Stream<Page, E, R>,
168 pick: (page: Page) => readonly Item[] | undefined,
169) =>
170 stream.pipe(
171 Stream.flatMap((page) => Stream.fromIterable(pick(page) ?? [])),
172 Stream.runCollect,
173 Effect.map((chunk) => Array.from(chunk)),
174 );
175
176export const emptyOnMissing = <A, E extends { readonly _tag: string }, R>(
177 effect: Effect.Effect<A[], E, R>,
178) =>
179 effect.pipe(
180 Effect.catchIf(
181 (error) => error._tag === "NotFound",
182 () => Effect.succeed([] as A[]),
183 ),
184 );
185
186/**
187 * Wait for a Dataproc operation. Cluster creates and batch/session
188 * startups take several minutes (clusters up to ~20). An operation that
189 * finished with `ALREADY_EXISTS` (6) is a lost create race; `notFoundOk`
190 * also accepts `NOT_FOUND` (5).
191 */
192export const waitForOperation = (
193 operation: dataproc.Operation,
194 options?: { notFoundOk?: boolean },
195) =>
196 waitForGcpOperation(
197 operation,
198 (name) =>
199 (name.includes("/regions/")
200 ? dataproc.getProjectsRegionsOperations({ name })
201 : dataproc.getProjectsLocationsOperations({ name })
202 ).pipe(
203 Effect.catchTag("NotFound", (error) =>
204 options?.notFoundOk === true
205 ? Effect.succeed<dataproc.Operation>({ name, done: true })
206 : Effect.fail(error),
207 ),
208 ),
209 { budget: "30 minutes" },
210 ).pipe(
211 Effect.catchIf(
212 (error) =>
213 error._tag === "GCP.OperationFailed" &&
214 (error.code === 6 ||
215 (options?.notFoundOk === true && error.code === 5)),
216 () => Effect.void,
217 ),
218 Effect.asVoid,
219 );
220
221export const waitUntilExists = <A, E extends { readonly _tag: string }, R>(
222 get: Effect.Effect<A | undefined, E, R>,
223 name: string,
224) =>
225 get.pipe(
226 Effect.flatMap((value) =>
227 value
228 ? Effect.succeed(value)
229 : Effect.fail(new DataprocNotResolved({ name })),
230 ),
231 Effect.retry({
232 while: (error) => error._tag === "GCP.Dataproc.ResourceNotResolved",
233 times: 8,
234 schedule: Schedule.spaced("1 second"),
235 }),
236 );
237
238export const waitUntilGone = <A, E extends { readonly _tag: string }, R>(
239 get: Effect.Effect<A | undefined, E, R>,
240 name: string,
241) =>
242 get.pipe(
243 Effect.flatMap((value) =>
244 value === undefined
245 ? Effect.void
246 : Effect.fail(new DataprocStillExists({ name })),
247 ),
248 Effect.retry({
249 while: (error) => error._tag === "GCP.Dataproc.ResourceStillExists",
250 times: MAX_POLLS,
251 schedule: Schedule.spaced("5 seconds"),
252 }),
253 );
254
255export const defaultYarnConfig = (): dataproc.BasicYarnAutoscalingConfig => ({
256 gracefulDecommissionTimeout: "3600s",
257 scaleUpFactor: 0.5,
258 scaleDownFactor: 0,
259 scaleUpMinWorkerFraction: 0,
260 scaleDownMinWorkerFraction: 0,
261});
262
263export const defaultWorkerConfig =
264 (): dataproc.InstanceGroupAutoscalingPolicyConfig => ({
265 minInstances: 2,
266 maxInstances: 3,
267 weight: 1,
268 });
269
270export const defaultSparkBatch = (): dataproc.SparkBatch => ({
271 mainClass: "org.apache.spark.examples.SparkPi",
272 jarFileUris: ["file:///usr/lib/spark/examples/jars/spark-examples.jar"],
273 args: ["1"],
274});
275
276export const defaultWorkflowJobs = (): dataproc.OrderedJobList => [
277 {
278 stepId: "spark-pi",
279 sparkJob: defaultSparkBatch(),
280 },
281];
282
283export const defaultWorkflowPlacement = (
284 clusterName: string,
285): dataproc.WorkflowTemplatePlacement => ({
286 managedCluster: {
287 clusterName,
288 config: {
289 gceClusterConfig: { zoneUri: DEFAULT_ZONE },
290 masterConfig: {
291 numInstances: 1,
292 machineTypeUri: "n1-standard-2",
293 diskConfig: { bootDiskSizeGb: 30 },
294 },
295 workerConfig: { numInstances: 2 },
296 softwareConfig: { imageVersion: "2.2-debian12" },
297 },
298 },
299});