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 { Unowned } from "../../AdoptPolicy.ts";
6import { isResolved } from "../../Diff.ts";
7import * as Provider from "../../Provider.ts";
8import { Resource } from "../../Resource.ts";
9import { tagRecord } from "../../Tags.ts";
10import { GcpEnvironment } from "../Environment.ts";
11import { createInternalLabels, hasAlchemyLabels, toLabels } from "../Labels.ts";
12import type { Providers } from "../Providers.ts";
13import {
14 collectPages,
15 LIST_LOCATIONS,
16 MAX_POLLS,
17 MAX_WORKLOAD_ID_LENGTH,
18 defaultSparkBatch,
19 emptyOnMissing,
20 hasAlchemyLabelMap,
21 locationParent,
22 normalizeLocation,
23 parseResourceName,
24 toPhysicalId,
25 userLabels,
26 waitForOperation,
27 waitUntilExists,
28 waitUntilGone,
29} from "./internal.ts";
30
31const TERMINAL = new Set(["SUCCEEDED", "FAILED", "CANCELLED", "CANCELLING"]);
32
33export type BatchProps = {
34 /**
35 * Batch id (the `{batch}` segment of
36 * `projects/{project}/locations/{location}/batches/{batch}`). If omitted,
37 * a unique RFC1035 name is generated. 4-63 characters. Immutable —
38 * changing it replaces the batch.
39 */
40 batchId?: string;
41 /**
42 * Location (`us-central1`, …). Immutable — changing it replaces the
43 * batch. `US-CENTRAL1` is accepted and normalized to `us-central1`.
44 * @default the stack's GCP region (`GCP.Region`, profile region, `us-central1`)
45 */
46 location?: string;
47 /**
48 * User labels. Alchemy ownership labels are merged in automatically.
49 */
50 labels?: Record<string, string>;
51 /**
52 * Spark batch config. Defaults to SparkPi from the runtime image when
53 * no other workload kind is set.
54 */
55 sparkBatch?: dataproc.SparkBatch;
56 /**
57 * PySpark batch config.
58 */
59 pysparkBatch?: dataproc.PySparkBatch;
60 /**
61 * SparkR batch config.
62 */
63 sparkRBatch?: dataproc.SparkRBatch;
64 /**
65 * Spark SQL batch config.
66 */
67 sparkSqlBatch?: dataproc.SparkSqlBatch;
68 /**
69 * PySpark notebook batch config.
70 */
71 pysparkNotebookBatch?: dataproc.PySparkNotebookBatch;
72 /**
73 * Runtime configuration (version, properties, container image).
74 */
75 runtimeConfig?: dataproc.RuntimeConfig;
76 /**
77 * Environment configuration (execution and peripherals).
78 */
79 environmentConfig?: dataproc.EnvironmentConfig;
80};
81
82export type Batch = Resource<
83 "GCP.Dataproc.Batch",
84 BatchProps,
85 {
86 /** Full resource name `projects/{project}/locations/{location}/batches/{batch}`. */
87 name: string;
88 /** Batch id (last path segment). */
89 batchId: string;
90 /** Project id. */
91 project: string;
92 /** Location id. */
93 location: string;
94 /** User labels (Alchemy ownership labels stripped). */
95 labels: Record<string, string>;
96 /** Server-reported state (`PENDING`, `RUNNING`, `SUCCEEDED`, …). */
97 state: string | undefined;
98 /** Extra status text, if any. */
99 stateMessage: string | undefined;
100 /** Server-assigned uuid. */
101 uuid: string | undefined;
102 /** Associated long-running operation. */
103 operation: string | undefined;
104 /** Creator email. */
105 creator: string | undefined;
106 /** RFC3339 creation timestamp. */
107 createTime: string | undefined;
108 },
109 never,
110 Providers
111>;
112
113/**
114 * A Dataproc serverless batch workload.
115 *
116 * Batches are immutable after create. Changing identity or config
117 * replaces the batch (re-runs the workload). The resource exists as soon
118 * as the create operation returns; Spark completion is not required.
119 *
120 * ### Creating a Batch
121 * **Example:** SparkPi
122 * ```typescript
123 * const batch = yield* GCP.Dataproc.Batch("Pi", {
124 * sparkBatch: {
125 * mainClass: "org.apache.spark.examples.SparkPi",
126 * jarFileUris: ["file:///usr/lib/spark/examples/jars/spark-examples.jar"],
127 * args: ["1"],
128 * },
129 * });
130 * ```
131 *
132 * **Example:** PySpark
133 * ```typescript
134 * const batch = yield* GCP.Dataproc.Batch("Etl", {
135 * pysparkBatch: { mainPythonFileUri: "gs://bucket/job.py" },
136 * labels: { env: "prod" },
137 * });
138 * ```
139 *
140 * @resource
141 * @category Dataproc
142 */
143export const Batch = Resource<Batch>("GCP.Dataproc.Batch");
144
145export class BatchNotResolved extends Data.TaggedError(
146 "GCP.Dataproc.BatchNotResolved",
147)<{
148 name: string;
149}> {}
150
151export class BatchNotTerminal extends Data.TaggedError(
152 "GCP.Dataproc.BatchNotTerminal",
153)<{
154 name: string;
155 state: string | undefined;
156}> {}
157
158const resourceName = (project: string, location: string, batchId: string) =>
159 `${locationParent(project, location)}/batches/${batchId}`;
160
161const defaultSpark = (news: BatchProps): dataproc.SparkBatch | undefined => {
162 if (
163 news.pysparkBatch !== undefined ||
164 news.sparkRBatch !== undefined ||
165 news.sparkSqlBatch !== undefined ||
166 news.pysparkNotebookBatch !== undefined
167 ) {
168 return news.sparkBatch;
169 }
170 return news.sparkBatch ?? defaultSparkBatch();
171};
172
173const toAttrs = (batch: dataproc.Batch, project: string, location: string) => {
174 const name = batch.name ?? "";
175 const parsed = parseResourceName(name, "batches", location);
176 return {
177 name,
178 batchId: parsed.id,
179 project: parsed.project || project,
180 location: parsed.location || location,
181 labels: userLabels(batch.labels),
182 state: batch.state,
183 stateMessage: batch.stateMessage,
184 uuid: batch.uuid,
185 operation: batch.operation,
186 creator: batch.creator,
187 createTime: batch.createTime,
188 };
189};
190
191const desiredBody = (
192 news: BatchProps,
193 desiredLabels: Record<string, string>,
194): dataproc.Batch => ({
195 labels: desiredLabels,
196 sparkBatch: defaultSpark(news),
197 pysparkBatch: news.pysparkBatch,
198 sparkRBatch: news.sparkRBatch,
199 sparkSqlBatch: news.sparkSqlBatch,
200 pysparkNotebookBatch: news.pysparkNotebookBatch,
201 runtimeConfig: news.runtimeConfig,
202 environmentConfig: news.environmentConfig ?? {
203 executionConfig: { ttl: "600s" },
204 },
205});
206
207const getByName = (name: string) =>
208 dataproc
209 .getProjectsLocationsBatches({ name })
210 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
211
212const waitUntilTerminal = (name: string) =>
213 getByName(name).pipe(
214 Effect.flatMap((batch) =>
215 batch === undefined
216 ? Effect.succeed(undefined)
217 : TERMINAL.has(batch.state ?? "")
218 ? Effect.succeed(batch)
219 : Effect.fail(
220 new BatchNotTerminal({
221 name,
222 state: batch.state,
223 }),
224 ),
225 ),
226 Effect.retry({
227 while: (error) => error._tag === "GCP.Dataproc.BatchNotTerminal",
228 times: MAX_POLLS,
229 schedule: Schedule.spaced("5 seconds"),
230 }),
231 );
232
233const listLocation = (project: string, location: string) =>
234 emptyOnMissing(
235 collectPages(
236 dataproc.listProjectsLocationsBatches.pages({
237 parent: locationParent(project, location),
238 pageSize: 1000,
239 }),
240 (page) => page.batches,
241 ).pipe(
242 Effect.map((items) =>
243 items
244 .filter((batch) => hasAlchemyLabelMap(batch.labels))
245 .map((batch) => toAttrs(batch, project, location)),
246 ),
247 ),
248 );
249
250export const BatchProvider = () =>
251 Provider.succeed(Batch, {
252 stables: ["name", "batchId", "project", "location", "uuid", "createTime"],
253
254 diff: Effect.fn(function* ({ news, olds, output }) {
255 if (!isResolved(news)) return undefined;
256 const env = yield* GcpEnvironment.current;
257 const previousId = olds?.batchId ?? output?.batchId;
258 const nextId = news.batchId ?? previousId;
259 const previousLocation = normalizeLocation(
260 olds?.location ?? output?.location,
261 env.region,
262 );
263 const nextLocation = normalizeLocation(
264 news.location ?? olds?.location ?? output?.location,
265 env.region,
266 );
267 if (
268 (previousId !== undefined &&
269 nextId !== undefined &&
270 previousId !== nextId) ||
271 (output !== undefined && previousLocation !== nextLocation)
272 ) {
273 return {
274 action: "replace" as const,
275 deleteFirst:
276 previousLocation === nextLocation &&
277 previousId !== undefined &&
278 nextId === previousId,
279 };
280 }
281 return undefined;
282 }),
283
284 read: Effect.fn(function* ({ id, olds, output }) {
285 const env = yield* GcpEnvironment.current;
286 const location = normalizeLocation(
287 olds?.location ?? output?.location,
288 env.region,
289 );
290 const batchId = yield* toPhysicalId(
291 id,
292 olds?.batchId,
293 output?.batchId,
294 MAX_WORKLOAD_ID_LENGTH,
295 "batch",
296 );
297 const name = output?.name ?? resourceName(env.project, location, batchId);
298 const existing = yield* getByName(name);
299 if (existing === undefined) return undefined;
300 const attrs = toAttrs(existing, env.project, location);
301 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
302 ? attrs
303 : Unowned(attrs);
304 }),
305
306 list: () =>
307 Effect.gen(function* () {
308 const env = yield* GcpEnvironment.current;
309 const pages = yield* Effect.forEach(
310 [...new Set([env.region, ...LIST_LOCATIONS])],
311 (location) => listLocation(env.project, location),
312 { concurrency: 4 },
313 );
314 return pages.flat();
315 }),
316
317 reconcile: Effect.fn(function* ({ id, news, output }) {
318 const env = yield* GcpEnvironment.current;
319 const location = normalizeLocation(
320 news.location ?? output?.location,
321 env.region,
322 );
323 const batchId = yield* toPhysicalId(
324 id,
325 news.batchId,
326 output?.batchId,
327 MAX_WORKLOAD_ID_LENGTH,
328 "batch",
329 );
330 const name = resourceName(env.project, location, batchId);
331 const desiredLabels = {
332 ...toLabels(news.labels),
333 ...(yield* createInternalLabels(id)),
334 };
335 const desired = desiredBody(news, desiredLabels);
336
337 let current = yield* getByName(name);
338
339 if (current === undefined) {
340 const created = yield* dataproc
341 .createProjectsLocationsBatches({
342 parent: locationParent(env.project, location),
343 batchId,
344 body: desired,
345 })
346 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
347 if (created) {
348 yield* waitForOperation(created);
349 }
350 current = yield* waitUntilExists(getByName(name), name);
351 }
352
353 if (current === undefined) {
354 return yield* new BatchNotResolved({ name });
355 }
356
357 return toAttrs(current, env.project, location);
358 }),
359
360 delete: Effect.fn(function* ({ output }) {
361 if (!output.name) return;
362 yield* waitUntilTerminal(output.name);
363 yield* dataproc
364 .deleteProjectsLocationsBatches({ name: output.name })
365 .pipe(
366 Effect.retry({
367 while: (error) =>
368 error._tag === "Conflict" || error._tag === "BadRequest",
369 times: 8,
370 schedule: Schedule.spaced("5 seconds"),
371 }),
372 Effect.catchTag("NotFound", () => Effect.void),
373 );
374 yield* waitUntilGone(getByName(output.name), output.name);
375 }),
376 });