1import * as athena from "@distilled.cloud/aws/athena";
2import * as Effect from "effect/Effect";
3import * as Stream from "effect/Stream";
4import { isResolved } from "../../Diff.ts";
5import { createPhysicalName } from "../../PhysicalName.ts";
6import * as Provider from "../../Provider.ts";
7import { Resource } from "../../Resource.ts";
8import { createInternalTags, diffTags } from "../../Tags.ts";
9import { AWSEnvironment, type AccountID } from "../Environment.ts";
10import type { Providers } from "../Providers.ts";
11import type { RegionID } from "../Region.ts";
12
13export type WorkGroupName = string;
14export type WorkGroupArn =
15 `arn:aws:athena:${RegionID}:${AccountID}:workgroup/${WorkGroupName}`;
16
17export interface WorkGroupProps {
18 /**
19 * Name of the workgroup. If omitted, a unique name is generated.
20 * Changing this replaces the workgroup. Up to 128 characters.
21 */
22 workGroupName?: string;
23 /**
24 * Human-readable description of the workgroup.
25 */
26 description?: string;
27 /**
28 * S3 URI (`s3://bucket/prefix/`) where query results are written.
29 * Applied as the workgroup's `ResultConfiguration.OutputLocation`.
30 */
31 outputLocation?: string;
32 /**
33 * Server-side encryption for query results written to S3.
34 */
35 encryptionOption?: "SSE_S3" | "SSE_KMS" | "CSE_KMS";
36 /**
37 * KMS key ARN/ID — required when `encryptionOption` is `SSE_KMS` or `CSE_KMS`.
38 */
39 kmsKey?: string;
40 /**
41 * Force queries to use this workgroup's result configuration (output
42 * location, encryption) instead of client-side settings.
43 * @default true
44 */
45 enforceWorkGroupConfiguration?: boolean;
46 /**
47 * Per-query data-scanned cutoff in bytes. Queries scanning more are
48 * cancelled. Minimum is 10 MB (10_000_000).
49 */
50 bytesScannedCutoffPerQuery?: number;
51 /**
52 * Publish per-query CloudWatch metrics for this workgroup.
53 * @default false
54 */
55 publishCloudWatchMetricsEnabled?: boolean;
56 /**
57 * Charge S3 data-transfer/request costs to the query requester.
58 * @default false
59 */
60 requesterPaysEnabled?: boolean;
61 /**
62 * Selected Athena engine version (e.g. `"Athena engine version 3"`).
63 * If omitted, Athena picks the account default (AUTO).
64 */
65 engineVersion?: string;
66 /**
67 * Whether the workgroup accepts queries.
68 * @default "ENABLED"
69 */
70 state?: "ENABLED" | "DISABLED";
71 /**
72 * User-defined tags to apply to the workgroup.
73 */
74 tags?: Record<string, string>;
75}
76
77export interface WorkGroup extends Resource<
78 "AWS.Athena.WorkGroup",
79 WorkGroupProps,
80 {
81 /**
82 * Name of the workgroup.
83 */
84 workGroupName: WorkGroupName;
85 /**
86 * ARN of the workgroup.
87 */
88 workGroupArn: WorkGroupArn;
89 /**
90 * Whether the workgroup accepts queries.
91 */
92 state: "ENABLED" | "DISABLED";
93 /**
94 * S3 location query results are written to (`s3://bucket/prefix/`).
95 */
96 outputLocation: string | undefined;
97 /**
98 * Tags on the workgroup.
99 */
100 tags: Record<string, string>;
101 },
102 never,
103 Providers
104> {}
105
106/**
107 * An Amazon Athena workgroup — an isolation boundary for queries that pins the
108 * S3 result-output location, result encryption, a bytes-scanned cutoff, and
109 * whether that configuration is enforced over per-query client settings.
110 *
111 * ### Creating Workgroups
112 * **Example:** Workgroup with an enforced result location
113 * ```typescript
114 * const results = yield* AWS.S3.Bucket("AthenaResults", {});
115 * const wg = yield* AWS.Athena.WorkGroup("Analytics", {
116 * outputLocation: results.bucketName.pipe(
117 * Output.map((b) => `s3://${b}/results/`),
118 * ),
119 * enforceWorkGroupConfiguration: true,
120 * });
121 * ```
122 *
123 * **Example:** Workgroup with a bytes-scanned cost guardrail
124 * ```typescript
125 * const wg = yield* AWS.Athena.WorkGroup("Guarded", {
126 * outputLocation: "s3://my-results-bucket/prefix/",
127 * bytesScannedCutoffPerQuery: 10_000_000, // 10 MB per query
128 * publishCloudWatchMetricsEnabled: true,
129 * });
130 * ```
131 *
132 * @resource
133 */
134export const WorkGroup = Resource<WorkGroup>("AWS.Athena.WorkGroup");
135
136const observedTagsOf = (tags: readonly athena.Tag[] | undefined) =>
137 Object.fromEntries(
138 (tags ?? []).flatMap((t) =>
139 t.Key !== undefined && t.Value !== undefined ? [[t.Key, t.Value]] : [],
140 ),
141 );
142
143const outputLocationOf = (wg: athena.WorkGroup | undefined) =>
144 wg?.Configuration?.ResultConfiguration?.OutputLocation;
145
146export const WorkGroupProvider = () =>
147 Provider.effect(
148 WorkGroup,
149 Effect.gen(function* () {
150 const toName = (id: string, props: WorkGroupProps = {}) =>
151 props.workGroupName
152 ? Effect.succeed(props.workGroupName)
153 : createPhysicalName({ id, maxLength: 128 });
154
155 const getOne = (name: string) =>
156 athena.getWorkGroup({ WorkGroup: name }).pipe(
157 Effect.map((res) => res.WorkGroup),
158 Effect.catchTag("WorkGroupNotFound", () => Effect.succeed(undefined)),
159 );
160
161 const fetchTags = (arn: string) =>
162 athena.listTagsForResource.items({ ResourceARN: arn }).pipe(
163 Stream.runCollect,
164 Effect.map((chunk) => observedTagsOf(Array.from(chunk))),
165 Effect.catch(() => Effect.succeed({} as Record<string, string>)),
166 );
167
168 return {
169 stables: ["workGroupName", "workGroupArn"],
170 diff: Effect.fn(function* ({ id, olds, news }) {
171 if (!isResolved(news)) return;
172 if (
173 (yield* toName(id, olds ?? {})) !== (yield* toName(id, news ?? {}))
174 ) {
175 return { action: "replace" } as const;
176 }
177 }),
178 read: Effect.fn(function* ({ id, olds, output }) {
179 const { accountId, region } = yield* AWSEnvironment.current;
180 const name = output?.workGroupName ?? (yield* toName(id, olds ?? {}));
181 const wg = yield* getOne(name);
182 if (!wg) return undefined;
183 const arn =
184 `arn:aws:athena:${region}:${accountId}:workgroup/${name}` as WorkGroupArn;
185 return {
186 workGroupName: name,
187 workGroupArn: arn,
188 state: (wg.State ?? "ENABLED") as "ENABLED" | "DISABLED",
189 outputLocation: outputLocationOf(wg),
190 tags: yield* fetchTags(arn),
191 };
192 }),
193 list: () =>
194 Effect.gen(function* () {
195 const { accountId, region } = yield* AWSEnvironment.current;
196 const pages = yield* athena.listWorkGroups
197 .pages({})
198 .pipe(Stream.runCollect);
199 return Array.from(pages)
200 .flatMap((page) => page.WorkGroups ?? [])
201 .flatMap((wg) =>
202 // `primary` is the account's built-in default workgroup;
203 // DeleteWorkGroup always rejects it, so keep it out of
204 // enumeration for account-wide teardown (nuke).
205 wg.Name && wg.Name !== "primary"
206 ? [
207 {
208 workGroupName: wg.Name,
209 workGroupArn:
210 `arn:aws:athena:${region}:${accountId}:workgroup/${wg.Name}` as WorkGroupArn,
211 state: (wg.State ?? "ENABLED") as
212 | "ENABLED"
213 | "DISABLED",
214 outputLocation: undefined,
215 tags: {} as Record<string, string>,
216 },
217 ]
218 : [],
219 );
220 }),
221 reconcile: Effect.fn(function* ({ id, news = {}, output, session }) {
222 const { accountId, region } = yield* AWSEnvironment.current;
223 const name = output?.workGroupName ?? (yield* toName(id, news));
224 const arn =
225 `arn:aws:athena:${region}:${accountId}:workgroup/${name}` as WorkGroupArn;
226 const internalTags = yield* createInternalTags(id);
227 const desiredTags = { ...internalTags, ...news.tags };
228 const desiredState = news.state ?? "ENABLED";
229 const enforce = news.enforceWorkGroupConfiguration ?? true;
230 const encryption = news.encryptionOption
231 ? {
232 EncryptionOption: news.encryptionOption,
233 KmsKey: news.kmsKey,
234 }
235 : undefined;
236
237 // Observe — cloud state is authoritative.
238 let wg = yield* getOne(name);
239
240 // Ensure — create if missing.
241 if (!wg) {
242 yield* athena.createWorkGroup({
243 Name: name,
244 Description: news.description,
245 Configuration: {
246 ResultConfiguration: news.outputLocation
247 ? {
248 OutputLocation: news.outputLocation,
249 EncryptionConfiguration: encryption,
250 }
251 : undefined,
252 EnforceWorkGroupConfiguration: enforce,
253 PublishCloudWatchMetricsEnabled:
254 news.publishCloudWatchMetricsEnabled,
255 BytesScannedCutoffPerQuery: news.bytesScannedCutoffPerQuery,
256 RequesterPaysEnabled: news.requesterPaysEnabled,
257 EngineVersion: news.engineVersion
258 ? { SelectedEngineVersion: news.engineVersion }
259 : undefined,
260 },
261 Tags: Object.entries(desiredTags).map(([Key, Value]) => ({
262 Key,
263 Value,
264 })),
265 });
266 wg = yield* getOne(name);
267 } else {
268 // Sync — diff observed configuration against desired, apply the
269 // delta via UpdateWorkGroup (which takes a ConfigurationUpdates
270 // shape distinct from CreateWorkGroup's Configuration).
271 const cfg = wg.Configuration;
272 const updates: athena.WorkGroupConfigurationUpdates = {};
273 let dirty = false;
274
275 if ((cfg?.EnforceWorkGroupConfiguration ?? false) !== enforce) {
276 updates.EnforceWorkGroupConfiguration = enforce;
277 dirty = true;
278 }
279 if (
280 news.outputLocation !== undefined &&
281 cfg?.ResultConfiguration?.OutputLocation !== news.outputLocation
282 ) {
283 updates.ResultConfigurationUpdates = {
284 ...updates.ResultConfigurationUpdates,
285 OutputLocation: news.outputLocation,
286 };
287 dirty = true;
288 }
289 if (
290 encryption &&
291 cfg?.ResultConfiguration?.EncryptionConfiguration
292 ?.EncryptionOption !== encryption.EncryptionOption
293 ) {
294 updates.ResultConfigurationUpdates = {
295 ...updates.ResultConfigurationUpdates,
296 EncryptionConfiguration: encryption,
297 };
298 dirty = true;
299 }
300 if (
301 (cfg?.PublishCloudWatchMetricsEnabled ?? false) !==
302 (news.publishCloudWatchMetricsEnabled ?? false)
303 ) {
304 updates.PublishCloudWatchMetricsEnabled =
305 news.publishCloudWatchMetricsEnabled ?? false;
306 dirty = true;
307 }
308 if (
309 (cfg?.RequesterPaysEnabled ?? false) !==
310 (news.requesterPaysEnabled ?? false)
311 ) {
312 updates.RequesterPaysEnabled = news.requesterPaysEnabled ?? false;
313 dirty = true;
314 }
315 if (
316 (cfg?.BytesScannedCutoffPerQuery ?? undefined) !==
317 news.bytesScannedCutoffPerQuery
318 ) {
319 if (news.bytesScannedCutoffPerQuery === undefined) {
320 updates.RemoveBytesScannedCutoffPerQuery = true;
321 } else {
322 updates.BytesScannedCutoffPerQuery =
323 news.bytesScannedCutoffPerQuery;
324 }
325 dirty = true;
326 }
327 if (
328 news.engineVersion !== undefined &&
329 cfg?.EngineVersion?.SelectedEngineVersion !== news.engineVersion
330 ) {
331 updates.EngineVersion = {
332 SelectedEngineVersion: news.engineVersion,
333 };
334 dirty = true;
335 }
336
337 const stateChange =
338 (wg.State ?? "ENABLED") !== desiredState
339 ? desiredState
340 : undefined;
341 const descChange =
342 news.description !== undefined &&
343 wg.Description !== news.description
344 ? news.description
345 : undefined;
346
347 if (dirty || stateChange || descChange) {
348 yield* athena.updateWorkGroup({
349 WorkGroup: name,
350 Description: descChange,
351 State: stateChange,
352 ConfigurationUpdates: dirty ? updates : undefined,
353 });
354 }
355 }
356
357 // Sync tags — diff against OBSERVED cloud tags.
358 const observed = yield* fetchTags(arn);
359 const { upsert, removed } = diffTags(observed, desiredTags);
360 if (upsert.length > 0) {
361 yield* athena.tagResource({
362 ResourceARN: arn,
363 Tags: upsert.map((t) => ({ Key: t.Key, Value: t.Value })),
364 });
365 }
366 if (removed.length > 0) {
367 yield* athena.untagResource({ ResourceARN: arn, TagKeys: removed });
368 }
369
370 const final = (yield* getOne(name)) ?? wg;
371 yield* session.note(arn);
372 return {
373 workGroupName: name,
374 workGroupArn: arn,
375 state: (final?.State ?? desiredState) as "ENABLED" | "DISABLED",
376 outputLocation: outputLocationOf(final) ?? news.outputLocation,
377 tags: desiredTags,
378 };
379 }),
380 delete: Effect.fn(function* ({ output }) {
381 // Idempotent: DeleteWorkGroup succeeds (no error) when the workgroup
382 // is already gone. RecursiveDeleteOption tears down any named
383 // queries / prepared statements still saved in it.
384 yield* athena.deleteWorkGroup({
385 WorkGroup: output.workGroupName,
386 RecursiveDeleteOption: true,
387 });
388 }),
389 };
390 }),
391 );