Allow passing metadata as part of creating a bidi (googleapis#7514)

crwilcox · web-flow · commit 7f1cfee92344 · 2019-03-27T10:58:39.000-07:00
* allows providing rpc metadata for bidi streams
diff --git a/api_core/google/api_core/bidi.py b/api_core/google/api_core/bidi.py
@@ -147,7 +147,11 @@ class BidiRpc(object):
 
         initial_request = example_pb2.StreamingRpcRequest(
             setting='example')
-        rpc = BidiRpc(stub.StreamingRpc, initial_request=initial_request)
+        rpc = BidiRpc(
+            stub.StreamingRpc,
+            initial_request=initial_request,
+            metadata=[('name', 'value')]
+        )
 
         rpc.open()
 
@@ -165,11 +169,14 @@ class BidiRpc(object):
                 Callable[None, protobuf.Message]]): The initial request to
             yield. This is useful if an initial request is needed to start the
             stream.
+        metadata (Sequence[Tuple(str, str)]): RPC metadata to include in
+            the request.
     """
 
-    def __init__(self, start_rpc, initial_request=None):
+    def __init__(self, start_rpc, initial_request=None, metadata=None):
         self._start_rpc = start_rpc
         self._initial_request = initial_request
+        self._rpc_metadata = metadata
         self._request_queue = queue.Queue()
         self._request_generator = None
         self._is_active = False
@@ -200,7 +207,7 @@ def open(self):
         request_generator = _RequestQueueGenerator(
             self._request_queue, initial_request=self._initial_request
         )
-        call = self._start_rpc(iter(request_generator))
+        call = self._start_rpc(iter(request_generator), metadata=self._rpc_metadata)
 
         request_generator.call = call
 
@@ -288,10 +295,14 @@ def should_recover(exc):
         initial_request = example_pb2.StreamingRpcRequest(
             setting='example')
 
-        rpc = ResumeableBidiRpc(
+        metadata = [('header_name', 'value')]
+
+        rpc = ResumableBidiRpc(
             stub.StreamingRpc,
+            should_recover=should_recover,
             initial_request=initial_request,
-            should_recover=should_recover)
+            metadata=metadata
+        )
 
         rpc.open()
 
@@ -310,10 +321,12 @@ def should_recover(exc):
         should_recover (Callable[[Exception], bool]): A function that returns
             True if the stream should be recovered. This will be called
             whenever an error is encountered on the stream.
+        metadata Sequence[Tuple(str, str)]: RPC metadata to include in
+            the request.
     """
 
-    def __init__(self, start_rpc, should_recover, initial_request=None):
-        super(ResumableBidiRpc, self).__init__(start_rpc, initial_request)
+    def __init__(self, start_rpc, should_recover, initial_request=None, metadata=None):
+        super(ResumableBidiRpc, self).__init__(start_rpc, initial_request, metadata)
         self._should_recover = should_recover
         self._operational_lock = threading.RLock()
         self._finalized = False
diff --git a/api_core/noxfile.py b/api_core/noxfile.py
@@ -88,7 +88,7 @@ def lint_setup_py(session):
 def pytype(session):
     """Run type-checking."""
     session.install(
-        ".", "grpcio >= 1.8.2", "grpcio-gcp >= 0.2.2", "pytype >= 2018.9.26"
+        ".", "grpcio >= 1.8.2", "grpcio-gcp >= 0.2.2", "pytype >= 2019.3.21"
     )
     session.run("pytype")
 
diff --git a/api_core/tests/unit/test_bidi.py b/api_core/tests/unit/test_bidi.py
@@ -125,9 +125,10 @@ def make_rpc():
     call = mock.create_autospec(_CallAndFuture, instance=True)
     rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True)
 
-    def rpc_side_effect(request):
+    def rpc_side_effect(request, metadata=None):
         call.is_active.return_value = True
         call.request = request
+        call.metadata = metadata
         return call
 
     rpc.side_effect = rpc_side_effect
@@ -172,6 +173,15 @@ def test_done_callbacks(self):
 
         callback.assert_called_once_with(mock.sentinel.future)
 
+    def test_metadata(self):
+        rpc, call = make_rpc()
+        bidi_rpc = bidi.BidiRpc(rpc, metadata=mock.sentinel.A)
+        assert bidi_rpc._rpc_metadata == mock.sentinel.A
+
+        bidi_rpc.open()
+        assert bidi_rpc.call == call
+        assert bidi_rpc.call.metadata == mock.sentinel.A
+
     def test_open(self):
         rpc, call = make_rpc()
         bidi_rpc = bidi.BidiRpc(rpc)
diff --git a/firestore/google/cloud/firestore_v1beta1/watch.py b/firestore/google/cloud/firestore_v1beta1/watch.py
@@ -213,9 +213,10 @@ def should_recover(exc):  # pragma: NO COVER
             ResumableBidiRpc = self.ResumableBidiRpc  # FBO unit tests
 
         self._rpc = ResumableBidiRpc(
-            self._api.transport._stubs["firestore_stub"].Listen,
+            self._api.transport.listen,
             initial_request=initial_request,
             should_recover=should_recover,
+            rpc_metadata=self._firestore._rpc_metadata,
         )
 
         self._rpc.add_done_callback(self._on_rpc_done)
diff --git a/firestore/tests/unit/v1beta1/test_cross_language.py b/firestore/tests/unit/v1beta1/test_cross_language.py
@@ -342,12 +342,13 @@ def convert_precondition(precond):
 
 
 class DummyRpc(object):  # pragma: NO COVER
-    def __init__(self, listen, initial_request, should_recover):
+    def __init__(self, listen, initial_request, should_recover, rpc_metadata=None):
         self.listen = listen
         self.initial_request = initial_request
         self.should_recover = should_recover
         self.closed = False
         self.callbacks = []
+        self._rpc_metadata = rpc_metadata
 
     def add_done_callback(self, callback):
         self.callbacks.append(callback)
diff --git a/firestore/tests/unit/v1beta1/test_watch.py b/firestore/tests/unit/v1beta1/test_watch.py
@@ -713,6 +713,7 @@ def _to_protobuf(self):
 class DummyFirestore(object):
     _firestore_api = DummyFirestoreClient()
     _database_string = "abc://bar/"
+    _rpc_metadata = None
 
     def document(self, *document_path):  # pragma: NO COVER
         if len(document_path) == 1:
@@ -781,12 +782,13 @@ def Thread(self, name, target, kwargs):
 
 
 class DummyRpc(object):
-    def __init__(self, listen, initial_request, should_recover):
+    def __init__(self, listen, initial_request, should_recover, rpc_metadata=None):
         self.listen = listen
         self.initial_request = initial_request
         self.should_recover = should_recover
         self.closed = False
         self.callbacks = []
+        self._rpc_metadata = rpc_metadata
 
     def add_done_callback(self, callback):
         self.callbacks.append(callback)

Original file line number	Diff line number	Diff line change
`@@ -88,7 +88,7 @@ def lint_setup_py(session):`
`88`	`88`	`def pytype(session):`
`89`	`89`	`"""Run type-checking."""`
`90`	`90`	`session.install(`
`91`		`- ".", "grpcio >= 1.8.2", "grpcio-gcp >= 0.2.2", "pytype >= 2018.9.26"`
	`91`	`+ ".", "grpcio >= 1.8.2", "grpcio-gcp >= 0.2.2", "pytype >= 2019.3.21"`
`92`	`92`	`)`
`93`	`93`	`session.run("pytype")`
`94`	`94`