starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-08-19 16:02:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(operator): refresh IAM RDS / Entra ID tokens in operator process (#9141)

Browse Source 
* fix(operator): refresh IAM RDS / Entra ID tokens in operator process

* fix: gate DEFAULT_MAX_CONNECTIONS_OPERATOR on operator feature
This commit is contained in:
Ruben Fiszel
2026-05-13 06:38:57 +00:00
committed by GitHub
parent d0f23cc523
commit 7ebb08133c
2 changed files with 99 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/src/db_connect.rs
+81 -54
View File
@@ -6,6 +6,8 @@ use windmill_common::{
pub const DEFAULT_MAX_CONNECTIONS_SERVER: u32 = 50;
pub const DEFAULT_MAX_CONNECTIONS_WORKER: u32 = 5;
pub const DEFAULT_MAX_CONNECTIONS_INDEXER: u32 = 5;
#[cfg(feature = "operator")]
pub const DEFAULT_MAX_CONNECTIONS_OPERATOR: u32 = 2;
pub async fn initial_connection() -> Result<sqlx::Pool<sqlx::Postgres>, error::Error> {
let connect_options = get_database_url().await?.connect_options().await?;
@@ -16,12 +18,35 @@ pub async fn initial_connection() -> Result<sqlx::Pool<sqlx::Postgres>, error::E
.map_err(|err| Error::ConnectingToDatabase(err.to_string()))
}
/// Connect to the database for the Kubernetes operator process.
///
/// Long-running operator pods need IAM RDS / Entra ID token refresh just like the server,
/// otherwise new pool connections start failing once the initial token expires (~15 min).
#[cfg(feature = "operator")]
pub async fn operator_connection(
#[cfg(all(feature = "enterprise", feature = "private"))]
killpill_rx: tokio::sync::broadcast::Receiver<()>,
) -> anyhow::Result<sqlx::Pool<sqlx::Postgres>> {
let database_url = get_database_url().await?;
let pool = connect(
database_url.clone(),
DEFAULT_MAX_CONNECTIONS_OPERATOR,
false,
)
.await?;
#[cfg(all(feature = "enterprise", feature = "private"))]
spawn_token_refresh_task(pool.clone(), database_url, killpill_rx);
Ok(pool)
}
pub async fn connect_db(
server_mode: bool,
indexer_mode: bool,
worker_mode: bool,
num_workers: i32,
#[cfg(feature = "private")] mut killpill_rx: tokio::sync::broadcast::Receiver<()>,
#[cfg(feature = "private")] killpill_rx: tokio::sync::broadcast::Receiver<()>,
) -> anyhow::Result<sqlx::Pool<sqlx::Postgres>> {
use anyhow::Context;
@@ -43,70 +68,72 @@ pub async fn connect_db(
let pool = connect(database_url.clone(), max_connections, worker_mode).await?;
#[cfg(all(feature = "enterprise", feature = "private"))]
{
let needs_token_refresh = matches!(
database_url,
DatabaseUrl::IamRds(_) | DatabaseUrl::EntraId(_)
);
let label = match &database_url {
DatabaseUrl::IamRds(_) => "IAM RDS",
DatabaseUrl::EntraId(_) => "Entra ID",
DatabaseUrl::Static(_) => "",
};
if needs_token_refresh {
let pool2 = pool.clone();
let database_url2 = database_url.clone();
tokio::spawn(async move {
loop {
tokio::select! {
_ = killpill_rx.recv() => {
break;
}
_ = tokio::time::sleep(std::time::Duration::from_secs(10)) => {
if !database_url2.needs_refresh().await {
continue;
}
let new_url = tokio::time::timeout(
std::time::Duration::from_secs(10),
get_database_url(),
)
.await;
match new_url {
Ok(Ok(new_url)) => {
match new_url.connect_options().await {
Ok(connect_options) => {
pool2.set_connect_options(connect_options);
tracing::info!("Refreshed {label} URL successfully");
}
Err(e) => {
tracing::error!(
"Error getting {label} connect options, retrying in 10s: {e}"
);
continue;
}
}
}
Ok(Err(e)) => {
tracing::error!(
"Error refreshing {label} URL, trying again in 10s: {e}"
);
continue;
spawn_token_refresh_task(pool.clone(), database_url, killpill_rx);
Ok(pool)
}
/// Spawn a background task that refreshes IAM RDS / Entra ID tokens before they expire
/// and updates the pool's connect options so new connections use the fresh token.
/// No-op for static (password-based) database URLs.
#[cfg(all(feature = "enterprise", feature = "private"))]
pub fn spawn_token_refresh_task(
pool: sqlx::Pool<sqlx::Postgres>,
database_url: DatabaseUrl,
mut killpill_rx: tokio::sync::broadcast::Receiver<()>,
) {
let label = match &database_url {
DatabaseUrl::IamRds(_) => "IAM RDS",
DatabaseUrl::EntraId(_) => "Entra ID",
DatabaseUrl::Static(_) => return,
};
tokio::spawn(async move {
loop {
tokio::select! {
_ = killpill_rx.recv() => {
break;
}
_ = tokio::time::sleep(std::time::Duration::from_secs(10)) => {
if !database_url.needs_refresh().await {
continue;
}
let new_url = tokio::time::timeout(
std::time::Duration::from_secs(10),
get_database_url(),
)
.await;
match new_url {
Ok(Ok(new_url)) => {
match new_url.connect_options().await {
Ok(connect_options) => {
pool.set_connect_options(connect_options);
tracing::info!("Refreshed {label} URL successfully");
}
Err(e) => {
tracing::error!(
"Timeout after 10s refreshing {label} URL, trying again in 10s: {e}"
"Error getting {label} connect options, retrying in 10s: {e}"
);
continue;
}
}
}
Ok(Err(e)) => {
tracing::error!(
"Error refreshing {label} URL, trying again in 10s: {e}"
);
continue;
}
Err(e) => {
tracing::error!(
"Timeout after 10s refreshing {label} URL, trying again in 10s: {e}"
);
continue;
}
}
}
});
}
}
}
Ok(pool)
});
}
pub async fn connect(
backend/src/main.rs
+18 -1
View File
@@ -670,7 +670,24 @@ async fn windmill_main() -> anyhow::Result<()> {
tracing_subscriber::fmt::init();
tracing::info!("Starting Windmill Kubernetes operator...");
tracing::info!("Connecting to database...");
let db = crate::db_connect::initial_connection().await?;
#[cfg(all(feature = "enterprise", feature = "private"))]
let (operator_killpill_tx, operator_killpill_rx) =
tokio::sync::broadcast::channel::<()>(2);
let db = crate::db_connect::operator_connection(
#[cfg(all(feature = "enterprise", feature = "private"))]
operator_killpill_rx,
)
.await?;
#[cfg(all(feature = "enterprise", feature = "private"))]
tokio::spawn(async move {
if let Ok(()) = tokio::signal::ctrl_c().await {
let _ = operator_killpill_tx.send(());
}
});
tracing::info!("Database connected. Starting ConfigMap watcher...");
windmill_operator::run(db).await?;
return Ok(());