kalugin66/windmill
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add Entra ID (Azure Workload Identity) database auth (#8526)

Browse Source 
* feat: add Entra ID (Azure Workload Identity) support for database auth

Add support for Azure Workload Identity to authenticate to Azure Database
for PostgreSQL using short-lived Entra ID tokens. Mirrors the existing
AWS IAM RDS auth pattern.

- Extract shared DatabaseParams to db_params.rs for reuse across providers
- Add DatabaseUrl::EntraId variant with token refresh
- Detect "entraid" magic password in DATABASE_URL
- Unified background refresh task for both IAM RDS and Entra ID
- Support sovereign clouds via AZURE_AUTHORITY_HOST env var

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix: restore needs_refresh() check in background token refresh task

The unified refresh task was missing the needs_refresh() gate, causing
it to refresh tokens every 10 seconds instead of only when near expiry.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* chore: update ee-repo-ref.txt for Entra ID branch

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix: move entraid env var reads inside cfg(private) block

Fixes unused variable warnings in OSS and EE-without-private builds
where -D warnings is enabled.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* chore: update ee-repo-ref to 0e001bab643e449b3310b0692dd3598ee0902ecc

This commit updates the EE repository reference after PR #483 was merged in windmill-ee-private.

Previous ee-repo-ref: 44199013ed0c96680672e718f35124aa34a5d010

New ee-repo-ref: 0e001bab643e449b3310b0692dd3598ee0902ecc

Automated by sync-ee-ref workflow.

* refactor: add needs_refresh() and refresh_if_needed() to DatabaseUrl

Simplify duplicated refresh logic per Claude review suggestion.
Background task and get_database_url() now use shared methods
instead of matching on each variant individually.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: windmill-internal-app[bot] <windmill-internal-app[bot]@users.noreply.github.com>
This commit is contained in:
hugocasa
2026-04-02 16:00:33 +02:00
committed by GitHub
parent 8c3c97f7a6
commit 6a5cfbc159
4 changed files with 185 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
backend/ee-repo-ref.txt
View File

@@ -1 +1 @@
e08a87450627bef9013498e40ee93a47bedda7ee
0e001bab643e449b3310b0692dd3598ee0902ecc

63
backend/src/db_connect.rs
View File

@@ -41,50 +41,69 @@ pub async fn connect_db(
};
let pool = connect(database_url.clone(), max_connections, worker_mode).await?;
#[cfg(all(feature = "enterprise", feature = "private"))]
let pool2 = pool.clone();
#[cfg(all(feature = "enterprise", feature = "private"))]
if let DatabaseUrl::IamRds(database_url) = database_url {
tokio::spawn(async move {
loop {
tokio::select! {
_ = killpill_rx.recv() => {
break;
}
_ = tokio::time::sleep(std::time::Duration::from_secs(10)) => {
let needs_refresh = {
let read_guard = database_url.read().await;
read_guard.needs_refresh()
};
if needs_refresh {
let new_url = tokio::time::timeout(std::time::Duration::from_secs(10), get_database_url()).await;
{
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 IAM RDS URL successfully");
tracing::info!("Refreshed {label} URL successfully");
}
Err(e) => {
tracing::error!("Error getting IAM RDS connect options, retrying in 10s: {}", e);
tracing::error!(
"Error getting {label} connect options, retrying in 10s: {e}"
);
continue;
}
}
}
Ok(Err(e)) => {
tracing::error!("Error refreshing IAM RDS URL, trying again in 10s: {}", e);
tracing::error!(
"Error refreshing {label} URL, trying again in 10s: {e}"
);
continue;
}
Err(e) => {
tracing::error!("Timeout after 10s refreshing IAM RDS URL, trying again in 10 seconds: {}", e);
tracing::error!(
"Timeout after 10s refreshing {label} URL, trying again in 10s: {e}"
);
continue;
}
}
}
}
}
}
});
});
}
}
Ok(pool)

45
backend/windmill-common/src/db_params.rs Normal file
View File

