Add a metadata_transformer to the Python stub.

3 files changed, 28 insertions, 5 deletions

diff --git a/src/python/src/grpc/_adapter/_links_test.py b/src/python/src/grpc/_adapter/_links_test.py
index cfdcc2c4bc..4987be389a 100644
--- a/src/python/src/grpc/_adapter/_links_test.py
+++ b/src/python/src/grpc/_adapter/_links_test.py
@@ -43,6 +43,14 @@ _IDENTITY = lambda x: x
 _TIMEOUT = 2
 
 
+# TODO(nathaniel): End-to-end metadata testing.
+def _transform_metadata(unused_metadata):
+  return (
+      ('one unused key', 'one unused value'),
+      ('another unused key', 'another unused value'),
+)
+
+
 class RoundTripTest(unittest.TestCase):
 
   def setUp(self):
@@ -76,7 +84,8 @@ class RoundTripTest(unittest.TestCase):
 
     rear_link = rear.RearLink(
         'localhost', port, self.rear_link_pool, {test_method: None},
-        {test_method: None}, False, None, None, None)
+        {test_method: None}, False, None, None, None,
+        metadata_transformer=_transform_metadata)
     rear_link.join_fore_link(test_fore_link)
     test_fore_link.join_rear_link(rear_link)
     rear_link.start()
diff --git a/src/python/src/grpc/_adapter/rear.py b/src/python/src/grpc/_adapter/rear.py
index f19321c426..2b93aa6331 100644
--- a/src/python/src/grpc/_adapter/rear.py
+++ b/src/python/src/grpc/_adapter/rear.py
@@ -93,7 +93,7 @@ class RearLink(base_interfaces.RearLink, activated.Activated):
   def __init__(
       self, host, port, pool, request_serializers, response_deserializers,
       secure, root_certificates, private_key, certificate_chain,
-      server_host_override=None):
+      metadata_transformer=None, server_host_override=None):
     """Constructor.
 
     Args:
@@ -111,6 +111,9 @@ class RearLink(base_interfaces.RearLink, activated.Activated):
         key should be used.
       certificate_chain: The PEM-encoded certificate chain to use or None if
         no certificate chain should be used.
+      metadata_transformer: A function that given a metadata object produces
+        another metadata to be used in the underlying communication on the
+        wire.
       server_host_override: (For testing only) the target name used for SSL
         host name checking.
     """
@@ -134,6 +137,7 @@ class RearLink(base_interfaces.RearLink, activated.Activated):
     self._root_certificates = root_certificates
     self._private_key = private_key
     self._certificate_chain = certificate_chain
+    self._metadata_transformer = metadata_transformer
     self._server_host_override = server_host_override
 
   def _on_write_event(self, operation_id, event, rpc_state):
@@ -243,6 +247,10 @@ class RearLink(base_interfaces.RearLink, activated.Activated):
     """
     request_serializer = self._request_serializers[name]
     call = _low.Call(self._channel, name, self._host, time.time() + timeout)
+    if self._metadata_transformer is not None:
+      metadata = self._metadata_transformer([])
+      for metadata_key, metadata_value in metadata:
+        call.add_metadata(metadata_key, metadata_value)
     call.invoke(self._completion_queue, operation_id, operation_id)
     outstanding = set(_INVOCATION_EVENT_KINDS)
 
diff --git a/src/python/src/grpc/early_adopter/implementations.py b/src/python/src/grpc/early_adopter/implementations.py
index 7d3d29f06c..35456d38c6 100644
--- a/src/python/src/grpc/early_adopter/implementations.py
+++ b/src/python/src/grpc/early_adopter/implementations.py
@@ -114,7 +114,7 @@ class _Stub(interfaces.Stub):
 
   def __init__(
       self, breakdown, host, port, secure, root_certificates, private_key,
-      certificate_chain, server_host_override=None):
+      certificate_chain, metadata_transformer=None, server_host_override=None):
     self._lock = threading.Lock()
     self._breakdown = breakdown
     self._host = host
@@ -123,6 +123,7 @@ class _Stub(interfaces.Stub):
     self._root_certificates = root_certificates
     self._private_key = private_key
     self._certificate_chain = certificate_chain
+    self._metadata_transformer = metadata_transformer
     self._server_host_override = server_host_override
 
     self._pool = None
@@ -141,6 +142,7 @@ class _Stub(interfaces.Stub):
             self._breakdown.request_serializers,
             self._breakdown.response_deserializers, self._secure,
             self._root_certificates, self._private_key, self._certificate_chain,
+            metadata_transformer=self._metadata_transformer,
             server_host_override=self._server_host_override)
         self._front.join_rear_link(self._rear_link)
         self._rear_link.join_fore_link(self._front)
@@ -189,8 +191,9 @@ class _Stub(interfaces.Stub):
 
 
 def stub(
-    service_name, methods, host, port, secure=False, root_certificates=None,
-    private_key=None, certificate_chain=None, server_host_override=None):
+    service_name, methods, host, port, metadata_transformer=None, secure=False,
+    root_certificates=None, private_key=None, certificate_chain=None,
+    server_host_override=None):
   """Constructs an interfaces.Stub.
 
   Args:
@@ -201,6 +204,9 @@ def stub(
       not qualified by the service name or decorated in any other way.
     host: The host to which to connect for RPC service.
     port: The port to which to connect for RPC service.
+    metadata_transformer: A callable that given a metadata object produces
+      another metadata object to be used in the underlying communication on the
+      wire.
     secure: Whether or not to construct the stub with a secure connection.
     root_certificates: The PEM-encoded root certificates or None to ask for
       them to be retrieved from a default location.