From d89564d72e9f520714f82a8b047bf55e6d226289 Mon Sep 17 00:00:00 2001
From: Han Xiao <artex.xh@gmail.com>
Date: Sun, 11 Apr 2021 18:55:29 +0800
Subject: [PATCH] feat(rest): expose reload api to rest interface (#2301)

---
 jina/peapods/runtimes/asyncio/rest/app.py    | 21 ++++++++++++++++++++
 jina/peapods/runtimes/asyncio/rest/models.py | 17 +++++++++++-----
 tests/unit/flow/test_flow_reload.py          | 11 ++++++++++
 3 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/jina/peapods/runtimes/asyncio/rest/app.py b/jina/peapods/runtimes/asyncio/rest/app.py
index 39c1b6f3ad716..b70cd0a6f75ca 100644
--- a/jina/peapods/runtimes/asyncio/rest/app.py
+++ b/jina/peapods/runtimes/asyncio/rest/app.py
@@ -40,6 +40,7 @@ def get_fastapi_app(args: 'argparse.Namespace', logger: 'JinaLogger'):
             JinaDeleteRequestModel,
             JinaUpdateRequestModel,
             JinaSearchRequestModel,
+            JinaReloadRequestModel,
         )
 
     app = FastAPI(
@@ -199,6 +200,26 @@ async def delete_api(body: JinaDeleteRequestModel):
             result_in_stream(request_generator(**bd)), media_type='application/json'
         )
 
+    @app.post(
+        path='/reload', summary='Reload the executor of certain Peas/Pods in the Flow'
+    )
+    async def reload_api(body: JinaReloadRequestModel):
+        """
+        Reload the executor of certain peas/pods in the Flow
+
+        :param body: reload request.
+        :return: Response of the results.
+        """
+        from .....clients import BaseClient
+
+        bd = body.dict()
+        bd['mode'] = RequestType.CONTROL
+        bd['command'] = 'RELOAD'
+        return StreamingResponse(
+            result_in_stream(request_generator(data=[], **bd)),
+            media_type='application/json',
+        )
+
     async def result_in_stream(req_iter):
         """
         Streams results from AsyncPrefetchCall as json
diff --git a/jina/peapods/runtimes/asyncio/rest/models.py b/jina/peapods/runtimes/asyncio/rest/models.py
index 9af51df639afe..cbe072f51350b 100644
--- a/jina/peapods/runtimes/asyncio/rest/models.py
+++ b/jina/peapods/runtimes/asyncio/rest/models.py
@@ -1,15 +1,12 @@
-from enum import Enum
+from collections import defaultdict
 from datetime import datetime
+from enum import Enum
 from types import SimpleNamespace
-from collections import defaultdict
 from typing import Callable, Dict, Any, Optional, List, Union
 
-from pydantic import Field, BaseModel, BaseConfig, create_model, root_validator
 from google.protobuf.descriptor import Descriptor, FieldDescriptor
 from google.protobuf.pyext.cpp_message import GeneratedProtocolMessageType
-
 from jina.enums import DataInputType
-from jina.types.document import Document
 from jina.parsers import set_client_cli_parser
 from jina.proto.jina_pb2 import (
     DenseNdArrayProto,
@@ -24,6 +21,8 @@
     RequestProto,
     QueryLangProto,
 )
+from jina.types.document import Document
+from pydantic import Field, BaseModel, BaseConfig, create_model, root_validator
 
 DEFAULT_REQUEST_SIZE = set_client_cli_parser().parse_args([]).request_size
 PROTO_TO_PYDANTIC_MODELS = SimpleNamespace()
@@ -238,6 +237,14 @@ class JinaStatusModel(BaseModel):
     used_memory: str
 
 
+class JinaReloadRequestModel(BaseModel):
+    """
+    Jina control request model.
+    """
+
+    targets: Union[str, List[str]]
+
+
 class JinaRequestModel(BaseModel):
     """
     Jina request model.
diff --git a/tests/unit/flow/test_flow_reload.py b/tests/unit/flow/test_flow_reload.py
index 96c19a02cfdf2..287868cc7d75c 100644
--- a/tests/unit/flow/test_flow_reload.py
+++ b/tests/unit/flow/test_flow_reload.py
@@ -1,11 +1,22 @@
 import os
 
 import pytest
+import requests
 
 from jina import Flow
 from jina.executors import BaseExecutor
 
 
+def test_flow_rest_reload():
+    f = Flow().add()
+    f.use_rest_gateway()
+    with f:
+        r = requests.post(
+            f'http://0.0.0.0:{f.port_expose}/reload', json={'targets': ['pod0']}
+        )
+        assert r.status_code == 200
+
+
 @pytest.mark.skipif(
     'GITHUB_WORKFLOW' in os.environ,
     reason='skip the test on github as it will hang the whole CI, locally is fine',