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 fingerprint,
24 hasAlchemyLabelMap,
25 lastSegment,
26 normalizeLocation,
27 parseResourceName,
28 specifiedEquals,
29 toPhysicalSnake,
30 userLabels,
31} from "./helpers.ts";
32
33const MAX_NAME_LENGTH = 60;
34
35export type FeatureViewBigQuerySource = {
36 /** BigQuery view URI materialized on each sync. */
37 uri?: string;
38 /** Columns used as entity ids. */
39 entityIdColumns?: string[];
40};
41
42export type FeatureViewFeatureRegistrySource = {
43 /** Feature groups and feature ids to sync. */
44 featureGroups?: Array<{
45 featureGroupId?: string;
46 featureIds?: string[];
47 }>;
48 /** Parent project number of the Feature Groups. */
49 projectNumber?: string;
50};
51
52export type FeatureViewVertexRagSource = {
53 /** BigQuery view/table URI. */
54 uri?: string;
55 /** RAG corpus id. */
56 ragCorpusId?: string;
57};
58
59export type FeatureViewSyncConfig = {
60 /** Cron schedule. */
61 cron?: string;
62 /** Continuous sync. */
63 continuous?: boolean;
64};
65
66export type FeatureViewIndexConfig = {
67 embeddingColumn?: string;
68 embeddingDimension?: number;
69 crowdingColumn?: string;
70 filterColumns?: string[];
71 distanceMeasureType?:
72 | aiplatform.GoogleCloudAiplatformV1FeatureViewIndexConfigDistanceMeasureTypeEnum
73 | (string & {});
74 treeAhConfig?: { leafNodeEmbeddingCount?: string };
75 bruteForceConfig?: Record<string, never>;
76};
77
78export type FeatureOnlineStoresFeatureViewProps = {
79 /**
80 * Parent Feature Online Store resource name or id. Immutable.
81 */
82 featureOnlineStore: string;
83 /**
84 * Feature View id. Valid characters `[a-z0-9_]`, must start with a
85 * letter, max 60. Immutable.
86 */
87 featureViewId?: string;
88 /**
89 * Region used when `featureOnlineStore` is a bare id.
90 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
91 */
92 location?: string;
93 /**
94 * User labels. Alchemy ownership labels are merged in automatically.
95 */
96 labels?: Record<string, string>;
97 /**
98 * BigQuery source. Mutually exclusive with `featureRegistrySource`
99 * and `vertexRagSource`.
100 */
101 bigQuerySource?: FeatureViewBigQuerySource;
102 /**
103 * Feature Registry source.
104 */
105 featureRegistrySource?: FeatureViewFeatureRegistrySource;
106 /**
107 * Vertex RAG source.
108 */
109 vertexRagSource?: FeatureViewVertexRagSource;
110 /**
111 * Sync schedule.
112 */
113 syncConfig?: FeatureViewSyncConfig;
114 /**
115 * Vector index config.
116 */
117 indexConfig?: FeatureViewIndexConfig;
118 /**
119 * Optimized-store replica config.
120 */
121 optimizedConfig?: {
122 automaticResources?: { minReplicaCount?: number; maxReplicaCount?: number };
123 };
124 /**
125 * Service agent type used during data sync.
126 */
127 serviceAgentType?:
128 | aiplatform.GoogleCloudAiplatformV1FeatureViewServiceAgentTypeEnum
129 | (string & {});
130 /**
131 * Run an on-demand sync immediately at create.
132 * @default false
133 */
134 runSyncImmediately?: boolean;
135};
136
137export type FeatureOnlineStoresFeatureView = Resource<
138 "GCP.AIPlatform.FeatureOnlineStoresFeatureView",
139 FeatureOnlineStoresFeatureViewProps,
140 {
141 /** Full resource name. */
142 name: string;
143 /** Feature view id. */
144 featureViewId: string;
145 /** Parent online store resource name. */
146 featureOnlineStore: string;
147 /** Parent online store id. */
148 featureOnlineStoreId: string;
149 /** Project id. */
150 project: string;
151 /** Location id. */
152 location: string;
153 /** User labels (Alchemy ownership labels stripped). */
154 labels: Record<string, string>;
155 /** BigQuery source. */
156 bigQuerySource:
157 | aiplatform.GoogleCloudAiplatformV1FeatureViewBigQuerySource
158 | undefined;
159 /** Feature Registry source. */
160 featureRegistrySource:
161 | aiplatform.GoogleCloudAiplatformV1FeatureViewFeatureRegistrySource
162 | undefined;
163 /** Vertex RAG source. */
164 vertexRagSource:
165 | aiplatform.GoogleCloudAiplatformV1FeatureViewVertexRagSource
166 | undefined;
167 /** Sync config. */
168 syncConfig:
169 | aiplatform.GoogleCloudAiplatformV1FeatureViewSyncConfig
170 | undefined;
171 /** Vector index config. */
172 indexConfig:
173 | aiplatform.GoogleCloudAiplatformV1FeatureViewIndexConfig
174 | undefined;
175 /** Optimized replica config. */
176 optimizedConfig:
177 | aiplatform.GoogleCloudAiplatformV1FeatureViewOptimizedConfig
178 | undefined;
179 /** Service agent type. */
180 serviceAgentType:
181 | aiplatform.GoogleCloudAiplatformV1FeatureViewServiceAgentTypeEnum
182 | (string & {})
183 | undefined;
184 /** Unique service account email. */
185 serviceAccountEmail: string | undefined;
186 /** RFC3339 creation timestamp. */
187 createTime: string | undefined;
188 /** RFC3339 last-update timestamp. */
189 updateTime: string | undefined;
190 /** Server etag. */
191 etag: string | undefined;
192 },
193 never,
194 Providers
195>;
196
197/**
198 * A Feature View — the serving projection of a Feature Online Store.
199 *
200 * Parent store, view id, and location are identity. Labels, sources, sync
201 * schedule, and optimized replica counts update in place.
202 *
203 * ### Creating a Feature View
204 * **Example:** Registry-backed view
205 * ```typescript
206 * const view = yield* GCP.AIPlatform.FeatureOnlineStoresFeatureView("Users", {
207 * featureOnlineStore: store.name,
208 * featureRegistrySource: {
209 * featureGroups: [{ featureGroupId: group.featureGroupId, featureIds: ["age"] }],
210 * },
211 * syncConfig: { cron: "0 * * * *" },
212 * });
213 * ```
214 *
215 * @resource
216 * @category AIPlatform
217 */
218export const FeatureOnlineStoresFeatureView =
219 Resource<FeatureOnlineStoresFeatureView>(
220 "GCP.AIPlatform.FeatureOnlineStoresFeatureView",
221 );
222
223export class FeatureOnlineStoresFeatureViewNotResolved extends Data.TaggedError(
224 "GCP.AIPlatform.FeatureOnlineStoresFeatureViewNotResolved",
225)<{
226 name: string;
227}> {}
228
229export class FeatureOnlineStoresFeatureViewStillExists extends Data.TaggedError(
230 "GCP.AIPlatform.FeatureOnlineStoresFeatureViewStillExists",
231)<{
232 name: string;
233}> {}
234
235const parentOf = (project: string, location: string, store: string) =>
236 expandParent(store, project, location, "featureOnlineStores");
237
238const resourceName = (parent: string, viewId: string) =>
239 `${parent}/featureViews/${viewId}`;
240
241const toAttrs = (
242 view: aiplatform.GoogleCloudAiplatformV1FeatureView,
243 project: string,
244) => {
245 const name = view.name ?? "";
246 const parsed = parseResourceName(name, "featureViews");
247 const store = parseResourceName(parsed.parent, "featureOnlineStores");
248 return {
249 name,
250 featureViewId: parsed.id,
251 featureOnlineStore: parsed.parent,
252 featureOnlineStoreId: store.id,
253 project: parsed.project || project,
254 location: parsed.location,
255 labels: userLabels(view.labels),
256 bigQuerySource: view.bigQuerySource,
257 featureRegistrySource: view.featureRegistrySource,
258 vertexRagSource: view.vertexRagSource,
259 syncConfig: view.syncConfig,
260 indexConfig: view.indexConfig,
261 optimizedConfig: view.optimizedConfig,
262 serviceAgentType: view.serviceAgentType,
263 serviceAccountEmail: view.serviceAccountEmail,
264 createTime: view.createTime,
265 updateTime: view.updateTime,
266 etag: view.etag,
267 };
268};
269
270const getByName = (name: string) =>
271 aiplatform
272 .getProjectsLocationsFeatureOnlineStoresFeatureViews({ name })
273 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
274
275const waitUntilExists = (name: string) =>
276 getByName(name).pipe(
277 Effect.flatMap((view) =>
278 view
279 ? Effect.succeed(view)
280 : Effect.fail(new FeatureOnlineStoresFeatureViewNotResolved({ name })),
281 ),
282 Effect.retry({
283 while: (error) =>
284 error._tag ===
285 "GCP.AIPlatform.FeatureOnlineStoresFeatureViewNotResolved",
286 times: 8,
287 schedule: Schedule.spaced("1 second"),
288 }),
289 );
290
291const waitUntilGone = (name: string) =>
292 getByName(name).pipe(
293 Effect.flatMap((view) =>
294 view === undefined
295 ? Effect.void
296 : Effect.fail(new FeatureOnlineStoresFeatureViewStillExists({ name })),
297 ),
298 Effect.retry({
299 while: (error) =>
300 error._tag ===
301 "GCP.AIPlatform.FeatureOnlineStoresFeatureViewStillExists",
302 times: 10,
303 schedule: Schedule.spaced("2 seconds"),
304 }),
305 );
306
307const listViewsUnder = (parent: string, project: string) =>
308 aiplatform.listProjectsLocationsFeatureOnlineStoresFeatureViews
309 .pages({ parent, pageSize: 100 })
310 .pipe(
311 Stream.flatMap((page) => Stream.fromIterable(page.featureViews ?? [])),
312 Stream.filter((view) => hasAlchemyLabelMap(view.labels)),
313 Stream.map((view) => toAttrs(view, project)),
314 Stream.runCollect,
315 Effect.map((chunk) => Array.from(chunk)),
316 Effect.catchTag("NotFound", () => Effect.succeed([])),
317 );
318
319export const FeatureOnlineStoresFeatureViewProvider = () =>
320 Provider.succeed(FeatureOnlineStoresFeatureView, {
321 stables: [
322 "name",
323 "featureViewId",
324 "featureOnlineStore",
325 "featureOnlineStoreId",
326 "project",
327 "location",
328 "createTime",
329 ],
330
331 diff: Effect.fn(function* ({ news, olds, output }) {
332 const env = yield* GcpEnvironment.current;
333 if (!isResolved(news)) return undefined;
334 const previousId = olds?.featureViewId ?? output?.featureViewId;
335 const nextId = news.featureViewId ?? previousId;
336 const previousParent =
337 olds?.featureOnlineStore ?? output?.featureOnlineStore ?? "";
338 const nextParent = news.featureOnlineStore ?? previousParent;
339 const previousLocation = normalizeLocation(
340 olds?.location ?? output?.location,
341 env.region,
342 );
343 const nextLocation = normalizeLocation(
344 news.location ?? output?.location,
345 env.region,
346 );
347 const parentChanged =
348 previousParent.length > 0 &&
349 lastSegment(nextParent) !== lastSegment(previousParent);
350 const replace =
351 (previousId !== undefined &&
352 nextId !== undefined &&
353 nextId !== previousId) ||
354 parentChanged ||
355 previousLocation !== nextLocation;
356 if (!replace) return undefined;
357 return { action: "replace" as const, deleteFirst: false };
358 }),
359
360 read: Effect.fn(function* ({ id, olds, output }) {
361 const env = yield* GcpEnvironment.current;
362 const location = normalizeLocation(
363 olds?.location ?? output?.location,
364 env.region,
365 );
366 const parentRef = olds?.featureOnlineStore ?? output?.featureOnlineStore;
367 // A create interrupted before its parent resolved has nothing to find.
368 if (output?.name === undefined && typeof parentRef !== "string") {
369 return undefined;
370 }
371 const parent = parentOf(env.project, location, parentRef ?? "");
372 const viewId = yield* toPhysicalSnake(
373 id,
374 olds?.featureViewId,
375 output?.featureViewId,
376 MAX_NAME_LENGTH,
377 );
378 const name = output?.name ?? resourceName(parent, viewId);
379 const existing = yield* getByName(name);
380 if (existing === undefined) return undefined;
381 const attrs = toAttrs(existing, env.project);
382 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
383 ? attrs
384 : Unowned(attrs);
385 }),
386
387 list: () =>
388 Effect.gen(function* () {
389 const env = yield* GcpEnvironment.current;
390 const stores = yield* Stream.fromIterable(listLocations(env.region))
391 .pipe(
392 Stream.flatMap((location) =>
393 aiplatform.listProjectsLocationsFeatureOnlineStores.pages({
394 parent: `projects/${env.project}/locations/${location}`,
395 pageSize: 100,
396 }),
397 ),
398 )
399 .pipe(
400 Stream.flatMap((page) =>
401 Stream.fromIterable(page.featureOnlineStores ?? []),
402 ),
403 Stream.filter((store) => hasAlchemyLabelMap(store.labels)),
404 Stream.runCollect,
405 Effect.map((chunk) => Array.from(chunk)),
406 Effect.catchTag("NotFound", () => Effect.succeed([])),
407 );
408 const nested = yield* Effect.forEach(
409 stores,
410 (store) =>
411 store.name
412 ? listViewsUnder(store.name, env.project)
413 : Effect.succeed([]),
414 { concurrency: 4 },
415 );
416 return nested.flat();
417 }),
418
419 reconcile: Effect.fn(function* ({ id, news, output }) {
420 const env = yield* GcpEnvironment.current;
421 const location = normalizeLocation(
422 news.location ?? output?.location,
423 env.region,
424 );
425 const parent = parentOf(env.project, location, news.featureOnlineStore);
426 const viewId = yield* toPhysicalSnake(
427 id,
428 news.featureViewId,
429 output?.featureViewId,
430 MAX_NAME_LENGTH,
431 );
432 const name = resourceName(parent, viewId);
433 const desiredLabels = {
434 ...toLabels(news.labels),
435 ...(yield* createInternalLabels(id)),
436 };
437
438 let current = yield* getByName(output?.name ?? name);
439
440 if (current === undefined) {
441 const created = yield* aiplatform
442 .createProjectsLocationsFeatureOnlineStoresFeatureViews({
443 parent,
444 featureViewId: viewId,
445 runSyncImmediately: news.runSyncImmediately === true,
446 body: {
447 labels: desiredLabels,
448 bigQuerySource: news.bigQuerySource,
449 featureRegistrySource: news.featureRegistrySource,
450 vertexRagSource: news.vertexRagSource,
451 syncConfig: news.syncConfig,
452 indexConfig: news.indexConfig,
453 optimizedConfig: news.optimizedConfig,
454 serviceAgentType: news.serviceAgentType,
455 },
456 })
457 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
458 if (created !== undefined) {
459 yield* waitForOperation(created);
460 }
461 current = yield* waitUntilExists(name);
462 }
463
464 if (current === undefined) {
465 return yield* new FeatureOnlineStoresFeatureViewNotResolved({ name });
466 }
467
468 const observedLabels = tagRecord(current.labels);
469 const { upsert, removed } = diffLabels(observedLabels, desiredLabels);
470 const labelsChanged = upsert.length > 0 || removed.length > 0;
471 const bqChanged =
472 news.bigQuerySource !== undefined &&
473 !specifiedEquals(news.bigQuerySource, current.bigQuerySource);
474 const registryChanged =
475 news.featureRegistrySource !== undefined &&
476 !specifiedEquals(
477 news.featureRegistrySource,
478 current.featureRegistrySource,
479 );
480 const syncChanged =
481 news.syncConfig !== undefined &&
482 fingerprint(current.syncConfig) !== fingerprint(news.syncConfig);
483 const optimizedChanged =
484 news.optimizedConfig !== undefined &&
485 !specifiedEquals(news.optimizedConfig, current.optimizedConfig);
486 const agentChanged =
487 news.serviceAgentType !== undefined &&
488 (current.serviceAgentType ?? "") !== news.serviceAgentType;
489
490 if (
491 labelsChanged ||
492 bqChanged ||
493 registryChanged ||
494 syncChanged ||
495 optimizedChanged ||
496 agentChanged
497 ) {
498 const updateMask = [
499 labelsChanged ? "labels" : undefined,
500 bqChanged ? "big_query_source" : undefined,
501 registryChanged ? "feature_registry_source" : undefined,
502 syncChanged ? "sync_config" : undefined,
503 optimizedChanged ? "optimized_config.automatic_resources" : undefined,
504 agentChanged ? "service_agent_type" : undefined,
505 ].filter((field): field is string => field !== undefined);
506
507 const patched = yield* aiplatform
508 .patchProjectsLocationsFeatureOnlineStoresFeatureViews({
509 name,
510 updateMask: updateMask.join(","),
511 body: {
512 name,
513 labels: desiredLabels,
514 bigQuerySource: news.bigQuerySource,
515 featureRegistrySource: news.featureRegistrySource,
516 syncConfig: news.syncConfig,
517 optimizedConfig: news.optimizedConfig,
518 serviceAgentType: news.serviceAgentType,
519 },
520 })
521 .pipe(
522 Effect.retry({
523 while: (error) => error._tag === "Conflict",
524 times: 8,
525 schedule: Schedule.spaced("5 seconds"),
526 }),
527 );
528 yield* waitForOperation(patched);
529 current = yield* waitUntilExists(name);
530 }
531
532 return toAttrs(current, env.project);
533 }),
534
535 delete: Effect.fn(function* ({ output }) {
536 const operation = yield* aiplatform
537 .deleteProjectsLocationsFeatureOnlineStoresFeatureViews({
538 name: output.name,
539 })
540 .pipe(
541 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
542 Effect.retry({
543 while: (error) => error._tag === "Conflict",
544 times: 8,
545 schedule: Schedule.spaced("5 seconds"),
546 }),
547 );
548 if (operation !== undefined) {
549 yield* waitForOperation(operation, { notFoundOk: true });
550 }
551 yield* waitUntilGone(output.name);
552 }),
553 });