1import * as aiplatform from "@distilled.cloud/gcp/aiplatform_v1";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
4import * as Schedule from "effect/Schedule";
5import * as Stream from "effect/Stream";
6import { Unowned } from "../../AdoptPolicy.ts";
7import { isResolved } from "../../Diff.ts";
8import * as Provider from "../../Provider.ts";
9import { Resource } from "../../Resource.ts";
10import { tagRecord } from "../../Tags.ts";
11import { GcpEnvironment } from "../Environment.ts";
12import { hasAlchemyLabels } from "../Labels.ts";
13import type { Providers } from "../Providers.ts";
14import {
15 createInternalLabels,
16 hasAlchemyPrefix,
17 labelsDiffer,
18 listLocations,
19 locationOf,
20 lastSegment,
21 projectOf,
22 stableJson,
23 toDisplayName,
24 toLabels,
25 toResourceId,
26 userLabels,
27} from "./names.ts";
28import { waitForOperation } from "./operations.ts";
29
30const DEFAULT_SCHEMA_TITLE = "system.Experiment";
31
32export type MetadataStoresContextProps = {
33 /**
34 * Parent MetadataStore resource name
35 * `projects/{project}/locations/{location}/metadataStores/{metadatastore}`.
36 * Immutable — changing it replaces the context.
37 */
38 metadataStore: string;
39 /**
40 * Context id. If omitted, a unique id is generated. Must be 4-128
41 * characters of `a-z` and `-`. Immutable.
42 */
43 contextId?: string;
44 /**
45 * Display name (max 128 Unicode characters).
46 */
47 displayName?: string;
48 /**
49 * Human-readable description.
50 */
51 description?: string;
52 /**
53 * User labels. Alchemy ownership labels are merged in automatically.
54 */
55 labels?: Record<string, string>;
56 /**
57 * Schema title registered in the MetadataStore. Vertex AI rejects
58 * entries without one.
59 * @default "system.Experiment"
60 */
61 schemaTitle?: string;
62 /**
63 * Schema version registered in the MetadataStore.
64 */
65 schemaVersion?: string;
66 /**
67 * Arbitrary metadata properties.
68 */
69 metadata?: aiplatform.DocumentMap;
70};
71
72export type MetadataStoresContext = Resource<
73 "GCP.AIPlatform.MetadataStoresContext",
74 MetadataStoresContextProps,
75 {
76 /** Full resource name. */
77 name: string;
78 /** Context id (last path segment). */
79 contextId: string;
80 /** Parent MetadataStore resource name. */
81 metadataStore: string;
82 /** Project id. */
83 project: string;
84 /** Location id. */
85 location: string;
86 /** Display name. */
87 displayName: string | undefined;
88 /** User-provided description. */
89 description: string | undefined;
90 /** User labels (Alchemy ownership labels stripped). */
91 labels: Record<string, string>;
92 /** Schema title. */
93 schemaTitle: string | undefined;
94 /** Schema version. */
95 schemaVersion: string | undefined;
96 /** Parent context resource names. */
97 parentContexts: string[];
98 /** RFC3339 creation timestamp. */
99 createTime: string | undefined;
100 /** RFC3339 last-update timestamp. */
101 updateTime: string | undefined;
102 },
103 never,
104 Providers
105>;
106
107/**
108 * A Vertex ML Metadata Context associated with a MetadataStore.
109 *
110 * Changing `metadataStore` or `contextId` replaces the context. Display
111 * name, description, labels, schema, and metadata update in place.
112 *
113 * ### Creating a Context
114 * **Example:** Experiment context
115 * ```typescript
116 * const context = yield* GCP.AIPlatform.MetadataStoresContext("Experiment", {
117 * metadataStore: store.name,
118 * displayName: "training-run",
119 * labels: { env: "prod" },
120 * });
121 * ```
122 *
123 * @resource
124 * @category AIPlatform
125 */
126export const MetadataStoresContext = Resource<MetadataStoresContext>(
127 "GCP.AIPlatform.MetadataStoresContext",
128);
129
130export class MetadataStoresContextNotResolved extends Data.TaggedError(
131 "GCP.AIPlatform.MetadataStoresContextNotResolved",
132)<{
133 name: string;
134}> {}
135
136export class MetadataStoresContextStillExists extends Data.TaggedError(
137 "GCP.AIPlatform.MetadataStoresContextStillExists",
138)<{
139 name: string;
140}> {}
141
142const parentOf = (name: string) => {
143 const at = name.lastIndexOf("/contexts/");
144 return at >= 0 ? name.slice(0, at) : name;
145};
146
147const resourceName = (metadataStore: string, contextId: string) =>
148 `${metadataStore}/contexts/${contextId}`;
149
150const toAttrs = (
151 context: aiplatform.GoogleCloudAiplatformV1Context,
152 project: string,
153) => {
154 const name = context.name ?? "";
155 return {
156 name,
157 contextId: lastSegment(name),
158 metadataStore: parentOf(name),
159 project: projectOf(name, project),
160 location: locationOf(name),
161 displayName: context.displayName,
162 description: context.description,
163 labels: userLabels(context.labels),
164 schemaTitle: context.schemaTitle,
165 schemaVersion: context.schemaVersion,
166 parentContexts: context.parentContexts ?? [],
167 createTime: context.createTime,
168 updateTime: context.updateTime,
169 };
170};
171
172const getByName = (name: string) =>
173 aiplatform
174 .getProjectsLocationsMetadataStoresContexts({ name })
175 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
176
177const listContexts = (parent: string) =>
178 aiplatform.listProjectsLocationsMetadataStoresContexts
179 .pages({ parent, pageSize: 100 })
180 .pipe(
181 Stream.runCollect,
182 Effect.map((pages) =>
183 Array.from(pages).flatMap((page) => page.contexts ?? []),
184 ),
185 Effect.catchTag("NotFound", () =>
186 Effect.succeed([] as aiplatform.GoogleCloudAiplatformV1Context[]),
187 ),
188 );
189
190const listStores = (parent: string) =>
191 aiplatform.listProjectsLocationsMetadataStores
192 .pages({ parent, pageSize: 100 })
193 .pipe(
194 Stream.runCollect,
195 Effect.map((pages) =>
196 Array.from(pages).flatMap((page) => page.metadataStores ?? []),
197 ),
198 Effect.catchTag("NotFound", () =>
199 Effect.succeed([] as aiplatform.GoogleCloudAiplatformV1MetadataStore[]),
200 ),
201 );
202
203const waitUntilGone = (name: string) =>
204 getByName(name).pipe(
205 Effect.filterOrFail(
206 (context) => context === undefined,
207 () => new MetadataStoresContextStillExists({ name }),
208 ),
209 Effect.asVoid,
210 Effect.retry({
211 while: (error) =>
212 error._tag === "GCP.AIPlatform.MetadataStoresContextStillExists",
213 times: 10,
214 schedule: Schedule.spaced("1 second"),
215 }),
216 );
217
218export const MetadataStoresContextProvider = () =>
219 Provider.succeed(MetadataStoresContext, {
220 stables: [
221 "name",
222 "contextId",
223 "metadataStore",
224 "project",
225 "location",
226 "createTime",
227 ],
228
229 diff: Effect.fn(function* ({ news, olds, output }) {
230 if (!isResolved(news)) return undefined;
231 const previousParent = olds?.metadataStore ?? output?.metadataStore;
232 const nextParent = news.metadataStore ?? previousParent;
233 const previousId = olds?.contextId ?? output?.contextId;
234 const nextId = news.contextId ?? previousId;
235 if (
236 (previousParent !== undefined &&
237 nextParent !== undefined &&
238 nextParent !== previousParent) ||
239 (previousId !== undefined &&
240 nextId !== undefined &&
241 nextId !== previousId)
242 ) {
243 return {
244 action: "replace" as const,
245 deleteFirst:
246 previousParent === nextParent &&
247 previousId !== undefined &&
248 nextId === previousId,
249 };
250 }
251 return undefined;
252 }),
253
254 read: Effect.fn(function* ({ id, olds, output }) {
255 const env = yield* GcpEnvironment.current;
256 const contextId = yield* toResourceId(
257 id,
258 olds?.contextId,
259 output?.contextId,
260 );
261 const metadataStore = olds?.metadataStore ?? output?.metadataStore;
262 const name =
263 output?.name ??
264 (metadataStore !== undefined
265 ? resourceName(metadataStore, contextId)
266 : "");
267 if (name.length === 0) return undefined;
268 const existing = yield* getByName(name);
269 if (existing === undefined) return undefined;
270 const attrs = toAttrs(existing, env.project);
271 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
272 ? attrs
273 : Unowned(attrs);
274 }),
275
276 list: () =>
277 Effect.gen(function* () {
278 const env = yield* GcpEnvironment.current;
279 const stores = yield* Effect.forEach(
280 listLocations(env.region),
281 (location) =>
282 listStores(`projects/${env.project}/locations/${location}`),
283 { concurrency: 4 },
284 );
285 const contexts = yield* Effect.forEach(
286 stores.flat().filter((store) => store.name !== undefined),
287 (store) => listContexts(store.name!),
288 { concurrency: 4 },
289 );
290 return contexts
291 .flat()
292 .filter((context) => hasAlchemyPrefix(context.labels))
293 .map((context) => toAttrs(context, env.project));
294 }),
295
296 reconcile: Effect.fn(function* ({ id, news, output }) {
297 const env = yield* GcpEnvironment.current;
298 const contextId = yield* toResourceId(
299 id,
300 news.contextId,
301 output?.contextId,
302 );
303 const metadataStore = news.metadataStore;
304 const name = resourceName(metadataStore, contextId);
305 const displayName = yield* toDisplayName(
306 id,
307 news.displayName,
308 output?.displayName,
309 );
310 const desiredLabels = {
311 ...toLabels(news.labels),
312 ...(yield* createInternalLabels(id)),
313 };
314
315 let current = yield* getByName(output?.name ?? name);
316
317 if (current === undefined) {
318 const created = yield* aiplatform
319 .createProjectsLocationsMetadataStoresContexts({
320 parent: metadataStore,
321 contextId,
322 body: {
323 displayName,
324 description: news.description,
325 labels: desiredLabels,
326 schemaTitle: news.schemaTitle ?? DEFAULT_SCHEMA_TITLE,
327 schemaVersion: news.schemaVersion,
328 metadata: news.metadata,
329 },
330 })
331 .pipe(Effect.catchTag("Conflict", () => getByName(name)));
332 current = created ?? undefined;
333 }
334
335 if (current === undefined) {
336 return yield* new MetadataStoresContextNotResolved({ name });
337 }
338
339 const descriptionChanged =
340 (current.description ?? "") !== (news.description ?? "");
341 const displayChanged = (current.displayName ?? "") !== displayName;
342 const labelsChanged = labelsDiffer(current.labels, desiredLabels);
343 const schemaTitleChanged =
344 (current.schemaTitle ?? "") !==
345 (news.schemaTitle ?? DEFAULT_SCHEMA_TITLE);
346 const schemaVersionChanged =
347 (current.schemaVersion ?? "") !== (news.schemaVersion ?? "");
348 const metadataChanged =
349 news.metadata !== undefined &&
350 stableJson(current.metadata) !== stableJson(news.metadata);
351
352 if (
353 descriptionChanged ||
354 displayChanged ||
355 labelsChanged ||
356 schemaTitleChanged ||
357 schemaVersionChanged ||
358 metadataChanged
359 ) {
360 current =
361 yield* aiplatform.patchProjectsLocationsMetadataStoresContexts({
362 name,
363 updateMask: [
364 descriptionChanged ? "description" : undefined,
365 displayChanged ? "display_name" : undefined,
366 labelsChanged ? "labels" : undefined,
367 schemaTitleChanged ? "schema_title" : undefined,
368 schemaVersionChanged ? "schema_version" : undefined,
369 metadataChanged ? "metadata" : undefined,
370 ]
371 .filter((field): field is string => field !== undefined)
372 .join(","),
373 body: {
374 name,
375 displayName,
376 description: news.description,
377 labels: desiredLabels,
378 schemaTitle: news.schemaTitle ?? DEFAULT_SCHEMA_TITLE,
379 schemaVersion: news.schemaVersion,
380 metadata: news.metadata,
381 etag: current.etag,
382 },
383 });
384 }
385
386 return toAttrs(current, env.project);
387 }),
388
389 delete: Effect.fn(function* ({ output }) {
390 const operation = yield* aiplatform
391 .deleteProjectsLocationsMetadataStoresContexts({
392 name: output.name,
393 })
394 .pipe(
395 Effect.retry({
396 while: (error) => error._tag === "Conflict",
397 times: 8,
398 schedule: Schedule.spaced("2 seconds"),
399 }),
400 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
401 );
402 if (operation !== undefined) {
403 yield* waitForOperation(operation, { notFoundOk: true });
404 }
405 yield* waitUntilGone(output.name);
406 }),
407 });