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 fingerprint,
23 hasAlchemyLabelMap,
24 normalizeLocation,
25 parseResourceName,
26 specifiedEquals,
27 toPhysicalSnake,
28 userLabels,
29} from "./helpers.ts";
30
31const MAX_NAME_LENGTH = 60;
32
33export type FeatureGroupBigQueryTimeSeries = {
34 /** Timestamp column. Defaults to `feature_timestamp`. */
35 timestampColumn?: string;
36};
37
38export type FeatureGroupBigQuery = {
39 /**
40 * BigQuery table or view URI (`bq://project.dataset.table`). Immutable.
41 */
42 inputUri?: string;
43 /** Convenience alias for `bigQuerySource.inputUri`. Immutable. */
44 bigQuerySource?: { inputUri?: string };
45 /** Treat the source as a single snapshot rather than a time series. */
46 staticDataSource?: boolean;
47 /** Collapse rows per entity including nulls. */
48 dense?: boolean;
49 /**
50 * Columns used as entity ids. Defaults to `entity_id`.
51 */
52 entityIdColumns?: string[];
53 /** Time-series options. */
54 timeSeries?: FeatureGroupBigQueryTimeSeries;
55};
56
57export type FeatureGroupProps = {
58 /**
59 * Feature group id. Valid characters `[a-z0-9_]`, must start with a
60 * letter, max 128. Immutable — changing it replaces the group.
61 */
62 featureGroupId?: string;
63 /**
64 * Region. Immutable — changing it replaces the group.
65 * @default the stack's GCP region (`GCP.Region`, else the profile region, else `us-central1`)
66 */
67 location?: string;
68 /**
69 * Description of the feature group.
70 */
71 description?: string;
72 /**
73 * User labels. Alchemy ownership labels are merged in automatically.
74 */
75 labels?: Record<string, string>;
76 /**
77 * BigQuery source. The table or view must include entity-id columns
78 * and (unless `staticDataSource`) a `feature_timestamp` column.
79 */
80 bigQuery?: FeatureGroupBigQuery;
81 /**
82 * Service agent used by jobs under this group.
83 */
84 serviceAgentType?:
85 | aiplatform.GoogleCloudAiplatformV1FeatureGroupServiceAgentTypeEnum
86 | (string & {});
87};
88
89export type FeatureGroup = Resource<
90 "GCP.AIPlatform.FeatureGroup",
91 FeatureGroupProps,
92 {
93 /** Full resource name. */
94 name: string;
95 /** Feature group id. */
96 featureGroupId: string;
97 /** Project id. */
98 project: string;
99 /** Location id. */
100 location: string;
101 /** Description. */
102 description: string | undefined;
103 /** User labels (Alchemy ownership labels stripped). */
104 labels: Record<string, string>;
105 /** BigQuery source. */
106 bigQuery: FeatureGroupBigQuery | undefined;
107 /** Service agent type. */
108 serviceAgentType: string | undefined;
109 /** Unique service account email. */
110 serviceAccountEmail: string | undefined;
111 /** RFC3339 creation timestamp. */
112 createTime: string | undefined;
113 /** RFC3339 last-update timestamp. */
114 updateTime: string | undefined;
115 /** Server etag. */
116 etag: string | undefined;
117 },
118 never,
119 Providers
120>;
121
122/**
123 * A Vertex AI Feature Registry Feature Group backed by a BigQuery table
124 * or view.
125 *
126 * Id, location, and BigQuery URI are identity. Description, labels,
127 * entity-id columns, and service-agent type update in place.
128 *
129 * ### Creating a Feature Group
130 * **Example:** Point at a BigQuery table
131 * ```typescript
132 * const group = yield* GCP.AIPlatform.FeatureGroup("Users", {
133 * bigQuery: {
134 * inputUri: "bq://my-project.features.users",
135 * entityIdColumns: ["entity_id"],
136 * },
137 * labels: { env: "prod" },
138 * });
139 * ```
140 *
141 * @resource
142 * @category AIPlatform
143 */
144export const FeatureGroup = Resource<FeatureGroup>(
145 "GCP.AIPlatform.FeatureGroup",
146);
147
148export class FeatureGroupNotResolved extends Data.TaggedError(
149 "GCP.AIPlatform.FeatureGroupNotResolved",
150)<{
151 name: string;
152}> {}
153
154export class FeatureGroupStillExists extends Data.TaggedError(
155 "GCP.AIPlatform.FeatureGroupStillExists",
156)<{
157 name: string;
158}> {}
159
160const resourceName = (project: string, location: string, groupId: string) =>
161 `projects/${project}/locations/${location}/featureGroups/${groupId}`;
162
163const inputUriOf = (bigQuery: FeatureGroupBigQuery | undefined) =>
164 bigQuery?.inputUri ?? bigQuery?.bigQuerySource?.inputUri;
165
166const toBigQuery = (
167 source:
168 | aiplatform.GoogleCloudAiplatformV1FeatureGroupBigQuery
169 | FeatureGroupBigQuery
170 | undefined,
171): FeatureGroupBigQuery | undefined => {
172 if (source === undefined) return undefined;
173 const inputUri =
174 "inputUri" in source
175 ? (source.inputUri ??
176 (source as FeatureGroupBigQuery).bigQuerySource?.inputUri)
177 : source.bigQuerySource?.inputUri;
178 return {
179 inputUri,
180 bigQuerySource: inputUri ? { inputUri } : undefined,
181 staticDataSource: source.staticDataSource,
182 dense: source.dense,
183 entityIdColumns: source.entityIdColumns,
184 timeSeries: source.timeSeries,
185 };
186};
187
188const toBigQueryBody = (
189 source: FeatureGroupBigQuery | undefined,
190): aiplatform.GoogleCloudAiplatformV1FeatureGroupBigQuery | undefined => {
191 if (source === undefined) return undefined;
192 const inputUri = inputUriOf(source);
193 return {
194 staticDataSource: source.staticDataSource,
195 dense: source.dense,
196 entityIdColumns: source.entityIdColumns,
197 timeSeries: source.timeSeries,
198 bigQuerySource: inputUri ? { inputUri } : source.bigQuerySource,
199 };
200};
201
202const toAttrs = (
203 group: aiplatform.GoogleCloudAiplatformV1FeatureGroup,
204 project: string,
205) => {
206 const name = group.name ?? "";
207 const parsed = parseResourceName(name, "featureGroups");
208 return {
209 name,
210 featureGroupId: parsed.id,
211 project: parsed.project || project,
212 location: parsed.location,
213 description: group.description,
214 labels: userLabels(group.labels),
215 bigQuery: toBigQuery(group.bigQuery),
216 serviceAgentType: group.serviceAgentType,
217 serviceAccountEmail: group.serviceAccountEmail,
218 createTime: group.createTime,
219 updateTime: group.updateTime,
220 etag: group.etag,
221 };
222};
223
224const getByName = (name: string) =>
225 aiplatform
226 .getProjectsLocationsFeatureGroups({ name })
227 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
228
229const waitUntilExists = (name: string) =>
230 getByName(name).pipe(
231 Effect.flatMap((group) =>
232 group
233 ? Effect.succeed(group)
234 : Effect.fail(new FeatureGroupNotResolved({ name })),
235 ),
236 Effect.retry({
237 while: (error) => error._tag === "GCP.AIPlatform.FeatureGroupNotResolved",
238 times: 8,
239 schedule: Schedule.spaced("1 second"),
240 }),
241 );
242
243const waitUntilGone = (name: string) =>
244 getByName(name).pipe(
245 Effect.flatMap((group) =>
246 group === undefined
247 ? Effect.void
248 : Effect.fail(new FeatureGroupStillExists({ name })),
249 ),
250 Effect.retry({
251 while: (error) => error._tag === "GCP.AIPlatform.FeatureGroupStillExists",
252 times: 10,
253 schedule: Schedule.spaced("2 seconds"),
254 }),
255 );
256
257export const FeatureGroupProvider = () =>
258 Provider.succeed(FeatureGroup, {
259 stables: ["name", "featureGroupId", "project", "location", "createTime"],
260
261 diff: Effect.fn(function* ({ news, olds, output }) {
262 const env = yield* GcpEnvironment.current;
263 if (!isResolved(news)) return undefined;
264 const previousId = olds?.featureGroupId ?? output?.featureGroupId;
265 const nextId = news.featureGroupId ?? previousId;
266 const previousLocation = normalizeLocation(
267 olds?.location ?? output?.location,
268 env.region,
269 );
270 const nextLocation = normalizeLocation(
271 news.location ?? output?.location,
272 env.region,
273 );
274 const previousUri =
275 inputUriOf(olds?.bigQuery) ?? inputUriOf(output?.bigQuery) ?? "";
276 const nextUri = inputUriOf(news.bigQuery) ?? previousUri;
277 const replace =
278 (previousId !== undefined &&
279 nextId !== undefined &&
280 nextId !== previousId) ||
281 previousLocation !== nextLocation ||
282 (news.bigQuery !== undefined && nextUri !== previousUri);
283 if (!replace) return undefined;
284 return {
285 action: "replace" as const,
286 deleteFirst:
287 previousLocation === nextLocation &&
288 previousId !== undefined &&
289 nextId === previousId,
290 };
291 }),
292
293 read: Effect.fn(function* ({ id, olds, output }) {
294 const env = yield* GcpEnvironment.current;
295 const featureGroupId = yield* toPhysicalSnake(
296 id,
297 olds?.featureGroupId,
298 output?.featureGroupId,
299 MAX_NAME_LENGTH,
300 );
301 const location = normalizeLocation(
302 olds?.location ?? output?.location,
303 env.region,
304 );
305 const name =
306 output?.name ?? resourceName(env.project, location, featureGroupId);
307 const existing = yield* getByName(name);
308 if (existing === undefined) return undefined;
309 const attrs = toAttrs(existing, env.project);
310 return (yield* hasAlchemyLabels(id, tagRecord(existing.labels)))
311 ? attrs
312 : Unowned(attrs);
313 }),
314
315 list: () =>
316 Effect.gen(function* () {
317 const env = yield* GcpEnvironment.current;
318 return yield* Stream.fromIterable(listLocations(env.region))
319 .pipe(
320 Stream.flatMap((location) =>
321 aiplatform.listProjectsLocationsFeatureGroups.pages({
322 parent: `projects/${env.project}/locations/${location}`,
323 pageSize: 100,
324 }),
325 ),
326 )
327 .pipe(
328 Stream.flatMap((page) =>
329 Stream.fromIterable(page.featureGroups ?? []),
330 ),
331 Stream.filter((group) => hasAlchemyLabelMap(group.labels)),
332 Stream.map((group) => toAttrs(group, env.project)),
333 Stream.runCollect,
334 Effect.map((chunk) => Array.from(chunk)),
335 Effect.catchTag("NotFound", () => Effect.succeed([])),
336 );
337 }),
338
339 reconcile: Effect.fn(function* ({ id, news, output }) {
340 const env = yield* GcpEnvironment.current;
341 const featureGroupId = yield* toPhysicalSnake(
342 id,
343 news.featureGroupId,
344 output?.featureGroupId,
345 MAX_NAME_LENGTH,
346 );
347 const location = normalizeLocation(
348 news.location ?? output?.location,
349 env.region,
350 );
351 const name = resourceName(env.project, location, featureGroupId);
352 const desiredLabels = {
353 ...toLabels(news.labels),
354 ...(yield* createInternalLabels(id)),
355 };
356 const bigQuery = toBigQueryBody(news.bigQuery);
357
358 let current = yield* getByName(output?.name ?? name);
359
360 if (current === undefined) {
361 const created = yield* aiplatform
362 .createProjectsLocationsFeatureGroups({
363 parent: `projects/${env.project}/locations/${location}`,
364 featureGroupId,
365 body: {
366 description: news.description,
367 labels: desiredLabels,
368 bigQuery,
369 serviceAgentType: news.serviceAgentType,
370 },
371 })
372 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
373 if (created !== undefined) {
374 yield* waitForOperation(created);
375 }
376 current = yield* waitUntilExists(name);
377 }
378
379 if (current === undefined) {
380 return yield* new FeatureGroupNotResolved({ name });
381 }
382
383 const observedLabels = tagRecord(current.labels);
384 const { upsert, removed } = diffLabels(observedLabels, desiredLabels);
385 const labelsChanged = upsert.length > 0 || removed.length > 0;
386 const descriptionChanged =
387 (current.description ?? "") !== (news.description ?? "");
388 const agentChanged =
389 news.serviceAgentType !== undefined &&
390 (current.serviceAgentType ?? "") !== news.serviceAgentType;
391 const entityColumnsChanged =
392 news.bigQuery?.entityIdColumns !== undefined &&
393 fingerprint(current.bigQuery?.entityIdColumns ?? []) !==
394 fingerprint(news.bigQuery.entityIdColumns);
395 const staticChanged =
396 news.bigQuery?.staticDataSource !== undefined &&
397 (current.bigQuery?.staticDataSource === true) !==
398 (news.bigQuery.staticDataSource === true);
399 const denseChanged =
400 news.bigQuery?.dense !== undefined &&
401 (current.bigQuery?.dense === true) !== (news.bigQuery.dense === true);
402 const timeSeriesChanged =
403 news.bigQuery?.timeSeries !== undefined &&
404 !specifiedEquals(
405 news.bigQuery.timeSeries,
406 current.bigQuery?.timeSeries,
407 );
408
409 if (
410 labelsChanged ||
411 descriptionChanged ||
412 agentChanged ||
413 entityColumnsChanged ||
414 staticChanged ||
415 denseChanged ||
416 timeSeriesChanged
417 ) {
418 const updateMask = [
419 labelsChanged ? "labels" : undefined,
420 descriptionChanged ? "description" : undefined,
421 agentChanged ? "service_agent_type" : undefined,
422 entityColumnsChanged ? "big_query.entity_id_columns" : undefined,
423 staticChanged ? "big_query.static_data_source" : undefined,
424 denseChanged ? "big_query.dense" : undefined,
425 timeSeriesChanged ? "big_query.time_series" : undefined,
426 ].filter((field): field is string => field !== undefined);
427
428 const patched = yield* aiplatform
429 .patchProjectsLocationsFeatureGroups({
430 name,
431 updateMask: updateMask.join(","),
432 body: {
433 name,
434 labels: desiredLabels,
435 description: news.description,
436 serviceAgentType: news.serviceAgentType,
437 bigQuery,
438 },
439 })
440 .pipe(
441 Effect.retry({
442 while: (error) => error._tag === "Conflict",
443 times: 8,
444 schedule: Schedule.spaced("5 seconds"),
445 }),
446 );
447 yield* waitForOperation(patched);
448 current = yield* waitUntilExists(name);
449 }
450
451 return toAttrs(current, env.project);
452 }),
453
454 delete: Effect.fn(function* ({ output }) {
455 const operation = yield* aiplatform
456 .deleteProjectsLocationsFeatureGroups({
457 name: output.name,
458 force: true,
459 })
460 .pipe(
461 Effect.catchTag("NotFound", () => Effect.succeed(undefined)),
462 Effect.retry({
463 while: (error) => error._tag === "Conflict",
464 times: 8,
465 schedule: Schedule.spaced("5 seconds"),
466 }),
467 );
468 if (operation !== undefined) {
469 yield* waitForOperation(operation, { notFoundOk: true });
470 }
471 yield* waitUntilGone(output.name);
472 }),
473 });