1import * as Lambda from "@/AWS/Lambda";
2import * as SageMaker from "@/AWS/SageMaker";
3import * as Duration from "effect/Duration";
4import * as Effect from "effect/Effect";
5import * as Layer from "effect/Layer";
6import * as Stream from "effect/Stream";
7import { HttpServerRequest } from "effect/http/HttpServerRequest";
8import * as HttpServerResponse from "effect/http/HttpServerResponse";
9import path from "pathe";
10
11const main = path.resolve(import.meta.dirname, "handler.ts");
12
13export class SageMakerTestFunction extends Lambda.Function<Lambda.Function>()(
14 "SageMakerTestFunction",
15) {}
16
17export default SageMakerTestFunction.make(
18 {
19 main,
20 functionUrl: true,
21 timeout: Duration.seconds(30),
22 },
23 Effect.gen(function* () {
24 const featureGroup = yield* SageMaker.FeatureGroup("BindingsFeatures", {
25 recordIdentifierFeatureName: "user_id",
26 eventTimeFeatureName: "event_time",
27 featureDefinitions: [
28 { FeatureName: "user_id", FeatureType: "String" },
29 { FeatureName: "event_time", FeatureType: "String" },
30 { FeatureName: "clicks", FeatureType: "Integral" },
31 ],
32 onlineStoreConfig: { EnableOnlineStore: true },
33 });
34
35 // Event source: subscribe the host to SageMaker feature-group state
36 // changes. The deploy proves the EventBridge rule + invoke permission
37 // wiring.
38 yield* SageMaker.consumeSageMakerEvents(
39 { kinds: ["feature-group"] },
40 (events) =>
41 Stream.runForEach(events, (event) =>
42 Effect.log(
43 `feature group ${event.detail.FeatureGroupName} -> ${event.detail.FeatureGroupStatus}`,
44 ),
45 ),
46 );
47
48 const putRecord = yield* SageMaker.PutRecord(featureGroup);
49 const getRecord = yield* SageMaker.GetRecord(featureGroup);
50 const deleteRecord = yield* SageMaker.DeleteRecord(featureGroup);
51 const batchGetRecord = yield* SageMaker.BatchGetRecord(featureGroup);
52 const batchWriteRecord = yield* SageMaker.BatchWriteRecord(featureGroup);
53 const listRecords = yield* SageMaker.ListRecords(featureGroup);
54
55 const bound = {
56 putRecord,
57 getRecord,
58 deleteRecord,
59 batchGetRecord,
60 batchWriteRecord,
61 listRecords,
62 };
63
64 const record = (userId: string, clicks: number) => [
65 { FeatureName: "user_id", ValueAsString: userId },
66 { FeatureName: "event_time", ValueAsString: new Date().toISOString() },
67 { FeatureName: "clicks", ValueAsString: String(clicks) },
68 ];
69
70 return {
71 fetch: Effect.gen(function* () {
72 const request = yield* HttpServerRequest;
73 const url = new URL(request.originalUrl);
74 const pathname = url.pathname;
75
76 if (request.method === "GET" && pathname === "/bindings") {
77 return yield* HttpServerResponse.json({ bound: Object.keys(bound) });
78 }
79
80 if (request.method === "POST" && pathname === "/put-record") {
81 const body = (yield* request.json) as unknown as {
82 userId: string;
83 clicks: number;
84 };
85 yield* putRecord({ Record: record(body.userId, body.clicks) });
86 return yield* HttpServerResponse.json({ success: true });
87 }
88
89 if (request.method === "GET" && pathname === "/get-record") {
90 const userId = url.searchParams.get("userId") ?? "";
91 const result = yield* getRecord({
92 RecordIdentifierValueAsString: userId,
93 });
94 return yield* HttpServerResponse.json({
95 record: result.Record ?? [],
96 });
97 }
98
99 if (request.method === "POST" && pathname === "/delete-record") {
100 const body = (yield* request.json) as unknown as { userId: string };
101 yield* deleteRecord({
102 RecordIdentifierValueAsString: body.userId,
103 EventTime: new Date().toISOString(),
104 });
105 return yield* HttpServerResponse.json({ success: true });
106 }
107
108 if (request.method === "POST" && pathname === "/batch-write-record") {
109 const body = (yield* request.json) as unknown as {
110 records: { userId: string; clicks: number }[];
111 };
112 const result = yield* batchWriteRecord({
113 Entries: body.records.map((r) => ({
114 Record: record(r.userId, r.clicks),
115 })),
116 });
117 return yield* HttpServerResponse.json({
118 errors: result.Errors ?? [],
119 unprocessed: result.UnprocessedEntries ?? [],
120 });
121 }
122
123 if (request.method === "POST" && pathname === "/batch-get-record") {
124 const body = (yield* request.json) as unknown as {
125 userIds: string[];
126 };
127 const result = yield* batchGetRecord({
128 RecordIdentifiersValueAsString: body.userIds,
129 });
130 return yield* HttpServerResponse.json({
131 records: (result.Records ?? []).map((r) => ({
132 userId: r.RecordIdentifierValueAsString,
133 record: r.Record ?? [],
134 })),
135 errors: result.Errors ?? [],
136 });
137 }
138
139 if (request.method === "GET" && pathname === "/list-records") {
140 const result = yield* listRecords({ MaxResults: 100 });
141 return yield* HttpServerResponse.json({
142 identifiers: result.RecordIdentifiers ?? [],
143 });
144 }
145
146 if (request.method === "GET" && pathname === "/health") {
147 return yield* HttpServerResponse.json({ ok: true });
148 }
149
150 return yield* HttpServerResponse.json(
151 { error: "Not found", pathname },
152 { status: 404 },
153 );
154 }).pipe(Effect.orDie),
155 };
156 }).pipe(
157 Effect.provide(
158 Layer.mergeAll(
159 Lambda.EventSource,
160 SageMaker.PutRecordHttp,
161 SageMaker.GetRecordHttp,
162 SageMaker.DeleteRecordHttp,
163 SageMaker.BatchGetRecordHttp,
164 SageMaker.BatchWriteRecordHttp,
165 SageMaker.ListRecordsHttp,
166 ),
167 ),
168 ),
169);