1import * as Lambda from "@/AWS/Lambda";
2import * as Timestream from "@/AWS/Timestream";
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, "sink-handler.ts");
12
13export class TimestreamSinkFunction extends Lambda.Function<Lambda.Function>()(
14 "TimestreamSinkFunction",
15) {}
16
17export default TimestreamSinkFunction.make(
18 {
19 main,
20 functionUrl: true,
21 // The sink drains fully (multiple sequential WriteRecords calls behind
22 // endpoint discovery) before the handler returns.
23 timeout: Duration.seconds(60),
24 },
25 Effect.gen(function* () {
26 const database = yield* Timestream.Database("SinkMetrics");
27 const table = yield* Timestream.Table("SinkCpu", {
28 databaseName: database.databaseName,
29 retentionProperties: {
30 memoryStoreRetention: "6 hours",
31 magneticStoreRetention: "30 days",
32 },
33 });
34 const DatabaseName = yield* database.databaseName;
35 const TableName = yield* table.tableName;
36
37 // `commonAttributes` exercises the CommonAttributes path: the measure
38 // shape is shared, each record carries only Dimensions/MeasureValue/Time.
39 const sink = yield* Timestream.RecordsSink(table, {
40 commonAttributes: {
41 MeasureName: "cpu",
42 MeasureValueType: "DOUBLE",
43 TimeUnit: "MILLISECONDS",
44 },
45 });
46 const query = yield* Timestream.Query(table);
47
48 const record = (host: string, value: number, time: number) => ({
49 Dimensions: [{ Name: "host", Value: host }],
50 MeasureValue: `${value}`,
51 Time: `${time}`,
52 });
53
54 return {
55 fetch: Effect.gen(function* () {
56 const request = yield* HttpServerRequest;
57 const pathname = new URL(request.originalUrl).pathname;
58
59 if (request.method === "GET" && pathname === "/ready") {
60 return yield* HttpServerResponse.json({ ok: true });
61 }
62
63 // Stream `count` records through the sink. count > 100 proves the
64 // sink splits the stream into <=100-record WriteRecords calls.
65 if (request.method === "POST" && pathname === "/sink") {
66 const body = (yield* request.json) as unknown as {
67 host: string;
68 count: number;
69 };
70 const base = Date.now() - body.count * 1_000;
71 yield* Stream.fromIterable(
72 Array.from({ length: body.count }, (_, i) =>
73 record(body.host, i, base + i * 1_000),
74 ),
75 ).pipe(Stream.run(sink));
76 return yield* HttpServerResponse.json({
77 ok: true,
78 count: body.count,
79 });
80 }
81
82 // Partial failure: one record with a Time outside the 6h memory-store
83 // retention window (magnetic writes are disabled by default) is
84 // permanently rejected via RejectedRecordsException; the sink drops
85 // it and lands the two valid records around it.
86 if (request.method === "POST" && pathname === "/sink-rejects") {
87 const body = (yield* request.json) as unknown as { host: string };
88 const now = Date.now();
89 yield* Stream.fromIterable([
90 record(body.host, 1, now - 2_000),
91 record(body.host, 2, now - 24 * 60 * 60 * 1_000),
92 record(body.host, 3, now - 1_000),
93 ]).pipe(Stream.run(sink));
94 return yield* HttpServerResponse.json({ ok: true });
95 }
96
97 if (request.method === "GET" && pathname === "/count") {
98 const host = new URL(request.originalUrl).searchParams.get("host");
99 const result = yield* query({
100 QueryString: `SELECT COUNT(*) AS c FROM "${DatabaseName}"."${TableName}" WHERE host = '${host}'`,
101 });
102 return yield* HttpServerResponse.json({ rows: result.Rows });
103 }
104
105 return yield* HttpServerResponse.json(
106 { error: "Not found", method: request.method, pathname },
107 { status: 404 },
108 );
109 }).pipe(Effect.orDie),
110 };
111 }).pipe(
112 Effect.provide(
113 Layer.provideMerge(
114 Layer.mergeAll(Timestream.RecordsSinkHttp, Timestream.QueryHttp),
115 Layer.mergeAll(Timestream.WriteRecordsHttp),
116 ),
117 ),
118 ),
119);