Rename metadata to headers

Author: NarekA

jina/constants.py

@@ -1,7 +1,7 @@
+import datetime as _datetime
 import os as _os
 import sys as _sys
 from pathlib import Path as _Path
-import datetime as _datetime
 
 __windows__ = _sys.platform == 'win32'
 __uptime__ = _datetime.datetime.now().isoformat()
@@ -53,6 +53,7 @@
 __args_executor_func__ = {
     'docs',
     'parameters',
+    'headers',
     'docs_matrix',
 }
 __args_executor_init__ = {'metas', 'requests', 'runtime_args'}

jina/serve/runtimes/gateway/http_fastapi_app_docarrayv2.py

@@ -187,8 +187,10 @@ def add_post_route(
         )
         app_kwargs['response_class'] = DocArrayResponse
 
+        from fastapi import Request
+
         @app.api_route(**app_kwargs)
-        async def post(body: input_model, response: Response):
+        async def post(body: input_model, response: Response, request: Request):
             target_executor = None
             req_id = None
             if body.header is not None:
@@ -208,6 +210,7 @@ async def post(body: input_model, response: Response):
                     docs,
                     exec_endpoint=endpoint_path,
                     parameters=body.parameters,
+                    headers=request.headers,
                     target_executor=target_executor,
                     request_id=req_id,
                     return_results=True,

jina/serve/runtimes/gateway/request_handling.py

@@ -157,48 +157,37 @@ 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:
-                        # Create a StreamResponse with the same headers and status as the target response
-                        stream_response = web.StreamResponse(
-                            status=response.status,
-                            headers=response.headers,
-                        )
-
-                        # Prepare the response to send headers
-                        await stream_response.prepare(request)
-
-                        # Stream the response from the target server to the client
-                        async for chunk in response.content.iter_any():
-                            await stream_response.write(chunk)
-
-                        # 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:
-                        content = await response.read()
-                        return web.Response(
-                            body=content,
-                            status=response.status,
-                            content_type=response.content_type,
-                        )
+                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.request(
+                    request.method,
+                    url=target_url,
+                    auto_decompress=False,
+                    **request_kwargs,
+                ) as response:
+                    # Create a StreamResponse with the same headers and status as the target response
+                    stream_response = web.StreamResponse(
+                        status=response.status,
+                        headers=response.headers,
+                    )
+
+                    # Prepare the response to send headers
+                    await stream_response.prepare(request)
+
+                    # Stream the response from the target server to the client
+                    async for chunk in response.content.iter_any():
+                        await stream_response.write(chunk)
+
+                    # Close the stream response once all chunks are sent
+                    await stream_response.write_eof()
+                    return stream_response
+
         except aiohttp.ClientError as e:
             return web.Response(text=f'Error: {str(e)}', status=500)
 

jina/serve/runtimes/gateway/streamer.py

@@ -6,6 +6,7 @@
     AsyncIterator,
     Dict,
     List,
+    Mapping,
     Optional,
     Sequence,
     Tuple,
@@ -209,6 +210,7 @@ async def stream(
         exec_endpoint: Optional[str] = None,
         target_executor: Optional[str] = None,
         parameters: Optional[Dict] = None,
+        headers: Optional[Mapping[str, str]] = None,
         results_in_order: bool = False,
         return_type: Type[DocumentArray] = DocumentArray,
     ) -> AsyncIterator[Tuple[Union[DocumentArray, 'Request'], 'ExecutorError']]:
