1import * as Athena from "@/AWS/Athena";
2import * as Glue from "@/AWS/Glue";
3import * as Lambda from "@/AWS/Lambda";
4import * as S3 from "@/AWS/S3";
5import * as Output from "@/Output";
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
17// Fixed lowercase Glue identifiers so the SQL below is deterministic. The
18// bucket and workgroup are engine-named (unique); the database/table/catalog
19// are scoped to this single Athena e2e suite.
20const DATABASE = "alchemy_athena_e2e";
21const TABLE = "people";
22const CATALOG = "alchemy_athena_e2e_catalog";
23const STATEMENT = "alchemy_athena_e2e_stmt";
24const CSV = "1,alice\n2,bob\n3,carol\n";
25
26export class AthenaTestFunction extends Lambda.Function<Lambda.Function>()(
27 "AthenaTestFunction",
28) {}
29
30export default AthenaTestFunction.make(
31 {
32 main,
33 functionUrl: true,
34 // startQueryExecution → poll getQueryExecution → getQueryResults fans out
35 // several SDK calls; AWS's 3s default would intermittently time out.
36 timeout: Duration.seconds(60),
37 },
38 Effect.gen(function* () {
39 // Single bucket holds the CSV source data (under data/), the Athena query
40 // results (under results/), and the event-source markers (under events/).
41 const bucket = yield* S3.Bucket("AthenaBucket", { forceDestroy: true });
42
43 const database = yield* Glue.Database("AthenaDatabase", {
44 databaseName: DATABASE,
45 });
46 yield* Glue.Table("AthenaTable", {
47 databaseName: database.databaseName,
48 tableName: TABLE,
49 tableType: "EXTERNAL_TABLE",
50 storageDescriptor: {
51 location: Output.interpolate`s3://${bucket.bucketName}/data/`,
52 columns: [
53 { name: "id", type: "string" },
54 { name: "name", type: "string" },
55 ],
56 inputFormat: "org.apache.hadoop.mapred.TextInputFormat",
57 outputFormat:
58 "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat",
59 serdeInfo: {
60 serializationLibrary:
61 "org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe",
62 parameters: { "field.delim": "," },
63 },
64 },
65 });
66
67 const workGroup = yield* Athena.WorkGroup("AthenaWorkGroup", {
68 outputLocation: Output.interpolate`s3://${bucket.bucketName}/results/`,
69 enforceWorkGroupConfiguration: true,
70 });
71
72 // A GLUE-type data catalog federating to this account's own Glue Data
73 // Catalog (catalog-id = the account id, derived from the workgroup ARN)
74 // — drives the catalog-metadata bindings against the fixture database.
75 const dataCatalog = yield* Athena.DataCatalog("AthenaGlueCatalog", {
76 name: CATALOG,
77 type: "GLUE",
78 parameters: {
79 "catalog-id": workGroup.workGroupArn.pipe(
80 Output.map((arn) => arn.split(":")[4]!),
81 ),
82 },
83 });
84
85 // Saved-query + prepared-statement fixtures in the workgroup — the
86 // ListNamedQueries/ListPreparedStatements bindings enumerate them.
87 yield* Athena.NamedQuery("AthenaNamedQuery", {
88 queryString: `SELECT * FROM ${TABLE} LIMIT 1`,
89 database: DATABASE,
90 workGroup: workGroup.workGroupName,
91 });
92 yield* Athena.PreparedStatement("AthenaPreparedStatement", {
93 statementName: STATEMENT,
94 queryStatement: `SELECT * FROM ${DATABASE}.${TABLE} WHERE id = ?`,
95 workGroup: workGroup.workGroupName,
96 });
97
98 const putObject = yield* S3.PutObject(bucket);
99 const getObject = yield* S3.GetObject(bucket);
100 const query = yield* Athena.Query(workGroup, bucket);
101 const runQuery = Effect.fn((request: Athena.RunQueryRequest) =>
102 query(request).pipe(
103 Effect.retry({
104 while: (error) => error._tag === "AccessDeniedException",
105 schedule: Schedule.spaced("3 seconds"),
106 times: 8,
107 }),
108 ),
109 );
110
111 // --- workgroup-scoped query-execution bindings ---
112 const getQueryExecution = yield* Athena.GetQueryExecution(workGroup);
113 const getQueryResults = yield* Athena.GetQueryResults(workGroup);
114 const stopQueryExecution = yield* Athena.StopQueryExecution(workGroup);
115 const batchGetQueryExecution =
116 yield* Athena.BatchGetQueryExecution(workGroup);
117 const listQueryExecutions = yield* Athena.ListQueryExecutions(workGroup);
118 const getQueryRuntimeStatistics =
119 yield* Athena.GetQueryRuntimeStatistics(workGroup);
120 const listNamedQueries = yield* Athena.ListNamedQueries(workGroup);
121 const listPreparedStatements =
122 yield* Athena.ListPreparedStatements(workGroup);
123
124 // --- catalog-metadata bindings ---
125 const listDatabases = yield* Athena.ListDatabases(dataCatalog);
126 const getDatabase = yield* Athena.GetDatabase(dataCatalog);
127 const listTableMetadata = yield* Athena.ListTableMetadata(dataCatalog);
128 const getTableMetadata = yield* Athena.GetTableMetadata(dataCatalog);
129
130 // --- event source ---
131 // Athena publishes every query state transition to the default bus; write
132 // a marker object per terminal transition so /events/probe can observe
133 // the delivery out-of-band (the event may arrive on another instance).
134 yield* Athena.consumeQueryStateChanges(
135 { states: ["SUCCEEDED", "FAILED", "CANCELLED"] },
136 (events) =>
137 Stream.runForEach(events, (event) =>
138 putObject({
139 Key: `events/${event.detail.queryExecutionId}`,
140 Body: JSON.stringify(event.detail),
141 ContentType: "application/json",
142 }).pipe(Effect.orDie, Effect.asVoid),
143 ),
144 );
145
146 return {
147 fetch: Effect.gen(function* () {
148 const request = yield* HttpServerRequest;
149 const url = new URL(request.originalUrl);
150 const pathname = url.pathname;
151 const param = (name: string) => url.searchParams.get(name) ?? "";
152
153 // Seed the CSV source data into s3://bucket/data/.
154 if (request.method === "POST" && pathname === "/seed") {
155 yield* putObject({
156 Key: "data/rows.csv",
157 Body: CSV,
158 ContentType: "text/csv",
159 });
160 return yield* HttpServerResponse.json({ seeded: true });
161 }
162
163 // Flagship: run a Glue-backed query end-to-end and read the result.
164 if (request.method === "GET" && pathname === "/count") {
165 const result = yield* runQuery({
166 QueryString: `SELECT COUNT(*) AS c FROM ${DATABASE}.${TABLE}`,
167 });
168 return yield* HttpServerResponse.json({
169 state: result.state,
170 columns: result.columns,
171 rows: result.rows,
172 });
173 }
174
175 // Glue-free path: proves start+poll+results without any table.
176 if (request.method === "GET" && pathname === "/select-one") {
177 const result = yield* runQuery({ QueryString: "SELECT 1" });
178 return yield* HttpServerResponse.json({
179 state: result.state,
180 columns: result.columns,
181 rows: result.rows,
182 });
183 }
184
185 // Run a trivial query to completion and hand back its execution id
186 // for the granular query-execution bindings below.
187 if (request.method === "GET" && pathname === "/exec/run") {
188 const result = yield* runQuery({ QueryString: "SELECT 1" });
189 return yield* HttpServerResponse.json({
190 id: result.queryExecutionId,
191 });
192 }
193
194 if (request.method === "GET" && pathname === "/exec/get") {
195 const res = yield* getQueryExecution({
196 QueryExecutionId: param("id"),
197 });
198 return yield* HttpServerResponse.json({
199 state: res.QueryExecution?.Status?.State,
200 workGroup: res.QueryExecution?.WorkGroup,
201 });
202 }
203
204 if (request.method === "GET" && pathname === "/exec/results") {
205 const res = yield* getQueryResults({
206 QueryExecutionId: param("id"),
207 });
208 return yield* HttpServerResponse.json({
209 rows: res.ResultSet?.Rows?.length ?? 0,
210 columns: (res.ResultSet?.ResultSetMetadata?.ColumnInfo ?? []).map(
211 (c) => c.Name,
212 ),
213 });
214 }
215
216 if (request.method === "GET" && pathname === "/exec/stats") {
217 // Runtime statistics materialize asynchronously shortly after the
218 // query completes — retry briefly.
219 const res = yield* getQueryRuntimeStatistics({
220 QueryExecutionId: param("id"),
221 }).pipe(
222 Effect.retry({
223 while: (e): boolean => e._tag === "InvalidRequestException",
224 schedule: Schedule.spaced("1 second"),
225 times: 8,
226 }),
227 );
228 return yield* HttpServerResponse.json({
229 totalMillis:
230 res.QueryRuntimeStatistics?.Timeline?.TotalExecutionTimeInMillis,
231 });
232 }
233
234 if (request.method === "GET" && pathname === "/exec/batch") {
235 const res = yield* batchGetQueryExecution({
236 QueryExecutionIds: [param("id")],
237 });
238 return yield* HttpServerResponse.json({
239 states: (res.QueryExecutions ?? []).map((qe) => qe.Status?.State),
240 });
241 }
242
243 if (request.method === "GET" && pathname === "/exec/list") {
244 const res = yield* listQueryExecutions({ MaxResults: 25 });
245 return yield* HttpServerResponse.json({
246 ids: res.QueryExecutionIds ?? [],
247 });
248 }
249
250 if (request.method === "POST" && pathname === "/exec/stop") {
251 // Stopping an already-completed query is a documented no-op, which
252 // makes this deterministic without racing a long-running query.
253 yield* stopQueryExecution({ QueryExecutionId: param("id") });
254 return yield* HttpServerResponse.json({ stopped: true });
255 }
256
257 if (request.method === "GET" && pathname === "/named/list") {
258 const res = yield* listNamedQueries({ MaxResults: 50 });
259 return yield* HttpServerResponse.json({
260 ids: res.NamedQueryIds ?? [],
261 });
262 }
263
264 if (request.method === "GET" && pathname === "/prepared/list") {
265 const res = yield* listPreparedStatements({ MaxResults: 50 });
266 return yield* HttpServerResponse.json({
267 names: (res.PreparedStatements ?? []).map((s) => s.StatementName),
268 });
269 }
270
271 if (request.method === "GET" && pathname === "/catalog/databases") {
272 const res = yield* listDatabases({});
273 return yield* HttpServerResponse.json({
274 names: (res.DatabaseList ?? []).map((db) => db.Name),
275 });
276 }
277
278 if (request.method === "GET" && pathname === "/catalog/database") {
279 const res = yield* getDatabase({ DatabaseName: param("name") });
280 return yield* HttpServerResponse.json({
281 name: res.Database?.Name,
282 });
283 }
284
285 if (request.method === "GET" && pathname === "/catalog/tables") {
286 const res = yield* listTableMetadata({ DatabaseName: param("db") });
287 return yield* HttpServerResponse.json({
288 names: (res.TableMetadataList ?? []).map((t) => t.Name),
289 });
290 }
291
292 if (request.method === "GET" && pathname === "/catalog/table") {
293 const res = yield* getTableMetadata({
294 DatabaseName: param("db"),
295 TableName: param("name"),
296 });
297 return yield* HttpServerResponse.json({
298 columns: (res.TableMetadata?.Columns ?? []).map((c) => c.Name),
299 });
300 }
301
302 // Run a query, then poll for the event-source marker the state-change
303 // handler wrote to S3 (bounded well under the 60s function timeout).
304 if (request.method === "GET" && pathname === "/events/probe") {
305 const result = yield* runQuery({ QueryString: "SELECT 1" });
306 const id = result.queryExecutionId;
307 const seen = yield* getObject({ Key: `events/${id}` }).pipe(
308 Effect.retry({
309 while: (e): boolean => e._tag === "NoSuchKey",
310 schedule: Schedule.spaced("2 seconds"),
311 times: 12,
312 }),
313 Effect.map(() => true),
314 Effect.catchTag("NoSuchKey", () => Effect.succeed(false)),
315 );
316 return yield* HttpServerResponse.json({ seen, id });
317 }
318
319 return yield* HttpServerResponse.json(
320 { error: "Not found", method: request.method, pathname },
321 { status: 404 },
322 );
323 }).pipe(Effect.orDie),
324 };
325 }).pipe(
326 Effect.provide(
327 Layer.mergeAll(
328 Lambda.EventSource,
329 Athena.QueryHttp,
330 Athena.GetQueryExecutionHttp,
331 Athena.GetQueryResultsHttp,
332 Athena.StopQueryExecutionHttp,
333 Athena.BatchGetQueryExecutionHttp,
334 Athena.ListQueryExecutionsHttp,
335 Athena.GetQueryRuntimeStatisticsHttp,
336 Athena.ListNamedQueriesHttp,
337 Athena.ListPreparedStatementsHttp,
338 Athena.ListDatabasesHttp,
339 Athena.GetDatabaseHttp,
340 Athena.ListTableMetadataHttp,
341 Athena.GetTableMetadataHttp,
342 S3.PutObjectHttp,
343 S3.GetObjectHttp,
344 ),
345 ),
346 ),
347);