Add support for gRPC binary metadata values (#1070)

Author: jazev-stripe

temporalio/bridge/client.py

@@ -7,7 +7,7 @@
 
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import Mapping, Optional, Tuple, Type, TypeVar
+from typing import Mapping, Optional, Tuple, Type, TypeVar, Union
 
 import google.protobuf.message
 
@@ -59,7 +59,7 @@ class ClientConfig:
     """Python representation of the Rust struct for configuring the client."""
 
     target_url: str
-    metadata: Mapping[str, str]
+    metadata: Mapping[str, Union[str, bytes]]
     api_key: Optional[str]
     identity: str
     tls_config: Optional[ClientTlsConfig]
@@ -77,7 +77,7 @@ class RpcCall:
     rpc: str
     req: bytes
     retry: bool
-    metadata: Mapping[str, str]
+    metadata: Mapping[str, Union[str, bytes]]
     timeout_millis: Optional[int]
 
 
@@ -108,7 +108,7 @@ def __init__(
         self._runtime = runtime
         self._ref = ref
 
-    def update_metadata(self, metadata: Mapping[str, str]) -> None:
+    def update_metadata(self, metadata: Mapping[str, Union[str, bytes]]) -> None:
         """Update underlying metadata on Core client."""
         self._ref.update_metadata(metadata)
 
@@ -124,7 +124,7 @@ async def call(
         req: google.protobuf.message.Message,
         resp_type: Type[ProtoMessage],
         retry: bool,
-        metadata: Mapping[str, str],
+        metadata: Mapping[str, Union[str, bytes]],
         timeout: Optional[timedelta],
     ) -> ProtoMessage:
         """Make RPC call using SDK Core."""

temporalio/bridge/sdk-core

@@ -8,7 +8,9 @@ use temporal_client::{
     ConfiguredClient, HealthService, HttpConnectProxyOptions, RetryClient, RetryConfig,
     TemporalServiceClientWithMetrics, TestService, TlsConfig, WorkflowService,
 };
-use tonic::metadata::MetadataKey;
+use tonic::metadata::{
+    AsciiMetadataKey, AsciiMetadataValue, BinaryMetadataKey, BinaryMetadataValue,
+};
 use url::Url;
 
 use crate::runtime;
@@ -28,7 +30,7 @@ pub struct ClientConfig {
     target_url: String,
     client_name: String,
     client_version: String,
-    metadata: HashMap<String, String>,
+    metadata: HashMap<String, RpcMetadataValue>,
     api_key: Option<String>,
     identity: String,
     tls_config: Option<ClientTlsConfig>,
@@ -72,10 +74,18 @@ struct RpcCall {
     rpc: String,
     req: Vec<u8>,
     retry: bool,
-    metadata: HashMap<String, String>,
+    metadata: HashMap<String, RpcMetadataValue>,
     timeout_millis: Option<u64>,
 }
 
+#[derive(FromPyObject)]
+enum RpcMetadataValue {
+    #[pyo3(transparent, annotation = "str")]
+    Str(String),
+    #[pyo3(transparent, annotation = "bytes")]
+    Bytes(Vec<u8>),
+}
+
 pub fn connect_client<'a>(
     py: Python<'a>,
     runtime_ref: &runtime::RuntimeRef,
@@ -116,8 +126,19 @@ macro_rules! rpc_call_on_trait {
 
 #[pymethods]
 impl ClientRef {
-    fn update_metadata(&self, headers: HashMap<String, String>) {
-        self.retry_client.get_client().set_headers(headers);
+    fn update_metadata(&self, headers: HashMap<String, RpcMetadataValue>) -> PyResult<()> {
+        let (ascii_headers, binary_headers) = partition_headers(headers);
+
+        self.retry_client
+            .get_client()
+            .set_headers(ascii_headers)
+            .map_err(|err| PyValueError::new_err(err.to_string()))?;
+        self.retry_client
+            .get_client()
+            .set_binary_headers(binary_headers)
+            .map_err(|err| PyValueError::new_err(err.to_string()))?;
+
+        Ok(())
     }
 
     fn update_api_key(&self, api_key: Option<String>) {
@@ -536,12 +557,32 @@ fn rpc_req<P: prost::Message + Default>(call: RpcCall) -> PyResult<tonic::Reques
         .map_err(|err| PyValueError::new_err(format!("Invalid proto: {err}")))?;
     let mut req = tonic::Request::new(proto);
     for (k, v) in call.metadata {
-        req.metadata_mut().insert(
-            MetadataKey::from_str(k.as_str())
-                .map_err(|err| PyValueError::new_err(format!("Invalid metadata key: {err}")))?,
-            v.parse()
-                .map_err(|err| PyValueError::new_err(format!("Invalid metadata value: {err}")))?,
-        );
+        if let Ok(binary_key) = BinaryMetadataKey::from_str(&k) {
+            let RpcMetadataValue::Bytes(bytes) = v else {
+                return Err(PyValueError::new_err(format!(
+                    "Invalid metadata value for binary key {k}: expected bytes"
+                )));
+            };
+
+            req.metadata_mut()
+                .insert_bin(binary_key, BinaryMetadataValue::from_bytes(&bytes));
+        } else {
+            let ascii_key = AsciiMetadataKey::from_str(&k)
+                .map_err(|err| PyValueError::new_err(format!("Invalid metadata key: {err}")))?;
+
+            let RpcMetadataValue::Str(string) = v else {
+                return Err(PyValueError::new_err(format!(
+                    "Invalid metadata value for ASCII key {k}: expected str"
+                )));
+            };
+
+            req.metadata_mut().insert(
+                ascii_key,
+                AsciiMetadataValue::from_str(&string).map_err(|err| {
+                    PyValueError::new_err(format!("Invalid metadata value: {err}"))
+                })?,
+            );
+        }
     }
     if let Some(timeout_millis) = call.timeout_millis {
         req.set_timeout(Duration::from_millis(timeout_millis));
@@ -568,11 +609,41 @@ where
     }
 }
 
+fn partition_headers(
+    headers: HashMap<String, RpcMetadataValue>,
+) -> (HashMap<String, String>, HashMap<String, Vec<u8>>) {
+    let (ascii_enum_headers, binary_enum_headers): (HashMap<_, _>, HashMap<_, _>) = headers
+        .into_iter()
+        .partition(|(_, v)| matches!(v, RpcMetadataValue::Str(_)));
+
+    let ascii_headers = ascii_enum_headers
+        .into_iter()
+        .map(|(k, v)| {
+            let RpcMetadataValue::Str(s) = v else {
+                unreachable!();
+            };
+            (k, s)
+        })
+        .collect();
+    let binary_headers = binary_enum_headers
+        .into_iter()
+        .map(|(k, v)| {
+            let RpcMetadataValue::Bytes(b) = v else {
+                unreachable!();
+            };
+            (k, b)
+        })
+        .collect();
+
+    (ascii_headers, binary_headers)
+}
+
 impl TryFrom<ClientConfig> for ClientOptions {
     type Error = PyErr;
 
     fn try_from(opts: ClientConfig) -> PyResult<Self> {
         let mut gateway_opts = ClientOptionsBuilder::default();
+        let (ascii_headers, binary_headers) = partition_headers(opts.metadata);
         gateway_opts
             .target_url(
                 Url::parse(&opts.target_url)
@@ -587,7 +658,8 @@ impl TryFrom<ClientConfig> for ClientOptions {
             )
             .keep_alive(opts.keep_alive_config.map(Into::into))
             .http_connect_proxy(opts.http_connect_proxy_config.map(Into::into))
-            .headers(Some(opts.metadata))
+            .headers(Some(ascii_headers))
+            .binary_headers(Some(binary_headers))
             .api_key(opts.api_key);
         // Builder does not allow us to set option here, so we have to make
         // a conditional to even call it