diff --git a/export/httpQuery.ts b/export/httpQuery.ts index a7b56ee..6fa9e90 100644 --- a/export/httpQuery.ts +++ b/export/httpQuery.ts @@ -48,6 +48,10 @@ interface HTTPQueryOptions { rows: any, opts: any, ) => void; + + // JWT auth token to be passed as the Bearer token in the Authorization + // header + authToken?: string | (() => string); } interface HTTPTransactionOptions extends HTTPQueryOptions { @@ -124,6 +128,7 @@ export function neon( deferrable: neonOptDeferrable, queryCallback, resultCallback, + authToken, }: HTTPTransactionOptions = {}, ) { // check the connection string @@ -218,7 +223,9 @@ export function neon( const url = typeof fetchEndpoint === 'function' - ? fetchEndpoint(hostname, port) + ? fetchEndpoint(hostname, port, { + jwtAuth: Boolean(authToken), + }) : fetchEndpoint; const bodyData = Array.isArray(parameterizedQuery) @@ -272,6 +279,12 @@ export function neon( 'Neon-Array-Mode': 'true', // this saves data and post-processing even if we return objects, not arrays }; + if (typeof authToken === 'string') { + headers['Authorization'] = `Bearer ${authToken}`; + } else if (typeof authToken === 'function') { + headers['Authorization'] = `Bearer ${authToken()}`; + } + if (Array.isArray(parameterizedQuery)) { // only send these headers for batch queries, where they matter if (resolvedIsolationLevel !== undefined) diff --git a/shims/net/index.ts b/shims/net/index.ts index ca65810..90cb8a3 100644 --- a/shims/net/index.ts +++ b/shims/net/index.ts @@ -52,10 +52,20 @@ export function isIP(input: string) { return 0; } +interface FetchEndpointOptions { + jwtAuth?: boolean; +} + export interface SocketDefaults { // these options relate to the fetch transport and take effect *only* when set globally poolQueryViaFetch: boolean; - fetchEndpoint: string | ((host: string, port: number | string) => string); + fetchEndpoint: + | string + | (( + host: string, + port: number | string, + options?: FetchEndpointOptions, + ) => string); fetchConnectionCache: boolean; fetchFunction: any; // these options relate to the WebSocket transport @@ -77,15 +87,23 @@ type GlobalOnlyDefaults = | 'fetchConnectionCache' | 'fetchFunction'; -const transformHost = (host: string): string => { - return host.replace(/^[^.]+\./, 'api.'); -}; - export class Socket extends EventEmitter { static defaults: SocketDefaults = { // these options relate to the fetch transport and take effect *only* when set globally poolQueryViaFetch: false, - fetchEndpoint: (host) => 'https://' + transformHost(host) + '/sql', + fetchEndpoint: (host, _port, options) => { + let newHost; + if (options?.jwtAuth) { + // If the caller sends in a JWT, we need to use the Neon Authorize API + // endpoint instead (this goes to the Auth Broker instead of the Neon + // Proxy). + newHost = host.replace(/^[^.]+\./, 'apiauth.'); + } else { + newHost = host.replace(/^[^.]+\./, 'api.'); + } + + return 'https://' + newHost + '/sql'; + }, fetchConnectionCache: true, fetchFunction: undefined, // these options relate to the WebSocket transport