do the thing again

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)

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