From f9efb3bc6d4b987ffe5336cb82e881a83b54f72e Mon Sep 17 00:00:00 2001
From: Alonso Vega <lavb91@gmail.com>
Date: Thu, 18 Jan 2024 16:13:16 -0600
Subject: [PATCH] feat(client): allow basic and custom auth, pass extra headers

---
 presto-client/src/client.ts       | 35 ++++++++++++++++++++++++++++---
 presto-client/src/client.types.ts |  6 ++++++
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/presto-client/src/client.ts b/presto-client/src/client.ts
index ef8071b..890f906 100644
--- a/presto-client/src/client.ts
+++ b/presto-client/src/client.ts
@@ -15,7 +15,12 @@ export class PrestoClient {
   /**
    * Creates an instance of PrestoClient.
    * @param {PrestoClientConfig} config - Configuration object for the PrestoClient.
+   * @param {Object} config.basicAuthorization - Optional object for basic authorization.
+   * @param {Object} config.basicAuthorization.user - The basic auth user name.
+   * @param {Object} config.basicAuthorization.password - The basic auth password.
+   * @param {string} config.authorizationToken - An optional token to be sent in the authorization header. Takes precedence over the basic auth.
    * @param {string} config.catalog - The default catalog to be used.
+   * @param {Record<string, string>} config.extraHeaders - Any extra headers to include in the API requests. Optional.
    * @param {string} config.host - The host address of the Presto server.
    * @param {number} config.interval - The polling interval in milliseconds for query status checks.
    * @param {number} config.port - The port number on which the Presto server is listening.
@@ -24,7 +29,19 @@ export class PrestoClient {
    * @param {string} [config.timezone] - The timezone to be used for the session. Optional.
    * @param {string} config.user - The username to be used for the Presto session.
    */
-  constructor({ catalog, host, interval, port, schema, source, timezone, user }: PrestoClientConfig) {
+  constructor({
+    basicAuthorization,
+    authorizationToken,
+    catalog,
+    extraHeaders,
+    host,
+    interval,
+    port,
+    schema,
+    source,
+    timezone,
+    user,
+  }: PrestoClientConfig) {
     this.baseUrl = `${host || 'http://localhost'}:${port || 8080}/v1/statement`
     this.catalog = catalog
     this.interval = interval
@@ -46,7 +63,19 @@ export class PrestoClient {
       this.headers['X-Presto-Time-Zone'] = this.timezone
     }
 
-    // TODO: Set up auth
+    if (authorizationToken) {
+      this.headers['Authorization'] = authorizationToken
+    } else if (basicAuthorization) {
+      // Note this is only available for Node.js
+      this.headers['Authorization'] = `Bearer ${Buffer.from(
+        `${basicAuthorization.user}:${basicAuthorization.password}`,
+      ).toString('base64')}`
+    }
+
+    this.headers = {
+      ...extraHeaders,
+      ...this.headers,
+    }
   }
 
   /**
@@ -197,7 +226,7 @@ export class PrestoClient {
     const data = []
 
     do {
-      const response = await this.request({ method: 'GET', url: nextUri })
+      const response = await this.request({ headers, method: 'GET', url: nextUri })
 
       // Server is overloaded, wait a bit
       if (response.status === 503) {
diff --git a/presto-client/src/client.types.ts b/presto-client/src/client.types.ts
index c13945d..6931d14 100644
--- a/presto-client/src/client.types.ts
+++ b/presto-client/src/client.types.ts
@@ -1,5 +1,11 @@
 export interface PrestoClientConfig {
+  authorizationToken?: string
+  basicAuthorization?: {
+    user: string
+    password: string
+  }
   catalog?: string
+  extraHeaders?: Record<string, string>
   host?: string
   interval?: number
   port?: number