From c3197beaf45dbe8db501496424d46ad99d4fc88d Mon Sep 17 00:00:00 2001
From: Corey Zuares <251062+czuares@users.noreply.github.com>
Date: Fri, 13 Jan 2023 03:59:44 -0500
Subject: [PATCH] Extra headers for metadata API (#9321)

* Configurable Extra Authorization Header

- Adds two new header that can be set via env-var
  (`OMETA_HEADER_EXTRA_AUTH_NAME` and `OMETA_HEADER_EXTRA_AUTH_VALUE`)
- By default, it copies the Authorization header value to the extra auth header
- If a value is specified in the corresponding env-var, that value is
  used instead of the existing token

The primary use case is for Google Cloud's Identity Aware Proxy (IAP) which allows secure access to private services but removes the Authorization header in the process. The `Proxy-Authorization` header gets forwarded along to the service as the `Authorization` header for the internal service.

* Use openMetadataConnection schema instead of env vars

* Disable too-many-locals for _request function

* Fix dict merge
---
 ingestion/src/metadata/ingestion/ometa/client.py    | 13 +++++++++++++
 ingestion/src/metadata/ingestion/ometa/ometa_api.py |  2 ++
 .../metadata/openMetadataConnection.json            | 11 +++++++++++
 3 files changed, 26 insertions(+)

diff --git a/ingestion/src/metadata/ingestion/ometa/client.py b/ingestion/src/metadata/ingestion/ometa/client.py
index e3f233e25fa..43406c3065a 100644
--- a/ingestion/src/metadata/ingestion/ometa/client.py
+++ b/ingestion/src/metadata/ingestion/ometa/client.py
@@ -101,6 +101,7 @@ class ClientConfig(ConfigModel):
     access_token: Optional[str] = None
     expires_in: Optional[int] = None
     auth_header: Optional[str] = None
+    extra_headers: Optional[dict] = None
     raw_data: Optional[bool] = False
     allow_redirects: Optional[bool] = False
     auth_token_mode: Optional[str] = "Bearer"
@@ -135,6 +136,7 @@ class REST:
         api_version: str = None,
         headers: dict = None,
     ):
+        # pylint: disable=too-many-locals
         if not headers:
             headers = {"Content-type": "application/json"}
         base_url = base_url or self._base_url
@@ -157,6 +159,17 @@ class REST:
             self.config.auth_header
         ] = f"{self._auth_token_mode} {self.config.access_token}"
 
+        # Merge extra headers if provided.
+        # If a header value is provided in modulo string format and matches an existing header,
+        # the value will be set to that value.
+        # Example: "Proxy-Authorization": "%(Authorization)s"
+        # This will result in the Authorization value being set for the Proxy-Authorization Extra Header
+        if self.config.extra_headers:
+            extra_headers: dict[str, str] = self.config.extra_headers
+            extra_headers = {k: (v % headers) for k, v in extra_headers.items()}
+            logger.debug("Extra headers provided '%s'", extra_headers)
+            headers = {**headers, **extra_headers}
+
         opts = {
             "headers": headers,
             # Since we allow users to set endpoint URL via env var,
diff --git a/ingestion/src/metadata/ingestion/ometa/ometa_api.py b/ingestion/src/metadata/ingestion/ometa/ometa_api.py
index 6df879f70ec..b2233a78db2 100644
--- a/ingestion/src/metadata/ingestion/ometa/ometa_api.py
+++ b/ingestion/src/metadata/ingestion/ometa/ometa_api.py
@@ -22,6 +22,7 @@ try:
 except ImportError:
     from typing_compat import get_args
 
+
 from pydantic import BaseModel
 from requests.utils import quote
 
@@ -194,6 +195,7 @@ class OpenMetadata(
             base_url=self.config.hostPort,
             api_version=self.config.apiVersion,
             auth_header="Authorization",
+            extra_headers=self.config.extraHeaders,
             auth_token=self._auth_provider.get_access_token,
             verify=get_verify_ssl(self.config.sslConfig),
         )
diff --git a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/metadata/openMetadataConnection.json b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/metadata/openMetadataConnection.json
index 2f4042a647b..e4c36d895c4 100644
--- a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/metadata/openMetadataConnection.json
+++ b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/metadata/openMetadataConnection.json
@@ -11,6 +11,13 @@
       "type": "string",
       "enum": ["OpenMetadata"],
       "default": "OpenMetadata"
+    },
+    "extraHeaders": {
+      "description": "Additional headers to be sent to the API endpoint.",
+      "type": "object",
+      "additionalProperties": {
+        "type": "string"
+      }
     }
   },
   "properties": {
@@ -191,6 +198,10 @@
     },
     "supportsElasticSearchReindexingExtraction": {
       "$ref": "../connectionBasicType.json#/definitions/supportsElasticSearchReindexingExtraction"
+    },
+    "extraHeaders": {
+      "title": "Extra Headers",
+      "$ref": "#/definitions/extraHeaders"
     }
   },
   "additionalProperties": false,