Explicit API key option on client (#699)

cretz · web-flow · commit bcda5d9f73b0 · 2024-03-11T11:27:07.000-05:00
diff --git a/client/src/lib.rs b/client/src/lib.rs
@@ -67,7 +67,7 @@ use tonic::{
     body::BoxBody,
     client::GrpcService,
     codegen::InterceptedService,
-    metadata::{MetadataKey, MetadataValue},
+    metadata::{MetadataKey, MetadataMap, MetadataValue},
     service::Interceptor,
     transport::{Certificate, Channel, Endpoint, Identity},
     Code, Status,
@@ -133,6 +133,15 @@ pub struct ClientOptions {
     /// If set (which it is by default), HTTP2 gRPC keep alive will be enabled.
     #[builder(default = "Some(ClientKeepAliveConfig::default())")]
     pub keep_alive: Option<ClientKeepAliveConfig>,
+
+    /// HTTP headers to include on every RPC call.
+    #[builder(default)]
+    pub headers: Option<HashMap<String, String>>,
+
+    /// API key which is set as the "Authorization" header with "Bearer " prepended. This will only
+    /// be applied if the headers don't already have an "Authorization" header.
+    #[builder(default)]
+    pub api_key: Option<String>,
 }
 
 /// Configuration options for TLS
@@ -279,7 +288,7 @@ pub enum ClientInitError {
 pub struct ConfiguredClient<C> {
     client: C,
     options: Arc<ClientOptions>,
-    headers: Arc<RwLock<HashMap<String, String>>>,
+    headers: Arc<RwLock<ClientHeaders>>,
     /// Capabilities as read from the `get_system_info` RPC call made on client connection
     capabilities: Option<get_system_info_response::Capabilities>,
     workers: Arc<SlotManager>,
@@ -288,8 +297,12 @@ pub struct ConfiguredClient<C> {
 impl<C> ConfiguredClient<C> {
     /// Set HTTP request headers overwriting previous headers
     pub fn set_headers(&self, headers: HashMap<String, String>) {
-        let mut guard = self.headers.write();
-        *guard = headers;
+        self.headers.write().user_headers = headers;
+    }
+
+    /// Set API key, overwriting previous
+    pub fn set_api_key(&self, api_key: Option<String>) {
+        self.headers.write().api_key = api_key;
     }
 
     /// Returns the options the client is configured with
@@ -309,6 +322,34 @@ impl<C> ConfiguredClient<C> {
     }
 }
 
+#[derive(Debug)]
+struct ClientHeaders {
+    user_headers: HashMap<String, String>,
+    api_key: Option<String>,
+}
+
+impl ClientHeaders {
+    fn apply_to_metadata(&self, metadata: &mut MetadataMap) {
+        for (key, val) in self.user_headers.iter() {
+            // Only if not already present
+            if !metadata.contains_key(key) {
+                // Ignore invalid keys/values
+                if let (Ok(key), Ok(val)) = (MetadataKey::from_str(key), val.parse()) {
+                    metadata.insert(key, val);
+                }
+            }
+        }
+        if let Some(api_key) = &self.api_key {
+            // Only if not already present
+            if !metadata.contains_key("authorization") {
+                if let Ok(val) = format!("Bearer {}", api_key).parse() {
+                    metadata.insert("authorization", val);
+                }
+            }
+        }
+    }
+}
+
 // The configured client is effectively a "smart" (dumb) pointer
 impl<C> Deref for ConfiguredClient<C> {
     type Target = C;
@@ -331,12 +372,8 @@ impl ClientOptions {
         &self,
         namespace: impl Into<String>,
         metrics_meter: Option<TemporalMeter>,
-        headers: Option<Arc<RwLock<HashMap<String, String>>>>,
     ) -> Result<RetryClient<Client>, ClientInitError> {
-        let client = self
-            .connect_no_namespace(metrics_meter, headers)
-            .await?
-            .into_inner();
+        let client = self.connect_no_namespace(metrics_meter).await?.into_inner();
         let client = Client::new(client, namespace.into());
         let retry_client = RetryClient::new(client, self.retry_config.clone());
         Ok(retry_client)
@@ -349,7 +386,6 @@ impl ClientOptions {
     pub async fn connect_no_namespace(
         &self,
         metrics_meter: Option<TemporalMeter>,
-        headers: Option<Arc<RwLock<HashMap<String, String>>>>,
     ) -> Result<RetryClient<ConfiguredClient<TemporalServiceClientWithMetrics>>, ClientInitError>
     {
         let channel = Channel::from_shared(self.target_url.to_string())?;
@@ -374,7 +410,10 @@ impl ClientOptions {
                 metrics: metrics_meter.clone().map(MetricsContext::new),
             })
             .service(channel);
-        let headers = headers.unwrap_or_default();
+        let headers = Arc::new(RwLock::new(ClientHeaders {
+            user_headers: self.headers.clone().unwrap_or_default(),
+            api_key: self.api_key.clone(),
+        }));
         let interceptor = ServiceCallInterceptor {
             opts: self.clone(),
             headers: headers.clone(),
@@ -442,7 +481,7 @@ impl ClientOptions {
 pub struct ServiceCallInterceptor {
     opts: ClientOptions,
     /// Only accessed as a reader
-    headers: Arc<RwLock<HashMap<String, String>>>,
+    headers: Arc<RwLock<ClientHeaders>>,
 }
 
 impl Interceptor for ServiceCallInterceptor {
@@ -468,16 +507,7 @@ impl Interceptor for ServiceCallInterceptor {
                     .unwrap_or_else(|_| MetadataValue::from_static("")),
             );
         }
-        let headers = &*self.headers.read();
-        for (k, v) in headers {
-            if metadata.contains_key(k) {
-                // Don't overwrite per-request specified headers
-                continue;
-            }
-            if let (Ok(k), Ok(v)) = (MetadataKey::from_str(k), v.parse()) {
-                metadata.insert(k, v);
-            }
-        }
+        self.headers.read().apply_to_metadata(metadata);
         if !metadata.contains_key("grpc-timeout") {
             request.set_timeout(OTHER_CALL_TIMEOUT);
         }
@@ -1559,7 +1589,7 @@ mod tests {
     use super::*;
 
     #[test]
-    fn respects_per_call_headers() {
+    fn applies_headers() {
         let opts = ClientOptionsBuilder::default()
             .identity("enchicat".to_string())
             .target_url(Url::parse("https://smolkitty").unwrap())
@@ -1568,16 +1598,55 @@ mod tests {
             .build()
             .unwrap();
 
-        let mut static_headers = HashMap::new();
-        static_headers.insert("enchi".to_string(), "kitty".to_string());
-        let mut iceptor = ServiceCallInterceptor {
+        // Initial header set
+        let headers = Arc::new(RwLock::new(ClientHeaders {
+            user_headers: HashMap::new(),
+            api_key: Some("my-api-key".to_owned()),
+        }));
+        headers
+            .clone()
+            .write()
+            .user_headers
+            .insert("my-meta-key".to_owned(), "my-meta-val".to_owned());
+        let mut interceptor = ServiceCallInterceptor {
             opts,
-            headers: Arc::new(RwLock::new(static_headers)),
+            headers: headers.clone(),
         };
+
+        // Confirm on metadata
+        let req = interceptor.call(tonic::Request::new(())).unwrap();
+        assert_eq!(req.metadata().get("my-meta-key").unwrap(), "my-meta-val");
+        assert_eq!(
+            req.metadata().get("authorization").unwrap(),
+            "Bearer my-api-key"
+        );
+
+        // Overwrite at request time
         let mut req = tonic::Request::new(());
-        req.metadata_mut().insert("enchi", "cat".parse().unwrap());
-        let next_req = iceptor.call(req).unwrap();
-        assert_eq!(next_req.metadata().get("enchi").unwrap(), "cat");
+        req.metadata_mut()
+            .insert("my-meta-key", "my-meta-val2".parse().unwrap());
+        req.metadata_mut()
+            .insert("authorization", "my-api-key2".parse().unwrap());
+        let req = interceptor.call(req).unwrap();
+        assert_eq!(req.metadata().get("my-meta-key").unwrap(), "my-meta-val2");
+        assert_eq!(req.metadata().get("authorization").unwrap(), "my-api-key2");
+
+        // Overwrite auth on header
+        headers
+            .clone()
+            .write()
+            .user_headers
+            .insert("authorization".to_owned(), "my-api-key3".to_owned());
+        let req = interceptor.call(tonic::Request::new(())).unwrap();
+        assert_eq!(req.metadata().get("my-meta-key").unwrap(), "my-meta-val");
+        assert_eq!(req.metadata().get("authorization").unwrap(), "my-api-key3");
+
+        // Remove headers and auth and confirm gone
+        headers.clone().write().user_headers.clear();
+        headers.clone().write().api_key.take();
+        let req = interceptor.call(tonic::Request::new(())).unwrap();
+        assert!(!req.metadata().contains_key("my-meta-key"));
+        assert!(!req.metadata().contains_key("authorization"));
     }
 
     #[test]
diff --git a/client/src/raw.rs b/client/src/raw.rs
@@ -1015,7 +1015,7 @@ mod tests {
     #[allow(dead_code)]
     async fn raw_client_retry_compiles() {
         let opts = ClientOptionsBuilder::default().build().unwrap();
-        let raw_client = opts.connect_no_namespace(None, None).await.unwrap();
+        let raw_client = opts.connect_no_namespace(None).await.unwrap();
         let mut retry_client = RetryClient::new(raw_client, opts.retry_config);
 
         let list_ns_req = ListNamespacesRequest::default();
diff --git a/core/src/ephemeral_server/mod.rs b/core/src/ephemeral_server/mod.rs
@@ -307,11 +307,7 @@ impl EphemeralServer {
             .build()?;
         for _ in 0..50 {
             sleep(Duration::from_millis(100)).await;
-            if client_options
-                .connect_no_namespace(None, None)
-                .await
-                .is_ok()
-            {
+            if client_options.connect_no_namespace(None).await.is_ok() {
                 return success;
             }
         }
diff --git a/sdk/src/lib.rs b/sdk/src/lib.rs
@@ -17,7 +17,7 @@
 //! async fn main() -> Result<(), Box<dyn std::error::Error>> {
 //!     let server_options = sdk_client_options(Url::from_str("http://localhost:7233")?).build()?;
 //!
-//!     let client = server_options.connect("default", None, None).await?;
+//!     let client = server_options.connect("default", None).await?;
 //!
 //!     let telemetry_options = TelemetryOptionsBuilder::default().build()?;
 //!     let runtime = CoreRuntime::new_assume_tokio(telemetry_options)?;
diff --git a/test-utils/src/histfetch.rs b/test-utils/src/histfetch.rs
@@ -11,7 +11,7 @@ use temporal_sdk_core_test_utils::get_integ_server_options;
 #[tokio::main]
 async fn main() -> Result<(), anyhow::Error> {
     let gw_opts = get_integ_server_options();
-    let client = gw_opts.connect("default", None, None).await?;
+    let client = gw_opts.connect("default", None).await?;
     let wf_id = std::env::args()
         .nth(1)
         .expect("must provide workflow id as only argument");
diff --git a/test-utils/src/lib.rs b/test-utils/src/lib.rs
@@ -378,7 +378,7 @@ impl CoreWfStarter {
                     .expect("Worker config must be valid");
                 let client = Arc::new(
                     get_integ_server_options()
-                        .connect(cfg.namespace.clone(), None, None)
+                        .connect(cfg.namespace.clone(), None)
                         .await
                         .expect("Must connect"),
                 );
diff --git a/tests/integ_tests/client_tests.rs b/tests/integ_tests/client_tests.rs
@@ -17,7 +17,7 @@ async fn can_use_retry_client() {
 #[tokio::test]
 async fn can_use_retry_raw_client() {
     let opts = get_integ_server_options();
-    let raw_client = opts.connect_no_namespace(None, None).await.unwrap();
+    let raw_client = opts.connect_no_namespace(None).await.unwrap();
     let mut retry_client = RetryClient::new(raw_client, opts.retry_config);
     retry_client
         .describe_namespace(DescribeNamespaceRequest {
@@ -31,6 +31,6 @@ async fn can_use_retry_raw_client() {
 #[tokio::test]
 async fn calls_get_system_info() {
     let opts = get_integ_server_options();
-    let raw_client = opts.connect_no_namespace(None, None).await.unwrap();
+    let raw_client = opts.connect_no_namespace(None).await.unwrap();
     assert!(raw_client.get_client().capabilities().is_some());
 }
diff --git a/tests/integ_tests/ephemeral_server_tests.rs b/tests/integ_tests/ephemeral_server_tests.rs
@@ -124,7 +124,7 @@ async fn assert_ephemeral_server(server: &EphemeralServer) {
         .client_version("0.1.0".to_string())
         .build()
         .unwrap()
-        .connect_no_namespace(None, None)
+        .connect_no_namespace(None)
         .await
         .unwrap();
     let resp = client
diff --git a/tests/integ_tests/metrics_tests.rs b/tests/integ_tests/metrics_tests.rs
@@ -72,7 +72,7 @@ async fn prometheus_metrics_exported() {
     let rt = CoreRuntime::new_assume_tokio(telemopts).unwrap();
     let opts = get_integ_server_options();
     let mut raw_client = opts
-        .connect_no_namespace(rt.telemetry().get_temporal_metric_meter(), None)
+        .connect_no_namespace(rt.telemetry().get_temporal_metric_meter())
         .await
         .unwrap();
     assert!(raw_client.get_client().capabilities().is_some());
@@ -125,7 +125,7 @@ async fn one_slot_worker_reports_available_slot() {
 
     let client = Arc::new(
         get_integ_server_options()
-            .connect(worker_cfg.namespace.clone(), None, None)
+            .connect(worker_cfg.namespace.clone(), None)
             .await
             .expect("Must connect"),
     );
@@ -453,7 +453,7 @@ fn runtime_new() {
     let opts = get_integ_server_options();
     handle.block_on(async {
         let mut raw_client = opts
-            .connect_no_namespace(rt.telemetry().get_temporal_metric_meter(), None)
+            .connect_no_namespace(rt.telemetry().get_temporal_metric_meter())
             .await
             .unwrap();
         assert!(raw_client.get_client().capabilities().is_some());
diff --git a/tests/integ_tests/visibility_tests.rs b/tests/integ_tests/visibility_tests.rs
@@ -91,7 +91,7 @@ async fn client_list_open_closed_workflow_executions() {
 async fn client_create_namespace() {
     let client = Arc::new(
         get_integ_server_options()
-            .connect(NAMESPACE.to_owned(), None, None)
+            .connect(NAMESPACE.to_owned(), None)
             .await
             .expect("Must connect"),
     );
@@ -138,7 +138,7 @@ async fn client_create_namespace() {
 async fn client_describe_namespace() {
     let client = Arc::new(
         get_integ_server_options()
-            .connect(NAMESPACE.to_owned(), None, None)
+            .connect(NAMESPACE.to_owned(), None)
             .await
             .expect("Must connect"),
     );
diff --git a/tests/integ_tests/workflow_tests/eager.rs b/tests/integ_tests/workflow_tests/eager.rs
@@ -33,7 +33,7 @@ async fn eager_wf_start_different_clients() {
     worker.register_wf(wf_name.to_owned(), eager_wf);
 
     let client = get_integ_server_options()
-        .connect(NAMESPACE.to_string(), None, None)
+        .connect(NAMESPACE.to_string(), None)
         .await
         .expect("Should connect");
     let w = starter.get_worker().await;
diff --git a/tests/main.rs b/tests/main.rs
@@ -38,7 +38,7 @@ mod integ_tests {
         let opts = get_integ_server_options();
         let runtime = CoreRuntime::new_assume_tokio(get_integ_telem_options()).unwrap();
         let mut retrying_client = opts
-            .connect_no_namespace(runtime.telemetry().get_temporal_metric_meter(), None)
+            .connect_no_namespace(runtime.telemetry().get_temporal_metric_meter())
             .await
             .unwrap();
 
@@ -85,7 +85,7 @@ mod integ_tests {
             .build()
             .unwrap();
         let con = sgo
-            .connect("spencer.temporal-dev".to_string(), None, None)
+            .connect("spencer.temporal-dev".to_string(), None)
             .await
             .unwrap();
         dbg!(con

Original file line number	Diff line number	Diff line change
`@@ -307,11 +307,7 @@ impl EphemeralServer {`
`307`	`307`	`.build()?;`
`308`	`308`	`for _ in 0..50 {`
`309`	`309`	`sleep(Duration::from_millis(100)).await;`
`310`		`- if client_options`
`311`		`- .connect_no_namespace(None, None)`
`312`		`- .await`
`313`		`- .is_ok()`
`314`		`- {`
	`310`	`+ if client_options.connect_no_namespace(None).await.is_ok() {`
`315`	`311`	`return success;`
`316`	`312`	`}`
`317`	`313`	`}`