1import * as athena from "@distilled.cloud/aws/athena";
2import * as Effect from "effect/Effect";
3import * as Layer from "effect/Layer";
4import * as Schedule from "effect/Schedule";
5import * as Binding from "../../Binding.ts";
6import * as Output from "../../Output.ts";
7import { isBindingHost } from "../Lambda/Function.ts";
8import type { Bucket } from "../S3/Bucket.ts";
9import type { Output as OutputType } from "../../Output.ts";
10import {
11 AthenaQueryFailed,
12 Query,
13 type QueryResult,
14 type RunQueryRequest,
15} from "./Query.ts";
16import type { WorkGroup } from "./WorkGroup.ts";
17
18// The workgroup ARN is `arn:aws:athena:{region}:{account}:workgroup/{name}` —
19// the Glue Data Catalog ARNs share that partition/region/account.
20const glueBase = (workGroupArn: OutputType<string>) =>
21 workGroupArn.pipe(
22 Output.map((arn) => {
23 const [, partition, , region, account] = arn.split(":");
24 return `arn:${partition}:glue:${region}:${account}`;
25 }),
26 );
27const glueCatalogArn = (workGroupArn: OutputType<string>) =>
28 glueBase(workGroupArn).pipe(Output.map((base) => `${base}:catalog`));
29const glueDatabaseArn = (workGroupArn: OutputType<string>) =>
30 glueBase(workGroupArn).pipe(Output.map((base) => `${base}:database/*`));
31const glueTableArn = (workGroupArn: OutputType<string>) =>
32 glueBase(workGroupArn).pipe(Output.map((base) => `${base}:table/*/*`));
33
34export const QueryHttp = Layer.effect(
35 Query,
36 Effect.gen(function* () {
37 const startQueryExecution = yield* athena.startQueryExecution;
38 const getQueryExecution = yield* athena.getQueryExecution;
39 const getQueryResults = yield* athena.getQueryResults;
40
41 return Effect.fn(function* (workGroup: WorkGroup, resultsBucket: Bucket) {
42 const WorkGroupName = yield* workGroup.workGroupName;
43 if (!globalThis.__ALCHEMY_RUNTIME__) {
44 const host = yield* Binding.Host;
45 if (isBindingHost(host)) {
46 yield* host.bind`Allow(${host}, AWS.Athena.Query(${workGroup}, ${resultsBucket}))`(
47 {
48 policyStatements: [
49 {
50 Effect: "Allow",
51 Action: [
52 "athena:StartQueryExecution",
53 "athena:GetQueryExecution",
54 "athena:GetQueryResults",
55 "athena:StopQueryExecution",
56 ],
57 Resource: [workGroup.workGroupArn],
58 },
59 {
60 Effect: "Allow",
61 Action: ["s3:GetBucketLocation", "s3:ListBucket"],
62 Resource: [resultsBucket.bucketArn],
63 },
64 {
65 Effect: "Allow",
66 Action: [
67 "s3:GetObject",
68 "s3:PutObject",
69 "s3:AbortMultipartUpload",
70 ],
71 Resource: [Output.interpolate`${resultsBucket.bucketArn}/*`],
72 },
73 {
74 // Athena resolves table metadata through the Glue Data
75 // Catalog on the caller's behalf — reads are required for any
76 // query over a Glue-backed table. Scoped to the workgroup's
77 // account/region catalog (derived from its ARN).
78 Effect: "Allow",
79 Action: [
80 "glue:GetDatabase",
81 "glue:GetDatabases",
82 "glue:GetTable",
83 "glue:GetTables",
84 "glue:GetPartition",
85 "glue:GetPartitions",
86 ],
87 Resource: [
88 glueCatalogArn(workGroup.workGroupArn),
89 glueDatabaseArn(workGroup.workGroupArn),
90 glueTableArn(workGroup.workGroupArn),
91 ],
92 },
93 ],
94 },
95 );
96 }
97 }
98 return Effect.fn(`AWS.Athena.Query(${workGroup.LogicalId})`)(function* (
99 request: RunQueryRequest,
100 ) {
101 const workGroupName = yield* WorkGroupName;
102 const started = yield* startQueryExecution({
103 ...request,
104 WorkGroup: workGroupName,
105 });
106 const queryExecutionId = started.QueryExecutionId!;
107
108 // Poll GetQueryExecution until terminal — bounded to ~90s so a hung
109 // query fails fast instead of hitting the caller's timeout.
110 const exec = yield* getQueryExecution({
111 QueryExecutionId: queryExecutionId,
112 }).pipe(
113 Effect.map((r) => r.QueryExecution),
114 Effect.repeat({
115 schedule: Schedule.spaced("1 second"),
116 until: (qe) => {
117 const s = qe?.Status?.State;
118 return s === "SUCCEEDED" || s === "FAILED" || s === "CANCELLED";
119 },
120 times: 90,
121 }),
122 );
123
124 const state = exec?.Status?.State ?? "FAILED";
125 const reason = exec?.Status?.StateChangeReason;
126 if (state !== "SUCCEEDED") {
127 return yield* Effect.fail(
128 new AthenaQueryFailed({ queryExecutionId, state, reason }),
129 );
130 }
131
132 const results = yield* getQueryResults({
133 QueryExecutionId: queryExecutionId,
134 });
135 const columns = (
136 results.ResultSet?.ResultSetMetadata?.ColumnInfo ?? []
137 ).map((c) => c.Name);
138 const rows = (results.ResultSet?.Rows ?? []).map((row) =>
139 (row.Data ?? []).map((d) => d.VarCharValue ?? ""),
140 );
141 return {
142 queryExecutionId,
143 state,
144 stateChangeReason: reason,
145 columns,
146 rows,
147 } satisfies QueryResult;
148 });
149 });
150 }),
151);