@@ -221,6 +223,7 @@ async def stream(
         :param exec_endpoint: The Executor endpoint to which to send the Documents
         :param target_executor: A regex expression indicating the Executors that should receive the Request
         :param parameters: Parameters to be attached to the Requests
+        :param headers: Http request headers
         :param results_in_order: return the results in the same order as the request_iterator
         :param return_type: the DocumentArray type to be returned. By default, it is `DocumentArray`.
         :yield: tuple of Documents or Responses and unpacked error from Executors if any
@@ -232,6 +235,7 @@ async def stream(
             exec_endpoint=exec_endpoint,
             target_executor=target_executor,
             parameters=parameters,
+            headers=headers,
             results_in_order=results_in_order,
             return_type=return_type,
         ):
@@ -256,6 +260,7 @@ async def stream_doc(
         exec_endpoint: Optional[str] = None,
         target_executor: Optional[str] = None,
         parameters: Optional[Dict] = None,
+        headers: Optional[Mapping[str, str]] = None,
         request_id: Optional[str] = None,
         return_type: Type[DocumentArray] = DocumentArray,
     ) -> AsyncIterator[Tuple[Union[DocumentArray, 'Request'], 'ExecutorError']]:
@@ -267,6 +272,7 @@ async def stream_doc(
         :param exec_endpoint: The Executor endpoint to which to send the Documents
         :param target_executor: A regex expression indicating the Executors that should receive the Request
         :param parameters: Parameters to be attached to the Requests
+        :param headers: Http request headers
         :param request_id: Request ID to add to the request streamed to Executor. Only applicable if request_size is equal or less to the length of the docs
         :param return_type: the DocumentArray type to be returned. By default, it is `DocumentArray`.
         :yield: tuple of Documents or Responses and unpacked error from Executors if any
@@ -282,6 +288,8 @@ async def stream_doc(
             req.header.target_executor = target_executor
         if parameters:
             req.parameters = parameters
+        if headers:
+            req.headers = headers
 
         async for result in self.rpc_stream_doc(request=req, return_type=return_type):
             error = None
@@ -306,6 +314,7 @@ async def stream_docs(
         exec_endpoint: Optional[str] = None,
         target_executor: Optional[str] = None,
         parameters: Optional[Dict] = None,
+        headers: Optional[Mapping[str, str]] = None,
         results_in_order: bool = False,
         request_id: Optional[str] = None,
         return_type: Type[DocumentArray] = DocumentArray,
@@ -319,6 +328,7 @@ async def stream_docs(
         :param exec_endpoint: The Executor endpoint to which to send the Documents
         :param target_executor: A regex expression indicating the Executors that should receive the Request
         :param parameters: Parameters to be attached to the Requests
+        :param headers: Http request headers
         :param results_in_order: return the results in the same order as the request_iterator
         :param request_id: Request ID to add to the request streamed to Executor. Only applicable if request_size is equal or less to the length of the docs
         :param return_type: the DocumentArray type to be returned. By default, it is `DocumentArray`.
@@ -339,6 +349,8 @@ def _req_generator():
                         req.header.target_executor = target_executor
                     if parameters:
                         req.parameters = parameters
+                    if headers:
+                        req.headers = headers
                     yield req
             else:
                 from docarray import BaseDoc
@@ -361,6 +373,8 @@ def batch(iterable, n=1):
                             req.header.target_executor = target_executor
                         if parameters:
                             req.parameters = parameters
+                        if headers:
+                            req.headers = headers
                         yield req
                 else:
                     req = DataRequest()
@@ -374,6 +388,8 @@ def batch(iterable, n=1):
                         req.header.target_executor = target_executor
                     if parameters:
                         req.parameters = parameters
+                    if headers:
+                        req.headers = headers
                     yield req
 
         async for resp in self.rpc_stream(
@@ -438,6 +454,7 @@ async def post(
         request_size: int = 100,
         on: Optional[str] = None,
         parameters: Optional[Dict] = None,
+        headers: Optional[Mapping[str, str]] = None,
         return_type: Type[DocumentArray] = DocumentArray,
         **kwargs,
     ):
@@ -505,6 +522,7 @@ async def stream_doc(
         inputs: 'Document',
         on: Optional[str] = None,
         parameters: Optional[Dict] = None,
+        headers: Optional[Mapping[str, str]] = None,
         **kwargs,
     ):
         req: SingleDocumentRequest = SingleDocumentRequest(inputs.to_protobuf())

jina/serve/runtimes/worker/http_fastapi_app.py

@@ -86,15 +86,18 @@ def add_post_route(
 
             app_kwargs['response_class'] = DocArrayResponse
 
+        from fastapi import Request
+
         @app.api_route(**app_kwargs)
-        async def post(body: input_model, response: Response):
+        async def post(body: input_model, response: Response, request: Request):
 
             req = DataRequest()
             if body.header is not None:
                 req.header.request_id = body.header.request_id
 
             if body.parameters is not None:
                 req.parameters = body.parameters
+            req.headers = request.headers
             req.header.exec_endpoint = endpoint_path
             data = body.data
             if isinstance(data, list):
@@ -149,6 +152,7 @@ async def streaming_get(request: Request = None, body: input_doc_model = None):
                     body = Document.from_pydantic_model(body)
             req = DataRequest()
             req.header.exec_endpoint = endpoint_path
+            req.headers = request.headers
             if not docarray_v2:
                 req.data.docs = DocumentArray([body])
             else: