Skip to content

Commit

Permalink
fix: loadbalance right method
Browse files Browse the repository at this point in the history
  • Loading branch information
NarekA committed Dec 4, 2023
1 parent ed70797 commit af64005
Showing 1 changed file with 16 additions and 20 deletions.
36 changes: 16 additions & 20 deletions jina/serve/runtimes/gateway/request_handling.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import itertools
from typing import TYPE_CHECKING, AsyncIterator, Dict

from aiohttp.client import _RequestContextManager

from jina.enums import ProtocolType
from jina.helper import get_full_version
from jina.proto import jina_pb2
Expand Down Expand Up @@ -157,18 +159,19 @@ async def _load_balance(self, request):
try:
async with aiohttp.ClientSession() as session:

if request.method == 'GET':
request_kwargs = {}
try:
payload = await request.json()
if payload:
request_kwargs['json'] = payload
except Exception:
self.logger.debug('No JSON payload found in request')

async with session.get(
url=target_url, **request_kwargs
) as response:
request_kwargs = {}
try:
payload = await request.json()
if payload:
request_kwargs['json'] = payload
except Exception:
self.logger.debug('No JSON payload found in request')

Check warning on line 168 in jina/serve/runtimes/gateway/request_handling.py

View check run for this annotation

Codecov / codecov/patch

jina/serve/runtimes/gateway/request_handling.py#L162-L168

Added lines #L162 - L168 were not covered by tests

async with _RequestContextManager(

Check warning on line 170 in jina/serve/runtimes/gateway/request_handling.py

View check run for this annotation

Codecov / codecov/patch

jina/serve/runtimes/gateway/request_handling.py#L170

Added line #L170 was not covered by tests
session._request(request.method, target_url, **request_kwargs)
) as response:
if request.content_type.endswith('stream'):

Check warning on line 173 in jina/serve/runtimes/gateway/request_handling.py

View check run for this annotation

Codecov / codecov/patch

jina/serve/runtimes/gateway/request_handling.py#L173

Added line #L173 was not covered by tests

# Create a StreamResponse with the same headers and status as the target response
stream_response = web.StreamResponse(
status=response.status,
Expand All @@ -185,14 +188,7 @@ async def _load_balance(self, request):
# Close the stream response once all chunks are sent
await stream_response.write_eof()
return stream_response

elif request.method == 'POST':
d = await request.read()
import json

async with session.post(
url=target_url, json=json.loads(d.decode())
) as response:
else:
content = await response.read()
return web.Response(
body=content,
Expand Down

0 comments on commit af64005

Please sign in to comment.