1import * as AWS from "@/AWS";
2import * as Core from "@/Test/Core";
3import * as Test from "@/Test/Alchemy";
4import { describe, expect } from "alchemy-test";
5import * as Data from "effect/Data";
6import * as Effect from "effect/Effect";
7import * as Schedule from "effect/Schedule";
8import * as HttpClient from "effect/http/HttpClient";
9import * as HttpClientRequest from "effect/http/HttpClientRequest";
10import * as athena from "@distilled.cloud/aws/athena";
11import * as glue from "@distilled.cloud/aws/glue";
12
13import AthenaTestFunctionLive, { AthenaTestFunction } from "./handler";
14
15const testOptions = { providers: AWS.providers() };
16const { test, beforeAll, afterAll } = Test.make(testOptions);
17const sharedStack = Core.scratchStack(testOptions, "AthenaQuery");
18
19let baseUrl: string;
20
21class TransientUpstream extends Data.TaggedError("TransientUpstream")<{
22 readonly status: number;
23 readonly body: string;
24}> {
25 // Without this the runner prints a bare `TransientUpstream:` when the retry
26 // budget is exhausted, which is unactionable — the status and the handler's
27 // body are the whole diagnosis.
28 override get message() {
29 const body = this.body.trim();
30 return `upstream returned ${this.status}: ${
31 body.length > 800 ? `${body.slice(0, 800)}… (${body.length} bytes)` : body
32 }`;
33 }
34}
35
36// Fresh Lambda role + Athena/S3 permissions propagate eventually — the first
37// queries can 500 with AccessDenied under the handler's `Effect.orDie`.
38// Retry 5xx only; a genuine 4xx fails immediately.
39const send = (request: HttpClientRequest.HttpClientRequest) =>
40 HttpClient.execute(request).pipe(
41 Effect.flatMap((response) =>
42 response.status >= 500
43 ? response.text.pipe(
44 Effect.flatMap((body) =>
45 Effect.fail(
46 new TransientUpstream({ status: response.status, body }),
47 ),
48 ),
49 )
50 : Effect.succeed(response),
51 ),
52 Effect.retry({
53 while: (e) => e._tag === "TransientUpstream",
54 // Bounded well under the 180s test timeout (~63s of sleeps) so a
55 // persistent 500 surfaces its body instead of an opaque timeout.
56 // 31s was not enough under a loaded full-suite run (`--concurrency 64`),
57 // where the emulator's Lambda invokes queue behind everything else.
58 schedule: Schedule.max([
59 Schedule.exponential("1 second"),
60 Schedule.recurs(6),
61 ]),
62 }),
63 );
64
65describe(
66 "Athena Query",
67 {
68 tags: [
69 "provider:aws",
70 "provider:aws:athena",
71 "provider:aws:batch",
72 "provider:aws:glue",
73 "provider:aws:lambda",
74 "provider:aws:s3",
75 "live",
76 ],
77 },
78 () => {
79 beforeAll(
80 Effect.gen(function* () {
81 yield* Effect.logInfo(
82 "Athena Query setup: destroying previous resources",
83 );
84 yield* sharedStack.destroy();
85
86 yield* Effect.logInfo(
87 "Athena Query setup: deploying bucket -> Glue db/table -> workgroup -> Lambda",
88 );
89 const { functionUrl } = yield* sharedStack.deploy(
90 Effect.gen(function* () {
91 return yield* AthenaTestFunction;
92 }).pipe(Effect.provide(AthenaTestFunctionLive)),
93 );
94
95 expect(functionUrl).toBeTruthy();
96 baseUrl = functionUrl!.replace(/\/+$/, "");
97
98 // Readiness probe — fresh function URLs take seconds to serve 200s.
99 yield* HttpClient.get(`${baseUrl}/nope`).pipe(
100 Effect.flatMap((response) =>
101 response.status === 404
102 ? Effect.succeed(response)
103 : Effect.fail(
104 new Error(`Function not ready: ${response.status}`),
105 ),
106 ),
107 Effect.retry({
108 schedule: Schedule.max([
109 Schedule.fixed("2 seconds"),
110 Schedule.recurs(75),
111 ]),
112 }),
113 );
114 }),
115 { timeout: 300_000 },
116 );
117 // No NO_DESTROY escape hatch here: scratch-stack state is in-memory per
118 // process, so a skipped destroy would orphan the whole stack forever.
119 afterAll(
120 Effect.gen(function* () {
121 yield* sharedStack.destroy();
122
123 // Assert-gone out-of-band via the fixture's fixed-name resources —
124 // engine-named resources (bucket/workgroup/Lambda) hang off the same
125 // stack, so these two vanishing proves the destroy actually ran.
126 // `afterAll` effects don't get the providers layer automatically the
127 // way `test.provider` bodies do, so wrap the distilled calls in
128 // `Core.withProviders` to supply AWS credentials.
129 yield* Core.withProviders(
130 Effect.gen(function* () {
131 const db = yield* glue
132 .getDatabase({ Name: "alchemy_athena_e2e" })
133 .pipe(
134 Effect.map((res) => res.Database),
135 Effect.catchTag("EntityNotFoundException", () =>
136 Effect.succeed(undefined),
137 ),
138 );
139 expect(db).toBeUndefined();
140
141 const catalog = yield* athena
142 .getDataCatalog({
143 Name: "alchemy_athena_e2e_catalog",
144 WorkGroup: "primary",
145 })
146 .pipe(
147 Effect.map((res) => res.DataCatalog),
148 Effect.catchTag("DataCatalogNotFound", () =>
149 Effect.succeed(undefined),
150 ),
151 );
152 expect(catalog).toBeUndefined();
153 }),
154 testOptions,
155 sharedStack.name,
156 );
157 }),
158 { timeout: 300_000 },
159 );
160
161 test.provider(
162 "runs SELECT 1 through the binding (execute + poll + results)",
163 () =>
164 Effect.gen(function* () {
165 const response = yield* send(
166 HttpClientRequest.get(`${baseUrl}/select-one`),
167 );
168 expect(response.status).toBe(200);
169 const body = (yield* response.json) as {
170 state: string;
171 columns: string[];
172 rows: string[][];
173 };
174 expect(body.state).toBe("SUCCEEDED");
175 // SELECT 1 → header row then the value row.
176 expect(body.rows.at(-1)?.[0]).toBe("1");
177 }),
178 { timeout: 180_000 },
179 );
180
181 test.provider(
182 "runs SELECT COUNT(*) over a Glue table on S3 CSV data",
183 () =>
184 Effect.gen(function* () {
185 // Seed the CSV source data first.
186 const seeded = yield* send(HttpClientRequest.post(`${baseUrl}/seed`));
187 expect(seeded.status).toBe(200);
188
189 const response = yield* send(
190 HttpClientRequest.get(`${baseUrl}/count`),
191 );
192 expect(response.status).toBe(200);
193 const body = (yield* response.json) as {
194 state: string;
195 columns: string[];
196 rows: string[][];
197 };
198 expect(body.state).toBe("SUCCEEDED");
199 // Three CSV rows → COUNT(*) = 3 (last row is the value, first is header).
200 expect(body.rows.at(-1)?.[0]).toBe("3");
201 }),
202 { timeout: 180_000 },
203 );
204
205 // One completed SELECT 1 execution shared by every query-execution binding
206 // below (each vitest file runs its tests sequentially).
207 let execId: string | undefined;
208 const ensureExecId = Effect.gen(function* () {
209 if (execId) return execId;
210 const response = yield* send(
211 HttpClientRequest.get(`${baseUrl}/exec/run`),
212 );
213 expect(response.status).toBe(200);
214 const body = (yield* response.json) as { id: string };
215 execId = body.id;
216 return body.id;
217 });
218
219 describe("GetQueryExecution", () => {
220 test.provider(
221 "reads the execution's terminal state",
222 () =>
223 Effect.gen(function* () {
224 const id = yield* ensureExecId;
225 const response = yield* send(
226 HttpClientRequest.get(`${baseUrl}/exec/get?id=${id}`),
227 );
228 expect(response.status).toBe(200);
229 const body = (yield* response.json) as {
230 state: string;
231 workGroup: string;
232 };
233 expect(body.state).toBe("SUCCEEDED");
234 expect(body.workGroup).toBeTruthy();
235 }),
236 { timeout: 120_000 },
237 );
238 });
239
240 describe("GetQueryResults", () => {
241 test.provider(
242 "reads the raw result page",
243 () =>
244 Effect.gen(function* () {
245 const id = yield* ensureExecId;
246 const response = yield* send(
247 HttpClientRequest.get(`${baseUrl}/exec/results?id=${id}`),
248 );
249 expect(response.status).toBe(200);
250 const body = (yield* response.json) as {
251 rows: number;
252 columns: string[];
253 };
254 // SELECT 1 → header row + value row.
255 expect(body.rows).toBe(2);
256 expect(body.columns.length).toBe(1);
257 }),
258 { timeout: 120_000 },
259 );
260 });
261
262 describe("GetQueryRuntimeStatistics", () => {
263 test.provider(
264 "reads the execution timeline",
265 () =>
266 Effect.gen(function* () {
267 const id = yield* ensureExecId;
268 const response = yield* send(
269 HttpClientRequest.get(`${baseUrl}/exec/stats?id=${id}`),
270 );
271 expect(response.status).toBe(200);
272 const body = (yield* response.json) as { totalMillis: number };
273 expect(body.totalMillis).toBeGreaterThan(0);
274 }),
275 { timeout: 120_000 },
276 );
277 });
278
279 describe("BatchGetQueryExecution", () => {
280 test.provider(
281 "reads executions in bulk",
282 () =>
283 Effect.gen(function* () {
284 const id = yield* ensureExecId;
285 const response = yield* send(
286 HttpClientRequest.get(`${baseUrl}/exec/batch?id=${id}`),
287 );
288 expect(response.status).toBe(200);
289 const body = (yield* response.json) as { states: string[] };
290 expect(body.states).toEqual(["SUCCEEDED"]);
291 }),
292 { timeout: 120_000 },
293 );
294 });
295
296 describe("ListQueryExecutions", () => {
297 test.provider(
298 "lists the workgroup's executions (workgroup injected)",
299 () =>
300 Effect.gen(function* () {
301 const id = yield* ensureExecId;
302 const response = yield* send(
303 HttpClientRequest.get(`${baseUrl}/exec/list`),
304 );
305 expect(response.status).toBe(200);
306 const body = (yield* response.json) as { ids: string[] };
307 expect(body.ids).toContain(id);
308 }),
309 { timeout: 120_000 },
310 );
311 });
312
313 describe("StopQueryExecution", () => {
314 test.provider(
315 "stop on a completed execution is an idempotent no-op",
316 () =>
317 Effect.gen(function* () {
318 const id = yield* ensureExecId;
319 const stopped = yield* send(
320 HttpClientRequest.post(`${baseUrl}/exec/stop?id=${id}`),
321 );
322 expect(stopped.status).toBe(200);
323 // Still SUCCEEDED — stop did not flip a terminal state.
324 const after = yield* send(
325 HttpClientRequest.get(`${baseUrl}/exec/get?id=${id}`),
326 );
327 const body = (yield* after.json) as { state: string };
328 expect(body.state).toBe("SUCCEEDED");
329 }),
330 { timeout: 120_000 },
331 );
332 });
333
334 describe("ListNamedQueries", () => {
335 test.provider(
336 "lists the workgroup's saved queries (workgroup injected)",
337 () =>
338 Effect.gen(function* () {
339 const response = yield* send(
340 HttpClientRequest.get(`${baseUrl}/named/list`),
341 );
342 expect(response.status).toBe(200);
343 const body = (yield* response.json) as { ids: string[] };
344 // The fixture saves exactly one named query in the (fresh) workgroup.
345 expect(body.ids.length).toBeGreaterThanOrEqual(1);
346 }),
347 { timeout: 120_000 },
348 );
349 });
350
351 describe("ListPreparedStatements", () => {
352 test.provider(
353 "lists the workgroup's prepared statements (workgroup injected)",
354 () =>
355 Effect.gen(function* () {
356 const response = yield* send(
357 HttpClientRequest.get(`${baseUrl}/prepared/list`),
358 );
359 expect(response.status).toBe(200);
360 const body = (yield* response.json) as { names: string[] };
361 expect(body.names).toContain("alchemy_athena_e2e_stmt");
362 }),
363 { timeout: 120_000 },
364 );
365 });
366
367 describe("ListDatabases", () => {
368 test.provider(
369 "lists databases through the GLUE-backed catalog",
370 () =>
371 Effect.gen(function* () {
372 const response = yield* send(
373 HttpClientRequest.get(`${baseUrl}/catalog/databases`),
374 );
375 expect(response.status).toBe(200);
376 const body = (yield* response.json) as { names: string[] };
377 expect(body.names).toContain("alchemy_athena_e2e");
378 }),
379 { timeout: 120_000 },
380 );
381 });
382
383 describe("GetDatabase", () => {
384 test.provider(
385 "reads a single database's metadata",
386 () =>
387 Effect.gen(function* () {
388 const response = yield* send(
389 HttpClientRequest.get(
390 `${baseUrl}/catalog/database?name=alchemy_athena_e2e`,
391 ),
392 );
393 expect(response.status).toBe(200);
394 const body = (yield* response.json) as { name: string };
395 expect(body.name).toBe("alchemy_athena_e2e");
396 }),
397 { timeout: 120_000 },
398 );
399 });
400
401 describe("ListTableMetadata", () => {
402 test.provider(
403 "lists the database's tables",
404 () =>
405 Effect.gen(function* () {
406 const response = yield* send(
407 HttpClientRequest.get(
408 `${baseUrl}/catalog/tables?db=alchemy_athena_e2e`,
409 ),
410 );
411 expect(response.status).toBe(200);
412 const body = (yield* response.json) as { names: string[] };
413 expect(body.names).toContain("people");
414 }),
415 { timeout: 120_000 },
416 );
417 });
418
419 describe("GetTableMetadata", () => {
420 test.provider(
421 "reads a table's columns",
422 () =>
423 Effect.gen(function* () {
424 const response = yield* send(
425 HttpClientRequest.get(
426 `${baseUrl}/catalog/table?db=alchemy_athena_e2e&name=people`,
427 ),
428 );
429 expect(response.status).toBe(200);
430 const body = (yield* response.json) as { columns: string[] };
431 expect(body.columns).toEqual(["id", "name"]);
432 }),
433 { timeout: 120_000 },
434 );
435 });
436
437 describe("QueryStateChangeEventSource", () => {
438 test.provider(
439 "delivers query state-change events to the handler",
440 () =>
441 Effect.gen(function* () {
442 // The probe runs a query and polls S3 for the marker the event
443 // handler writes; retry the whole probe a few times to ride out
444 // fresh-rule propagation on the default bus.
445 const seen = yield* Effect.gen(function* () {
446 const response = yield* send(
447 HttpClientRequest.get(`${baseUrl}/events/probe`),
448 );
449 expect(response.status).toBe(200);
450 const body = (yield* response.json) as {
451 seen: boolean;
452 id: string;
453 };
454 return body.seen;
455 }).pipe(
456 Effect.repeat({
457 schedule: Schedule.spaced("2 seconds"),
458 until: (s): boolean => s,
459 times: 3,
460 }),
461 );
462 expect(seen).toBe(true);
463 }),
464 { timeout: 180_000 },
465 );
466 });
467 },
468);