1import * as AWS from "@/AWS";
2import * as Output from "@/Output";
3import * as EC2 from "@distilled.cloud/aws/ec2";
4import * as Effect from "effect/Effect";
5import * as Layer from "effect/Layer";
6import { HttpServerRequest } from "effect/http/HttpServerRequest";
7import * as HttpServerResponse from "effect/http/HttpServerResponse";
8import { getDefaultVpc } from "../DefaultVpc.ts";
9
10export class MWAATestFunction extends AWS.Lambda.Function<AWS.Lambda.Function>()(
11 "MWAATestFunction",
12) {}
13
14/**
15 * Resolve two default-for-AZ subnets (distinct AZs) plus the default security
16 * group from the account's default VPC. Runtime-guarded: the Lambda runtime
17 * re-executes this props effect on cold start with no ec2:Describe*
18 * permission, so return empty values there (network config is deploy-time
19 * only). NOTE: a REAL MWAA environment requires PRIVATE subnets with outbound
20 * egress — an entitled/patient runner supplies that out of band, which is why
21 * the fixture is only deployed behind the AWS_TEST_SLOW gate.
22 */
23const resolveNetwork = Effect.gen(function* () {
24 if (globalThis.__ALCHEMY_RUNTIME__) {
25 return { subnetIds: [] as string[], securityGroupIds: [] as string[] };
26 }
27 const vpc = yield* getDefaultVpc;
28 const subnets = yield* EC2.describeSubnets({
29 Filters: [
30 { Name: "vpc-id", Values: [vpc.vpcId] },
31 { Name: "default-for-az", Values: ["true"] },
32 ],
33 });
34 const byAz = new Map<string, string>();
35 for (const s of subnets.Subnets ?? []) {
36 if (s.AvailabilityZone && s.SubnetId && !byAz.has(s.AvailabilityZone)) {
37 byAz.set(s.AvailabilityZone, s.SubnetId);
38 }
39 }
40 const sgs = yield* EC2.describeSecurityGroups({
41 Filters: [
42 { Name: "vpc-id", Values: [vpc.vpcId] },
43 { Name: "group-name", Values: ["default"] },
44 ],
45 });
46 const securityGroupId = sgs.SecurityGroups?.[0]?.GroupId;
47 return {
48 subnetIds: [...byAz.values()].slice(0, 2),
49 securityGroupIds: securityGroupId ? [securityGroupId] : [],
50 };
51}).pipe(
52 // Deploy-time lookup only; a failure here is a fixture defect, not a typed
53 // error the Function impl contract can carry.
54 Effect.orDie,
55);
56
57/**
58 * Every route answers `{ …fields }` on success or `{ errorTag }` when the
59 * operation fails with a TYPED error — the test asserts on concrete fields
60 * (or a typed tag), which proves the binding wiring, the environment-name
61 * injection, and the IAM grants (environment-ARN-scoped for
62 * CliToken/GetEnvironment, Airflow-role-ARN-scoped for WebLoginToken and
63 * InvokeRestApi). An untyped error crashes into a 500.
64 */
65const errorTagged = <A, E extends { _tag: string }, R>(
66 effect: Effect.Effect<A, E, R>,
67): Effect.Effect<A | { errorTag: string }, never, R> =>
68 effect.pipe(
69 Effect.map((a): A | { errorTag: string } => a),
70 Effect.catch((e) => Effect.succeed({ errorTag: e._tag })),
71 );
72
73/**
74 * Environment-scoped binding fixture: deploys a real MWAA environment
75 * (~20-30 minutes to provision, billed hourly while it exists — gated behind
76 * AWS_TEST_SLOW) plus its supporting cast (versioned DAGs bucket, execution
77 * role) and a Lambda bound to all four MWAA bindings.
78 */
79export default MWAATestFunction.make(
80 { main: import.meta.url, functionUrl: true },
81 Effect.gen(function* () {
82 const { subnetIds, securityGroupIds } = yield* resolveNetwork;
83
84 const bucket = yield* AWS.S3.Bucket("MWAABindingsDags", {
85 versioning: "Enabled",
86 forceDestroy: true,
87 });
88
89 // The execution role Amazon MWAA and Airflow assume. Fixture-only wide
90 // grants; the interesting IAM (the scoped airflow:* statements) is what
91 // the bindings attach to the Lambda itself.
92 const role = yield* AWS.IAM.Role("MWAABindingsExecutionRole", {
93 assumeRolePolicyDocument: {
94 Version: "2012-10-17",
95 Statement: [
96 {
97 Effect: "Allow",
98 Principal: {
99 Service: ["airflow.amazonaws.com", "airflow-env.amazonaws.com"],
100 },
101 Action: ["sts:AssumeRole"],
102 },
103 ],
104 },
105 inlinePolicies: {
106 "mwaa-execution": {
107 Version: "2012-10-17",
108 Statement: [
109 {
110 Effect: "Allow",
111 Action: ["airflow:PublishMetrics"],
112 Resource: ["*"],
113 },
114 {
115 Effect: "Allow",
116 Action: ["s3:GetObject*", "s3:GetBucket*", "s3:List*"],
117 Resource: [
118 bucket.bucketArn,
119 Output.interpolate`${bucket.bucketArn}/*`,
120 ],
121 },
122 {
123 Effect: "Allow",
124 Action: [
125 "logs:CreateLogStream",
126 "logs:CreateLogGroup",
127 "logs:PutLogEvents",
128 "logs:GetLogEvents",
129 "logs:GetLogRecord",
130 "logs:GetLogGroupFields",
131 "logs:GetQueryResults",
132 "logs:DescribeLogGroups",
133 ],
134 Resource: ["arn:aws:logs:*:*:log-group:airflow-*"],
135 },
136 {
137 Effect: "Allow",
138 Action: ["cloudwatch:PutMetricData"],
139 Resource: ["*"],
140 },
141 {
142 Effect: "Allow",
143 Action: [
144 "sqs:ChangeMessageVisibility",
145 "sqs:DeleteMessage",
146 "sqs:GetQueueAttributes",
147 "sqs:GetQueueUrl",
148 "sqs:ReceiveMessage",
149 "sqs:SendMessage",
150 ],
151 Resource: ["arn:aws:sqs:*:*:airflow-celery-*"],
152 },
153 ],
154 },
155 },
156 });
157
158 const environment = yield* AWS.MWAA.Environment("BindingsAirflow", {
159 executionRoleArn: role.roleArn,
160 sourceBucketArn: bucket.bucketArn,
161 dagS3Path: "dags",
162 subnetIds,
163 securityGroupIds:
164 securityGroupIds.length > 0 ? securityGroupIds : undefined,
165 environmentClass: "mw1.small",
166 maxWorkers: 2,
167 minWorkers: 1,
168 webserverAccessMode: "PUBLIC_ONLY",
169 tags: { fixture: "mwaa-bindings" },
170 });
171
172 const getEnvironment = yield* AWS.MWAA.GetEnvironment(environment);
173 const createCliToken = yield* AWS.MWAA.CreateCliToken(environment);
174 const createWebLoginToken =
175 yield* AWS.MWAA.CreateWebLoginToken(environment);
176 const invokeRestApi = yield* AWS.MWAA.InvokeRestApi(environment);
177
178 const bound = {
179 getEnvironment,
180 createCliToken,
181 createWebLoginToken,
182 invokeRestApi,
183 };
184
185 return {
186 fetch: Effect.gen(function* () {
187 const request = yield* HttpServerRequest;
188 const url = new URL(request.originalUrl ?? request.url, "http://x");
189 const pathname = url.pathname;
190
191 if (pathname === "/bindings") {
192 return yield* HttpServerResponse.json({ bound: Object.keys(bound) });
193 }
194
195 if (pathname === "/environment") {
196 const result = yield* errorTagged(getEnvironment());
197 return yield* HttpServerResponse.json(
198 "errorTag" in result
199 ? result
200 : {
201 status: result.Environment?.Status,
202 webserverUrl: result.Environment?.WebserverUrl,
203 },
204 );
205 }
206
207 if (pathname === "/cli-token") {
208 const result = yield* errorTagged(createCliToken());
209 return yield* HttpServerResponse.json(
210 "errorTag" in result
211 ? result
212 : {
213 // Never echo the Redacted token itself.
214 hasToken: result.CliToken !== undefined,
215 hostname: result.WebServerHostname,
216 },
217 );
218 }
219
220 if (pathname === "/web-login-token") {
221 const result = yield* errorTagged(createWebLoginToken());
222 return yield* HttpServerResponse.json(
223 "errorTag" in result
224 ? result
225 : {
226 hasToken: result.WebToken !== undefined,
227 hostname: result.WebServerHostname,
228 airflowIdentity: result.AirflowIdentity,
229 },
230 );
231 }
232
233 if (pathname === "/dags") {
234 const result = yield* errorTagged(
235 invokeRestApi({ Method: "GET", Path: "/dags" }),
236 );
237 return yield* HttpServerResponse.json(
238 "errorTag" in result
239 ? result
240 : {
241 statusCode: result.RestApiStatusCode,
242 totalEntries: (
243 result.RestApiResponse as
244 | { total_entries?: number }
245 | undefined
246 )?.total_entries,
247 },
248 );
249 }
250
251 return yield* HttpServerResponse.json(
252 { error: "Not found", pathname },
253 { status: 404 },
254 );
255 }).pipe(Effect.orDie),
256 };
257 }).pipe(
258 Effect.provide(
259 Layer.mergeAll(
260 AWS.MWAA.GetEnvironmentHttp,
261 AWS.MWAA.CreateCliTokenHttp,
262 AWS.MWAA.CreateWebLoginTokenHttp,
263 AWS.MWAA.InvokeRestApiHttp,
264 ),
265 ),
266 ),
267);