From 168c7f2c02aab2caa3ac595c2dd136e2ed7609a0 Mon Sep 17 00:00:00 2001
From: z0al <ahmed.ali@contentful.com>
Date: Wed, 14 Apr 2021 11:59:24 +0200
Subject: [PATCH] feat: add custom headers support (#698)

---
 README.md                        |  1 +
 src/bin/lib/config.ts            | 50 ++++++++++++++++++++++++++++++--
 src/bin/usage-params.ts          |  5 ++++
 test/unit/bin/lib/config.spec.ts | 10 +++++++
 4 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index b7f807b38..f0608a7fc 100644
--- a/README.md
+++ b/README.md
@@ -186,6 +186,7 @@ module.exports = function (migration, context) {
 | accessToken       |            | string  | The access token to use                                     | true     |
 | yes               | false      | boolean | Skips any confirmation before applying the migration,script | false    |
 | requestBatchSize  | 100        | number  | Limit for every single request                              | false    |
+| headers           |            | object  | Additional headers to attach to the requests                | false    |
 
 ### Chaining vs Object notation
 
diff --git a/src/bin/lib/config.ts b/src/bin/lib/config.ts
index 45755f572..5707b0ecd 100644
--- a/src/bin/lib/config.ts
+++ b/src/bin/lib/config.ts
@@ -12,6 +12,7 @@ interface ClientConfig {
   proxy?: string,
   rawProxy?: boolean
   requestBatchSize?: number
+  headers?: Record<string, unknown>
 }
 
 function getFileConfig (): ClientConfig {
@@ -32,14 +33,15 @@ function getEnvConfig (): ClientConfig {
     {}
 }
 
-function getArgvConfig ({ spaceId, environmentId = 'master', accessToken, proxy, rawProxy, requestBatchSize }): ClientConfig {
+function getArgvConfig ({ spaceId, environmentId = 'master', accessToken, proxy, rawProxy, requestBatchSize, headers }): ClientConfig {
   const config = {
     spaceId,
     environmentId,
     accessToken,
     proxy,
     rawProxy,
-    requestBatchSize
+    requestBatchSize,
+    headers
   }
 
   if (!config.accessToken) {
@@ -49,10 +51,52 @@ function getArgvConfig ({ spaceId, environmentId = 'master', accessToken, proxy,
   return config
 }
 
+/**
+ * Turn header option into an object. Invalid header values
+ * are ignored.
+ *
+ * @example
+ * getHeadersConfig('Accept: Any')
+ * // -> {Accept: 'Any'}
+ *
+ * @example
+ * getHeadersConfig(['Accept: Any', 'X-Version: 1'])
+ * // -> {Accept: 'Any', 'X-Version': '1'}
+ */
+function getHeadersConfig (value?: string | string[]) {
+  if (!value) {
+    return {}
+  }
+
+  const values = Array.isArray(value) ? value : [value]
+
+  return values.reduce((headers, value) => {
+    value = value.trim()
+
+    const separatorIndex = value.indexOf(':')
+
+    // Invalid header format
+    if (separatorIndex === -1) {
+      return headers
+    }
+
+    const headerKey = value.slice(0, separatorIndex).trim()
+    const headerValue = value.slice(separatorIndex + 1).trim()
+
+    return {
+      ...headers,
+      [headerKey]: headerValue
+    }
+  }, {})
+}
+
 function getConfig (argv) {
   const fileConfig = getFileConfig()
   const envConfig = getEnvConfig()
-  const argvConfig = getArgvConfig(argv || {})
+  const argvConfig = getArgvConfig({
+    ...argv,
+    headers: argv?.headers || getHeadersConfig(argv?.header)
+  })
 
   return Object.assign(fileConfig, envConfig, argvConfig)
 }
diff --git a/src/bin/usage-params.ts b/src/bin/usage-params.ts
index dcf32c926..729420f80 100644
--- a/src/bin/usage-params.ts
+++ b/src/bin/usage-params.ts
@@ -56,6 +56,11 @@ export default yargs
     type: 'number',
     default: 100
   })
+  .option('header', {
+    alias: 'H',
+    type: 'string',
+    describe: 'Pass an additional HTTP Header'
+  })
   .demandOption(['space-id'], 'Please provide a space ID')
   .help('h')
   .alias('h', 'help')
diff --git a/test/unit/bin/lib/config.spec.ts b/test/unit/bin/lib/config.spec.ts
index 18de9ca63..1d487b446 100644
--- a/test/unit/bin/lib/config.spec.ts
+++ b/test/unit/bin/lib/config.spec.ts
@@ -41,4 +41,14 @@ describe('Config', function () {
     const config = getConfig({ requestBatchSize: 99 })
     expect(config.requestBatchSize).to.eql(99)
   })
+
+  it('exposes headers from argv', function () {
+    const config = getConfig({ headers: { test: true } })
+    expect(config.headers).to.eql({ test: true })
+  })
+
+  it('parses argv.header if provided', function () {
+    const config = getConfig({ header: ['Accept   : application/json ', ' X-Header: 1'] })
+    expect(config.headers).to.eql({ Accept: 'application/json', 'X-Header': '1' })
+  })
 })