Implement response handlers (#54)

* Implement resopnse handlers

Signed-off-by: Daniel Widdis <widdis@gmail.com>

* Add missing test which uncovered missing return statement

Signed-off-by: Daniel Widdis <widdis@gmail.com>

* Ignore unreachable protobuf code branch

Signed-off-by: Daniel Widdis <widdis@gmail.com>

* Fix event loop deprecation warning

Signed-off-by: Daniel Widdis <widdis@gmail.com>

---------

Signed-off-by: Daniel Widdis <widdis@gmail.com>

.github/workflows/test.yml

@@ -30,6 +30,7 @@ jobs:
           poetry install
           protoc -I=$PROTOBUF_SRC_DIR --python_out=$PROTOBUF_SRC_DIR $PROTOBUF_SRC_DIR/*.proto
           2to3 --no-diff -n -w $PROTOBUF_SRC_DIR
+          find $PROTOBUF_SRC_DIR -type f -name "*_pb2.py" | xargs sed -i.bak -E 's/(USE_C_DESCRIPTORS == False:)/\1  # pragma: no cover/g'
           poetry run coverage run --source=src -m pytest -v
           poetry run coverage xml
       - name: Upload Coverage Report

src/opensearch_sdk_py/actions/internal/discovery_extensions_request_handler.py

@@ -9,36 +9,58 @@
 
 import logging
 
+from opensearch_sdk_py.actions.internal.register_rest_actions_response_handler import RegisterRestActionsResponseHandler
 from opensearch_sdk_py.actions.request_handler import RequestHandler
+from opensearch_sdk_py.actions.response_handlers import ResponseHandlers
 from opensearch_sdk_py.api.action_extension import ActionExtension
 from opensearch_sdk_py.transport.initialize_extension_request import InitializeExtensionRequest
+from opensearch_sdk_py.transport.initialize_extension_response import InitializeExtensionResponse
 from opensearch_sdk_py.transport.outbound_message_request import OutboundMessageRequest
+from opensearch_sdk_py.transport.outbound_message_response import OutboundMessageResponse
 from opensearch_sdk_py.transport.register_rest_actions_request import RegisterRestActionsRequest
 from opensearch_sdk_py.transport.stream_input import StreamInput
 from opensearch_sdk_py.transport.stream_output import StreamOutput
 
 
 class DiscoveryExtensionsRequestHandler(RequestHandler):
-    def __init__(self, extension: ActionExtension) -> None:
+    def __init__(self, extension: ActionExtension, response_handlers: ResponseHandlers) -> None:
         super().__init__("internal:discovery/extensions", extension)
+        self.response_handlers = response_handlers
 
     def handle(self, request: OutboundMessageRequest, input: StreamInput) -> StreamOutput:
         initialize_extension_request = InitializeExtensionRequest().read_from(input)
         logging.debug(f"< {initialize_extension_request}")
 
+        # Create the response message preserving the request id, but don't send it yet.
+        # This will be sent when response handler calls send()
+        self.response = OutboundMessageResponse(
+            request.thread_context_struct,
+            request.features,
+            InitializeExtensionResponse(self.extension.name, self.extension.implemented_interfaces),
+            request.version,
+            request.request_id,
+            request.is_handshake,
+            request.is_compress,
+        )
+
         # Sometime between tcp and transport handshakes and the eventual response,
         # the uniqueId gets added to the thread context.
         # request.thread_context_struct.request_headers["extension_unique_id"] = self.extension.name
-        self.extension.init_response_request_id = request.request_id
-
-        return self.send(
-            OutboundMessageRequest(
-                thread_context=request.thread_context_struct,
-                features=request.features,
-                message=RegisterRestActionsRequest(self.extension.name, self.extension.named_routes),
-                version=request.version,
-                action="internal:discovery/registerrestactions",
-                is_handshake=False,
-                is_compress=False,
-            )
+
+        # Now send our own initialization requests.
+
+        # Create the request, this gets us an auto-increment request id
+        register_request = OutboundMessageRequest(
+            thread_context=request.thread_context_struct,
+            features=request.features,
+            message=RegisterRestActionsRequest(self.extension.name, self.extension.named_routes),
+            version=request.version,
+            action="internal:discovery/registerrestactions",
+            is_handshake=False,
+            is_compress=False,
         )
+        # Register response handler to handle this request ID invoking this class's send()
+        register_response_handler = RegisterRestActionsResponseHandler(self)
+        self.response_handlers.register(register_request.request_id, register_response_handler)
+        # Now send the request
+        return register_response_handler.send(register_request)

src/opensearch_sdk_py/actions/internal/extension_rest_request_handler.py

@@ -30,21 +30,20 @@ def handle(self, request: OutboundMessageRequest, input: StreamInput) -> StreamO
         route = f"{extension_rest_request.method.name} {extension_rest_request.path}"
         response = self.extension.handle(route, extension_rest_request)
 
-        return self.send(
-            OutboundMessageResponse(
-                request.thread_context_struct,
-                request.features,
-                RestExecuteOnExtensionResponse(
-                    status=response.status,
-                    content_type=response.content_type,
-                    content=response.content,
-                    headers=response.headers,
-                    consumed_params=response.consumed_params,
-                    content_consumed=response.content_consumed,
-                ),
-                request.version,
-                request.request_id,
-                request.is_handshake,
-                request.is_compress,
-            )
+        self.response = OutboundMessageResponse(
+            request.thread_context_struct,
+            request.features,
+            RestExecuteOnExtensionResponse(
+                status=response.status,
+                content_type=response.content_type,
+                content=response.content,
+                headers=response.headers,
+                consumed_params=response.consumed_params,
+                content_consumed=response.content_consumed,
+            ),
+            request.version,
+            request.request_id,
+            request.is_handshake,
+            request.is_compress,
         )
+        return self.send()

src/opensearch_sdk_py/actions/internal/register_rest_actions_response_handler.py

@@ -0,0 +1,31 @@
+#
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+#
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+#
+
+import logging
+
+from opensearch_sdk_py.actions.request_response_handler import RequestResponseHandler
+from opensearch_sdk_py.actions.response_handler import ResponseHandler
+from opensearch_sdk_py.transport.acknowledged_response import AcknowledgedResponse
+from opensearch_sdk_py.transport.outbound_message_request import OutboundMessageRequest
+from opensearch_sdk_py.transport.stream_input import StreamInput
+from opensearch_sdk_py.transport.stream_output import StreamOutput
+
+
+class RegisterRestActionsResponseHandler(ResponseHandler):
+    def __init__(self, next_handler: RequestResponseHandler) -> None:
+        self.next_handler = next_handler
+
+    def handle(self, request: OutboundMessageRequest, input: StreamInput) -> StreamOutput:
+        ack_response = AcknowledgedResponse().read_from(input)
+        logging.debug(f"< {ack_response}")
+        if ack_response.status:
+            return self.next_handler.send()
+        else:
+            # TODO error handling
+            return None

src/opensearch_sdk_py/actions/internal/request_error_handler.py

@@ -31,18 +31,17 @@ def __init__(
         super().__init__("internal:error", extension)
 
     def handle(self, request: OutboundMessageRequest, input: StreamInput = None) -> StreamOutput:
-        return self.send(
-            OutboundMessageResponse(
-                request.thread_context_struct,
-                request.features,
-                RestExecuteOnExtensionResponse(
-                    status=self.status,
-                    content_type=self.content_type,
-                    content=self.content,
-                ),
-                request.version,
-                request.request_id,
-                request.is_handshake,
-                request.is_compress,
-            )
+        self.response = OutboundMessageResponse(
+            request.thread_context_struct,
+            request.features,
+            RestExecuteOnExtensionResponse(
+                status=self.status,
+                content_type=self.content_type,
+                content=self.content,
+            ),
+            request.version,
+            request.request_id,
+            request.is_handshake,
+            request.is_compress,
         )
+        return self.send()

src/opensearch_sdk_py/actions/internal/tcp_handshake_request_handler.py

@@ -26,14 +26,13 @@ def __init__(self, extension: Extension) -> None:
     def handle(self, request: OutboundMessageRequest, input: StreamInput) -> StreamOutput:
         tcp_handshake = TransportHandshakerHandshakeRequest().read_from(input)
         logging.debug(f"< {tcp_handshake}")
-        return self.send(
-            OutboundMessageResponse(
-                request.thread_context_struct,
-                request.features,
-                TransportHandshakerHandshakeResponse(request.version),
-                request.version,
-                request.request_id,
-                request.is_handshake,
-                request.is_compress,
-            )
+        self.response = OutboundMessageResponse(
+            request.thread_context_struct,
+            request.features,
+            TransportHandshakerHandshakeResponse(request.version),
+            request.version,
+            request.request_id,
+            request.is_handshake,
+            request.is_compress,
         )
+        return self.send()

src/opensearch_sdk_py/actions/internal/transport_handshake_request_handler.py

@@ -30,26 +30,25 @@ def handle(self, request: OutboundMessageRequest, input: StreamInput) -> StreamO
         transport_handshake = TransportServiceHandshakeRequest().read_from(input)
         logging.debug(f"< {transport_handshake}")
 
-        return self.send(
-            OutboundMessageResponse(
-                request.thread_context_struct,
-                request.features,
-                TransportServiceHandshakeResponse(
-                    DiscoveryNode(
-                        node_name="hello-world",
-                        node_id="hello-world",
-                        address=TransportAddress("127.0.0.1", 1234),
-                        roles={
-                            DiscoveryNodeRole.CLUSTER_MANAGER_ROLE,
-                            DiscoveryNodeRole.DATA_ROLE,
-                            DiscoveryNodeRole.INGEST_ROLE,
-                            DiscoveryNodeRole.REMOTE_CLUSTER_CLIENT_ROLE,
-                        },
-                    )
-                ),
-                request.version,
-                request.request_id,
-                request.is_handshake,
-                request.is_compress,
-            )
+        self.response = OutboundMessageResponse(
+            request.thread_context_struct,
+            request.features,
+            TransportServiceHandshakeResponse(
+                DiscoveryNode(
+                    node_name="hello-world",
+                    node_id="hello-world",
+                    address=TransportAddress("127.0.0.1", 1234),
+                    roles={
+                        DiscoveryNodeRole.CLUSTER_MANAGER_ROLE,
+                        DiscoveryNodeRole.DATA_ROLE,
+                        DiscoveryNodeRole.INGEST_ROLE,
+                        DiscoveryNodeRole.REMOTE_CLUSTER_CLIENT_ROLE,
+                    },
+                )
+            ),
+            request.version,
+            request.request_id,
+            request.is_handshake,
+            request.is_compress,
         )
+        return self.send()

src/opensearch_sdk_py/actions/request_handler.py

@@ -8,29 +8,31 @@
 #
 
 import logging
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from typing import Optional
 
+from opensearch_sdk_py.actions.request_response_handler import RequestResponseHandler
 from opensearch_sdk_py.extension import Extension
-from opensearch_sdk_py.transport.outbound_message import OutboundMessage
 from opensearch_sdk_py.transport.outbound_message_request import OutboundMessageRequest
+from opensearch_sdk_py.transport.outbound_message_response import OutboundMessageResponse
 from opensearch_sdk_py.transport.stream_input import StreamInput
 from opensearch_sdk_py.transport.stream_output import StreamOutput
 
 
-class RequestHandler(ABC):
+class RequestHandler(RequestResponseHandler):
     def __init__(self, action: str, extension: Extension):
         self.action = action
         self.extension = extension
+        self.response: OutboundMessageResponse = None
 
     @abstractmethod
     def handle(self, request: OutboundMessageRequest, input: StreamInput = None) -> Optional[bytes]:
         pass  # pragma: no cover
 
-    def send(self, message: OutboundMessage) -> StreamOutput:
+    def send(self) -> StreamOutput:
         output = StreamOutput()
-        message.write_to(output)
+        self.response.write_to(output)
         raw_out = output.getvalue()
-        logging.info(f"> {message.__str__()}, size={len(raw_out)} byte(s)")
+        logging.info(f"> {self.response.__str__()}, size={len(raw_out)} byte(s)")
         logging.debug(f"> #{raw_out}")
         return output

src/opensearch_sdk_py/actions/request_response_handler.py

@@ -0,0 +1,18 @@
+#
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+#
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+#
+
+from abc import ABC, abstractmethod
+
+from opensearch_sdk_py.transport.stream_output import StreamOutput
+
+
+class RequestResponseHandler(ABC):
+    @abstractmethod
+    def send(self) -> StreamOutput:
+        pass  # pragma: no cover

src/opensearch_sdk_py/actions/response_handler.py

@@ -0,0 +1,32 @@
+#
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+#
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+#
+
+import logging
+from abc import abstractmethod
+from typing import Optional
+
+from opensearch_sdk_py.actions.request_response_handler import RequestResponseHandler
+from opensearch_sdk_py.transport.outbound_message_request import OutboundMessageRequest
+from opensearch_sdk_py.transport.outbound_message_response import OutboundMessageResponse
+from opensearch_sdk_py.transport.stream_input import StreamInput
+from opensearch_sdk_py.transport.stream_output import StreamOutput
+
+
+class ResponseHandler(RequestResponseHandler):
+    @abstractmethod
+    def handle(self, response: OutboundMessageResponse, input: StreamInput = None) -> Optional[bytes]:
+        pass  # pragma: no cover
+
+    def send(self, request: OutboundMessageRequest) -> StreamOutput:
+        output = StreamOutput()
+        request.write_to(output)
+        raw_out = output.getvalue()
+        logging.info(f"> {request.__str__()}, size={len(raw_out)} byte(s)")
+        logging.debug(f"> #{raw_out}")
+        return output

src/opensearch_sdk_py/actions/response_handlers.py

@@ -0,0 +1,28 @@
+#
+# Copyright OpenSearch Contributors
+# SPDX-License-Identifier: Apache-2.0
+#
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+#
+
+from typing import Dict, Optional
+
+from opensearch_sdk_py.actions.response_handler import ResponseHandler
+from opensearch_sdk_py.extension import Extension
+from opensearch_sdk_py.transport.outbound_message_response import OutboundMessageResponse
+from opensearch_sdk_py.transport.stream_input import StreamInput
+
+
+class ResponseHandlers(Dict[int, ResponseHandler]):
+    def __init__(self, extension: Extension) -> None:
+        self.extension = extension
+
+    def register(self, request_id: int, handler: ResponseHandler) -> None:
+        self[request_id] = handler
+
+    def handle(self, response: OutboundMessageResponse, input: StreamInput = None) -> Optional[bytes]:
+        handler = self[response.request_id]
+        del self[response.request_id]
+        return handler.handle(response, input) if handler else None