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 { listLocations } from "./names.ts";
7import { Unowned } from "../../AdoptPolicy.ts";
8import { isResolved } from "../../Diff.ts";
9import * as Provider from "../../Provider.ts";
10import { Resource } from "../../Resource.ts";
11import { tagRecord } from "../../Tags.ts";
12import { GcpEnvironment } from "../Environment.ts";
13import {
14 createInternalLabels,
15 diffLabels,
16 hasAlchemyLabels,
17 toLabels,
18} from "../Labels.ts";
19import type { Providers } from "../Providers.ts";
20import { waitForOperation } from "./operations.ts";
21import {
22 expandParent,
23 hasAlchemyLabelMap,
24 lastSegment,
25 normalizeLocation,
26 parseResourceName,
27 toPhysicalSnake,
28 userLabels,
29} from "./helpers.ts";
30
31const MAX_NAME_LENGTH = 63;
32
33export type FeatureGroupsFeatureProps = {
34 /**
35 * Parent Feature Group resource name
36 * `projects/{project}/locations/{location}/featureGroups/{featureGroup}`
37 * or the feature-group id (uses `location`). Immutable.
38 */
39 featureGroup: string;
40 /**
41 * Feature id. Valid characters `[a-z0-9_]`, must start with a letter.
42 * Immutable — changing it replaces the feature.
43 */
44 featureId?: string;
45 /**
46 * Region used when `featureGroup` is a bare id.
47 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
48 */
49 location?: string;
50 /**
51 * Description of the feature.
52 */
53 description?: string;
54 /**
55 * User labels. Alchemy ownership labels are merged in automatically.
56 */
57 labels?: Record<string, string>;
58 /**
59 * BigQuery column hosting values for this feature. Defaults to the
60 * feature id.
61 */
62 versionColumnName?: string;
63 /**
64 * Entity responsible for maintaining this feature (emails or URIs).
65 */
66 pointOfContact?: string;
67};
68
69export type FeatureGroupsFeature = Resource<
70 "GCP.AIPlatform.FeatureGroupsFeature",
71 FeatureGroupsFeatureProps,
72 {
73 /** Full resource name. */
74 name: string;
75 /** Feature id. */
76 featureId: string;
77 /** Parent feature group resource name. */
78 featureGroup: string;
79 /** Feature group id. */
80 featureGroupId: string;
81 /** Project id. */
82 project: string;
83 /** Location id. */
84 location: string;
85 /** Description. */
86 description: string | undefined;
87 /** User labels (Alchemy ownership labels stripped). */
88 labels: Record<string, string>;
89 /** Version column name. */
90 versionColumnName: string | undefined;
91 /** Point of contact. */
92 pointOfContact: string | undefined;
93 /** Value type (legacy Feature Store only). */
94 valueType: string | undefined;
95 /** RFC3339 creation timestamp. */
96 createTime: string | undefined;
97 /** RFC3339 last-update timestamp. */
98 updateTime: string | undefined;
99 /** Server etag. */
100 etag: string | undefined;
101 },
102 never,
103 Providers
104>;
105
106/**
107 * A Feature Registry Feature belonging to a Feature Group.
108 *
109 * Parent group, feature id, and location are identity. Description,
110 * labels, version column, and point of contact update in place.
111 *
112 * ### Creating a Feature
113 * **Example:** Feature under a group
114 * ```typescript
115 * const feature = yield* GCP.AIPlatform.FeatureGroupsFeature("Age", {
116 * featureGroup: group.name,
117 * versionColumnName: "age",
118 * labels: { env: "prod" },
119 * });
120 * ```
121 *
122 * @resource
123 * @category AIPlatform
124 */
125export const FeatureGroupsFeature = Resource<FeatureGroupsFeature>(
126 "GCP.AIPlatform.FeatureGroupsFeature",
127);
128
129export class FeatureGroupsFeatureNotResolved extends Data.TaggedError(
130 "GCP.AIPlatform.FeatureGroupsFeatureNotResolved",
131)<{
132 name: string;
133}> {}
134
135export class FeatureGroupsFeatureStillExists extends Data.TaggedError(
136 "GCP.AIPlatform.FeatureGroupsFeatureStillExists",
137)<{
138 name: string;
139}> {}
140
141const parentOf = (project: string, location: string, featureGroup: string) =>
142 expandParent(featureGroup, project, location, "featureGroups");
143
144const resourceName = (parent: string, featureId: string) =>
145 `${parent}/features/${featureId}`;
146
147const toAttrs = (
148 feature: aiplatform.GoogleCloudAiplatformV1Feature,
149 project: string,
150) => {
151 const name = feature.name ?? "";
152 const parsed = parseResourceName(name, "features");
153 const group = parseResourceName(parsed.parent, "featureGroups");
154 return {
155 name,
156 featureId: parsed.id,
157 featureGroup: parsed.parent,
158 featureGroupId: group.id,
159 project: parsed.project || project,
160 location: parsed.location,
161 description: feature.description,
162 labels: userLabels(feature.labels),
163 versionColumnName: feature.versionColumnName,
164 pointOfContact: feature.pointOfContact,
165 valueType: feature.valueType,
166 createTime: feature.createTime,
167 updateTime: feature.updateTime,
168 etag: feature.etag,
169 };
170};
171
172const getByName = (name: string) =>
173 aiplatform
174 .getProjectsLocationsFeatureGroupsFeatures({ name })
175 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
176
177const waitUntilExists = (name: string) =>
178 getByName(name).pipe(
179 Effect.flatMap((feature) =>
180 feature
181 ? Effect.succeed(feature)
182 : Effect.fail(new FeatureGroupsFeatureNotResolved({ name })),
183 ),
184 Effect.retry({
185 while: (error) =>
186 error._tag === "GCP.AIPlatform.FeatureGroupsFeatureNotResolved",
187 times: 8,
188 schedule: Schedule.spaced("1 second"),
189 }),
190 );
191
192const waitUntilGone = (name: string) =>
193 getByName(name).pipe(
194 Effect.flatMap((feature) =>
195 feature === undefined
196 ? Effect.void
197 : Effect.fail(new FeatureGroupsFeatureStillExists({ name })),
198 ),
199 Effect.retry({
200 while: (error) =>
201 error._tag === "GCP.AIPlatform.FeatureGroupsFeatureStillExists",
202 times: 10,
203 schedule: Schedule.spaced("2 seconds"),
204 }),
205 );
206
207const listFeaturesUnder = (parent: string, project: string) =>
208 aiplatform.listProjectsLocationsFeatureGroupsFeatures
209 .pages({ parent, pageSize: 100 })
210 .pipe(
211 Stream.flatMap((page) => Stream.fromIterable(page.features ?? [])),
212 Stream.filter((feature) => hasAlchemyLabelMap(feature.labels)),
213 Stream.map((feature) => toAttrs(feature, project)),
214 Stream.runCollect,
215 Effect.map((chunk) => Array.from(chunk)),
216 Effect.catchTag("NotFound", () => Effect.succeed([])),
217 );
218
219export const FeatureGroupsFeatureProvider = () =>
220 Provider.succeed(FeatureGroupsFeature, {
221 stables: [
222 "name",
223 "featureId",
224 "featureGroup",
225 "featureGroupId",
226 "project",
227 "location",
228 "createTime",
229 ],
230
231 diff: Effect.fn(function* ({ news, olds, output }) {
232 const env = yield* GcpEnvironment.current;
233 if (!isResolved(news)) return undefined;
234 const previousId = olds?.featureId ?? output?.featureId;
235 const nextId = news.featureId ?? previousId;
236 const previousParent = olds?.featureGroup ?? output?.featureGroup ?? "";
237 const nextParent = news.featureGroup
238 ? lastSegment(news.featureGroup) === lastSegment(previousParent) ||
239 news.featureGroup === previousParent
240 ? previousParent
241 : news.featureGroup
242 : previousParent;
243 const previousLocation = normalizeLocation(
244 olds?.location ?? output?.location,
245 env.region,
246 );
247 const nextLocation = normalizeLocation(
248 news.location ?? output?.location,
249 env.region,
250 );
251 const parentChanged =
252 previousParent.length > 0 &&
253 lastSegment(nextParent) !== lastSegment(previousParent);
254 const replace =
255 (previousId !== undefined &&
256 nextId !== undefined &&
257 nextId !== previousId) ||
258 parentChanged ||
259 previousLocation !== nextLocation;
260 if (!replace) return undefined;
261 return { action: "replace" as const, deleteFirst: false };
262 }),
263
264 read: Effect.fn(function* ({ id, olds, output }) {
265 const env = yield* GcpEnvironment.current;
266 const location = normalizeLocation(
267 olds?.location ?? output?.location,
268 env.region,
269 );
270 const parentRef = olds?.featureGroup ?? output?.featureGroup;
271 // A create interrupted before its parent resolved has nothing to find.
272 if (output?.name === undefined && typeof parentRef !== "string") {
273 return undefined;
274 }
275 const parent = parentOf(env.project, location, parentRef ?? "");
276 const featureId = yield* toPhysicalSnake(
277 id,
278 olds?.featureId,
279 output?.featureId,
280 MAX_NAME_LENGTH,
281 );
282 const name = output?.name ?? resourceName(parent, featureId);
283 const existing = yield* getByName(name);
284 if (existing === undefined) return undefined;
285 const attrs = toAttrs(existing, env.project);
286 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
287 ? attrs
288 : Unowned(attrs);
289 }),
290
291 list: () =>
292 Effect.gen(function* () {
293 const env = yield* GcpEnvironment.current;
294 const groups = yield* Stream.fromIterable(listLocations(env.region))
295 .pipe(
296 Stream.flatMap((location) =>
297 aiplatform.listProjectsLocationsFeatureGroups.pages({
298 parent: `projects/${env.project}/locations/${location}`,
299 pageSize: 100,
300 }),
301 ),
302 )
303 .pipe(
304 Stream.flatMap((page) =>
305 Stream.fromIterable(page.featureGroups ?? []),
306 ),
307 Stream.filter((group) => hasAlchemyLabelMap(group.labels)),
308 Stream.runCollect,
309 Effect.map((chunk) => Array.from(chunk)),
310 Effect.catchTag("NotFound", () => Effect.succeed([])),
311 );
312 const nested = yield* Effect.forEach(
313 groups,
314 (group) =>
315 group.name
316 ? listFeaturesUnder(group.name, env.project)
317 : Effect.succeed([]),
318 { concurrency: 4 },
319 );
320 return nested.flat();
321 }),
322
323 reconcile: Effect.fn(function* ({ id, news, output }) {
324 const env = yield* GcpEnvironment.current;
325 const location = normalizeLocation(
326 news.location ?? output?.location,
327 env.region,
328 );
329 const parent = parentOf(env.project, location, news.featureGroup);
330 const featureId = yield* toPhysicalSnake(
331 id,
332 news.featureId,
333 output?.featureId,
334 MAX_NAME_LENGTH,
335 );
336 const name = resourceName(parent, featureId);
337 const desiredLabels = {
338 ...toLabels(news.labels),
339 ...(yield* createInternalLabels(id)),
340 };
341
342 let current = yield* getByName(output?.name ?? name);
343
344 if (current === undefined) {
345 const created = yield* aiplatform
346 .createProjectsLocationsFeatureGroupsFeatures({
347 parent,
348 featureId,
349 body: {
350 description: news.description,
351 labels: desiredLabels,
352 versionColumnName: news.versionColumnName,
353 pointOfContact: news.pointOfContact,
354 },
355 })
356 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
357 if (created !== undefined) {
358 yield* waitForOperation(created);
359 }
360 current = yield* waitUntilExists(name);
361 }
362
363 if (current === undefined) {
364 return yield* new FeatureGroupsFeatureNotResolved({ name });
365 }
366
367 const observedLabels = tagRecord(current.labels);
368 const { upsert, removed } = diffLabels(observedLabels, desiredLabels);
369 const labelsChanged = upsert.length > 0 || removed.length > 0;
370 const descriptionChanged =
371 (current.description ?? "") !== (news.description ?? "");
372 const columnChanged =
373 news.versionColumnName !== undefined &&
374 (current.versionColumnName ?? "") !== news.versionColumnName;
375 const contactChanged =
376 news.pointOfContact !== undefined &&
377 (current.pointOfContact ?? "") !== news.pointOfContact;
378
379 if (
380 labelsChanged ||
381 descriptionChanged ||
382 columnChanged ||
383 contactChanged
384 ) {
385 const updateMask = [
386 labelsChanged ? "labels" : undefined,
387 descriptionChanged ? "description" : undefined,
388 columnChanged ? "version_column_name" : undefined,
389 contactChanged ? "point_of_contact" : undefined,
390 ].filter((field): field is string => field !== undefined);
391
392 const patched = yield* aiplatform
393 .patchProjectsLocationsFeatureGroupsFeatures({
394 name,
395 updateMask: updateMask.join(","),
396 body: {
397 name,
398 labels: desiredLabels,
399 description: news.description,
400 versionColumnName: news.versionColumnName,
401 pointOfContact: news.pointOfContact,
402 },
403 })
404 .pipe(
405 Effect.retry({
406 while: (error) => error._tag === "Conflict",
407 times: 8,
408 schedule: Schedule.spaced("5 seconds"),
409 }),
410 );
411 yield* waitForOperation(patched);
412 current = yield* waitUntilExists(name);
413 }
414
415 return toAttrs(current, env.project);
416 }),
417
418 delete: Effect.fn(function* ({ output }) {
419 const operation = yield* aiplatform
420 .deleteProjectsLocationsFeatureGroupsFeatures({ name: output.name })
421 .pipe(
422 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
423 Effect.retry({
424 while: (error) => error._tag === "Conflict",
425 times: 8,
426 schedule: Schedule.spaced("5 seconds"),
427 }),
428 );
429 if (operation !== undefined) {
430 yield* waitForOperation(operation, { notFoundOk: true });
431 }
432 yield* waitUntilGone(output.name);
433 }),
434 });