1import * as DataBrew from "@/AWS/DataBrew";
2import * as IAM from "@/AWS/IAM";
3import * as Lambda from "@/AWS/Lambda";
4import * as S3 from "@/AWS/S3";
5import type * as databrew from "@distilled.cloud/aws/databrew";
6import * as Duration from "effect/Duration";
7import * as Effect from "effect/Effect";
8import * as Layer from "effect/Layer";
9import * as Schedule from "effect/Schedule";
10import * as Stream from "effect/Stream";
11import { HttpServerRequest } from "effect/http/HttpServerRequest";
12import * as HttpServerResponse from "effect/http/HttpServerResponse";
13import path from "pathe";
14
15const main = path.resolve(import.meta.dirname, "handler.ts");
16
17export const SOURCE_KEY = "raw/data.csv";
18
19// DataBrew can observe a fresh role's grants at different times per operation.
20const authorizationPolicy = {
21 while: (
22 error:
23 | databrew.StartJobRunError
24 | databrew.PublishRecipeError
25 | databrew.StartProjectSessionError
26 | databrew.SendProjectSessionActionError,
27 ) => error._tag === "AccessDeniedException",
28 schedule: Schedule.spaced("3 seconds"),
29 times: 8,
30};
31
32/**
33 * Shared foundation for the bindings fixture: bucket + DataBrew service role
34 * + dataset + published recipe. The Bindings test deploys this FIRST, seeds
35 * the CSV source object, and only then deploys the full fixture — DataBrew
36 * validates the dataset's source object with the job role when the job and
37 * project are created.
38 */
39export const foundation = Effect.gen(function* () {
40 const bucket = yield* S3.Bucket("DataBrewBindingsBucket", {
41 forceDestroy: true,
42 });
43 const role = yield* IAM.Role("DataBrewBindingsRole", {
44 assumeRolePolicyDocument: {
45 Version: "2012-10-17",
46 Statement: [
47 {
48 Effect: "Allow",
49 Principal: { Service: "databrew.amazonaws.com" },
50 Action: ["sts:AssumeRole"],
51 },
52 ],
53 },
54 managedPolicyArns: [
55 "arn:aws:iam::aws:policy/service-role/AwsGlueDataBrewServiceRole",
56 "arn:aws:iam::aws:policy/AmazonS3FullAccess",
57 ],
58 });
59 const dataset = yield* DataBrew.Dataset("DataBrewBindingsSource", {
60 format: "CSV",
61 formatOptions: { csv: { delimiter: ",", headerRow: true } },
62 input: {
63 s3InputDefinition: { bucket: bucket.bucketName, key: SOURCE_KEY },
64 },
65 });
66 const recipe = yield* DataBrew.Recipe("DataBrewBindingsRecipe", {
67 publish: true,
68 steps: [
69 {
70 action: {
71 operation: "UPPER_CASE",
72 parameters: { sourceColumn: "name" },
73 },
74 },
75 ],
76 });
77 return { bucket, role, dataset, recipe };
78});
79
80export class DataBrewTestFunction extends Lambda.Function<Lambda.Function>()(
81 "DataBrewTestFunction",
82) {}
83
84/**
85 * Every route answers `{ …fields }` on success or `{ errorTag }` when the
86 * operation fails with a TYPED error — the test asserts the tag is a typed,
87 * non-authorization tag, which proves both the binding wiring and the IAM
88 * grant. An untyped error crashes into a 500 instead.
89 */
90const errorTagged = <A, E extends { _tag: string }, R>(
91 effect: Effect.Effect<A, E, R>,
92): Effect.Effect<A | { errorTag: string }, never, R> =>
93 effect.pipe(
94 Effect.map((a): A | { errorTag: string } => a),
95 Effect.catch((e) => Effect.succeed({ errorTag: e._tag })),
96 );
97
98export default DataBrewTestFunction.make(
99 {
100 main,
101 functionUrl: true,
102 // Session start + action retries fan out SDK calls — AWS's 3s default
103 // intermittently times out under cold starts.
104 timeout: Duration.seconds(90),
105 },
106 Effect.gen(function* () {
107 const base = yield* foundation;
108
109 const profileJob = yield* DataBrew.Job("DataBrewBindingsProfileJob", {
110 type: "PROFILE",
111 datasetName: base.dataset.datasetName,
112 role: base.role.roleArn,
113 outputLocation: { bucket: base.bucket.bucketName, key: "profiles/" },
114 jobSample: { mode: "CUSTOM_ROWS", size: 100 },
115 maxCapacity: 2,
116 timeout: "30 minutes",
117 });
118
119 const project = yield* DataBrew.Project("DataBrewBindingsProject", {
120 datasetName: base.dataset.datasetName,
121 recipeName: base.recipe.recipeName,
122 sample: { type: "FIRST_N", size: 100 },
123 role: base.role.roleArn,
124 });
125
126 // Job-run plane
127 const startJobRun = yield* DataBrew.StartJobRun(profileJob);
128 const stopJobRun = yield* DataBrew.StopJobRun(profileJob);
129 const describeJobRun = yield* DataBrew.DescribeJobRun(profileJob);
130 const listJobRuns = yield* DataBrew.ListJobRuns(profileJob);
131 // Recipe plane
132 const publishRecipe = yield* DataBrew.PublishRecipe(base.recipe);
133 // Interactive-session plane
134 const startProjectSession = yield* DataBrew.StartProjectSession(project);
135 const sendProjectSessionAction =
136 yield* DataBrew.SendProjectSessionAction(project);
137
138 // Deploy-time: creates the EventBridge rule (default bus, source
139 // aws.databrew) targeting this Function. Runtime firing rides on the
140 // real job runs the suite starts; the test verifies the rule deploys.
141 yield* DataBrew.consumeJobEvents({}, (events) =>
142 Stream.runForEach(events, (event) =>
143 Effect.log(
144 `databrew event: ${event.detail.jobRunId} of ${event.detail.jobName} -> ${event.detail.state}`,
145 ),
146 ),
147 );
148
149 return {
150 fetch: Effect.gen(function* () {
151 const request = yield* HttpServerRequest;
152 const url = new URL(request.originalUrl);
153 const pathname = url.pathname;
154 const route = `${request.method} ${pathname}`;
155 const param = (name: string) => url.searchParams.get(name)!;
156
157 switch (route) {
158 // ---- job-run plane ----
159 case "POST /run/start": {
160 const result = yield* errorTagged(
161 startJobRun().pipe(Effect.retry(authorizationPolicy)),
162 );
163 return yield* HttpServerResponse.json(
164 "errorTag" in result ? result : { runId: result.RunId },
165 );
166 }
167 case "GET /run/list": {
168 const result = yield* errorTagged(listJobRuns());
169 return yield* HttpServerResponse.json(
170 "errorTag" in result
171 ? result
172 : { runIds: (result.JobRuns ?? []).map((r) => r.RunId) },
173 );
174 }
175 case "GET /run/get": {
176 const result = yield* errorTagged(
177 describeJobRun({ RunId: param("id") }),
178 );
179 return yield* HttpServerResponse.json(
180 "errorTag" in result ? result : { state: result.State },
181 );
182 }
183 case "POST /run/stop": {
184 const body = (yield* request.json) as unknown as { id: string };
185 const result = yield* errorTagged(stopJobRun({ RunId: body.id }));
186 return yield* HttpServerResponse.json(
187 "errorTag" in result ? result : { runId: result.RunId },
188 );
189 }
190
191 // ---- recipe plane ----
192 case "POST /recipe/publish": {
193 const result = yield* errorTagged(
194 publishRecipe({
195 Description: "published by bindings fixture",
196 }).pipe(Effect.retry(authorizationPolicy)),
197 );
198 return yield* HttpServerResponse.json(
199 "errorTag" in result ? result : { name: result.Name },
200 );
201 }
202
203 // ---- interactive-session plane ----
204 case "POST /session/run": {
205 const started = yield* errorTagged(
206 startProjectSession({ AssumeControl: true }).pipe(
207 Effect.retry(authorizationPolicy),
208 ),
209 );
210 if ("errorTag" in started) {
211 return yield* HttpServerResponse.json({ started });
212 }
213 // Session readiness and IAM propagation share one bounded retry.
214 const action = yield* errorTagged(
215 sendProjectSessionAction({
216 Preview: true,
217 // The Redacted session token flows straight back in.
218 ClientSessionId: started.ClientSessionId,
219 RecipeStep: {
220 Action: {
221 Operation: "UPPER_CASE",
222 Parameters: { sourceColumn: "name" },
223 },
224 },
225 }).pipe(
226 Effect.retry({
227 while: (e): boolean => e._tag === "ConflictException",
228 schedule: Schedule.max([
229 Schedule.fixed("5 seconds"),
230 Schedule.recurs(6),
231 ]),
232 }),
233 ),
234 );
235 return yield* HttpServerResponse.json({
236 started: {
237 name: started.Name,
238 hasSessionId: started.ClientSessionId !== undefined,
239 },
240 action:
241 "errorTag" in action
242 ? action
243 : { actionId: action.ActionId ?? null },
244 });
245 }
246
247 default:
248 return yield* HttpServerResponse.json(
249 { error: "Not found", route },
250 { status: 404 },
251 );
252 }
253 }).pipe(Effect.orDie),
254 };
255 }).pipe(
256 Effect.provide(
257 Layer.mergeAll(
258 Lambda.EventSource,
259 DataBrew.StartJobRunHttp,
260 DataBrew.StopJobRunHttp,
261 DataBrew.DescribeJobRunHttp,
262 DataBrew.ListJobRunsHttp,
263 DataBrew.PublishRecipeHttp,
264 DataBrew.StartProjectSessionHttp,
265 DataBrew.SendProjectSessionActionHttp,
266 ),
267 ),
268 ),
269);