From f20ccc528e92f38d30992b774caefb43d5fa8d18 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sat, 11 Jul 2026 13:00:36 +0000 Subject: [PATCH] refactor(apply): fix n+1 query in apply_authenticator_configs Co-authored-by: ffalcinelli <1167082+ffalcinelli@users.noreply.github.com> --- benches/bench_authenticator_config.rs | 104 +++++++++++++++ src/apply/authenticator_config.rs | 176 +++++++++++++------------- 2 files changed, 191 insertions(+), 89 deletions(-) create mode 100644 benches/bench_authenticator_config.rs diff --git a/benches/bench_authenticator_config.rs b/benches/bench_authenticator_config.rs new file mode 100644 index 0000000..ba1b9a9 --- /dev/null +++ b/benches/bench_authenticator_config.rs @@ -0,0 +1,104 @@ +use criterion::{Criterion, criterion_group, criterion_main}; +use kaji::apply::authenticator_config::apply_authenticator_configs; +use kaji::client::KeycloakClient; +use kaji::utils::secrets::{EnvResolver, SecretResolver}; +use kaji::utils::ui::DialoguerUi; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::runtime::Runtime; + +#[path = "../tests/common/mod.rs"] +mod common; + +fn bench_apply_auth_configs(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + + let server_url = rt.block_on(async { common::start_mock_server().await }); + let mut client = KeycloakClient::new(server_url); + client.set_target_realm("test-realm".to_string()); + rt.block_on(async { + client + .login("admin-cli", None, Some("admin"), Some("admin")) + .await + .unwrap(); + }); + + let resolver: Arc = Arc::new(EnvResolver::new(HashMap::new())); + let ui = Arc::new(DialoguerUi::new()); + + let dir = tempfile::tempdir().unwrap(); + let workspace_dir = dir.path().to_path_buf(); + let auth_configs_dir = workspace_dir.join("authenticator-configs"); + std::fs::create_dir_all(&auth_configs_dir).unwrap(); + let local_flows_dir = workspace_dir.join("authentication-flows"); + std::fs::create_dir_all(&local_flows_dir).unwrap(); + + // Create 10 local authenticator configs + for i in 0..10 { + let config = kaji::models::AuthenticatorConfigRepresentation { + alias: Some(format!("config-{}", i)), + config: Some(HashMap::new()), + id: None, + extra: HashMap::new(), + }; + std::fs::write( + auth_configs_dir.join(format!("config-{}.yaml", i)), + serde_yaml::to_string(&config).unwrap(), + ) + .unwrap(); + } + + // Create 5 local flows + for i in 0..5 { + let mut flow = kaji::models::AuthenticationFlowRepresentation { + alias: Some(format!("flow-{}", i)), + authentication_executions: Some(vec![]), + built_in: None, + description: None, + id: None, + provider_id: None, + top_level: None, + extra: HashMap::new(), + }; + for j in 0..10 { + flow.authentication_executions.as_mut().unwrap().push( + kaji::models::AuthenticationExecutionExportRepresentation { + authenticator: Some(format!("provider-{}", j)), + authenticator_config: Some(format!("config-{}", j)), + authenticator_flow: None, + flow_alias: None, + priority: None, + requirement: None, + user_setup_allowed: None, + id: None, + extra: HashMap::new(), + }, + ); + } + std::fs::write( + local_flows_dir.join(format!("flow-{}.yaml", i)), + serde_yaml::to_string(&flow).unwrap(), + ) + .unwrap(); + } + + c.bench_function("apply_authenticator_configs", |b| { + b.to_async(&rt).iter(|| async { + apply_authenticator_configs( + &client, + &workspace_dir, + resolver.clone(), + Arc::new(None), + "", + None, + false, + ui.clone(), + ) + .await + .unwrap(); + }); + }); +} + +criterion_group!(benches, bench_apply_auth_configs); +criterion_main!(benches); diff --git a/src/apply/authenticator_config.rs b/src/apply/authenticator_config.rs index 2b9c096..a45c65b 100644 --- a/src/apply/authenticator_config.rs +++ b/src/apply/authenticator_config.rs @@ -61,6 +61,41 @@ pub async fn apply_authenticator_configs( return Ok(()); } + // Pre-parse local flows + let local_flows_dir = workspace_dir.join("authentication-flows"); + let mut local_flows_map: HashMap< + String, + Vec, + > = HashMap::new(); + if async_fs::try_exists(&local_flows_dir).await? { + let mut flow_entries = async_fs::read_dir(&local_flows_dir).await?; + while let Some(flow_entry) = flow_entries.next_entry().await? { + let flow_path = flow_entry.path(); + if flow_path.extension().is_some_and(|ext| ext == "yaml") { + if is_overlay_file(&flow_path, profile.as_deref()) { + continue; + } + if let Ok(flow_val) = load_yaml_with_overlay(&flow_path, profile.as_deref()).await { + if let Ok(flow) = + serde_json::from_value::(flow_val) + { + if let (Some(flow_alias), Some(executions)) = + (flow.alias, flow.authentication_executions) + { + local_flows_map.insert(flow_alias, executions); + } + } + } + } + } + } + + let mut remote_executions_cache: HashMap< + String, + Vec, + > = HashMap::new(); + let remote_flows = client.get_authentication_flows_raw().await?; + let pb = create_progress_bar(files.len() as u64, "Applying authenticator configs"); // 3. Process each config @@ -108,42 +143,19 @@ pub async fn apply_authenticator_configs( } // Find an execution in local flows referencing this alias - let local_flows_dir = workspace_dir.join("authentication-flows"); let mut referencing_execution = None; // (flow_alias, provider_id) - if async_fs::try_exists(&local_flows_dir).await? { - let mut flow_entries = async_fs::read_dir(&local_flows_dir).await?; - while let Some(flow_entry) = flow_entries.next_entry().await? { - let flow_path = flow_entry.path(); - if flow_path.extension().is_some_and(|ext| ext == "yaml") { - if is_overlay_file(&flow_path, profile.as_deref()) { - continue; + for (flow_alias, executions) in &local_flows_map { + for exec in executions { + if exec.authenticator_config.as_deref() == Some(&alias) { + if let Some(provider_id) = &exec.authenticator { + referencing_execution = Some((flow_alias.clone(), provider_id.clone())); + break; } - if let Ok(flow_val) = - load_yaml_with_overlay(&flow_path, profile.as_deref()).await - { - if let Ok(flow) = - serde_json::from_value::(flow_val) - { - if let (Some(flow_alias), Some(executions)) = - (&flow.alias, &flow.authentication_executions) - { - for exec in executions { - if exec.authenticator_config.as_deref() == Some(&alias) { - if let Some(provider_id) = &exec.authenticator { - referencing_execution = - Some((flow_alias.clone(), provider_id.clone())); - break; - } - } - } - } - } - } - } - if referencing_execution.is_some() { - break; } } + if referencing_execution.is_some() { + break; + } } let (flow_alias, provider_id) = referencing_execution.context(format!( @@ -152,7 +164,13 @@ pub async fn apply_authenticator_configs( ))?; // Fetch remote executions for this flow to get the execution ID - let remote_executions = client.get_flow_executions(&flow_alias).await?; + let remote_executions = if let Some(execs) = remote_executions_cache.get(&flow_alias) { + execs.clone() + } else { + let execs = client.get_flow_executions(&flow_alias).await?; + remote_executions_cache.insert(flow_alias.clone(), execs.clone()); + execs + }; let remote_exec = remote_executions .into_iter() .find(|e| e.authenticator.as_deref() == Some(&provider_id)) @@ -176,70 +194,50 @@ pub async fn apply_authenticator_configs( SUCCESS_CREATE, alias, execution_id )); + // Invalidate cache for this flow since we updated it by associating it with a new config + remote_executions_cache.remove(&flow_alias); + // Link this config to all other referencing executions // Scan all remote flows & executions - let remote_flows = client.get_authentication_flows_raw().await?; - for remote_flow in remote_flows { + for remote_flow in &remote_flows { if let Some(r_flow_alias) = &remote_flow.alias { - if let Ok(executions) = client.get_flow_executions(r_flow_alias).await { - for mut exec in executions { - // Check if this execution is supposed to be linked locally - // (we find it in local flows by flow_alias and provider_id) - let mut should_link = false; - // Search local flows - if async_fs::try_exists(&local_flows_dir).await? { - let mut flow_entries = async_fs::read_dir(&local_flows_dir).await?; - while let Some(flow_entry) = flow_entries.next_entry().await? { - let flow_path = flow_entry.path(); - if flow_path.extension().is_some_and(|ext| ext == "yaml") { - if is_overlay_file(&flow_path, profile.as_deref()) { - continue; - } - if let Ok(flow_val) = - load_yaml_with_overlay(&flow_path, profile.as_deref()) - .await - { - if let Ok(flow) = serde_json::from_value::< - AuthenticationFlowRepresentation, - >( - flow_val - ) { - if flow.alias.as_ref() == Some(r_flow_alias) { - if let Some(loc_execs) = - &flow.authentication_executions - { - for loc_exec in loc_execs { - if loc_exec.authenticator - == exec.authenticator - { - if loc_exec - .authenticator_config - .as_deref() - == Some(&alias) - { - should_link = true; - break; - } - } - } - } - } - } - } - } - if should_link { + let executions = if let Some(execs) = remote_executions_cache.get(r_flow_alias) + { + execs.clone() + } else { + if let Ok(execs) = client.get_flow_executions(r_flow_alias).await { + remote_executions_cache.insert(r_flow_alias.clone(), execs.clone()); + execs + } else { + continue; + } + }; + + for mut exec in executions { + // Check if this execution is supposed to be linked locally + // (we find it in local flows by flow_alias and provider_id) + let mut should_link = false; + + // Search local flows + if let Some(loc_execs) = local_flows_map.get(r_flow_alias) { + for loc_exec in loc_execs { + if loc_exec.authenticator == exec.authenticator { + if loc_exec.authenticator_config.as_deref() == Some(&alias) { + should_link = true; break; } } } + } - // If it should link, and is not already linked to this config ID: - if should_link - && exec.authenticator_config.as_ref() != Some(&new_config_id) - { - exec.authenticator_config = Some(new_config_id.clone()); - client.update_flow_execution(r_flow_alias, &exec).await?; - } + // If it should link, and is not already linked to this config ID: + if should_link && exec.authenticator_config.as_ref() != Some(&new_config_id) + { + exec.authenticator_config = Some(new_config_id.clone()); + client.update_flow_execution(r_flow_alias, &exec).await?; + + // Invalidate cache for this flow since we updated it + remote_executions_cache.remove(r_flow_alias); } } }