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 {
5 : TRINO_CHANNEL_READY,
6 : TRINO_CHANNEL_SETUP_CANCELLED,
7 : TRINO_CHANNEL_SETUP_FAILED,
8 : TRINO_CHANNEL_SETUP_STARTED,
9 : TrinoConnectorAction,
10 : OAUTH_STARTED,
11 : OAUTH_CANCELLED,
12 : OAUTH_FAILED,
13 : GENERATING_PKCE_CHALLENGE,
14 : GENERATED_PKCE_CHALLENGE,
15 : OAUTH_BROWSER_OPENED,
16 : RECEIVED_OAUTH_CODE,
17 : REQUESTING_ACCESS_TOKEN,
18 : RECEIVED_ACCESS_TOKEN,
19 : } from './trino_connection_state.js';
20 : import { Dispatch } from '../../utils/variant.js';
21 : import { LoggableException, Logger, stringifyError } from '../../platform/logger/logger.js';
22 : import { RESET_CONNECTION } from '../connection_state.js';
23 : import { requestTrinoAccessToken, TrinoApiClientInterface, TrinoApiEndpoint } from './trino_api_client.js';
24 : import { TrinoChannel, TrinoChannelInterface } from './trino_channel.js';
25 : import { trinoSessionDefaults } from './trino_catalogs.js';
26 : import { TrinoConnectorConfig } from '../connector_configs.js';
27 : import { generatePKCEChallenge } from '../../utils/pkce.js';
28 : import { PlatformType } from '../../platform/platform_type.js';
29 : import { HttpClient } from '../../platform/http/http_client.js';
30 : import { dateToTimestamp } from '../proto_helper.js';
31 :
32 1 : const LOG_CTX = "trino_setup";
33 :
34 : // Default OAuth callback URL for Trino
35 1 : const DEFAULT_OAUTH_CALLBACK_URL = "http://localhost:56512/Callback";
36 : // Default OAuth scopes if the connection does not configure any
37 1 : const DEFAULT_OAUTH_SCOPES = ["openid", "email", "profile", "offline_access"];
38 :
39 : /// Setup Trino connection with basic auth (username/password)
40 1 : async function setupTrinoConnectionBasic(
41 1 : modifyState: Dispatch<TrinoConnectorAction>,
42 1 : logger: Logger,
43 1 : params: connection.TrinoConnectionParams,
44 1 : client: TrinoApiClientInterface,
45 1 : abortSignal: AbortSignal
46 1 : ): Promise<TrinoChannelInterface> {
47 1 : let channel: TrinoChannelInterface;
48 1 : try {
49 : // Start the channel setup
50 1 : modifyState({
51 1 : type: TRINO_CHANNEL_SETUP_STARTED,
52 1 : value: params,
53 1 : });
54 1 : abortSignal.throwIfAborted();
55 :
56 : // Create the channel
57 1 : const auth = params.auth ?? ({
58 0 : authType: "AUTH_BASIC",
59 0 : basic: ({
60 0 : username: "",
61 0 : secret: "",
62 0 : })
63 0 : });
64 1 : const endpoint = new TrinoApiEndpoint(params.endpoint, auth);
65 1 : channel = new TrinoChannel(logger, client, endpoint, trinoSessionDefaults(params));
66 :
67 : // Mark the channel as ready
68 1 : modifyState({
69 1 : type: TRINO_CHANNEL_READY,
70 1 : value: channel,
71 1 : });
72 1 : abortSignal.throwIfAborted();
73 :
74 1 : } catch (error: any) {
75 1 : if (error.name === 'AbortError') {
76 1 : logger.warn("Cancelled setup", {}, LOG_CTX);
77 1 : modifyState({
78 1 : type: TRINO_CHANNEL_SETUP_CANCELLED,
79 1 : value: error.message,
80 1 : });
81 1 : } else {
82 0 : logger.warn("Setup failed", { "message": error.message, "details": error.data }, LOG_CTX);
83 0 : modifyState({
84 0 : type: TRINO_CHANNEL_SETUP_FAILED,
85 0 : value: error,
86 0 : });
87 0 : }
88 1 : throw error;
89 1 : }
90 :
91 1 : return channel;
92 1 : }
93 :
94 : /// Wait for the OAuth redirect to reach the native loopback server
95 1 : async function waitForOAuthCallback(bridge: DashQLElectronBridge, loopbackId: number, abortSignal: AbortSignal): Promise<DashQLOAuthLoopbackResult> {
96 1 : const onAbort = () => void bridge.oauthLoopback.close(loopbackId);
97 1 : abortSignal.addEventListener('abort', onAbort, { once: true });
98 1 : try {
99 1 : return await bridge.oauthLoopback.wait(loopbackId);
100 1 : } catch (e: any) {
101 0 : abortSignal.throwIfAborted();
102 0 : throw e;
103 1 : } finally {
104 1 : abortSignal.removeEventListener('abort', onAbort);
105 1 : }
106 1 : }
107 :
108 : /// Setup Trino connection with OAuth
109 1 : async function setupTrinoConnectionOAuth(
110 1 : modifyState: Dispatch<TrinoConnectorAction>,
111 1 : logger: Logger,
112 1 : params: connection.TrinoConnectionParams,
113 1 : _config: TrinoConnectorConfig,
114 1 : client: TrinoApiClientInterface,
115 1 : httpClient: HttpClient,
116 1 : forceReLogin: boolean,
117 1 : abortSignal: AbortSignal
118 1 : ): Promise<TrinoChannelInterface> {
119 1 : let channel: TrinoChannelInterface;
120 :
121 1 : if (params.auth?.authType != "AUTH_OAUTH" || !params.auth?.oauth) {
122 0 : throw new LoggableException("OAuth configuration is required for OAuth authentication", {}, LOG_CTX);
123 : }
124 1 : const oauthConfig = params.auth.oauth;
125 1 : let loopbackId: number | null = null;
126 :
127 1 : try {
128 : // Start OAuth flow
129 1 : logger.debug("Starting Trino OAuth flow", {}, LOG_CTX);
130 1 : modifyState({ type: OAUTH_STARTED, value: params });
131 1 : abortSignal.throwIfAborted();
132 :
133 : // Generate PKCE challenge
134 1 : logger.debug("Generating PKCE challenge", {}, LOG_CTX);
135 1 : modifyState({ type: GENERATING_PKCE_CHALLENGE, value: null });
136 1 : const pkceChallenge = await generatePKCEChallenge();
137 1 : abortSignal.throwIfAborted();
138 1 : modifyState({ type: GENERATED_PKCE_CHALLENGE, value: pkceChallenge });
139 :
140 : // The redirect is received by a loopback server of the desktop app
141 1 : const bridge = globalThis.dashqlElectron;
142 1 : if (!bridge) {
143 0 : throw new LoggableException("Trino OAuth requires the DashQL desktop app", {}, LOG_CTX);
144 : }
145 1 : const callbackUrl = oauthConfig.callbackUrl || DEFAULT_OAUTH_CALLBACK_URL;
146 1 : loopbackId = await bridge.oauthLoopback.listen(callbackUrl);
147 1 : abortSignal.throwIfAborted();
148 :
149 : // Build the authorization URL
150 1 : const state = crypto.randomUUID();
151 1 : const scopes = oauthConfig.scopes?.length ? oauthConfig.scopes : DEFAULT_OAUTH_SCOPES;
152 1 : const authURLParams = new URLSearchParams({
153 1 : client_id: oauthConfig.clientId,
154 1 : redirect_uri: callbackUrl,
155 1 : response_type: 'code',
156 1 : code_challenge: pkceChallenge.value,
157 1 : code_challenge_method: 'S256',
158 1 : scope: scopes.join(' '),
159 1 : state,
160 1 : });
161 1 : if (forceReLogin) {
162 : // Force the identity provider to re-authenticate instead of silently reusing an
163 : // existing session cookie in the system browser.
164 0 : authURLParams.set('prompt', 'login');
165 : }
166 1 : const authUrl = `${oauthConfig.authorizationUrl}?${authURLParams.toString()}`;
167 :
168 : // Open the browser for OAuth
169 1 : logger.debug("Opening OAuth URL", { url: authUrl }, LOG_CTX);
170 1 : await bridge.openExternal(authUrl);
171 1 : modifyState({
172 1 : type: OAUTH_BROWSER_OPENED,
173 1 : value: null,
174 1 : });
175 1 : abortSignal.throwIfAborted();
176 :
177 : // Wait for the OAuth callback
178 1 : const callbackData = await waitForOAuthCallback(bridge, loopbackId, abortSignal);
179 1 : abortSignal.throwIfAborted();
180 :
181 : // Check for errors
182 1 : if (callbackData.error) {
183 1 : throw new LoggableException(callbackData.errorDescription || callbackData.error, {}, LOG_CTX);
184 : }
185 1 : if (callbackData.state !== state) {
186 1 : throw new LoggableException("OAuth callback state does not match the authorization request", {}, LOG_CTX);
187 : }
188 1 : if (!callbackData.code) {
189 0 : throw new LoggableException("No authorization code received", {}, LOG_CTX);
190 : }
191 1 : const authCode = {
192 1 : token: callbackData.code,
193 1 : createdAt: dateToTimestamp(new Date())!
194 1 : };
195 :
196 1 : logger.debug("Received authorization code", {}, LOG_CTX);
197 1 : modifyState({
198 1 : type: RECEIVED_OAUTH_CODE,
199 1 : value: {
200 1 : token: authCode.token,
201 1 : createdAt: authCode.createdAt,
202 1 : },
203 1 : });
204 :
205 : // Exchange the code for an access token
206 1 : modifyState({
207 1 : type: REQUESTING_ACCESS_TOKEN,
208 1 : value: null,
209 1 : });
210 1 : logger.debug("Exchanging authorization code for token", { tokenEndpoint: oauthConfig.tokenUrl }, LOG_CTX);
211 1 : const accessToken = await requestTrinoAccessToken(httpClient, oauthConfig.tokenUrl, new URLSearchParams({
212 1 : grant_type: 'authorization_code',
213 1 : client_id: oauthConfig.clientId,
214 1 : code: callbackData.code,
215 1 : redirect_uri: callbackUrl,
216 1 : code_verifier: pkceChallenge.verifier,
217 1 : }));
218 1 : abortSignal.throwIfAborted();
219 :
220 1 : logger.debug("Received access token", {}, LOG_CTX);
221 1 : modifyState({
222 1 : type: RECEIVED_ACCESS_TOKEN,
223 1 : value: accessToken,
224 1 : });
225 :
226 : // Now setup the channel with the access token
227 1 : modifyState({
228 1 : type: TRINO_CHANNEL_SETUP_STARTED,
229 1 : value: params,
230 1 : });
231 1 : abortSignal.throwIfAborted();
232 :
233 : // Create the channel
234 1 : const endpoint = new TrinoApiEndpoint(params.endpoint, params.auth);
235 1 : endpoint.oauthState = ({
236 1 : oauthPkce: pkceChallenge,
237 1 : authCode,
238 1 : accessToken
239 1 : });
240 1 : channel = new TrinoChannel(logger, client, endpoint, trinoSessionDefaults(params));
241 :
242 : // Mark the channel as ready
243 1 : modifyState({
244 1 : type: TRINO_CHANNEL_READY,
245 1 : value: channel,
246 1 : });
247 1 : abortSignal.throwIfAborted();
248 :
249 1 : } catch (error: any) {
250 1 : if (error.name === 'AbortError') {
251 0 : logger.info("Cancelled OAuth flow", {}, LOG_CTX);
252 0 : modifyState({
253 0 : type: OAUTH_CANCELLED,
254 0 : value: {
255 0 : message: error.message || "OAuth flow was aborted",
256 0 : },
257 0 : });
258 1 : } else {
259 1 : logger.warn("Failed OAuth flow", { "error": stringifyError(error) }, LOG_CTX);
260 1 : modifyState({
261 1 : type: OAUTH_FAILED,
262 1 : value: {
263 1 : message: error.message || "OAuth flow failed",
264 1 : },
265 1 : });
266 1 : }
267 1 : throw error;
268 1 : } finally {
269 1 : if (loopbackId !== null) {
270 1 : void globalThis.dashqlElectron?.oauthLoopback.close(loopbackId).catch(() => { });
271 : }
272 1 : }
273 :
274 1 : return channel;
275 1 : }
276 :
277 : /// Main setup function that routes to the appropriate auth method
278 1 : export async function setupTrinoConnection(
279 1 : modifyState: Dispatch<TrinoConnectorAction>,
280 1 : logger: Logger,
281 1 : params: connection.TrinoConnectionParams,
282 1 : config: TrinoConnectorConfig,
283 1 : client: TrinoApiClientInterface,
284 1 : httpClient: HttpClient,
285 1 : _platformType: PlatformType,
286 1 : forceReLogin: boolean,
287 1 : abortSignal: AbortSignal
288 1 : ): Promise<TrinoChannelInterface> {
289 : // Determine auth type
290 1 : params = structuredClone(params);
291 1 : delete params.draftRows;
292 1 : const authType = params.auth?.authType ?? "AUTH_BASIC";
293 1 : switch (authType) {
294 1 : case "AUTH_OAUTH":
295 1 : return await setupTrinoConnectionOAuth(modifyState, logger, params, config, client, httpClient, forceReLogin, abortSignal);
296 : case "AUTH_BASIC":
297 1 : default:
298 1 : return await setupTrinoConnectionBasic(modifyState, logger, params, client, abortSignal);
299 1 : }
300 1 : }
301 :
302 : export interface TrinoSetupApi {
303 : setup(dispatch: Dispatch<TrinoConnectorAction>, params: connection.TrinoConnectionParams, abortSignal: AbortSignal): Promise<TrinoChannelInterface | null>
304 : reset(dispatch: Dispatch<TrinoConnectorAction>): Promise<void>
305 : }
306 :
307 1 : export function createTrinoSetup(
308 1 : trinoClient: TrinoApiClientInterface,
309 1 : config: TrinoConnectorConfig,
310 1 : logger: Logger,
311 1 : httpClient: HttpClient,
312 1 : platformType: PlatformType,
313 1 : forceReLogin: boolean
314 1 : ): (TrinoSetupApi | null) {
315 1 : const setup = async (modifyState: Dispatch<TrinoConnectorAction>, params: connection.TrinoConnectionParams, abort: AbortSignal) => {
316 1 : return await setupTrinoConnection(modifyState, logger, params, config, trinoClient, httpClient, platformType, forceReLogin, abort);
317 1 : };
318 1 : const reset = async (updateState: Dispatch<TrinoConnectorAction>) => {
319 1 : updateState({
320 1 : type: RESET_CONNECTION,
321 1 : value: null,
322 1 : })
323 1 : };
324 1 : return { setup, reset };
325 1 : };
|
