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