mirror of
https://github.com/bartvdbraak/keyweave.git
synced 2025-04-28 23:31:20 +00:00
feat: multi-threading secret lookups
This commit is contained in:
parent
dc1d8a7a1f
commit
cd7f7f59ab
1 changed files with 32 additions and 4 deletions
34
src/main.rs
34
src/main.rs
|
@ -4,6 +4,7 @@ use clap::Parser;
|
||||||
use futures::stream::StreamExt;
|
use futures::stream::StreamExt;
|
||||||
use std::fs::File;
|
use std::fs::File;
|
||||||
use std::io::Write;
|
use std::io::Write;
|
||||||
|
use tokio::sync::mpsc;
|
||||||
|
|
||||||
#[derive(Parser)]
|
#[derive(Parser)]
|
||||||
#[clap(author, version, about, long_about = None)]
|
#[clap(author, version, about, long_about = None)]
|
||||||
|
@ -46,15 +47,42 @@ async fn fetch_secrets_from_key_vault(
|
||||||
|
|
||||||
while let Some(page) = secret_pages.next().await {
|
while let Some(page) = secret_pages.next().await {
|
||||||
let page = page?;
|
let page = page?;
|
||||||
|
let (tx, mut rx) = mpsc::channel(32); // Channel for concurrent secret retrieval
|
||||||
|
|
||||||
for secret in &page.value {
|
for secret in &page.value {
|
||||||
if let Some(filter) = filter {
|
if let Some(filter) = filter {
|
||||||
if !secret.id.contains(filter) {
|
if !secret.id.contains(filter) {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
let secret_name = secret.id.split('/').last().unwrap_or_default();
|
let tx = tx.clone();
|
||||||
let secret_bundle = client.get(secret_name).await?;
|
|
||||||
secret_values.push((secret.id.clone(), secret_bundle.value));
|
// Clone necessary data before moving into the spawned task
|
||||||
|
let secret_id = secret.id.clone();
|
||||||
|
let client_clone = client.clone();
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let secret_name = secret_id.split('/').last().unwrap_or_default();
|
||||||
|
let secret_bundle = client_clone.get(secret_name).await;
|
||||||
|
|
||||||
|
// Handle the result and send it through the channel
|
||||||
|
match secret_bundle {
|
||||||
|
Ok(bundle) => {
|
||||||
|
tx.send((secret_id, bundle.value)).await.expect("Send error");
|
||||||
|
}
|
||||||
|
Err(err) => {
|
||||||
|
eprintln!("Error fetching secret: {}", err);
|
||||||
|
// You can decide to continue or not in case of an error.
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
drop(tx); // Drop the sender to signal the end of sending tasks
|
||||||
|
|
||||||
|
while let Some(result) = rx.recv().await {
|
||||||
|
let (key, value) = result;
|
||||||
|
secret_values.push((key, value));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
Loading…
Reference in a new issue