Support environment variable substitution

In the Kafka, Polling HTTP, WebSocket, and SSE sources, support environment variables in the passwords, secrets, and headers fields. This is done by setting the format of the field in the JSON schema to 'var-str', which tells typify to replace the type with the VarStr struct, which has a function for substituting environment variables. The VarStr struct serializes/deserializes from a JSON string.
ArroyoSystems · Dec 6, 2023 · 20efd7e · 20efd7e
1 parent 4d3ff50
commit 20efd7e
Show file tree

Hide file tree

Showing 21 changed files with 335 additions and 116 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/arroyo-connectors/src/kafka.rs b/arroyo-connectors/src/kafka.rs
@@ -1,21 +1,21 @@
 use anyhow::{anyhow, bail};
-use arroyo_rpc::OperatorConfig;
-use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
-use std::convert::Infallible;
-use typify::import_types;
-
 use arroyo_rpc::api_types::connections::{ConnectionProfile, ConnectionSchema, TestSourceMessage};
+use arroyo_rpc::OperatorConfig;
+use arroyo_types::VarStr;
 use axum::response::sse::Event;
 use rdkafka::{
     consumer::{BaseConsumer, Consumer},
     message::BorrowedMessage,
     ClientConfig, Offset, TopicPartitionList,
 };
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::convert::Infallible;
 use std::time::{Duration, Instant};
 use tokio::sync::mpsc::Sender;
 use tonic::Status;
 use tracing::{error, info, warn};
+use typify::import_types;
 
 use crate::{pull_opt, Connection, ConnectionType};
 
@@ -25,7 +25,13 @@ const CONFIG_SCHEMA: &str = include_str!("../../connector-schemas/kafka/connecti
 const TABLE_SCHEMA: &str = include_str!("../../connector-schemas/kafka/table.json");
 const ICON: &str = include_str!("../resources/kafka.svg");
 
-import_types!(schema = "../connector-schemas/kafka/connection.json",);
+import_types!(
+    schema = "../connector-schemas/kafka/connection.json",
+    convert = {
+        {type = "string", format = "var-str"} = VarStr,
+        {type = "string", format = "var-str", isSensitive = true} = VarStr
+    }
+);
 import_types!(schema = "../connector-schemas/kafka/table.json");
 
 pub struct KafkaConnector {}
@@ -150,14 +156,16 @@ impl Connector for KafkaConnector {
                         mechanism: pull_opt("auth.mechanism", options)?,
                         protocol: pull_opt("auth.protocol", options)?,
                         username: pull_opt("auth.username", options)?,
-                        password: pull_opt("auth.password", options)?,
+                        password: VarStr::new(pull_opt("auth.password", options)?),
                     },
                     Some(other) => bail!("unknown auth type '{}'", other),
                 };
 
                 let schema_registry = options.remove("schema_registry.endpoint").map(|endpoint| {
                     let api_key = options.remove("schema_registry.api_key");
-                    let api_secret = options.remove("schema_registry.api_secret");
+                    let api_secret = options
+                        .remove("schema_registry.api_secret")
+                        .map(|secret| VarStr::new(secret));
                     SchemaRegistry::ConfluentSchemaRegistry {
                         endpoint,
                         api_key,
@@ -251,7 +259,7 @@ impl KafkaTester {
                 client_config.set("sasl.mechanism", mechanism);
                 client_config.set("security.protocol", protocol);
                 client_config.set("sasl.username", username);
-                client_config.set("sasl.password", password);
+                client_config.set("sasl.password", password.sub_env_vars()?);
             }
         };
 

diff --git a/arroyo-connectors/src/lib.rs b/arroyo-connectors/src/lib.rs
@@ -331,7 +331,7 @@ pub(crate) fn nullable_field(name: &str, field_type: SourceFieldType) -> SourceF
     }
 }
 