@@ -0,0 +1,45 @@
use anyhow::Result;
/// Parsed database connection parameters, shared across DB auth providers (IAM RDS, Entra ID, etc.)
#[derive(Debug, Clone)]
pub struct DatabaseParams {
pub hostname: String,
pub port: u64,
pub username: String,
pub database: String,
}
/// Extract database connection parameters from a PostgreSQL URL
pub fn extract_database_params(database_url: &str) -> Result<DatabaseParams> {
let url = url::Url::parse(database_url)
.map_err(|e| anyhow::anyhow!("Failed to parse database URL: {}", e))?;
let hostname = url
.host_str()
.ok_or_else(|| anyhow::anyhow!("Database URL missing hostname"))?
.to_string();
let port = url.port().unwrap_or(5432) as u64;
let username = if url.username().is_empty() {
return Err(anyhow::anyhow!("Database URL missing username"));
} else {
urlencoding::decode(url.username())?.to_string()
};
let database = url
.path()
.trim_start_matches('/')
.split('/')
.next()
.filter(|s| !s.is_empty())
.ok_or_else(|| anyhow::anyhow!("Database URL missing database name"))?
.to_string();
Ok(DatabaseParams {
hostname,
port,
username,
database: urlencoding::decode(&database)?.to_string(),
})
}

123
backend/windmill-common/src/lib.rs
View File

