Add ThriftBidiType for metadata

Summary:
- Add metadata type for Bidi
 - Bidi has 3 fields: firstResponse, streamType & sinkType
 - Fix thrift python metadata API accordingly

Reviewed By: sazonovkirill, tlj77

Differential Revision: D88336249

fbshipit-source-id: 3d2b5ae37776f4b282a8e8aadc353af01e253585

Author: Vishwa Prakash Gayasen

third-party/thrift/src/thrift/lib/python/metadata.py

@@ -22,6 +22,7 @@
 from typing import cast, Iterator, Mapping, Optional, Sequence, Tuple, Type, Union
 
 from apache.thrift.metadata.thrift_types import (
+    ThriftBidiType,
     ThriftConstStruct,
     ThriftConstValue,
     ThriftEnum,
@@ -59,6 +60,7 @@ class ThriftKind(enum.Enum):
     TYPEDEF = 7
     STREAM = 8
     SINK = 9
+    BIDI = 10
 
 
 class ThriftConstKind(enum.Enum):
@@ -80,6 +82,7 @@ class ThriftConstKind(enum.Enum):
     ThriftMapType,
     ThriftTypedefType,
     ThriftSinkType,
+    ThriftBidiType,
     ThriftStreamType,
 ]
 
@@ -103,6 +106,7 @@ def __init__(
                 ThriftSetType,
                 ThriftMapType,
                 ThriftTypedefType,
+                ThriftBidiType,
                 ThriftSinkType,
                 ThriftStreamType,
             ),
@@ -165,6 +169,10 @@ def _fbthrift_create(
             val = thriftType.value
             assert isinstance(val, ThriftSinkType)
             return ThriftSinkProxy(val, thriftMeta)
+        elif thriftType.type is ThriftType.Type.t_bidi:
+            val = thriftType.value
+            assert isinstance(val, ThriftBidiType)
+            return ThriftBidiProxy(val, thriftMeta)
         val = thriftType.value
         assert isinstance(val, ThriftPrimitiveType)
         specialType = ThriftTypeProxy(val, thriftMeta)
@@ -221,6 +229,11 @@ def as_sink(self) -> "ThriftSinkProxy":
             return cast(ThriftSinkProxy, self)
         raise TypeError("Type is not a sink")
 
+    def as_bidi(self) -> "ThriftBidiProxy":
+        if self.kind == ThriftKind.BIDI:
+            return cast(ThriftBidiProxy, self)
+        raise TypeError("Type is not a bidirectional stream")
+
 
 class ThriftSetProxy(ThriftTypeProxy):
     valueType: ThriftTypeProxy
@@ -313,6 +326,26 @@ def __init__(
             )
 
 
+class ThriftBidiProxy(ThriftTypeProxy):
+    streamElemType: ThriftTypeProxy
+    sinkElemType: ThriftTypeProxy
+    initialResponseType: ThriftTypeProxy
+
+    def __init__(self, thriftType: ThriftBidiType, thriftMeta: ThriftMetadata) -> None:
+        super().__init__(thriftType, thriftMeta)
+        self.kind: ThriftKind = ThriftKind.BIDI
+        self.streamElemType = ThriftTypeProxy._fbthrift_create(
+            thriftType.streamElemType, self.thriftMeta
+        )
+        self.sinkElemType = ThriftTypeProxy._fbthrift_create(
+            thriftType.sinkElemType, self.thriftMeta
+        )
+        if thriftType.initialResponseType is not None:
+            self.initialResponseType = ThriftTypeProxy._fbthrift_create(
+                thriftType.initialResponseType, self.thriftMeta
+            )
+
+
 class ThriftFieldProxy:
     type: ThriftTypeProxy
     thriftType: ThriftField

third-party/thrift/src/thrift/lib/python/metadata.pyi

@@ -28,6 +28,7 @@ from typing import (
 )
 
 from apache.thrift.metadata.thrift_types import (
+    ThriftBidiType,
     ThriftConstStruct,
     ThriftConstValue,
     ThriftEnum,
@@ -67,6 +68,7 @@ class ThriftKind(Enum):
     TYPEDEF: ThriftKind = ...
     STREAM: ThriftKind = ...
     SINK: ThriftKind = ...
+    BIDI: ThriftKind = ...
 
 class ThriftConstKind(Enum):
     CV_BOOL: ThriftConstKind = ...
@@ -90,6 +92,7 @@ class ThriftTypeProxy:
         ThriftEnum,
         ThriftStruct,
         ThriftSinkType,
+        ThriftBidiType,
         ThriftStreamType,
     ]
     thriftMeta: ThriftMetadata
@@ -104,6 +107,7 @@ class ThriftTypeProxy:
     def as_typedef(self) -> ThriftTypedefProxy: ...
     def as_stream(self) -> ThriftStreamProxy: ...
     def as_sink(self) -> ThriftSinkProxy: ...
+    def as_bidi(self) -> ThriftBidiProxy: ...
 
 class ThriftSetProxy(ThriftTypeProxy):
     thriftType: ThriftSetType
@@ -134,6 +138,11 @@ class ThriftStreamProxy(ThriftTypeProxy):
     elemType: ThriftTypeProxy
     initialResponseType: Optional[ThriftTypeProxy]
 
+class ThriftBidiProxy(ThriftTypeProxy):
+    streamElemType: ThriftTypeProxy
+    sinkElemType: ThriftTypeProxy
+    initialResponseType: ThriftTypeProxy
+
 class ThriftFieldProxy(Protocol):
     id: int
     name: str

third-party/thrift/src/thrift/lib/python/test/metadata.py

@@ -32,6 +32,7 @@
 )
 from testing.thrift_services import TestingServiceInterface
 from testing.thrift_types import Complex, hard, HardError, mixed, NestedError, Perm
