1import * as Axiom from "@distilled.cloud/axiom";
2import { Credentials } from "@distilled.cloud/axiom/Credentials";
3import * as Effect from "effect/Effect";
4import { Unowned } from "../AdoptPolicy.ts";
5import { isResolved } from "../Diff.ts";
6import * as Provider from "../Provider.ts";
7import { Resource } from "../Resource.ts";
8import { Stack } from "../Stack.ts";
9import { Stage } from "../Stage.ts";
10import type { Providers } from "./Providers.ts";
11
12export type DatasetKind =
13 | "otel:metrics:v1"
14 | "otel:traces:v1"
15 | "otel:logs:v1"
16 | "axiom:events:v1"
17 | (string & {});
18
19export type DatasetProps = {
20 /**
21 * Dataset name. Used as the dataset's stable identifier in Axiom — changing
22 * this triggers a replacement.
23 */
24 name: string;
25 /** Free-form description shown in the Axiom UI. */
26 description?: string;
27 /**
28 * Edge deployment where Axiom stores and processes the dataset's data.
29 * Defaults to the organization's default edge deployment.
30 *
31 * Cannot be changed after creation — triggers a replacement.
32 */
33 edgeDeployment?: string;
34 /**
35 * Dataset kind. Defaults to `axiom:events:v1`.
36 *
37 * For OTEL pipelines, choose:
38 * - `otel:traces:v1` for traces
39 * - `otel:logs:v1` for logs
40 * - `otel:metrics:v1` for metrics
41 *
42 * Cannot be changed after creation — triggers a replacement.
43 */
44 kind?: DatasetKind;
45 /** Retention in days. Plan-dependent — see Axiom docs. */
46 retentionDays?: number;
47 /** Whether to enforce the configured retention period. */
48 useRetentionPeriod?: boolean;
49};
50
51export type Dataset = Resource<
52 "Axiom.Dataset",
53 DatasetProps,
54 {
55 id: string;
56 name: string;
57 kind: DatasetKind;
58 description: string;
59 created: string;
60 /** Edge deployment where the dataset's data is stored and processed. */
61 edgeDeployment: string;
62 /** Base URL for ingesting and querying data in the dataset's edge deployment. */
63 edgeDeploymentUrl: string;
64 /** Base URL for Axiom's centralized management API. */
65 apiBaseUrl: string;
66 /** Root OTLP endpoint. Most exporters auto-append the signal path. */
67 otelEndpoint: string;
68 /** OTLP/HTTP traces endpoint. */
69 otelTracesEndpoint: string;
70 /** OTLP/HTTP logs endpoint. */
71 otelLogsEndpoint: string;
72 /** OTLP/HTTP metrics endpoint. */
73 otelMetricsEndpoint: string;
74 /**
75 * Headers required for OTLP shipping aside from the bearer token.
76 * Add `Authorization: Bearer <AXIOM_TOKEN>` separately at runtime so the
77 * secret is never persisted in resource state.
78 */
79 otelHeaders: Record<string, string>;
80 },
81 never,
82 Providers
83>;
84
85/**
86 * An Axiom dataset — the top-level container that stores events, logs,
87 * traces, or metrics. Pick a `kind` up-front: it determines schema and how
88 * the data is shown in the UI, and **cannot be changed** after creation
89 * (changing it triggers a replacement, which deletes the data).
90 *
91 * Datasets expose their edge deployment and Axiom's OTLP/HTTP endpoints
92 * (`otelTracesEndpoint`, `otelLogsEndpoint`, `otelMetricsEndpoint`) as output
93 * attributes so you can inject them into a Worker / Lambda's env vars for
94 * OpenTelemetry shipping. The bearer token is **not** stored in resource state
95 * — supply `Authorization: Bearer <AXIOM_TOKEN>` separately at runtime.
96 * @see https://axiom.co/docs/reference/datasets
97 * @see https://axiom.co/docs/send-data/opentelemetry — OTLP endpoint reference
98 *
99 * ### Creating a Dataset
100 * **Example:** Logs dataset with 30-day retention
101 * ```typescript
102 * const logs = yield* Axiom.Dataset("app-logs", {
103 * name: "my-app-logs",
104 * kind: "otel:logs:v1",
105 * description: "Application logs from prod workers",
106 * retentionDays: 30,
107 * useRetentionPeriod: true,
108 * });
109 * ```
110 *
111 * **Example:** Separate datasets per OTEL signal
112 * ```typescript
113 * const traces = yield* Axiom.Dataset("traces", { name: "app-traces", kind: "otel:traces:v1" });
114 * const logs = yield* Axiom.Dataset("logs", { name: "app-logs", kind: "otel:logs:v1" });
115 * const metrics = yield* Axiom.Dataset("metrics", { name: "app-metrics", kind: "otel:metrics:v1" });
116 * ```
117 *
118 * ### Shipping OTEL data
119 * **Example:** Wire OTEL env vars into a Cloudflare Worker
120 * ```typescript
121 * yield* Cloudflare.Worker("api", {
122 * vars: {
123 * OTEL_EXPORTER_OTLP_TRACES_ENDPOINT: traces.otelTracesEndpoint,
124 * OTEL_EXPORTER_OTLP_LOGS_ENDPOINT: logs.otelLogsEndpoint,
125 * // Bearer token must come from a secret store, not the dataset state.
126 * OTEL_EXPORTER_OTLP_HEADERS:
127 * `Authorization=Bearer ${env.AXIOM_TOKEN},X-Axiom-Dataset=${traces.name}`,
128 * },
129 * });
130 * ```
131 *
132 * @resource
133 * @product Dataset
134 */
135export const Dataset = Resource<Dataset>("Axiom.Dataset");
136
137/**
138 * Axiom has no tags/labels API. The only writable field we can use to mark
139 * ownership is `description`. We append a deterministic marker on create so
140 * that on a re-apply (e.g. state was wiped) we can safely identify a dataset
141 * that **we** previously created and adopt it idempotently — without
142 * accidentally hijacking a dataset created by someone else with the same name.
143 */
144const MARKER_RE = /\s*\[alchemy:stack=([^;]+);stage=([^;]+);id=([^\]]+)\]\s*$/;
145
146const buildMarker = (stack: string, stage: string, id: string) =>
147 `[alchemy:stack=${stack};stage=${stage};id=${id}]`;
148
149const augmentDescription = (
150 description: string | undefined,
151 marker: string,
152) => {
153 const base = stripMarker(description ?? "");
154 return base.length > 0 ? `${base}\n${marker}` : marker;
155};
156
157const stripMarker = (description: string): string =>
158 description.replace(MARKER_RE, "").trimEnd();
159
160const parseMarker = (
161 description: string | undefined,
162): { stack: string; stage: string; id: string } | undefined => {
163 if (!description) return undefined;
164 const m = description.match(MARKER_RE);
165 if (!m) return undefined;
166 return { stack: m[1], stage: m[2], id: m[3] };
167};
168
169export const DatasetProvider = () =>
170 Provider.effect(
171 Dataset,
172 Effect.gen(function* () {
173 const { apiBaseUrl } = yield* yield* Credentials;
174 const create = yield* Axiom.createDataset;
175 const update = yield* Axiom.updateDataset;
176 const get = yield* Axiom.getDataset;
177 const listDatasets = yield* Axiom.getDatasets;
178 const del = yield* Axiom.deleteDataset;
179
180 const toAttrs = Effect.fn(function* (dataset: Axiom.Dataset) {
181 if (!dataset.edgeDeployment || !dataset.edgeDeploymentUrl) {
182 return yield* Effect.fail(
183 new Error(
184 `Axiom dataset "${dataset.name}" is missing its edge deployment metadata`,
185 ),
186 );
187 }
188 const apiRoot = apiBaseUrl.replace(/\/$/, "");
189 const otelRoot = dataset.edgeDeploymentUrl.replace(/\/$/, "");
190 return {
191 id: dataset.id,
192 name: dataset.name,
193 kind: dataset.kind,
194 description: stripMarker(dataset.description),
195 created: dataset.created,
196 edgeDeployment: dataset.edgeDeployment,
197 edgeDeploymentUrl: dataset.edgeDeploymentUrl,
198 apiBaseUrl: apiRoot,
199 otelEndpoint: otelRoot,
200 otelTracesEndpoint: `${otelRoot}/v1/traces`,
201 otelLogsEndpoint: `${otelRoot}/v1/logs`,
202 otelMetricsEndpoint: `${otelRoot}/v1/metrics`,
203 otelHeaders: {
204 "X-Axiom-Dataset": dataset.name,
205 } as Record<string, string>,
206 };
207 });
208
209 return {
210 stables: ["id", "name", "kind", "edgeDeployment", "edgeDeploymentUrl"],
211 // Enumerate every dataset in the org. Axiom exposes a single
212 // account-wide `GET /v2/datasets` collection op (no pagination), so we
213 // fetch it once and hydrate each row into the exact `read`/`toAttrs`
214 // Attributes shape — directly usable by `delete` with no follow-up get.
215 list: () =>
216 Effect.gen(function* () {
217 const datasets = yield* listDatasets({});
218 return yield* Effect.forEach(datasets, toAttrs);
219 }),
220 diff: Effect.fn(function* ({ olds, news, output }) {
221 if (!isResolved(news)) return undefined;
222 if (output && news.name !== output.name) {
223 return { action: "replace" } as const;
224 }
225 if (news.kind && output && news.kind !== output.kind) {
226 return { action: "replace" } as const;
227 }
228 if (
229 news.edgeDeployment &&
230 output &&
231 news.edgeDeployment !== output.edgeDeployment
232 ) {
233 return { action: "replace" } as const;
234 }
235 if (
236 news.description !== olds?.description ||
237 news.retentionDays !== olds?.retentionDays ||
238 news.useRetentionPeriod !== olds?.useRetentionPeriod
239 ) {
240 return { action: "update" } as const;
241 }
242 return undefined;
243 }),
244 reconcile: Effect.fn(function* ({ id, news, output }) {
245 const stack = yield* Stack;
246 const stage = yield* Stage;
247 const marker = buildMarker(stack.name, stage, id);
248
249 // Observe — Axiom's dataset path identifier is `name`, which is
250 // also stable input. Probe for live state by name (preferring the
251 // cached `output.id` when present). `read` upstream has already
252 // surfaced foreign datasets as `Unowned`, so by the time we land
253 // here mutation is safe.
254 const datasetId = output?.id ?? news.name;
255 const observed = yield* get({ dataset_id: datasetId }).pipe(
256 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
257 );
258
259 // Ensure — POST creates the dataset. Tolerate Conflict/
260 // UnprocessableEntity as a race with a peer reconciler (or with
261 // upstream read↔create), falling through to the sync path.
262 let current = observed;
263 if (current === undefined) {
264 current = yield* (
265 create({
266 name: news.name,
267 description: augmentDescription(news.description, marker),
268 edgeDeployment: news.edgeDeployment,
269 kind: news.kind,
270 retentionDays: news.retentionDays,
271 useRetentionPeriod: news.useRetentionPeriod,
272 }) as Effect.Effect<
273 Axiom.Dataset,
274 { readonly _tag: string },
275 never
276 >
277 ).pipe(
278 Effect.catchIf(
279 (
280 e,
281 ): e is { readonly _tag: "Conflict" | "UnprocessableEntity" } =>
282 e._tag === "Conflict" || e._tag === "UnprocessableEntity",
283 () =>
284 update({
285 dataset_id: news.name,
286 description: augmentDescription(news.description, marker),
287 retentionDays: news.retentionDays,
288 useRetentionPeriod: news.useRetentionPeriod,
289 }),
290 ),
291 );
292 return yield* toAttrs(current);
293 }
294
295 // Sync — the dataset exists. Apply mutable aspects (description,
296 // retentionDays, useRetentionPeriod) via PATCH. `kind` and `name`
297 // are stable and replacement-only via diff above.
298 const desiredDescription = augmentDescription(
299 news.description,
300 marker,
301 );
302 const needsSync =
303 current.description !== desiredDescription ||
304 current.retentionDays !== news.retentionDays ||
305 current.useRetentionPeriod !== news.useRetentionPeriod;
306 if (!needsSync) {
307 return yield* toAttrs(current);
308 }
309 const updated = yield* update({
310 dataset_id: current.id,
311 description: desiredDescription,
312 retentionDays: news.retentionDays,
313 useRetentionPeriod: news.useRetentionPeriod,
314 });
315 return yield* toAttrs(updated);
316 }),
317 delete: Effect.fn(function* ({ output }) {
318 yield* del({ dataset_id: output.id }).pipe(
319 Effect.catchTag("NotFound", () => Effect.void),
320 );
321 }),
322 read: Effect.fn(function* ({ id, olds, output }) {
323 const stack = yield* Stack;
324 const stage = yield* Stage;
325 const datasetId = output?.id ?? olds?.name;
326 if (!datasetId) return undefined;
327 const existing = yield* get({ dataset_id: datasetId }).pipe(
328 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
329 );
330 if (!existing) return undefined;
331 const ownership = parseMarker(existing.description);
332 const isOurs =
333 ownership !== undefined &&
334 ownership.stack === stack.name &&
335 ownership.stage === stage &&
336 ownership.id === id;
337 const attrs = yield* toAttrs(existing);
338 return isOurs ? attrs : Unowned(attrs);
339 }),
340 };
341 }),
342 );