Skip to content

Commit

Permalink
chore: Refactor secrets loading to avoid use of futures::executor::bl…
Browse files Browse the repository at this point in the history
…ock_on

Since the caller was async no need to have the nested functions be sync.

Signed-off-by: Jesse Szwedko <[email protected]>
  • Loading branch information
jszwedko committed Aug 14, 2024
1 parent 61b1b18 commit 5523b8d
Show file tree
Hide file tree
Showing 7 changed files with 47 additions and 45 deletions.
1 change: 1 addition & 0 deletions src/config/loading/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,7 @@ pub async fn load_from_paths_with_provider_and_secrets(
debug!(message = "Secret placeholders found, retrieving secrets from configured backends.");
let resolved_secrets = secrets_backends_loader
.retrieve(&mut signal_handler.subscribe())
.await
.map_err(|e| vec![e])?;
load_builder_from_paths_with_secrets(config_paths, resolved_secrets)?
} else {
Expand Down
49 changes: 27 additions & 22 deletions src/config/loading/secret.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use std::{
io::Read,
};

use futures::TryFutureExt;
use indexmap::IndexMap;
use once_cell::sync::Lazy;
use regex::{Captures, Regex};
Expand Down Expand Up @@ -51,31 +52,35 @@ impl SecretBackendLoader {
}
}

pub(crate) fn retrieve(
pub(crate) async fn retrieve(
&mut self,
signal_rx: &mut signal::SignalRx,
) -> Result<HashMap<String, String>, String> {
let secrets = self.secret_keys.iter().flat_map(|(backend_name, keys)| {
match self.backends.get_mut(&ComponentKey::from(backend_name.clone())) {
None => {
vec![Err(format!("Backend \"{}\" is required for secret retrieval but was not found in config.", backend_name))]
},
Some(backend) => {
debug!(message = "Retrieving secret from a backend.", backend = ?backend_name);
match backend.retrieve(keys.clone(), signal_rx) {
Err(e) => {
vec![Err(format!("Error while retrieving secret from backend \"{}\": {}.", backend_name, e))]
},
Ok(s) => {
s.into_iter().map(|(k, v)| {
trace!(message = "Successfully retrieved a secret.", backend = ?backend_name, secret_key = ?k);
Ok((format!("{}.{}", backend_name, k), v))
}).collect::<Vec<Result<(String, String), String>>>()
}
}
},
}
}).collect::<Result<HashMap<String, String>, String>>()?;
let mut secrets: HashMap<String, String> = HashMap::new();

for (backend_name, keys) in &self.secret_keys {
let backend = self.backends.get_mut(&ComponentKey::from(backend_name.clone())).ok_or_else(|| format!("Backend \"{}\" is required for secret retrieval but was not found in config.", backend_name))?;

debug!(message = "Retrieving secrets from a backend.", backend = ?backend_name, keys = ?keys);
let backend_secrets = backend
.retrieve(keys.clone(), signal_rx)
.map_ok(|backend_secrets| {
backend_secrets.into_iter().map(|(k, v)| {
trace!(message = "Successfully retrieved a secret.", backend = ?backend_name, key = ?k);
(format!("{}.{}", backend_name, k), v)
}).collect::<HashMap<String, String>>()
})
.map_err(|e| {
format!(
"Error while retrieving secret from backend \"{}\": {}.",
backend_name, e
)
})
.await?;

secrets.extend(backend_secrets);
}

Ok(secrets)
}

Expand Down
2 changes: 1 addition & 1 deletion src/config/secret.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use crate::signal;
/// Generalized interface to a secret backend.
#[enum_dispatch]
pub trait SecretBackend: NamedComponent + core::fmt::Debug + Send + Sync {
fn retrieve(
async fn retrieve(
&mut self,
secret_keys: HashSet<String>,
signal_rx: &mut signal::SignalRx,
Expand Down
1 change: 1 addition & 0 deletions src/config/unit_test/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,7 @@ pub async fn build_unit_tests_main(
let config_builder = if secrets_backends_loader.has_secrets_to_retrieve() {
let resolved_secrets = secrets_backends_loader
.retrieve(&mut signal_handler.subscribe())
.await
.map_err(|e| vec![e])?;
loading::load_builder_from_paths_with_secrets(paths, resolved_secrets)?
} else {
Expand Down
35 changes: 15 additions & 20 deletions src/secrets/aws_secrets_manager.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
use std::collections::{HashMap, HashSet};

use aws_sdk_secretsmanager::{config, Client};
use futures::executor;
use vector_lib::configurable::{component::GenerateConfig, configurable_component};

use crate::aws::{create_client, AwsAuthentication, ClientBuilder, RegionOrEndpoint};
Expand Down Expand Up @@ -52,30 +51,26 @@ impl GenerateConfig for AwsSecretsManagerBackend {
}

impl SecretBackend for AwsSecretsManagerBackend {
fn retrieve(
async fn retrieve(
&mut self,
secret_keys: HashSet<String>,
_: &mut signal::SignalRx,
) -> crate::Result<HashMap<String, String>> {
let client = executor::block_on(async {
create_client::<SecretsManagerClientBuilder>(
&self.auth,
self.region.region(),
self.region.endpoint(),
&ProxyConfig::default(),
&self.tls,
&None,
)
.await
})?;
let client = create_client::<SecretsManagerClientBuilder>(
&self.auth,
self.region.region(),
self.region.endpoint(),
&ProxyConfig::default(),
&self.tls,
&None,
)
.await?;

let get_secret_value_response = executor::block_on(async {
client
.get_secret_value()
.secret_id(&self.secret_id)
.send()
.await
})?;
let get_secret_value_response = client
.get_secret_value()
.secret_id(&self.secret_id)
.send()
.await?;

let secret_string = get_secret_value_response
.secret_string
Expand Down
2 changes: 1 addition & 1 deletion src/secrets/exec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ struct ExecResponse {
}

impl SecretBackend for ExecBackend {
fn retrieve(
async fn retrieve(
&mut self,
secret_keys: HashSet<String>,
signal_rx: &mut signal::SignalRx,
Expand Down
2 changes: 1 addition & 1 deletion src/secrets/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ pub struct TestBackend {
impl_generate_config_from_default!(TestBackend);

impl SecretBackend for TestBackend {
fn retrieve(
async fn retrieve(
&mut self,
secret_keys: HashSet<String>,
_: &mut signal::SignalRx,
Expand Down

0 comments on commit 5523b8d

Please sign in to comment.