Line data Source code
1 : import * as connection from '@ankoh/dashql-jsonschema/connection.js';
2 : import * as auth from '@ankoh/dashql-jsonschema/auth.js';
3 :
4 : import { HttpClient, HttpFetchResult } from "../../platform/http/http_client.js";
5 : import { LoggableException, Logger } from "../../platform/logger/logger.js";
6 : import { dateToTimestamp } from '../proto_helper.js';
7 : import type { TrinoSessionDefaults } from './trino_catalogs.js';
8 :
9 1 : const LOG_CTX = "trino_api";
10 :
11 1 : export class TrinoApiEndpoint {
12 : // The endpoint url
13 1 : endpoint: string;
14 : // The auth params
15 1 : authParams: auth.TrinoAuthParams;
16 : // The auth state
17 1 : oauthState: connection.TrinoOAuthState | null;
18 : // The pending access token refresh
19 1 : oauthRefresh: Promise<void> | null;
20 :
21 1 : constructor(endpoint: string, params: auth.TrinoAuthParams) {
22 1 : this.endpoint = endpoint;
23 1 : this.authParams = params;
24 1 : this.oauthState = null;
25 1 : this.oauthRefresh = null;
26 1 : }
27 1 : }
28 :
29 : /// Request an access token from an OAuth token endpoint
30 1 : export async function requestTrinoAccessToken(httpClient: HttpClient, tokenUrl: string, body: URLSearchParams): Promise<connection.TrinoAccessToken> {
31 1 : const response = await httpClient.fetch(new URL(tokenUrl), {
32 1 : method: 'POST',
33 1 : headers: {
34 1 : 'Content-Type': 'application/x-www-form-urlencoded',
35 1 : },
36 1 : body: body.toString(),
37 1 : });
38 1 : if (response.status < 200 || response.status >= 300) {
39 0 : const errorText = await response.text();
40 0 : throw new Error(`Token request failed: ${response.status} ${errorText}`);
41 0 : }
42 1 : const tokenData = await response.json();
43 1 : const expiresAt = tokenData.expires_in
44 1 : ? new Date(Date.now() + (tokenData.expires_in * 1000))
45 0 : : null;
46 1 : return {
47 1 : createdAt: dateToTimestamp(new Date())!,
48 1 : accessToken: tokenData.access_token,
49 1 : tokenType: tokenData.token_type || 'Bearer',
50 1 : refreshToken: tokenData.refresh_token,
51 1 : expiresAt: dateToTimestamp(expiresAt),
52 1 : scope: tokenData.scope,
53 1 : };
54 1 : }
55 :
56 : /// Refresh access tokens a minute before they expire
57 1 : const ACCESS_TOKEN_REFRESH_MARGIN_MS = 60 * 1000;
58 :
59 : /// The size of the result pages that we ask Trino for.
60 : /// Trino returns pages of 1 MB by default and caps the size at 128 MB.
61 : /// A page also ends early when the coordinator has no more rows buffered, its exchange buffer holds 32 MB by default.
62 1 : export const TARGET_RESULT_SIZE = "32MB";
63 :
64 : export interface TrinoQueryFailureInfo {
65 : /// The failure type
66 : type: string;
67 : /// The failure message
68 : message: string;
69 : /// The suppressed error
70 : suppressed: string[];
71 : /// The error stack
72 : stack: string[];
73 : };
74 :
75 : export interface TrinoQueryError {
76 : /// The erro message
77 : message: string;
78 : /// The error code
79 : errorCode: number;
80 : /// The error name
81 : errorName: string;
82 : /// The error type
83 : errorType: string;
84 : /// The failure info
85 : failureInfo: TrinoQueryFailureInfo;
86 : };
87 :
88 : export interface TrinoQueryResultColumn {
89 : /// The column name
90 : name: string;
91 : /// The column type
92 : type: string
93 : };
94 :
95 : export type TrinoQueryData = any[];
96 :
97 : export interface TrinoQueryStage {
98 : /// The trino stage id
99 : stageId: string;
100 : /// The state
101 : state: string;
102 : /// Is the query done?
103 : done: boolean;
104 : /// The number of nodes in the stage
105 : nodes: number;
106 : /// The total number of splits
107 : totalSplits: number;
108 : /// The number of queued splits
109 : queuedSplits: number;
110 : /// The splits that are currently running
111 : runningSplits: number;
112 : /// The splits that are completed
113 : completedSplits: number;
114 : /// The cpu time in milliseconds
115 : cpuTimeMillis: number;
116 : /// The wall time in milliseconds
117 : wallTimeMillis: number;
118 : /// The processed rows
119 : processedRows: number;
120 : /// The processed bytes
121 : processedBytes: number;
122 : /// The physical input bytes
123 : physicalInputBytes: number;
124 : /// The number of failed tasks
125 : failedTasks: number;
126 : /// Coordinator-only query?
127 : coordinatorOnly: boolean;
128 : /// The number of substages
129 : subStages: TrinoQueryStage[];
130 : };
131 :
132 : export interface TrinoQueryStatistics {
133 : /// The state
134 : state: string;
135 : /// Is the query queued?
136 : queued: boolean;
137 : /// Is the query scheduled?
138 : scheduled: boolean;
139 : /// The number of nodes
140 : nodes: number;
141 : /// The total number of splits
142 : totalSplits: number;
143 : /// The queued number of splits
144 : queuedSplits: number;
145 : /// The currently running splits
146 : runningSplits: number;
147 : /// The number of completed splits
148 : completedSplits: number;
149 : /// The CPU time in milliseconds
150 : cpuTimeMillis: number;
151 : /// The wallclock time in milliseconds
152 : wallTimeMillis: number;
153 : /// The time that the query was queued
154 : queuedTimeMillis: number;
155 : /// The elapsed time
156 : elapsedTimeMillis: number;
157 : /// The number of processed rows
158 : processedRows: number;
159 : /// The number of processed bytes
160 : processedBytes: number;
161 : /// The number of input bytes
162 : physicalInputBytes: number;
163 : /// The peak memory size
164 : peakMemoryBytes: number;
165 : /// The number of spilled bytes
166 : spilledBytes: number;
167 : /// The root stage
168 : rootStage: TrinoQueryStage;
169 : /// The progess percentage
170 : progressPercentage: number;
171 : };
172 :
173 : export interface TrinoQueryResult {
174 : /// The query id
175 : id: string;
176 : /// The URI for the query info call
177 : infoUri?: string;
178 : /// The URI for fetching the next query batch
179 : nextUri?: string;
180 : /// The result columns
181 : columns?: TrinoQueryResultColumn[];
182 : /// The query data
183 : data?: TrinoQueryData[];
184 : /// The query statistics
185 : stats?: TrinoQueryStatistics;
186 : /// The warnings during the query execution
187 : warnings?: string[];
188 : /// The query error (if any)
189 : error?: TrinoQueryError;
190 : };
191 :
192 : export interface TrinoQueryInfo {
193 : /// The query id
194 : queryId: string;
195 : /// The query state
196 : state: string;
197 : /// The query text
198 : query: string;
199 : /// The failure info
200 : failureInfo?: TrinoQueryFailureInfo;
201 : /// Did Trino collect the final infos of all tasks?
202 : finalQueryInfo?: boolean;
203 : /// The stages, a flat list
204 : stages?: { stages?: TrinoStageInfo[] };
205 : };
206 :
207 : /// A stage of a QueryInfo, the progress model reads the rest
208 : export interface TrinoStageInfo {
209 : /// The stage id, `<queryId>.<fragmentId>`
210 : stageId: string;
211 : /// The stage state
212 : state: string;
213 : /// The statistics, with the operator summaries once the tasks of the stage are done
214 : stageStats?: { operatorSummaries?: object[] };
215 : }
216 :
217 : export interface TrinoApiClientInterface {
218 : /// Run a query
219 : runQuery(endpoint: TrinoApiEndpoint, defaults: TrinoSessionDefaults, text: string, abort?: AbortSignal): Promise<TrinoQueryResult>;
220 : /// Get a query result
221 : getQueryResult(endpoint: TrinoApiEndpoint, nextUri: string, abort?: AbortSignal): Promise<TrinoQueryResult>;
222 : /// Get the QueryInfo of a query: the stages with their plans and statistics
223 : getQueryInfo(endpoint: TrinoApiEndpoint, queryId: string, abort?: AbortSignal): Promise<TrinoQueryInfo>;
224 : /// Cancel a query by its next URI
225 : cancelQuery(endpoint: TrinoApiEndpoint, nextUri: string): Promise<void>;
226 : }
227 :
228 :
229 1 : export class TrinoApiClient implements TrinoApiClientInterface {
230 : /// The logger
231 1 : logger: Logger;
232 : /// The http client
233 1 : httpClient: HttpClient;
234 : /// The retries of failed requests of query results
235 1 : retryPolicy: TrinoRetryPolicy;
236 :
237 1 : constructor(logger: Logger, httpClient: HttpClient, retryPolicy: TrinoRetryPolicy = DEFAULT_TRINO_RETRY_POLICY) {
238 1 : this.logger = logger;
239 1 : this.httpClient = httpClient;
240 1 : this.retryPolicy = retryPolicy;
241 1 : }
242 :
243 : /// Refresh the OAuth access token if it is about to expire
244 1 : async refreshAccessToken(endpoint: TrinoApiEndpoint): Promise<void> {
245 1 : const oauth = endpoint.authParams.oauth;
246 1 : const token = endpoint.oauthState?.accessToken;
247 1 : if (endpoint.authParams.authType != "AUTH_OAUTH" || !oauth || !token?.refreshToken || !token.expiresAt) {
248 1 : return;
249 : }
250 1 : if (new Date(token.expiresAt).getTime() - Date.now() > ACCESS_TOKEN_REFRESH_MARGIN_MS) {
251 0 : return;
252 : }
253 1 : if (!endpoint.oauthRefresh) {
254 1 : this.logger.debug("Refreshing access token", {}, LOG_CTX);
255 1 : endpoint.oauthRefresh = requestTrinoAccessToken(this.httpClient, oauth.tokenUrl, new URLSearchParams({
256 1 : grant_type: 'refresh_token',
257 1 : client_id: oauth.clientId,
258 1 : refresh_token: token.refreshToken,
259 1 : })).then(refreshed => {
260 1 : endpoint.oauthState!.accessToken = { ...refreshed, refreshToken: refreshed.refreshToken ?? token.refreshToken };
261 1 : }).finally(() => {
262 1 : endpoint.oauthRefresh = null;
263 1 : });
264 1 : }
265 1 : await endpoint.oauthRefresh;
266 1 : }
267 :
268 : /// Build the request headers for an endpoint
269 1 : async getHeaders(endpoint: TrinoApiEndpoint): Promise<Headers> {
270 1 : await this.refreshAccessToken(endpoint);
271 1 : const headers = new Headers();
272 1 : addAuthHeaders(headers, endpoint);
273 1 : return headers;
274 1 : }
275 :
276 : /// Run a query
277 1 : async runQuery(endpoint: TrinoApiEndpoint, defaults: TrinoSessionDefaults, text: string, abort?: AbortSignal): Promise<TrinoQueryResult> {
278 1 : this.logger.debug("Running query", { "text": text }, LOG_CTX);
279 1 : const url = new URL(`${endpoint.endpoint}/v1/statement`);
280 1 : const headers = await this.getHeaders(endpoint);
281 : // Trino sees every catalog of the server, the defaults only resolve the names that statements leave out
282 1 : if (defaults.catalog != null) {
283 1 : headers.set('X-Trino-Catalog', defaults.catalog);
284 1 : if (defaults.schema != null) {
285 1 : headers.set('X-Trino-Schema', defaults.schema);
286 : }
287 1 : }
288 1 : const rawResponse = await this.httpClient.fetch(url, {
289 1 : method: 'POST',
290 1 : body: text,
291 1 : headers,
292 1 : signal: abort,
293 1 : });
294 1 : if (rawResponse.status != 200) {
295 1 : throw new LoggableException(`Query failed`, {
296 1 : status: rawResponse.status.toString(),
297 1 : message: rawResponse.statusText.toString()
298 1 : }, LOG_CTX);
299 : }
300 1 : const responseJson = await rawResponse.json() as TrinoQueryResult;
301 : const response = responseJson as TrinoQueryResult;
302 1 : return response;
303 1 : }
304 :
305 : /// Get the query result batch.
306 : /// Trino answers a repeated request of a next URI with the same page, so failed requests are retried.
307 1 : async getQueryResult(endpoint: TrinoApiEndpoint, nextUri: string, abort?: AbortSignal): Promise<TrinoQueryResult> {
308 1 : this.logger.debug("Getting query results", { "nextUri": nextUri }, LOG_CTX);
309 1 : const url = new URL(nextUri);
310 : // Only the URIs of executing queries return rows
311 1 : if (url.pathname.includes('/v1/statement/executing/')) {
312 1 : url.searchParams.set('targetResultSize', TARGET_RESULT_SIZE);
313 : }
314 1 : for (let attempt = 1; ; ++attempt) {
315 1 : const retriesLeft = attempt < this.retryPolicy.attempts;
316 1 : let rawResponse: HttpFetchResult;
317 1 : try {
318 1 : const headers = await this.getHeaders(endpoint);
319 1 : rawResponse = await this.httpClient.fetch(url, {
320 1 : method: 'GET',
321 1 : headers,
322 1 : signal: abort,
323 1 : });
324 1 : } catch (e: any) {
325 : // Retry requests that failed in the network
326 1 : if (!retriesLeft || abort?.aborted) throw e;
327 1 : this.logger.warn("Retrying a failed request of query results", { "attempt": attempt.toString(), "error": String(e) }, LOG_CTX);
328 1 : await sleepUnlessAborted(retryDelayMs(this.retryPolicy, attempt), abort);
329 1 : continue;
330 1 : }
331 1 : if (rawResponse.status == 200) {
332 1 : return parseTrinoJson(await rawResponse.text()) as TrinoQueryResult;
333 : }
334 : // Retry responses of a busy or restarting coordinator or gateway
335 1 : if (retriesLeft && RETRIED_STATUS_CODES.has(rawResponse.status)) {
336 1 : this.logger.warn("Retrying a rejected request of query results", { "attempt": attempt.toString(), "status": rawResponse.status.toString() }, LOG_CTX);
337 1 : await sleepUnlessAborted(retryDelayMs(this.retryPolicy, attempt), abort);
338 1 : continue;
339 1 : }
340 1 : throw new LoggableException(`Failed to fetch query results`, {
341 1 : status: rawResponse.status.toString(),
342 1 : message: rawResponse.statusText.toString()
343 1 : }, LOG_CTX);
344 1 : }
345 1 : }
346 :
347 : /// Get a query info
348 1 : async getQueryInfo(endpoint: TrinoApiEndpoint, queryId: string, abort?: AbortSignal): Promise<TrinoQueryInfo> {
349 1 : this.logger.debug("Getting query info", { "query": queryId }, LOG_CTX);
350 1 : const url = new URL(`${endpoint.endpoint}/v1/query/${encodeURIComponent(queryId)}`);
351 1 : const headers = await this.getHeaders(endpoint);
352 1 : const rawResponse = await this.httpClient.fetch(url, {
353 1 : method: 'GET',
354 1 : headers,
355 1 : signal: abort,
356 1 : });
357 1 : if (rawResponse.status != 200) {
358 1 : throw new LoggableException(`Failed to fetch query info`, {
359 1 : status: rawResponse.status.toString(),
360 1 : message: rawResponse.statusText.toString()
361 1 : }, LOG_CTX);
362 : }
363 1 : const responseJson = await rawResponse.json() as TrinoQueryInfo;
364 : const response = responseJson as TrinoQueryInfo;
365 1 : return response;
366 1 : }
367 :
368 : /// Cancel a query by deleting its next URI, Trino answers with 204 No Content
369 1 : async cancelQuery(endpoint: TrinoApiEndpoint, nextUri: string): Promise<void> {
370 1 : this.logger.debug("Cancelling query", { "nextUri": nextUri }, LOG_CTX);
371 1 : const url = new URL(nextUri);
372 1 : const headers = await this.getHeaders(endpoint);
373 1 : const rawResponse = await this.httpClient.fetch(url, {
374 1 : method: 'DELETE',
375 1 : headers
376 1 : });
377 1 : if (rawResponse.status < 200 || rawResponse.status >= 300) {
378 1 : throw new LoggableException(`Cancelling a query failed`, {
379 1 : status: rawResponse.status.toString(),
380 1 : message: rawResponse.statusText.toString()
381 1 : }, LOG_CTX);
382 : }
383 1 : }
384 1 : }
385 :
386 : /// The status codes of requests that are retried: too many requests, bad gateway, service unavailable, gateway timeout
387 1 : const RETRIED_STATUS_CODES = new Set([429, 502, 503, 504]);
388 :
389 : /// The retries of failed requests of query results
390 : export interface TrinoRetryPolicy {
391 : /// The number of attempts, including the first one
392 : attempts: number;
393 : /// The delay before the first retry, it doubles with every retry
394 : initialDelayMs: number;
395 : /// The maximum delay between retries
396 : maxDelayMs: number;
397 : }
398 :
399 : /// Retry for about 20 seconds before failing a query
400 1 : export const DEFAULT_TRINO_RETRY_POLICY: TrinoRetryPolicy = {
401 1 : attempts: 10,
402 1 : initialDelayMs: 100,
403 1 : maxDelayMs: 3000,
404 1 : };
405 :
406 : /// The delay after the failed attempt
407 1 : function retryDelayMs(policy: TrinoRetryPolicy, attempt: number): number {
408 1 : return Math.min(policy.initialDelayMs * 2 ** (attempt - 1), policy.maxDelayMs);
409 1 : }
410 :
411 : /// Sleep, an abort ends the sleep with the abort reason
412 1 : async function sleepUnlessAborted(ms: number, abort?: AbortSignal): Promise<void> {
413 1 : abort?.throwIfAborted();
414 1 : await new Promise<void>((resolve, reject) => {
415 1 : const onAbort = () => {
416 0 : clearTimeout(timer);
417 0 : reject(abort!.reason);
418 0 : };
419 1 : const timer = setTimeout(() => {
420 1 : abort?.removeEventListener('abort', onAbort);
421 1 : resolve();
422 1 : }, ms);
423 1 : abort?.addEventListener('abort', onAbort, { once: true });
424 1 : });
425 1 : }
426 :
427 : /// Integers with 16 or more digits may exceed the safe integers of JavaScript numbers
428 1 : const LONG_INTEGER_PATTERN = /\d{16}/;
429 : /// The source text of an integer
430 1 : const INTEGER_TEXT_PATTERN = /^-?\d+$/;
431 :
432 : /// Parse a JSON response of Trino.
433 : /// Trino writes bigint values as JSON numbers, integers beyond the safe integers are parsed from their source text as BigInt.
434 : /// Only texts with a long digit run pay for the reviver.
435 1 : export function parseTrinoJson(text: string): any {
436 1 : if (!LONG_INTEGER_PATTERN.test(text)) {
437 1 : return JSON.parse(text);
438 : }
439 1 : return JSON.parse(text, (_key: string, value: any, context?: { source?: string }) => {
440 1 : if (typeof value === 'number' && !Number.isSafeInteger(value) && context?.source != null && INTEGER_TEXT_PATTERN.test(context.source)) {
441 1 : return BigInt(context.source);
442 : }
443 1 : return value;
444 1 : });
445 1 : }
446 :
447 : /// Helper to add auth headers based on auth type
448 1 : function addAuthHeaders(headers: Headers, endpoint: TrinoApiEndpoint): void {
449 1 : const authParams = endpoint.authParams;
450 1 : const authType = authParams.authType ?? "AUTH_BASIC";
451 1 : switch (authType) {
452 1 : case "AUTH_OAUTH":
453 : // OAuth: Use Bearer token
454 1 : if (endpoint.oauthState?.accessToken?.accessToken) {
455 1 : headers.set('Authorization', `Bearer ${endpoint.oauthState?.accessToken?.accessToken}`);
456 : }
457 : // XXX X-Trino-User from user hint?
458 1 : break;
459 :
460 : case "AUTH_BASIC":
461 1 : default:
462 : // Basic auth: Use username:password
463 1 : if (authParams.basic?.username && authParams.basic.username.length > 0) {
464 1 : headers.set('Authorization', 'Basic ' + btoa(authParams.basic.username + ":" + (authParams.basic.secret ?? "")));
465 1 : headers.set('X-Trino-User', authParams.basic.username);
466 1 : }
467 : break;
468 1 : }
469 1 : }
|
