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 { GcpEnvironment } from "../Environment.ts";
11import { hasAlchemyLabels } from "../Labels.ts";
12import type { Providers } from "../Providers.ts";
13import {
14 createInternalLabels,
15 encodeDescription,
16 hasOwnershipMarker,
17 listLocations,
18 locationOf,
19 lastSegment,
20 normalizeLocation,
21 parseDescription,
22 projectOf,
23 toResourceId,
24} from "./names.ts";
25import { waitForOperation } from "./operations.ts";
26
27export type MetadataStoreProps = {
28 /**
29 * MetadataStore id (the `{metadatastore}` segment of
30 * `projects/{project}/locations/{location}/metadataStores/{metadatastore}`).
31 * If omitted, a unique name is generated. Must be 4-128 characters of
32 * `a-z` and `-`. Immutable — changing it replaces the store.
33 */
34 metadataStoreId?: string;
35 /**
36 * Vertex AI location. Immutable — changing it replaces the store.
37 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
38 */
39 location?: string;
40 /**
41 * Human-readable description. Metadata stores have no labels field, so
42 * Alchemy ownership is stored in a `[alchemy …]` prefix and stripped
43 * from attributes.
44 */
45 description?: string;
46 /**
47 * Customer-managed encryption key. Immutable.
48 */
49 encryptionSpec?: aiplatform.GoogleCloudAiplatformV1EncryptionSpec;
50 /**
51 * Dataplex integration settings.
52 */
53 dataplexConfig?: aiplatform.GoogleCloudAiplatformV1MetadataStoreDataplexConfig;
54};
55
56export type MetadataStore = Resource<
57 "GCP.AIPlatform.MetadataStore",
58 MetadataStoreProps,
59 {
60 /** Full resource name `projects/{project}/locations/{location}/metadataStores/{metadatastore}`. */
61 name: string;
62 /** MetadataStore id (last path segment). */
63 metadataStoreId: string;
64 /** Project id. */
65 project: string;
66 /** Location id. */
67 location: string;
68 /** User description with the Alchemy ownership prefix stripped. */
69 description: string | undefined;
70 /** Disk utilization of the store, in bytes. */
71 diskUtilizationBytes: string | undefined;
72 /** KMS key used for encryption, if any. */
73 kmsKeyName: string | undefined;
74 /** RFC3339 creation timestamp. */
75 createTime: string | undefined;
76 /** RFC3339 last-update timestamp. */
77 updateTime: string | undefined;
78 },
79 never,
80 Providers
81>;
82
83/**
84 * A Vertex ML Metadata store for artifacts, contexts, and executions.
85 *
86 * Metadata stores have no labels field, so Alchemy stamps ownership into
87 * the description for `list` / nuke. There is no update API — changing
88 * identity, location, encryption, or Dataplex config replaces the store.
89 *
90 * ### Creating a MetadataStore
91 * **Example:** Generated id
92 * ```typescript
93 * const store = yield* GCP.AIPlatform.MetadataStore("Mlmd", {
94 * description: "pipeline metadata",
95 * });
96 * ```
97 *
98 * @resource
99 * @category AIPlatform
100 */
101export const MetadataStore = Resource<MetadataStore>(
102 "GCP.AIPlatform.MetadataStore",
103);
104
105export class MetadataStoreNotResolved extends Data.TaggedError(
106 "GCP.AIPlatform.MetadataStoreNotResolved",
107)<{
108 name: string;
109}> {}
110
111export class MetadataStoreStillExists extends Data.TaggedError(
112 "GCP.AIPlatform.MetadataStoreStillExists",
113)<{
114 name: string;
115}> {}
116
117const resourceName = (
118 project: string,
119 location: string,
120 metadataStoreId: string,
121) =>
122 `projects/${project}/locations/${location}/metadataStores/${metadataStoreId}`;
123
124const toAttrs = (
125 store: aiplatform.GoogleCloudAiplatformV1MetadataStore,
126 project: string,
127) => {
128 const name = store.name ?? "";
129 const parsed = parseDescription(store.description);
130 return {
131 name,
132 metadataStoreId: lastSegment(name),
133 project: projectOf(name, project),
134 location: locationOf(name),
135 description: parsed.description,
136 diskUtilizationBytes: store.state?.diskUtilizationBytes,
137 kmsKeyName: store.encryptionSpec?.kmsKeyName,
138 createTime: store.createTime,
139 updateTime: store.updateTime,
140 };
141};
142
143const getByName = (name: string) =>
144 aiplatform
145 .getProjectsLocationsMetadataStores({ name })
146 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
147
148const listPage = (parent: string) =>
149 aiplatform.listProjectsLocationsMetadataStores
150 .pages({ parent, pageSize: 100 })
151 .pipe(
152 Stream.runCollect,
153 Effect.map((pages) =>
154 Array.from(pages).flatMap((page) => page.metadataStores ?? []),
155 ),
156 Effect.catchTag("NotFound", () =>
157 Effect.succeed([] as aiplatform.GoogleCloudAiplatformV1MetadataStore[]),
158 ),
159 );
160
161const waitUntilExists = (name: string) =>
162 getByName(name).pipe(
163 Effect.filterOrFail(
164 (store): store is aiplatform.GoogleCloudAiplatformV1MetadataStore =>
165 store !== undefined,
166 () => new MetadataStoreNotResolved({ name }),
167 ),
168 Effect.retry({
169 while: (error) =>
170 error._tag === "GCP.AIPlatform.MetadataStoreNotResolved",
171 times: 8,
172 schedule: Schedule.spaced("2 seconds"),
173 }),
174 );
175
176const waitUntilGone = (name: string) =>
177 getByName(name).pipe(
178 Effect.filterOrFail(
179 (store) => store === undefined,
180 () => new MetadataStoreStillExists({ name }),
181 ),
182 Effect.asVoid,
183 Effect.retry({
184 while: (error) =>
185 error._tag === "GCP.AIPlatform.MetadataStoreStillExists",
186 times: 10,
187 schedule: Schedule.spaced("2 seconds"),
188 }),
189 );
190
191export const MetadataStoreProvider = () =>
192 Provider.succeed(MetadataStore, {
193 stables: [
194 "name",
195 "metadataStoreId",
196 "project",
197 "location",
198 "kmsKeyName",
199 "createTime",
200 ],
201
202 diff: Effect.fn(function* ({ news, olds, output }) {
203 const env = yield* GcpEnvironment.current;
204 if (!isResolved(news)) return undefined;
205 const previousId = olds?.metadataStoreId ?? output?.metadataStoreId;
206 const nextId = news.metadataStoreId ?? previousId;
207 const previousLocation = normalizeLocation(
208 olds?.location ?? output?.location,
209 env.region,
210 );
211 const nextLocation = normalizeLocation(
212 news.location ?? olds?.location ?? output?.location,
213 env.region,
214 );
215 const previousKms =
216 olds?.encryptionSpec?.kmsKeyName ?? output?.kmsKeyName ?? "";
217 const nextKms = news.encryptionSpec?.kmsKeyName ?? previousKms;
218 if (
219 (previousId !== undefined &&
220 nextId !== undefined &&
221 nextId !== previousId) ||
222 previousLocation !== nextLocation ||
223 nextKms !== previousKms
224 ) {
225 return {
226 action: "replace" as const,
227 deleteFirst:
228 previousLocation === nextLocation &&
229 previousId !== undefined &&
230 nextId === previousId,
231 };
232 }
233 return undefined;
234 }),
235
236 read: Effect.fn(function* ({ id, olds, output }) {
237 const env = yield* GcpEnvironment.current;
238 const metadataStoreId = yield* toResourceId(
239 id,
240 olds?.metadataStoreId,
241 output?.metadataStoreId,
242 );
243 const location = normalizeLocation(
244 olds?.location ?? output?.location,
245 env.region,
246 );
247 const name =
248 output?.name ?? resourceName(env.project, location, metadataStoreId);
249 const existing = yield* getByName(name);
250 if (existing === undefined) return undefined;
251 const attrs = toAttrs(existing, env.project);
252 const { labels } = parseDescription(existing.description);
253 return (yield* hasAlchemyLabels(id, labels)) ? attrs : Unowned(attrs);
254 }),
255
256 list: () =>
257 Effect.gen(function* () {
258 const env = yield* GcpEnvironment.current;
259 const pages = yield* Effect.forEach(
260 listLocations(env.region),
261 (location) =>
262 listPage(`projects/${env.project}/locations/${location}`),
263 { concurrency: 4 },
264 );
265 return pages
266 .flat()
267 .filter((store) => hasOwnershipMarker(store.description))
268 .map((store) => toAttrs(store, env.project));
269 }),
270
271 reconcile: Effect.fn(function* ({ id, news, output }) {
272 const env = yield* GcpEnvironment.current;
273 const metadataStoreId = yield* toResourceId(
274 id,
275 news.metadataStoreId,
276 output?.metadataStoreId,
277 );
278 const location = normalizeLocation(
279 news.location ?? output?.location,
280 env.region,
281 );
282 const name = resourceName(env.project, location, metadataStoreId);
283 const ownership = yield* createInternalLabels(id);
284 const desiredDescription = encodeDescription(ownership, news.description);
285
286 let current = yield* getByName(output?.name ?? name);
287
288 if (current === undefined) {
289 const created = yield* aiplatform
290 .createProjectsLocationsMetadataStores({
291 parent: `projects/${env.project}/locations/${location}`,
292 metadataStoreId,
293 body: {
294 description: desiredDescription,
295 encryptionSpec: news.encryptionSpec,
296 dataplexConfig: news.dataplexConfig,
297 },
298 })
299 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
300 if (created !== undefined) {
301 yield* waitForOperation(created);
302 }
303 current = yield* waitUntilExists(name);
304 }
305
306 return toAttrs(current, env.project);
307 }),
308
309 delete: Effect.fn(function* ({ output }) {
310 const operation = yield* aiplatform
311 .deleteProjectsLocationsMetadataStores({
312 name: output.name,
313 force: true,
314 })
315 .pipe(
316 Effect.retry({
317 while: (error) => error._tag === "Conflict",
318 times: 8,
319 schedule: Schedule.spaced("2 seconds"),
320 }),
321 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
322 );
323 if (operation !== undefined) {
324 yield* waitForOperation(operation, { notFoundOk: true });
325 }
326 yield* waitUntilGone(output.name);
327 }),
328 });