@@ -43,7 +43,10 @@ pub mod cache;
pub mod client;
pub mod db;
#[cfg(all(feature = "enterprise", feature = "private"))]
mod db_entra_ee;
#[cfg(all(feature = "enterprise", feature = "private"))]
mod db_iam_ee;
pub mod db_params;
#[cfg(feature = "private")]
pub mod ee;
pub mod ee_oss;
@@ -633,12 +636,14 @@ impl PgDatabase {
pub enum DatabaseUrl {
#[cfg(all(feature = "enterprise", feature = "private"))]
IamRds(std::sync::Arc<tokio::sync::RwLock<db_iam_ee::IamRdsUrl>>),
#[cfg(all(feature = "enterprise", feature = "private"))]
EntraId(std::sync::Arc<tokio::sync::RwLock<db_entra_ee::EntraIdUrl>>),
Static(String),
}
impl DatabaseUrl {
/// Get the database URL as a string.
/// Note: For IAM RDS, this returns the original URL (for metadata extraction).
/// For token-based auth, this returns the original URL (for metadata extraction).
/// For actual database connections, use connect_options() instead.
pub async fn as_str(&self) -> String {
match self {
@@ -647,13 +652,18 @@ impl DatabaseUrl {
let guard = rds_url.read().await;
guard.as_str().to_string()
}
#[cfg(all(feature = "enterprise", feature = "private"))]
DatabaseUrl::EntraId(entra_url) => {
let guard = entra_url.read().await;
guard.as_str().to_string()
}
DatabaseUrl::Static(url) => url.clone(),
}
}
/// Get PgConnectOptions for this database URL.
/// For IAM RDS, this returns options built directly from the token to avoid double-encoding
/// issues with temporary credentials (IRSA/Pod Identity).
/// For token-based auth (IAM RDS, Entra ID), this returns options built directly from the
/// token to avoid double-encoding issues with temporary credentials.
/// For static URLs, this parses the URL string.
pub async fn connect_options(&self) -> Result<sqlx::postgres::PgConnectOptions, Error> {
match self {
@@ -662,6 +672,11 @@ impl DatabaseUrl {
let guard = rds_url.read().await;
Ok(guard.connect_options())
}
#[cfg(all(feature = "enterprise", feature = "private"))]
DatabaseUrl::EntraId(entra_url) => {
let guard = entra_url.read().await;
Ok(guard.connect_options())
}
DatabaseUrl::Static(url) => sqlx::postgres::PgConnectOptions::from_str(url)
.map_err(|e| Error::InternalErr(format!("Failed to parse database URL: {}", e))),
}
@@ -671,9 +686,31 @@ impl DatabaseUrl {
match self {
#[cfg(all(feature = "enterprise", feature = "private"))]
DatabaseUrl::IamRds(rds_url) => rds_url.write().await.refresh().await,
#[cfg(all(feature = "enterprise", feature = "private"))]
DatabaseUrl::EntraId(entra_url) => entra_url.write().await.refresh().await,
DatabaseUrl::Static(_) => Ok(()),
}
}
pub async fn needs_refresh(&self) -> bool {
match self {
#[cfg(all(feature = "enterprise", feature = "private"))]
DatabaseUrl::IamRds(rds_url) => rds_url.read().await.needs_refresh(),
#[cfg(all(feature = "enterprise", feature = "private"))]
DatabaseUrl::EntraId(entra_url) => entra_url.read().await.needs_refresh(),
DatabaseUrl::Static(_) => false,
}
}
/// Double-checked refresh: read-lock to check, then write-lock to refresh if still needed.
pub async fn refresh_if_needed(&self) -> Result<(), Error> {
if self.needs_refresh().await {
self.refresh().await.map_err(|e| {
Error::InternalErr(format!("Failed to refresh database token: {}", e))
})?;
}
Ok(())
}
}
static DATABASE_URL_CACHE: tokio::sync::OnceCell<DatabaseUrl> = tokio::sync::OnceCell::const_new();
@@ -701,7 +738,9 @@ pub async fn get_database_url() -> Result<DatabaseUrl, Error> {
let parsed_url = url::Url::parse(&url)?;
if parsed_url.password().is_some_and(|x| x == "iamrds") {
let password = parsed_url.password().unwrap_or_default();
if password == "iamrds" {
let region = var("AWS_REGION").map_err(|_| {
Error::BadConfig(
"AWS_REGION env var is required for IAM RDS authentication".to_string(),
@@ -731,34 +770,68 @@ pub async fn get_database_url() -> Result<DatabaseUrl, Error> {
"IAM RDS authentication is not enabled in OSS mode".to_string(),
));
}
} else if password == "entraid" {
let tenant_id = var("AZURE_TENANT_ID").map_err(|_| {
Error::BadConfig(
"AZURE_TENANT_ID env var is required for Entra ID authentication"
.to_string(),
)
})?;
tracing::info!(
"entraid mode detected, generating Entra ID URL for tenant: {tenant_id}"
);
#[cfg(all(feature = "enterprise", feature = "private"))]
{
let client_id = var("AZURE_CLIENT_ID").map_err(|_| {
Error::BadConfig(
"AZURE_CLIENT_ID env var is required for Entra ID authentication"
.to_string(),
)
})?;
let federated_token_file =
var("AZURE_FEDERATED_TOKEN_FILE").map_err(|_| {
Error::BadConfig(
"AZURE_FEDERATED_TOKEN_FILE env var is required for Entra ID authentication".to_string(),
)
})?;
let authority_host = var("AZURE_AUTHORITY_HOST")
.unwrap_or_else(|_| "login.microsoftonline.com".to_string());
let entra_url = db_entra_ee::generate_database_url(
&url,
&tenant_id,
&client_id,
&federated_token_file,
&authority_host,
)
.await
.map_err(|e| {
Error::InternalErr(format!(
"Failed to generate Entra ID database URL: {}",
e
))
})?;
tracing::info!("Entra ID URL generated successfully");
Ok::<DatabaseUrl, Error>(DatabaseUrl::EntraId(std::sync::Arc::new(
tokio::sync::RwLock::new(entra_url),
)))
}
#[cfg(not(all(feature = "enterprise", feature = "private")))]
{
return Err(Error::BadConfig(
"Entra ID authentication is not enabled in OSS mode".to_string(),
));
}
} else {
Ok::<DatabaseUrl, Error>(DatabaseUrl::Static(url.to_string()))
}
})
.await?;
// Check if we need to refresh and do so if necessary
#[cfg(all(feature = "enterprise", feature = "private"))]
if let DatabaseUrl::IamRds(ref rds_url_lock) = database_url {
// Check if refresh is needed
let needs_refresh = {
let read_guard = rds_url_lock.read().await;
read_guard.needs_refresh()
};
database_url.refresh_if_needed().await?;
// If refresh is needed, acquire write lock and refresh
if needs_refresh {
let mut write_guard = rds_url_lock.write().await;
// Double-check after acquiring write lock (another task might have refreshed)
if write_guard.needs_refresh() {
write_guard.refresh().await.map_err(|e| {
Error::InternalErr(format!("Failed to refresh IAM token: {}", e))
})?;
}
}
}
// Return the URL string
Ok(database_url.clone())
}