Always send grpc-status in trailers

public · Dec 30, 2021 · 0971bef · 0971bef
1 parent cfd0b2c
commit 0971bef
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 28 deletions.
diff --git a/sonora/asgi.py b/sonora/asgi.py
@@ -177,6 +177,7 @@ async def _do_streaming_response(
                 recv_task.cancel()
 
         trailers = [("grpc-status", str(context.code.value[0]))]
+
         if context.details:
             trailers.append(("grpc-message", quote(context.details)))
 
@@ -199,12 +200,6 @@ async def _do_unary_response(
 
         status = 200
 
-        headers.append((b"grpc-status", str(context.code.value[0]).encode()))
-        if context.details:
-            headers.append(
-                (b"grpc-message", quote(context.details.encode("utf8")).encode("ascii"))
-            )
-
         if context._initial_metadata:
             headers.extend(context._initial_metadata)
 
@@ -215,13 +210,18 @@ async def _do_unary_response(
         else:
             message_data = b""
 
+        trailers = [(b"grpc-status", str(context.code.value[0]).encode())]
+
+        if context.details:
+            trailers.append(
+                (b"grpc-message", quote(context.details.encode("utf8")).encode("ascii"))
+            )
+
         if context._trailing_metadata:
-            trailers = context._trailing_metadata
+            trailers.extend(context._trailing_metadata)
 
-            trailer_message = protocol.pack_trailers(trailers)
-            trailer_data = wrap_message(True, False, trailer_message)
-        else:
-            trailer_data = b""
+        trailer_message = protocol.pack_trailers(trailers)
+        trailer_data = wrap_message(True, False, trailer_message)
 
         content_length = len(message_data) + len(trailer_data)
 
@@ -230,9 +230,11 @@ async def _do_unary_response(
         await send(
             {"type": "http.response.start", "status": status, "headers": headers}
         )
+
         await send(
             {"type": "http.response.body", "body": message_data, "more_body": True}
         )
+
         await send(
             {"type": "http.response.body", "body": trailer_data, "more_body": False}
         )

diff --git a/sonora/wsgi.py b/sonora/wsgi.py
@@ -138,7 +138,7 @@ def _do_streaming_response(
             pass
 
         if context._initial_metadata:
-            headers.extend(protocol.encode_headers(context._initial_metadata))
+            headers.extend(context._initial_metadata)
 
         start_response("200 OK", headers)
 
@@ -158,7 +158,7 @@ def _do_streaming_response(
             trailers.append(("grpc-message", quote(context.details.encode("utf8"))))
 
         if context._trailing_metadata:
-            trailers.extend(protocol.encode_headers(context._trailing_metadata))
+            trailers.extend(context._trailing_metadata)
 
         trailer_message = protocol.pack_trailers(trailers)
 
@@ -174,28 +174,28 @@ def _do_unary_response(
         else:
             message_data = b""
 
-        if context._trailing_metadata:
-            trailers = protocol.encode_headers(context._trailing_metadata)
-            trailer_message = protocol.pack_trailers(trailers)
-            trailer_data = wrap_message(True, False, trailer_message)
-        else:
-            trailer_data = b""
+        if context._initial_metadata:
+            headers.extend(context._initial_metadata)
 
-        content_length = len(message_data) + len(trailer_data)
+        trailers = [("grpc-status", str(context.code.value[0]))]
 
-        headers.append(("content-length", str(content_length)))
+        if context.details:
+            trailers.append(("grpc-message", quote(context.details.encode("utf8"))))
 
-        headers.append(("grpc-status", str(context.code.value[0])))
+        if context._trailing_metadata:
+            trailers.extend(context._trailing_metadata)
 
-        if context.details:
-            headers.append(("grpc-message", quote(context.details.encode("utf8"))))
+        trailer_message = protocol.pack_trailers(trailers)
+        trailer_data = wrap_message(True, False, trailer_message)
 
-        if context._initial_metadata:
-            headers.extend(protocol.encode_headers(context._initial_metadata))
+        content_length = len(message_data) + len(trailer_data)
+
+        headers.append(("content-length", str(content_length)))
 
         start_response("200 OK", headers)
 
         yield message_data
+
         yield trailer_data
 
     def _do_cors_preflight(self, environ, start_response):
@@ -342,10 +342,10 @@ def invocation_metadata(self):
         return self._invocation_metadata
 
     def send_initial_metadata(self, initial_metadata):
-        self._initial_metadata = initial_metadata
+        self._initial_metadata = protocol.encode_headers(initial_metadata)
 
     def set_trailing_metadata(self, trailing_metadata):
-        self._trailing_metadata = trailing_metadata
+        self._trailing_metadata = protocol.encode_headers(trailing_metadata)
 
     def peer(self):
         raise NotImplementedError()