diff --git a/backend/.sqlx/query-1ba2e23d4ba816048ec1e88af9e342867fc0443cabea16d111afa2b91d3fe03b.json b/backend/.sqlx/query-1ba2e23d4ba816048ec1e88af9e342867fc0443cabea16d111afa2b91d3fe03b.json deleted file mode 100644 index 3552ba7a79..0000000000 --- a/backend/.sqlx/query-1ba2e23d4ba816048ec1e88af9e342867fc0443cabea16d111afa2b91d3fe03b.json +++ /dev/null @@ -1,65 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "SELECT client, refresh_token, grant_type, cc_client_id, cc_client_secret, cc_token_url, mcp_server_url, is_workspace_integration FROM account WHERE workspace_id = $1 AND id = $2", - "describe": { - "columns": [ - { - "ordinal": 0, - "name": "client", - "type_info": "Varchar" - }, - { - "ordinal": 1, - "name": "refresh_token", - "type_info": "Varchar" - }, - { - "ordinal": 2, - "name": "grant_type", - "type_info": "Varchar" - }, - { - "ordinal": 3, - "name": "cc_client_id", - "type_info": "Varchar" - }, - { - "ordinal": 4, - "name": "cc_client_secret", - "type_info": "Varchar" - }, - { - "ordinal": 5, - "name": "cc_token_url", - "type_info": "Varchar" - }, - { - "ordinal": 6, - "name": "mcp_server_url", - "type_info": "Text" - }, - { - "ordinal": 7, - "name": "is_workspace_integration", - "type_info": "Bool" - } - ], - "parameters": { - "Left": [ - "Text", - "Int4" - ] - }, - "nullable": [ - false, - false, - false, - true, - true, - true, - true, - false - ] - }, - "hash": "1ba2e23d4ba816048ec1e88af9e342867fc0443cabea16d111afa2b91d3fe03b" -} diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index 480dfc6747..f94cdbf754 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -62f6bd6238e313fe88c716bceefe9c94a19ba9eb +d842747738a2f10fc2fd0cd61f536efffcb45e41 diff --git a/backend/windmill-oauth/src/lib.rs b/backend/windmill-oauth/src/lib.rs index aa9524fe6c..e0c499c355 100644 --- a/backend/windmill-oauth/src/lib.rs +++ b/backend/windmill-oauth/src/lib.rs @@ -576,7 +576,18 @@ pub async fn exchange_token( Ok(token) } -/// Refresh an OAuth token and update the database +/// Pre-fetched account fields needed for token refresh. +pub struct OAuthAccountInfo { + pub client: String, + pub refresh_token: String, + pub grant_type: String, + pub cc_client_id: Option, + pub cc_client_secret: Option, + pub cc_token_url: Option, +} + +/// Refresh an OAuth token and update the database. +/// Fetches the account from DB, then delegates to `refresh_token_for_account`. pub async fn refresh_token<'c>( mut tx: Transaction<'c, Postgres>, path: &str, @@ -587,7 +598,8 @@ pub async fn refresh_token<'c>( http_client: &reqwest::Client, connect_configs_json: &str, ) -> error::Result { - let account = sqlx::query!( + let account = sqlx::query_as!( + OAuthAccountInfo, "SELECT client, refresh_token, grant_type, cc_client_id, cc_client_secret, cc_token_url FROM account WHERE workspace_id = $1 AND id = $2", w_id, id, @@ -595,6 +607,22 @@ pub async fn refresh_token<'c>( .fetch_optional(&mut *tx) .await?; let account = windmill_common::utils::not_found_if_none(account, "Account", &id.to_string())?; + + refresh_token_for_account(tx, path, w_id, id, db, account, oauth_clients, http_client, connect_configs_json).await +} + +/// Refresh an OAuth token given pre-fetched account info (no additional SELECT). +pub async fn refresh_token_for_account<'c>( + mut tx: Transaction<'c, Postgres>, + path: &str, + w_id: &str, + id: i32, + db: &DB, + account: OAuthAccountInfo, + oauth_clients: &AllClients, + http_client: &reqwest::Client, + connect_configs_json: &str, +) -> error::Result { let oauth_client_info = oauth_clients .connects .get(&account.client) diff --git a/backend/windmill-store/src/oauth_refresh_oss.rs b/backend/windmill-store/src/oauth_refresh_oss.rs index e2e1355cac..7402e66e49 100644 --- a/backend/windmill-store/src/oauth_refresh_oss.rs +++ b/backend/windmill-store/src/oauth_refresh_oss.rs @@ -8,8 +8,6 @@ #[cfg(feature = "private")] pub use crate::oauth_refresh_ee::_refresh_token; -#[cfg(feature = "private")] -pub use crate::oauth_refresh_ee::_refresh_workspace_integration_token; #[cfg(not(feature = "private"))] use sqlx::{Postgres, Transaction}; @@ -38,156 +36,3 @@ pub async fn _refresh_token<'c>( ) .await } - -#[cfg(not(feature = "private"))] -pub async fn _refresh_workspace_integration_token<'c>( - mut tx: Transaction<'c, Postgres>, - path: &str, - w_id: &str, - account_id: i32, - db: &DB, - client_name: &str, - refresh_token: &str, -) -> error::Result { - use windmill_common::global_settings::{ - get_instance_oauth_credentials, workspace_integration_auth_endpoint, - workspace_integration_oauth_key, workspace_integration_token_endpoint, - }; - use windmill_common::utils::now_from_db; - use windmill_common::variables::{build_crypt, encrypt}; - use windmill_oauth::{OClient, RefreshToken, Url, OAUTH_HTTP_CLIENT}; - - tracing::info!( - client = %client_name, - workspace_id = %w_id, - account_id = %account_id, - "Refreshing workspace integration OAuth token" - ); - - let oauth_data: serde_json::Value = sqlx::query_scalar( - "SELECT oauth_data FROM workspace_integrations \ - WHERE workspace_id = $1 AND service_name::text = $2", - ) - .bind(w_id) - .bind(client_name) - .fetch_optional(&mut *tx) - .await? - .ok_or_else(|| { - error::Error::NotFound(format!( - "Workspace integration for {} not found or not configured", - client_name - )) - })?; - - let is_instance_shared = oauth_data - .get("instance_shared") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - - let (client_id, client_secret, base_url); - if is_instance_shared { - let oauth_key = workspace_integration_oauth_key(client_name); - let (id, secret) = get_instance_oauth_credentials(db, oauth_key).await?; - client_id = id; - client_secret = secret; - base_url = String::new(); - } else { - client_id = oauth_data["client_id"] - .as_str() - .ok_or_else(|| { - error::Error::InternalErr("Missing client_id in workspace integration".into()) - })? - .to_string(); - client_secret = oauth_data["client_secret"] - .as_str() - .ok_or_else(|| { - error::Error::InternalErr( - "Missing client_secret in workspace integration".into(), - ) - })? - .to_string(); - base_url = oauth_data["base_url"].as_str().unwrap_or("").to_string(); - } - - let token_endpoint = workspace_integration_token_endpoint(client_name, &base_url); - let auth_endpoint = workspace_integration_auth_endpoint(client_name, &base_url); - - let auth_url = Url::parse(&auth_endpoint) - .map_err(|e| error::Error::InternalErr(format!("Invalid auth URL: {}", e)))?; - let token_url = Url::parse(&token_endpoint) - .map_err(|e| error::Error::InternalErr(format!("Invalid token URL: {}", e)))?; - - let mut client = OClient::new(client_id, auth_url, token_url); - client.set_client_secret(client_secret); - - let token = client - .exchange_refresh_token(&RefreshToken::from(refresh_token)) - .with_client(&*OAUTH_HTTP_CLIENT) - .execute::() - .await - .map_err(|e| { - error::Error::InternalErr(format!( - "Failed to refresh workspace integration token: {:?}", - e - )) - })?; - - #[derive(serde::Deserialize)] - struct WsTokenResponse { - access_token: String, - refresh_token: Option, - expires_in: Option, - } - - let token_result: WsTokenResponse = serde_json::from_value(token) - .map_err(|e| error::Error::InternalErr(format!("Failed to parse token response: {}", e)))?; - - let expires_at = now_from_db(&mut *tx).await? - + chrono::Duration::try_seconds( - token_result - .expires_in - .ok_or_else(|| { - error::Error::InternalErr("expires_in expected and not found".into()) - })? - .try_into() - .unwrap(), - ) - .unwrap_or_default(); - - sqlx::query( - "UPDATE account SET refresh_token = $1, expires_at = $2, refresh_error = NULL \ - WHERE workspace_id = $3 AND id = $4", - ) - .bind( - token_result - .refresh_token - .as_deref() - .unwrap_or(refresh_token), - ) - .bind(expires_at) - .bind(w_id) - .bind(account_id) - .execute(&mut *tx) - .await?; - tx.commit().await?; - - let token_str = &token_result.access_token; - let mc = build_crypt(db, w_id).await?; - let encrypted_token = encrypt(&mc, token_str); - - sqlx::query("UPDATE variable SET value = $1 WHERE workspace_id = $2 AND path = $3") - .bind(encrypted_token) - .bind(w_id) - .bind(path) - .execute(db) - .await?; - - tracing::info!( - client = %client_name, - workspace_id = %w_id, - account_id = %account_id, - "Workspace integration OAuth token refreshed successfully" - ); - - Ok(token_result.access_token) -}