feat(openapi): allow Bearer token

datahub-project · Jan 19, 2024 · b274151 · b274151
1 parent dadd51e
commit b274151
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git 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
@@ -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,14 +73,15 @@ 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 = "",
     proxies: Optional[dict] = None,
 ) -> 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