+from thrift.python.bidi_service.thrift_clients import TestBidiService
 from thrift.python.metadata import (
     gen_metadata,
     ThriftExceptionProxy,
@@ -344,3 +345,51 @@ def test_gen_metadata_with_module_returns_thrift_metadata(self) -> None:
         self.assertIsNotNone(meta.structs)
         self.assertIsNotNone(meta.services)
         self.assertIsNotNone(meta.exceptions)
+
+    def test_metadata_bidirectional_streaming(self) -> None:
+        bidi_service = gen_metadata(TestBidiService)
+        self.assertEqual(bidi_service.name, "bidi_service.TestBidiService")
+
+        # Get the echo function which has bidirectional streaming
+        echo_func = None
+        for func in bidi_service.functions:
+            if func.name == "echo":
+                echo_func = func
+                break
+
+        self.assertIsNotNone(echo_func, "echo function should exist")
+
+        # Test that the return type is a bidirectional stream
+        bidi_type = echo_func.return_type.as_bidi()
+
+        # Verify stream element type is string
+        self.assertEqual(
+            bidi_type.streamElemType.as_primitive(),
+            ThriftPrimitiveType.THRIFT_STRING_TYPE,
+        )
+
+        # Verify sink element type is string
+        self.assertEqual(
+            bidi_type.sinkElemType.as_primitive(),
+            ThriftPrimitiveType.THRIFT_STRING_TYPE,
+        )
+
+        # Test echoWithResponse which has an initial response
+        echo_with_response = None
+        for func in bidi_service.functions:
+            if func.name == "echoWithResponse":
+                echo_with_response = func
+                break
+
+        self.assertIsNotNone(
+            echo_with_response, "echoWithResponse function should exist"
+        )
+
+        bidi_with_response = echo_with_response.return_type.as_bidi()
+
+        # Verify it has an initial response type
+        self.assertIsNotNone(bidi_with_response.initialResponseType)
+        self.assertEqual(
+            bidi_with_response.initialResponseType.as_primitive(),
+            ThriftPrimitiveType.THRIFT_STRING_TYPE,
+        )

third-party/thrift/src/thrift/lib/thrift/metadata.thrift

@@ -147,6 +147,21 @@ struct ThriftSinkType {
   3: optional ThriftType initialResponseType;
 }
 
+struct ThriftBidiType {
+  @java.Recursive
+  @rust.Box
+  @cpp.Ref{type = cpp.RefType.Unique}
+  1: optional ThriftType initialResponseType;
+  @java.Recursive
+  @rust.Box
+  @cpp.Ref{type = cpp.RefType.Unique}
+  2: optional ThriftType streamElemType;
+  @java.Recursive
+  @rust.Box
+  @cpp.Ref{type = cpp.RefType.Unique}
+  3: optional ThriftType sinkElemType;
+}
+
 @hack.MigrationBlockingLegacyJSONSerialization
 union ThriftType {
   1: ThriftPrimitiveType t_primitive;
@@ -159,6 +174,7 @@ union ThriftType {
   8: ThriftTypedefType t_typedef;
   9: ThriftStreamType t_stream;
   10: ThriftSinkType t_sink;
+  11: ThriftBidiType t_bidi;
 }
 
 struct ThriftEnum {