1import { waitForOperation } from "@/GCP/DiscoveryEngine";
2import * as Category from "@distilled.cloud/core/category";
3import * as discoveryengine from "@distilled.cloud/gcp/discoveryengine_v1";
4import * as GcpRetry from "@distilled.cloud/gcp/Retry";
5import * as Effect from "effect/Effect";
6import * as Schedule from "effect/Schedule";
7
8/**
9 * The suite runs ~25 files at once against one project, which overruns the
10 * per-minute "Data store read requests" quota. Distilled's default backoff
11 * gives up before the minute rolls over, so the suite rides it out.
12 */
13export const quotaTolerant = GcpRetry.policy({
14 while: (error) =>
15 Category.isThrottling(error) || Category.isTransientError(error),
16 schedule: Schedule.max([Schedule.spaced("10 seconds"), Schedule.recurs(12)]),
17});
18
19export const DEFAULT_LOCATION = "global";
20export const DEFAULT_COLLECTION = "default_collection";
21
22export const collectionParent = (
23 project: string,
24 location = DEFAULT_LOCATION,
25) =>
26 `projects/${project}/locations/${location}/collections/${DEFAULT_COLLECTION}`;
27
28export const dataStoreName = (
29 project: string,
30 dataStoreId: string,
31 location = DEFAULT_LOCATION,
32) => `${collectionParent(project, location)}/dataStores/${dataStoreId}`;
33
34export const getDataStore = (name: string) =>
35 discoveryengine
36 .getProjectsLocationsCollectionsDataStores({ name })
37 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
38
39export const ensureDataStore = (
40 project: string,
41 dataStoreId: string,
42 options?: {
43 contentConfig?: string;
44 industryVertical?: string;
45 location?: string;
46 withEngine?: boolean;
47 skipDefaultSchemaCreation?: boolean;
48 },
49) =>
50 Effect.gen(function* () {
51 const location = options?.location ?? DEFAULT_LOCATION;
52 const parent = collectionParent(project, location);
53 const name = dataStoreName(project, dataStoreId, location);
54 let store = yield* getDataStore(name);
55 if (store === undefined) {
56 const created = yield* discoveryengine
57 .createProjectsLocationsCollectionsDataStores({
58 parent,
59 dataStoreId,
60 disableCmek: true,
61 skipDefaultSchemaCreation: options?.skipDefaultSchemaCreation,
62 body: {
63 displayName: dataStoreId,
64 industryVertical: options?.industryVertical ?? "GENERIC",
65 contentConfig: options?.contentConfig ?? "CONTENT_REQUIRED",
66 solutionTypes: ["SOLUTION_TYPE_SEARCH"],
67 },
68 })
69 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
70 if (created !== undefined) {
71 yield* waitForOperation(created);
72 }
73 store = yield* discoveryengine
74 .getProjectsLocationsCollectionsDataStores({ name })
75 .pipe(
76 Effect.retry({
77 while: (error) => error._tag === "NotFound",
78 times: 8,
79 schedule: Schedule.spaced("2 seconds"),
80 }),
81 );
82 }
83
84 if (options?.withEngine === true) {
85 const engineId = `${dataStoreId}eng`
86 .replace(/[^a-z0-9]/g, "")
87 .slice(0, 63);
88 const engineName = `${parent}/engines/${engineId}`;
89 const existingEngine = yield* discoveryengine
90 .getProjectsLocationsCollectionsEngines({ name: engineName })
91 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
92 if (existingEngine === undefined) {
93 const createdEngine = yield* discoveryengine
94 .createProjectsLocationsCollectionsEngines({
95 parent,
96 engineId,
97 body: {
98 displayName: engineId,
99 solutionType: "SOLUTION_TYPE_SEARCH",
100 dataStoreIds: [dataStoreId],
101 searchEngineConfig: { searchTier: "SEARCH_TIER_STANDARD" },
102 },
103 })
104 .pipe(Effect.catchTag("Conflict", () => Effect.succeed(undefined)));
105 if (createdEngine !== undefined) {
106 yield* waitForOperation(createdEngine);
107 }
108 }
109 }
110
111 return store;
112 });
113
114export const engineNameOf = (
115 project: string,
116 dataStoreId: string,
117 location = DEFAULT_LOCATION,
118) =>
119 `${collectionParent(project, location)}/engines/${dataStoreId.replace(/[^a-z0-9]/g, "").slice(0, 60)}eng`;
120
121export const deleteEngine = (name: string) =>
122 Effect.gen(function* () {
123 const existing = yield* discoveryengine
124 .getProjectsLocationsCollectionsEngines({ name })
125 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
126 if (existing === undefined) return;
127 const operation = yield* discoveryengine
128 .deleteProjectsLocationsCollectionsEngines({ name })
129 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
130 if (operation !== undefined) {
131 yield* waitForOperation(operation, { notFoundOk: true });
132 }
133 });
134
135export const deleteDataStore = (name: string) =>
136 Effect.gen(function* () {
137 const existing = yield* getDataStore(name);
138 if (existing === undefined) return;
139 const operation = yield* discoveryengine
140 .deleteProjectsLocationsCollectionsDataStores({ name })
141 .pipe(Effect.catchTag("NotFound", () => Effect.succeed(undefined)));
142 if (operation !== undefined) {
143 yield* waitForOperation(operation, { notFoundOk: true });
144 }
145 });
146
147export const deleteSearchParent = (
148 project: string,
149 dataStoreId: string,
150 location = DEFAULT_LOCATION,
151) =>
152 Effect.gen(function* () {
153 yield* deleteEngine(engineNameOf(project, dataStoreId, location));
154 yield* deleteDataStore(dataStoreName(project, dataStoreId, location));
155 });