Cache OAuth tokens in memory by default to avoid re-running OAuth flo…

…w on every request

Signed-off-by: Levko Kravets <[email protected]>

lib/connection/auth/DatabricksOAuth/OAuthPersistence.ts

@@ -5,3 +5,15 @@ export default interface OAuthPersistence {
 
   read(host: string): Promise<OAuthToken | undefined>;
 }
+
+export class OAuthPersistenceCache implements OAuthPersistence {
+  private tokens: Record<string, OAuthToken | undefined> = {};
+
+  async persist(host: string, token: OAuthToken) {
+    this.tokens[host] = token;
+  }
+
+  async read(host: string) {
+    return this.tokens[host];
+  }
+}

lib/connection/auth/DatabricksOAuth/index.ts

@@ -1,7 +1,7 @@
 import { HttpHeaders } from 'thrift';
 import IAuthentication from '../../contracts/IAuthentication';
 import IDBSQLLogger from '../../../contracts/IDBSQLLogger';
-import OAuthPersistence from './OAuthPersistence';
+import OAuthPersistence, { OAuthPersistenceCache } from './OAuthPersistence';
 import OAuthManager, { OAuthManagerOptions } from './OAuthManager';
 import { OAuthScopes, defaultOAuthScopes } from './OAuthScope';
 
@@ -19,22 +19,26 @@ export default class DatabricksOAuth implements IAuthentication {
 
   private readonly manager: OAuthManager;
 
+  private readonly defaultPersistence = new OAuthPersistenceCache();
+
   constructor(options: DatabricksOAuthOptions) {
     this.options = options;
     this.logger = options.logger;
     this.manager = OAuthManager.getManager(this.options);
   }
 
   public async authenticate(): Promise<HttpHeaders> {
-    const { host, scopes, headers, persistence } = this.options;
+    const { host, scopes, headers } = this.options;
+
+    const persistence = this.options.persistence ?? this.defaultPersistence;
 
-    let token = await persistence?.read(host);
+    let token = await persistence.read(host);
     if (!token) {
       token = await this.manager.getToken(scopes ?? defaultOAuthScopes);
     }
 
     token = await this.manager.refreshAccessToken(token);
-    await persistence?.persist(host, token);
+    await persistence.persist(host, token);
 
     return {
       ...headers,