-fn construct_http_client(endpoint: &str, headers: Option<&String>) -> anyhow::Result<Client> {
+fn construct_http_client(endpoint: &str, headers: Option<String>) -> anyhow::Result<Client> {
     if let Err(e) = reqwest::Url::parse(&endpoint) {
         bail!("invalid endpoint '{}': {:?}", endpoint, e)
     };

diff --git a/arroyo-connectors/src/polling_http.rs b/arroyo-connectors/src/polling_http.rs
@@ -3,7 +3,7 @@ use std::convert::Infallible;
 
 use anyhow::anyhow;
 use arroyo_rpc::OperatorConfig;
-use arroyo_types::string_to_map;
+use arroyo_types::{string_to_map, VarStr};
 use axum::response::sse::Event;
 use reqwest::{Client, Request};
 use tokio::sync::mpsc::Sender;
@@ -20,7 +20,10 @@ use super::Connector;
 
 const TABLE_SCHEMA: &str = include_str!("../../connector-schemas/polling_http/table.json");
 
-import_types!(schema = "../connector-schemas/polling_http/table.json");
+import_types!(
+    schema = "../connector-schemas/polling_http/table.json",
+    convert = { {type = "string", format = "var-str"} = VarStr }
+);
 const ICON: &str = include_str!("../resources/http.svg");
 
 pub struct PollingHTTPConnector {}
@@ -55,8 +58,14 @@ impl PollingHTTPConnector {
         config: &PollingHttpTable,
         tx: Sender<Result<Event, Infallible>>,
     ) -> anyhow::Result<()> {
-        let client =
-            construct_http_client(&config.endpoint, config.headers.as_ref().map(|t| &t.0))?;
+        let headers = config
+            .headers
+            .as_ref()
+            .map(|s| s.sub_env_vars())
+            .transpose()
+            .map_err(|e| anyhow!(e))?;
+
+        let client = construct_http_client(&config.endpoint, headers)?;
         let req = Self::construct_test_request(&client, config)?;
 
         tx.send(Ok(Event::default()
@@ -126,7 +135,7 @@ impl Connector for PollingHTTPConnector {
                 Err(err) => TestSourceMessage {
                     error: true,
                     done: true,
-                    message: format!("{:?}", err),
+                    message: format!("{:?}", err.root_cause()),
                 },
             };
 
@@ -166,7 +175,7 @@ impl Connector for PollingHTTPConnector {
             EmptyConfig {},
             PollingHttpTable {
                 endpoint,
-                headers: headers.map(Headers),
+                headers: headers.map(|s| VarStr::new(s)),
                 method,
                 body,
                 poll_interval_ms: interval,
@@ -187,7 +196,7 @@ impl Connector for PollingHTTPConnector {
         let description = format!("PollingHTTPSource<{}>", table.endpoint);
 
         if let Some(headers) = &table.headers {
-            string_to_map(headers).ok_or_else(|| {
+            string_to_map(&headers.sub_env_vars().map_err(|e| anyhow!(e))?).ok_or_else(|| {
                 anyhow!(
                     "Invalid format for headers; should be a \
                     comma-separated list of colon-separated key value pairs"

diff --git a/arroyo-connectors/src/sse.rs b/arroyo-connectors/src/sse.rs
@@ -4,7 +4,7 @@ use std::time::Duration;
 
 use anyhow::{anyhow, bail};
 use arroyo_rpc::OperatorConfig;
-use arroyo_types::string_to_map;
+use arroyo_types::{string_to_map, VarStr};
 use axum::response::sse::Event;
 use eventsource_client::Client;
 use futures::StreamExt;
@@ -22,7 +22,7 @@ use super::Connector;
 
 const TABLE_SCHEMA: &str = include_str!("../../connector-schemas/sse/table.json");
 
-import_types!(schema = "../connector-schemas/sse/table.json");
+import_types!(schema = "../connector-schemas/sse/table.json", convert = { {type = "string", format = "var-str"} = VarStr });
 const ICON: &str = include_str!("../resources/sse.svg");
 
 pub struct SSEConnector {}
@@ -79,7 +79,7 @@ impl Connector for SSEConnector {
         let description = format!("SSESource<{}>", table.endpoint);
 
         if let Some(headers) = &table.headers {
-            string_to_map(headers).ok_or_else(|| {
+            string_to_map(&headers.sub_env_vars().map_err(|e| anyhow!(e))?).ok_or_else(|| {
                 anyhow!(
                     "Invalid format for headers; should be a \
                     comma-separated list of colon-separated key value pairs"
@@ -134,7 +134,7 @@ impl Connector for SSEConnector {
             SseTable {
                 endpoint,
                 events,
-                headers: headers.map(Headers),
+                headers: headers.map(|s| VarStr::new(s)),
             },
             schema,
         )
@@ -174,11 +174,14 @@ impl SseTester {
             .map_err(|_| anyhow!("Endpoint URL is invalid"))?;
 
         let headers = string_to_map(
-            self.config
+            &self
+                .config
                 .headers
                 .as_ref()
-                .map(|t| t.0.as_str())
-                .unwrap_or(""),
+                .map(|s| s.sub_env_vars())
+                .transpose()
+                .map_err(|e| anyhow!(e))?
+                .unwrap_or("".to_string()),
         )
         .ok_or_else(|| anyhow!("Headers are invalid; should be comma-separated pairs"))?;
 

diff --git a/arroyo-connectors/src/webhook.rs b/arroyo-connectors/src/webhook.rs
@@ -4,6 +4,7 @@ use std::convert::Infallible;
 use anyhow::anyhow;
 use arroyo_rpc::OperatorConfig;
 
+use arroyo_types::VarStr;
 use axum::response::sse::Event;
 use reqwest::{Client, Request};
 use serde_json::json;
@@ -21,7 +22,7 @@ use super::Connector;
 
 const TABLE_SCHEMA: &str = include_str!("../../connector-schemas/webhook/table.json");
 
-import_types!(schema = "../connector-schemas/webhook/table.json");
+import_types!(schema = "../connector-schemas/webhook/table.json", convert = { {type = "string", format = "var-str"} = VarStr });
 const ICON: &str = include_str!("../resources/webhook.svg");
 
 pub struct WebhookConnector {}
@@ -47,8 +48,14 @@ impl WebhookConnector {
         config: &WebhookTable,
         tx: Sender<Result<Event, Infallible>>,
     ) -> anyhow::Result<()> {
-        let client =
-            construct_http_client(&config.endpoint, config.headers.as_ref().map(|t| &t.0))?;
+        let headers = config
+            .headers
+            .as_ref()
+            .map(|s| s.sub_env_vars())
+            .transpose()
+            .map_err(|e| anyhow!(e))?;
+
+        let client = construct_http_client(&config.endpoint, headers)?;
         let req = Self::construct_test_request(&client, config)?;
 
         tx.send(Ok(Event::default()
@@ -176,14 +183,19 @@ impl Connector for WebhookConnector {
     ) -> anyhow::Result<Connection> {
         let endpoint = pull_opt("endpoint", options)?;
 
-        let headers = options
-            .remove("headers")
-            .map(|s| s.try_into())
-            .transpose()
-            .map_err(|e| anyhow!("invalid value for 'headers' config: {:?}", e))?;
+        let headers = options.remove("headers").map(|s| VarStr::new(s));
 
         let table = WebhookTable { endpoint, headers };
-        let client = construct_http_client(&table.endpoint, table.headers.as_ref().map(|t| &t.0))?;
+
+        let client = construct_http_client(
+            &table.endpoint,
+            table
+                .headers
+                .as_ref()
+                .map(|s| s.sub_env_vars())
+                .transpose()
+                .map_err(|e| anyhow!(e))?,
+        )?;
         let _ = Self::construct_test_request(&client, &table)?;
 
         self.from_config(None, name, EmptyConfig {}, table, schema)

diff --git a/arroyo-connectors/src/websocket.rs b/arroyo-connectors/src/websocket.rs
@@ -8,7 +8,7 @@ use arroyo_rpc::api_types::connections::{
     ConnectionProfile, ConnectionSchema, ConnectionType, TestSourceMessage,
 };
 use arroyo_rpc::OperatorConfig;
-use arroyo_types::string_to_map;
+use arroyo_types::{string_to_map, VarStr};
 use axum::response::sse::Event;
 use futures::{SinkExt, StreamExt};
 use serde::{Deserialize, Serialize};
@@ -25,7 +25,7 @@ use super::Connector;
 
 const TABLE_SCHEMA: &str = include_str!("../../connector-schemas/websocket/table.json");
 
-import_types!(schema = "../connector-schemas/websocket/table.json");
+import_types!(schema = "../connector-schemas/websocket/table.json", convert = { {type = "string", format = "var-str"} = VarStr });
 const ICON: &str = include_str!("../resources/websocket.svg");
 
 pub struct WebsocketConnector {}
@@ -79,16 +79,29 @@ impl Connector for WebsocketConnector {
                 }
             };
 
-            let headers =
-                match string_to_map(table.headers.as_ref().map(|t| t.0.as_str()).unwrap_or(""))
-                    .ok_or_else(|| anyhow!("Headers are invalid; should be comma-separated pairs"))
-                {
-                    Ok(headers) => headers,
-                    Err(e) => {
-                        send(true, true, format!("Failed to parse headers: {:?}", e)).await;
-                        return;
-                    }
-                };
+            let headers_str = match table
+                .headers
+                .as_ref()
+                .map(|s| s.sub_env_vars())
+                .transpose()
+                .map_err(|e| anyhow!(e))
+            {
+                Ok(headers) => headers,
+                Err(e) => {
+                    send(true, true, format!("{}", e)).await;
+                    return;
+                }
+            };
+
+            let headers = match string_to_map(&headers_str.unwrap_or("".to_string()))
+                .ok_or_else(|| anyhow!("Headers are invalid; should be comma-separated pairs"))
+            {
+                Ok(headers) => headers,
+                Err(e) => {
+                    send(true, true, format!("Failed to parse headers: {:?}", e)).await;
+                    return;
+                }
+            };
 
             let uri = match Uri::from_str(&table.endpoint.to_string()) {
                 Ok(uri) => uri,
@@ -215,7 +228,7 @@ impl Connector for WebsocketConnector {
         let description = format!("WebsocketSource<{}>", table.endpoint);
 
         if let Some(headers) = &table.headers {
-            string_to_map(headers).ok_or_else(|| {
+            string_to_map(&headers.sub_env_vars().map_err(|e| anyhow!(e))?).ok_or_else(|| {
                 anyhow!(
                     "Invalid format for headers; should be a \
                     comma-separated list of colon-separated key value pairs"
@@ -289,7 +302,7 @@ impl Connector for WebsocketConnector {
             EmptyConfig {},
             WebsocketTable {
                 endpoint,
-                headers: headers.map(Headers),
+                headers: headers.map(|s| VarStr::new(s)),
                 subscription_message: None,
                 subscription_messages,
             },

diff --git a/arroyo-console/src/routes/connections/JsonForm.tsx b/arroyo-console/src/routes/connections/JsonForm.tsx
@@ -635,6 +635,7 @@ export function JsonForm({
 }) {
   let ajv = new Ajv();
   ajv.addKeyword('isSensitive');
+  ajv.addFormat('var-str', { validate: () => true });
   const memoAjv = useMemo(() => addFormats(ajv), [schema]);
 
   const formik = useFormik({