fix: Send custom headers on every request, not just the initial POST

README.md

@@ -102,7 +102,7 @@ Attributes of opts [object] are:
 * info [boolean :optional]
   * fetch query info (execution statistics) for success callback, or not (default false)
 * headers [object :optional]
-  * additional headers to be included in the request, check the full list for [Trino](https://trino.io/docs/current/develop/client-protocol.html#client-request-headers) and [Presto](https://prestodb.io/docs/current/develop/client-protocol.html#client-request-headers) engines
+  * additional headers to be included in every request of the query (the initial POST as well as the follow-up `nextUri` polls), check the full list for [Trino](https://trino.io/docs/current/develop/client-protocol.html#client-request-headers) and [Presto](https://prestodb.io/docs/current/develop/client-protocol.html#client-request-headers) engines
 * authorization [string: optional]
   * AUTHORIZATION header value, which overrides the client's configuration per request
 * timeout [integer :optional]
@@ -211,6 +211,8 @@ npm run test
 
 ## Versions
 
+* Unreleased:
+  * fix "headers" option so custom headers are sent on every request of a query (the `nextUri` polls, cancel and info requests), not just the initial POST
 * 1.1.0:
   * add automatic retries for server errors
   * follow redirects if servers simply redirect client's request

index.spec.js

@@ -488,3 +488,83 @@ describe('redirect tests', function(){
     }, 10000);
   });
 });
+
+describe('custom headers', function(){
+  var responses = {
+    '/v1/statement': {
+      "stats": {
+        "state": "QUEUED",
+      },
+      "nextUri": "http://localhost:8111/v1/statement/20140120_032523_00000_32v8g/1",
+      "infoUri": "http://localhost:8111/v1/query/20140120_032523_00000_32v8g",
+      "id": "20140120_032523_00000_32v8g",
+    },
+    '/v1/statement/20140120_032523_00000_32v8g/1': {
+      "stats": {
+        "state": "FINISHED",
+      },
+      "columns": [ { "type": "integer", "name": "col" } ],
+      "data": [ [ 1 ] ],
+      "infoUri": "http://localhost:8111/v1/query/20140120_032523_00000_32v8g",
+      "id": "20140120_032523_00000_32v8g"
+    }
+  };
+
+  var server;
+  var received;
+
+  beforeAll(function(done) {
+    server = http.createServer(function(req, res){
+      received.push({ method: req.method, url: req.url, headers: req.headers });
+      res.statusCode = 200;
+      res.setHeader('Content-Type', 'application/json');
+      res.write(JSON.stringify(responses[req.url]));
+      res.end();
+    });
+    server.listen(8111, function(){
+      done();
+    });
+  });
+
+  beforeEach(function(){
+    received = [];
+  });
+
+  afterAll(function(done){
+    server.close(function(){
+      done();
+    });
+  });
+
+  test('forwards custom headers on every request, including nextUri polls', function(done){
+    expect.assertions(6);
+    var client = new Client({
+      host: 'localhost',
+      port: 8111,
+    });
+    client.execute({
+      query: 'SELECT 1 AS col',
+      headers: {
+        'X-Custom-Header': 'custom-value',
+        'Proxy-Authorization': 'Basic dGVzdA==',
+      },
+      callback: function(error){
+        expect(error).toBeNull();
+
+        var post = received.find(function(r){ return r.method === 'POST'; });
+        var poll = received.find(function(r){ return r.method === 'GET'; });
+
+        // Custom headers have always been sent on the initial POST.
+        expect(post.headers['x-custom-header']).toBe('custom-value');
+        expect(post.headers['proxy-authorization']).toBe('Basic dGVzdA==');
+
+        // The nextUri poll must carry them too — this is what regressed.
+        expect(poll).toBeDefined();
+        expect(poll.headers['x-custom-header']).toBe('custom-value');
+        expect(poll.headers['proxy-authorization']).toBe('Basic dGVzdA==');
+
+        done();
+      },
+    });
+  });
+});

lib/presto-client/index.js

@@ -64,7 +64,7 @@ var Client = exports.Client = function(args){
     }
 };
 
-Client.prototype.request = function(opts, callback, authorization) {
+Client.prototype.request = function(opts, callback, authorization, extraHeaders) {
     var client = this;
     var contentBody = null;
 
@@ -94,6 +94,12 @@ Client.prototype.request = function(opts, callback, authorization) {
     }
     var adapter = adapters[opts.protocol];
 
+    if (extraHeaders) {
+        for (var headerName in extraHeaders) {
+            opts.headers[headerName] = extraHeaders[headerName];
+        }
+    }
+
     if (opts.user)
         opts.headers[client.headers.USER] = opts.user;
 
@@ -363,7 +369,7 @@ Client.prototype.statementResource = function(opts) {
                 } else {
                     error_callback({message: "query fetch canceled by operation"});
                 }
-            }, opts.authorization);
+            }, opts.authorization, opts.headers);
             return;
         }
         current_req = client.request(uri_obj, function(error, code, response){
@@ -451,12 +457,12 @@ Client.prototype.statementResource = function(opts) {
             if (fetch_info && response.infoUri) {
                 client.request(response.infoUri, function(error, code, response){
                     success_callback(null, finishedStats, response);
-                }, opts.authorization);
+                }, opts.authorization, opts.headers);
             }
             else {
                 success_callback(null, finishedStats);
             }
-        }, opts.authorization);
+        }, opts.authorization, opts.headers);
     };
     fetch({ method: 'POST', path: '/v1/statement', headers: header, body: opts.query, user: opts.user });
 };