1import * as bigquery from "@distilled.cloud/gcp/bigquery_v2";
2import * as Effect from "effect/Effect";
3import { bindGcpHost } from "../Host.ts";
4import { type BindingIam, grantFor } from "../HttpBinding.ts";
5import {
6 type ListTableResult,
7 QueryNotComplete,
8 type ReadTableClient,
9} from "./ReadTable.ts";
10import { decodeRows, encodeRow, queryParameter } from "./Rows.ts";
11import type { Table } from "./Table.ts";
12import { InsertRowsFailed, type WriteTableClient } from "./WriteTable.ts";
13
14// bigquery.jobs.create is only grantable on the project; data access is
15// granted on the table's own IAM policy.
16export const readTableIam: ReadonlyArray<BindingIam> = [
17 { role: "roles/bigquery.dataViewer", on: "bigquery.table" },
18 { role: "roles/bigquery.jobUser" },
19];
20export const writeTableIam: ReadonlyArray<BindingIam> = [
21 { role: "roles/bigquery.dataEditor", on: "bigquery.table" },
22];
23export const readWriteTableIam: ReadonlyArray<BindingIam> = [
24 { role: "roles/bigquery.dataEditor", on: "bigquery.table" },
25 { role: "roles/bigquery.jobUser" },
26];
27
28export interface TableRef {
29 project: string;
30 datasetId: string;
31 tableId: string;
32 location: string | undefined;
33}
34
35// `getQueryResults` long-polls up to `timeoutMs` per call, so this bounds
36// the total wait for a query job at ~2 minutes.
37const QUERY_WAIT_MS = 10_000;
38const QUERY_POLLS = 12;
39
40interface QueryPage {
41 jobComplete?: boolean;
42 schema?: bigquery.TableSchema;
43 rows?: bigquery.TableRow[];
44 pageToken?: string;
45}
46
47/** Fetch every result page after `pageToken`. */
48const remainingRows = <E>(
49 fetchPage: (pageToken: string) => Effect.Effect<QueryPage, E>,
50 pageToken: string | undefined,
51): Effect.Effect<ReadonlyArray<bigquery.TableRow>, E> =>
52 pageToken
53 ? fetchPage(pageToken).pipe(
54 Effect.flatMap((page) =>
55 remainingRows(fetchPage, page.pageToken).pipe(
56 Effect.map((rest) => [...(page.rows ?? []), ...rest]),
57 ),
58 ),
59 )
60 : Effect.succeed([]);
61
62/**
63 * Shared HTTP scaffolding for the BigQuery Read/Write/ReadWrite table
64 * bindings: resolves the distilled operations once at Layer construction
65 * and grants the level's roles at deploy time.
66 *
67 * NOT exported from `index.ts`.
68 */
69export const makeBigQueryTableHelpers = Effect.gen(function* () {
70 const getTable = yield* bigquery.getTables;
71 const listTabledata = yield* bigquery.listTabledata;
72 const queryJobs = yield* bigquery.queryJobs;
73 const getQueryResults = yield* bigquery.getQueryResultsJobs;
74 const insertAll = yield* bigquery.insertAllTabledata;
75
76 const makeRead = (table: Effect.Effect<TableRef>): ReadTableClient => ({
77 list: (
78 options,
79 ): Effect.Effect<
80 ListTableResult,
81 bigquery.GetTablesError | bigquery.ListTabledataError
82 > =>
83 Effect.gen(function* () {
84 const ref = yield* table;
85 const key = {
86 projectId: ref.project,
87 datasetId: ref.datasetId,
88 tableId: ref.tableId,
89 };
90 const metadata = yield* getTable({ ...key, view: "BASIC" });
91 const page = yield* listTabledata({
92 ...key,
93 maxResults: options?.maxResults,
94 pageToken: options?.pageToken,
95 "formatOptions.useInt64Timestamp": true,
96 });
97 return {
98 rows: decodeRows(metadata.schema, page.rows),
99 nextPageToken: page.pageToken || undefined,
100 };
101 }),
102 query: (
103 sql,
104 params,
105 ): Effect.Effect<
106 Record<string, unknown>[],
107 | bigquery.QueryJobsError
108 | bigquery.GetQueryResultsJobsError
109 | QueryNotComplete
110 > =>
111 Effect.gen(function* () {
112 const ref = yield* table;
113 const first = yield* queryJobs({
114 projectId: ref.project,
115 body: {
116 query: sql,
117 useLegacySql: false,
118 location: ref.location,
119 defaultDataset: {
120 projectId: ref.project,
121 datasetId: ref.datasetId,
122 },
123 timeoutMs: QUERY_WAIT_MS,
124 formatOptions: { useInt64Timestamp: true },
125 ...(params === undefined
126 ? {}
127 : {
128 parameterMode: "NAMED",
129 queryParameters: Object.entries(params).map(([name, value]) =>
130 queryParameter(name, value),
131 ),
132 }),
133 },
134 });
135 const jobId = first.jobReference?.jobId;
136 const location = first.jobReference?.location ?? ref.location;
137 const results = (
138 pageToken: string | undefined,
139 ): Effect.Effect<
140 bigquery.GetQueryResultsResponse,
141 bigquery.GetQueryResultsJobsError | QueryNotComplete
142 > =>
143 jobId === undefined
144 ? Effect.fail(new QueryNotComplete({ jobId }))
145 : getQueryResults({
146 projectId: ref.project,
147 jobId,
148 location,
149 pageToken,
150 timeoutMs: QUERY_WAIT_MS,
151 "formatOptions.useInt64Timestamp": true,
152 });
153
154 const settled: Effect.Effect<
155 QueryPage,
156 bigquery.GetQueryResultsJobsError | QueryNotComplete
157 > =
158 first.jobComplete === true
159 ? Effect.succeed(first)
160 : results(undefined).pipe(
161 Effect.repeat({
162 until: (r: bigquery.GetQueryResultsResponse) =>
163 r.jobComplete === true,
164 times: QUERY_POLLS,
165 }),
166 );
167 const complete = yield* settled;
168 if (complete.jobComplete !== true) {
169 return yield* new QueryNotComplete({ jobId });
170 }
171 const schema = complete.schema;
172 const rest: ReadonlyArray<bigquery.TableRow> = yield* remainingRows(
173 results,
174 complete.pageToken,
175 );
176 return decodeRows(schema, [...(complete.rows ?? []), ...rest]);
177 }),
178 });
179
180 const makeWrite = (table: Effect.Effect<TableRef>): WriteTableClient => ({
181 insert: (rows, options) =>
182 Effect.gen(function* () {
183 if (rows.length === 0) return;
184 const ref = yield* table;
185 const response = yield* insertAll({
186 projectId: ref.project,
187 datasetId: ref.datasetId,
188 tableId: ref.tableId,
189 body: {
190 rows: rows.map((row, index) => ({
191 json: encodeRow(row),
192 insertId: options?.insertIds?.[index],
193 })),
194 },
195 });
196 const insertErrors = response.insertErrors ?? [];
197 if (insertErrors.length > 0) {
198 return yield* new InsertRowsFailed({
199 insertErrors: insertErrors.map((item) => ({
200 index: item.index,
201 errors: item.errors ?? [],
202 })),
203 });
204 }
205 }),
206 });
207
208 return { makeRead, makeWrite };
209});
210
211/** Build a table binding that grants `iam` and returns `makeClient`'s client. */
212export const makeBigQueryTableBinding = <Client>(options: {
213 tag: string;
214 iam: ReadonlyArray<BindingIam>;
215 makeClient: (
216 helpers: Effect.Success<typeof makeBigQueryTableHelpers>,
217 table: Effect.Effect<TableRef>,
218 ) => Client;
219}) =>
220 Effect.gen(function* () {
221 const helpers = yield* makeBigQueryTableHelpers;
222 return Effect.fn(function* (table: Table) {
223 yield* bindGcpHost({
224 tag: options.tag,
225 resource: table,
226 iam: options.iam.map((iam) => grantFor(iam, table.name)),
227 });
228 const project = yield* table.project;
229 const datasetId = yield* table.datasetId;
230 const tableId = yield* table.tableId;
231 const location = yield* table.location;
232 const ref = Effect.all({
233 project,
234 datasetId,
235 tableId,
236 location,
237 });
238 return options.makeClient(helpers, ref);
239 });
240 });