From b27415166f96f0ebd0d67fa2f222b563f81b8854 Mon Sep 17 00:00:00 2001
From: Dimitri GRISARD <dgrisard-externe@bedrockstreaming.com>
Date: Fri, 19 Jan 2024 16:34:16 +0100
Subject: [PATCH] feat(openapi): allow Bearer token

---
 .../src/datahub/ingestion/source/openapi.py               | 8 +++++---
 .../src/datahub/ingestion/source/openapi_parser.py        | 8 ++++++--
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/metadata-ingestion/src/datahub/ingestion/source/openapi.py b/metadata-ingestion/src/datahub/ingestion/source/openapi.py
index ad62ef7362aebd..ba7f25aa3a15ab 100755
--- a/metadata-ingestion/src/datahub/ingestion/source/openapi.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/openapi.py
@@ -61,6 +61,7 @@ class OpenApiConfig(ConfigModel):
     )
     forced_examples: dict = Field(default={}, description="")
     token: Optional[str] = Field(default=None, description="")
+    bearer_token: Optional[bool] = Field(default=False, description="Use Bearer token")
     get_token: dict = Field(default={}, description="")
 
     def get_swagger(self) -> Dict:
@@ -99,6 +100,7 @@ def get_swagger(self) -> Dict:
             sw_dict = get_swag_json(
                 self.url,
                 token=self.token,
+                bearer_token=self.bearer_token,
                 swagger_file=self.swagger_file,
                 proxies=self.proxies,
             )  # load the swagger file
@@ -271,7 +273,7 @@ def get_workunits_internal(self) -> Iterable[ApiWorkUnit]:  # noqa: C901
 
                 if config.token:
                     response = request_call(
-                        tot_url, token=config.token, proxies=config.proxies
+                        tot_url, token=config.token, bearer_token=config.bearer_token, proxies=config.proxies
                     )
                 else:
                     response = request_call(
@@ -299,7 +301,7 @@ def get_workunits_internal(self) -> Iterable[ApiWorkUnit]:  # noqa: C901
                     tot_url = clean_url(config.url + self.url_basepath + url_guess)
                     if config.token:
                         response = request_call(
-                            tot_url, token=config.token, proxies=config.proxies
+                            tot_url, token=config.token, bearer_token=config.bearer_token, proxies=config.proxies
                         )
                     else:
                         response = request_call(
@@ -327,7 +329,7 @@ def get_workunits_internal(self) -> Iterable[ApiWorkUnit]:  # noqa: C901
                     tot_url = clean_url(config.url + self.url_basepath + composed_url)
                     if config.token:
                         response = request_call(
-                            tot_url, token=config.token, proxies=config.proxies
+                            tot_url, token=config.token, bearer_token=config.bearer_token, proxies=config.proxies
                         )
                     else:
                         response = request_call(
diff --git a/metadata-ingestion/src/datahub/ingestion/source/openapi_parser.py b/metadata-ingestion/src/datahub/ingestion/source/openapi_parser.py
index 84bb3ad4526117..075cd7b731a9d3 100755
--- a/metadata-ingestion/src/datahub/ingestion/source/openapi_parser.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/openapi_parser.py
@@ -49,17 +49,20 @@ def flatten2list(d: dict) -> list:
 def request_call(
     url: str,
     token: Optional[str] = None,
+    bearer_token: Optional[bool] = False,
     username: Optional[str] = None,
     password: Optional[str] = None,
     proxies: Optional[dict] = None,
 ) -> requests.Response:
     headers = {"accept": "application/json"}
-
     if username is not None and password is not None:
         return requests.get(
             url, headers=headers, auth=HTTPBasicAuth(username, password)
         )
 
+    elif token is not None and bearer_token is True:
+        headers["Authorization"] = f"Bearer {token}"
+        return requests.get(url, proxies=proxies, headers=headers)
     elif token is not None:
         headers["Authorization"] = f"{token}"
         return requests.get(url, proxies=proxies, headers=headers)
@@ -70,6 +73,7 @@ def request_call(
 def get_swag_json(
     url: str,
     token: Optional[str] = None,
+    bearer_token: Optional[bool] = False,
     username: Optional[str] = None,
     password: Optional[str] = None,
     swagger_file: str = "",
@@ -77,7 +81,7 @@ def get_swag_json(
 ) -> Dict:
     tot_url = url + swagger_file
     if token is not None:
-        response = request_call(url=tot_url, token=token, proxies=proxies)
+        response = request_call(url=tot_url, token=token, bearer_token=bearer_token, proxies=proxies)
     else:
         response = request_call(
             url=tot_url, username=username, password=password, proxies=proxies