From 09bd76552a6f4e3abb71b0f1aada95f6184381fa Mon Sep 17 00:00:00 2001 From: Joan Fontanals Date: Wed, 10 Jul 2024 18:48:30 +0200 Subject: [PATCH] fix: pass header kwargs to get openapi endpoints (#6177) --- jina/clients/base/http.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/jina/clients/base/http.py b/jina/clients/base/http.py index 44715ce732d6a..653a98f051629 100644 --- a/jina/clients/base/http.py +++ b/jina/clients/base/http.py @@ -24,7 +24,7 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._endpoints = [] - async def _get_endpoints_from_openapi(self): + async def _get_endpoints_from_openapi(self, **kwargs): def extract_paths_by_method(spec): paths_by_method = {} for path, methods in spec['paths'].items(): @@ -39,10 +39,15 @@ def extract_paths_by_method(spec): import aiohttp + session_kwargs = {} + if 'headers' in kwargs: + session_kwargs = {'headers': kwargs['headers']} + proto = 'https' if self.args.tls else 'http' target_url = f'{proto}://{self.args.host}:{self.args.port}/openapi.json' try: - async with aiohttp.ClientSession() as session: + + async with aiohttp.ClientSession(**session_kwargs) as session: async with session.get(target_url) as response: content = await response.read() openapi_response = json.loads(content.decode()) @@ -129,7 +134,7 @@ async def _get_results( request_iterator = self._get_requests(**kwargs) on = kwargs.get('on', '/post') if len(self._endpoints) == 0: - await self._get_endpoints_from_openapi() + await self._get_endpoints_from_openapi(**kwargs) async with AsyncExitStack() as stack: cm1 = ProgressBar(