1import * as AWS from "@/AWS";
2import { Dataset, Job, Recipe } from "@/AWS/DataBrew";
3import { Role } from "@/AWS/IAM";
4import { Bucket } from "@/AWS/S3";
5import * as Test from "@/Test/Alchemy";
6import * as databrew from "@distilled.cloud/aws/databrew";
7import * as s3 from "@distilled.cloud/aws/s3";
8import { expect } from "alchemy-test";
9import * as Effect from "effect/Effect";
10import * as Schedule from "effect/Schedule";
11
12const { test } = Test.make({ providers: AWS.providers() });
13
14const getJob = (name: string) =>
15 databrew
16 .describeJob({ Name: name })
17 .pipe(
18 Effect.catchTag("ResourceNotFoundException", () =>
19 Effect.succeed(undefined),
20 ),
21 );
22
23const brewRole = () =>
24 Role("DataBrewJobRole", {
25 assumeRolePolicyDocument: {
26 Version: "2012-10-17",
27 Statement: [
28 {
29 Effect: "Allow",
30 Principal: { Service: "databrew.amazonaws.com" },
31 Action: ["sts:AssumeRole"],
32 },
33 ],
34 },
35 managedPolicyArns: [
36 "arn:aws:iam::aws:policy/service-role/AwsGlueDataBrewServiceRole",
37 "arn:aws:iam::aws:policy/AmazonS3FullAccess",
38 ],
39 });
40
41// Shared foundation: bucket + role + dataset + published recipe.
42const foundation = Effect.gen(function* () {
43 const bucket = yield* Bucket("JobBucket", { forceDestroy: true });
44 const role = yield* brewRole();
45 const dataset = yield* Dataset("Source", {
46 format: "CSV",
47 formatOptions: { csv: { delimiter: ",", headerRow: true } },
48 input: {
49 s3InputDefinition: { bucket: bucket.bucketName, key: "raw/data.csv" },
50 },
51 });
52 const recipe = yield* Recipe("Transform", {
53 publish: true,
54 steps: [
55 {
56 action: {
57 operation: "UPPER_CASE",
58 parameters: { sourceColumn: "name" },
59 },
60 },
61 ],
62 });
63 return { bucket, role, dataset, recipe };
64});
65
66// CreateProfileJob/CreateRecipeJob validate the dataset's source object with
67// the job role, so the CSV is seeded between the foundation and job deploys.
68const seedSource = (bucketName: string) =>
69 s3.putObject({
70 Bucket: bucketName,
71 Key: "raw/data.csv",
72 Body: new TextEncoder().encode("id,name\n1,alice\n2,bob\n"),
73 ContentType: "text/csv",
74 });
75
76const withJobs = (maxCapacity: number) =>
77 Effect.gen(function* () {
78 const base = yield* foundation;
79 const profileJob = yield* Job("ProfileJob", {
80 type: "PROFILE",
81 datasetName: base.dataset.datasetName,
82 role: base.role.roleArn,
83 outputLocation: { bucket: base.bucket.bucketName, key: "profiles/" },
84 jobSample: { mode: "CUSTOM_ROWS", size: 100 },
85 maxCapacity,
86 timeout: "1 hour",
87 tags: { Environment: "test" },
88 });
89 const recipeJob = yield* Job("RecipeJob", {
90 type: "RECIPE",
91 datasetName: base.dataset.datasetName,
92 recipeReference: { name: base.recipe.recipeName },
93 role: base.role.roleArn,
94 outputs: [
95 {
96 location: { bucket: base.bucket.bucketName, key: "curated/" },
97 format: "CSV",
98 overwrite: true,
99 },
100 ],
101 maxCapacity,
102 tags: { Environment: "test" },
103 });
104 return { ...base, profileJob, recipeJob };
105 });
106
107test.provider(
108 "create, update, delete profile and recipe job definitions",
109 (stack) =>
110 Effect.gen(function* () {
111 yield* stack.destroy();
112
113 // stage 1: infrastructure the jobs depend on
114 const base = yield* stack.deploy(foundation);
115 yield* seedSource(base.bucket.bucketName);
116
117 // stage 2: the job definitions
118 const created = yield* stack.deploy(withJobs(2));
119
120 expect(created.profileJob.jobArn).toContain(
121 `:job/${created.profileJob.jobName}`,
122 );
123 expect(created.profileJob.type).toEqual("PROFILE");
124 expect(created.recipeJob.type).toEqual("RECIPE");
125
126 // out-of-band verification
127 const profile = yield* getJob(created.profileJob.jobName);
128 expect(profile?.Type).toEqual("PROFILE");
129 expect(profile?.DatasetName).toEqual(created.dataset.datasetName);
130 expect(profile?.JobSample?.Size).toEqual(100);
131 expect(profile?.MaxCapacity).toEqual(2);
132 // Duration.Input "1 hour" lands on the wire as 60 (minutes)
133 expect(profile?.Timeout).toEqual(60);
134 expect(profile?.Tags?.["alchemy::id"]).toBeDefined();
135
136 const recipeJob = yield* getJob(created.recipeJob.jobName);
137 expect(recipeJob?.Type).toEqual("RECIPE");
138 expect(recipeJob?.RecipeReference?.Name).toEqual(
139 created.recipe.recipeName,
140 );
141 expect(recipeJob?.Outputs?.[0]?.Location.Bucket).toEqual(
142 created.bucket.bucketName,
143 );
144 expect(recipeJob?.Outputs?.[0]?.Overwrite).toEqual(true);
145
146 // update: bump capacity on both jobs
147 yield* stack.deploy(withJobs(3));
148 const reProfile = yield* getJob(created.profileJob.jobName);
149 expect(reProfile?.MaxCapacity).toEqual(3);
150 const reRecipeJob = yield* getJob(created.recipeJob.jobName);
151 expect(reRecipeJob?.MaxCapacity).toEqual(3);
152
153 yield* stack.destroy();
154 expect(yield* getJob(created.profileJob.jobName)).toBeUndefined();
155 expect(yield* getJob(created.recipeJob.jobName)).toBeUndefined();
156 }),
157 {
158 tags: [
159 "provider:aws",
160 "provider:aws:databrew",
161 "provider:aws:iam",
162 "provider:aws:s3",
163 "live",
164 ],
165 timeout: 180_000,
166 },
167);
168
169// A live DataBrew job run spins up managed Spark capacity (billed per
170// node-hour, takes several minutes) — gated behind AWS_TEST_SLOW=1.
171test.provider.skipIf(!process.env.AWS_TEST_SLOW)(
172 "run a live recipe job to success",
173 (stack) =>
174 Effect.gen(function* () {
175 yield* stack.destroy();
176
177 const base = yield* stack.deploy(foundation);
178 yield* seedSource(base.bucket.bucketName);
179
180 const created = yield* stack.deploy(
181 Effect.gen(function* () {
182 const infra = yield* foundation;
183 const job = yield* Job("RunJob", {
184 type: "RECIPE",
185 datasetName: infra.dataset.datasetName,
186 recipeReference: { name: infra.recipe.recipeName },
187 role: infra.role.roleArn,
188 outputs: [
189 {
190 location: { bucket: infra.bucket.bucketName, key: "curated/" },
191 format: "CSV",
192 overwrite: true,
193 },
194 ],
195 maxCapacity: 2,
196 timeout: "20 minutes",
197 });
198 return { ...infra, job };
199 }),
200 );
201
202 const { RunId } = yield* databrew.startJobRun({
203 Name: created.job.jobName,
204 });
205
206 const run = yield* databrew
207 .describeJobRun({ Name: created.job.jobName, RunId })
208 .pipe(
209 Effect.repeat({
210 schedule: Schedule.spaced("10 seconds"),
211 until: (r) =>
212 r.State === "SUCCEEDED" ||
213 r.State === "FAILED" ||
214 r.State === "TIMEOUT" ||
215 r.State === "STOPPED",
216 times: 60,
217 }),
218 );
219 expect(run.State).toEqual("SUCCEEDED");
220
221 // the transformed output landed in S3
222 const listed = yield* s3.listObjectsV2({
223 Bucket: created.bucket.bucketName,
224 Prefix: "curated/",
225 });
226 expect((listed.KeyCount ?? 0) > 0).toBe(true);
227
228 yield* stack.destroy();
229 expect(yield* getJob(created.job.jobName)).toBeUndefined();
230 }),
231 {
232 tags: [
233 "provider:aws",
234 "provider:aws:databrew",
235 "provider:aws:iam",
236 "provider:aws:s3",
237 "live",
238 ],
239 timeout: 900_000,
240 },
241);