THRIFT-5331: Py: make THeader subprotocol configurable (#2302)

Client: py

This allows clients to choose which subprotocol (TBinary/TCompact) to frame with headers. The server will already accept either protocol and reply correctly.

2 files changed, 14 insertions, 7 deletions

diff --git a/lib/py/src/protocol/THeaderProtocol.py b/lib/py/src/protocol/THeaderProtocol.py
index 13982e898..4b58e639d 100644
--- a/lib/py/src/protocol/THeaderProtocol.py
+++ b/lib/py/src/protocol/THeaderProtocol.py
@@ -35,7 +35,9 @@ class THeaderProtocol(TProtocolBase):
 
     THeaderProtocol frames other Thrift protocols and adds support for optional
     out-of-band headers. The currently supported subprotocols are
-    TBinaryProtocol and TCompactProtocol.
+    TBinaryProtocol and TCompactProtocol. When used as a client, the
+    subprotocol to frame can be chosen with the `default_protocol` parameter to
+    the constructor.
 
     It's also possible to apply transforms to the encoded message payload. The
     only transform currently supported is to gzip.
@@ -53,14 +55,14 @@ class THeaderProtocol(TProtocolBase):
 
     """
 
-    def __init__(self, transport, allowed_client_types):
+    def __init__(self, transport, allowed_client_types, default_protocol=THeaderSubprotocolID.BINARY):
         # much of the actual work for THeaderProtocol happens down in
         # THeaderTransport since we need to do low-level shenanigans to detect
         # if the client is sending us headers or one of the headerless formats
         # we support. this wraps the real transport with the one that does all
         # the magic.
         if not isinstance(transport, THeaderTransport):
-            transport = THeaderTransport(transport, allowed_client_types)
+            transport = THeaderTransport(transport, allowed_client_types, default_protocol)
         super(THeaderProtocol, self).__init__(transport)
         self._set_protocol()
 
@@ -218,8 +220,13 @@ class THeaderProtocol(TProtocolBase):
 
 
 class THeaderProtocolFactory(TProtocolFactory):
-    def __init__(self, allowed_client_types=(THeaderClientType.HEADERS,)):
+    def __init__(
+        self,
+        allowed_client_types=(THeaderClientType.HEADERS,),
+        default_protocol=THeaderSubprotocolID.BINARY,
+    ):
         self.allowed_client_types = allowed_client_types
+        self.default_protocol = default_protocol
 
     def getProtocol(self, trans):
-        return THeaderProtocol(trans, self.allowed_client_types)
+        return THeaderProtocol(trans, self.allowed_client_types, self.default_protocol)
diff --git a/lib/py/src/transport/THeaderTransport.py b/lib/py/src/transport/THeaderTransport.py
index c0d564012..7c9827ba3 100644
--- a/lib/py/src/transport/THeaderTransport.py
+++ b/lib/py/src/transport/THeaderTransport.py
@@ -87,7 +87,7 @@ def _writeString(trans, value):
 
 
 class THeaderTransport(TTransportBase, CReadableTransport):
-    def __init__(self, transport, allowed_client_types):
+    def __init__(self, transport, allowed_client_types, default_protocol=THeaderSubprotocolID.BINARY):
         self._transport = transport
         self._client_type = THeaderClientType.HEADERS
         self._allowed_client_types = allowed_client_types
@@ -101,7 +101,7 @@ class THeaderTransport(TTransportBase, CReadableTransport):
 
         self.flags = 0
         self.sequence_id = 0
-        self._protocol_id = THeaderSubprotocolID.BINARY
+        self._protocol_id = default_protocol
         self._max_frame_size = HARD_MAX_FRAME_SIZE
 
     def isOpen(self):