diff --git a/backend/.sqlx/query-67170c7e1a3bfeab685716ec352271f41021a4be2e351e6ef96d7af70358f0aa.json b/backend/.sqlx/query-67170c7e1a3bfeab685716ec352271f41021a4be2e351e6ef96d7af70358f0aa.json new file mode 100644 index 0000000000..a67274af63 --- /dev/null +++ b/backend/.sqlx/query-67170c7e1a3bfeab685716ec352271f41021a4be2e351e6ef96d7af70358f0aa.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO usage (id, is_workspace, month_, usage)\n VALUES ($1, FALSE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), $2)\n ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + EXCLUDED.usage", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Varchar", + "Int4" + ] + }, + "nullable": [] + }, + "hash": "67170c7e1a3bfeab685716ec352271f41021a4be2e351e6ef96d7af70358f0aa" +} diff --git a/backend/.sqlx/query-124e67b0cee1baa6295846db4ad6242a39dd40186f1dbb48ad3018bd9f6913ec.json b/backend/.sqlx/query-6f660e55963ac74db95c44fc95da542a18812b403641104bbd24599ee4b9d187.json similarity index 51% rename from backend/.sqlx/query-124e67b0cee1baa6295846db4ad6242a39dd40186f1dbb48ad3018bd9f6913ec.json rename to backend/.sqlx/query-6f660e55963ac74db95c44fc95da542a18812b403641104bbd24599ee4b9d187.json index 51a1ec68cb..e623c0c1e0 100644 --- a/backend/.sqlx/query-124e67b0cee1baa6295846db4ad6242a39dd40186f1dbb48ad3018bd9f6913ec.json +++ b/backend/.sqlx/query-6f660e55963ac74db95c44fc95da542a18812b403641104bbd24599ee4b9d187.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "\n SELECT\n w.premium,\n COALESCE(cw.is_past_due, false) as \"is_past_due!\",\n cw.max_tolerated_executions\n FROM\n workspace w\n LEFT JOIN cloud_workspace_settings cw ON cw.workspace_id = w.id\n WHERE\n w.id = $1\n ", + "query": "\n SELECT\n w.premium,\n COALESCE(cw.is_past_due, false) as \"is_past_due!\",\n cw.max_tolerated_executions\n FROM\n workspace w\n LEFT JOIN cloud_workspace_settings cw ON cw.workspace_id = w.id\n WHERE\n w.id = $1\n ", "describe": { "columns": [ { @@ -30,5 +30,5 @@ true ] }, - "hash": "124e67b0cee1baa6295846db4ad6242a39dd40186f1dbb48ad3018bd9f6913ec" + "hash": "6f660e55963ac74db95c44fc95da542a18812b403641104bbd24599ee4b9d187" } diff --git a/backend/.sqlx/query-c41787a3efbb2c520a8fee93b2078c45e2a2615a8fdda3cc70778f8d037da2cc.json b/backend/.sqlx/query-c41787a3efbb2c520a8fee93b2078c45e2a2615a8fdda3cc70778f8d037da2cc.json deleted file mode 100644 index 88a6899759..0000000000 --- a/backend/.sqlx/query-c41787a3efbb2c520a8fee93b2078c45e2a2615a8fdda3cc70778f8d037da2cc.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "INSERT INTO usage (id, is_workspace, month_, usage)\n VALUES ($1, FALSE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), $2)\n ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + EXCLUDED.usage", - "describe": { - "columns": [], - "parameters": { - "Left": [ - "Varchar", - "Int4" - ] - }, - "nullable": [] - }, - "hash": "c41787a3efbb2c520a8fee93b2078c45e2a2615a8fdda3cc70778f8d037da2cc" -} diff --git a/backend/windmill-api/src/workspaces.rs b/backend/windmill-api/src/workspaces.rs index ff399ee7c6..345ff73086 100644 --- a/backend/windmill-api/src/workspaces.rs +++ b/backend/windmill-api/src/workspaces.rs @@ -439,7 +439,7 @@ async fn is_premium( require_admin(authed.is_admin, &authed.username)?; #[cfg(feature = "cloud")] let premium = windmill_common::workspaces::get_team_plan_status(&_db, &_w_id) - .await + .await? .premium; #[cfg(not(feature = "cloud"))] let premium = false; diff --git a/backend/windmill-common/src/workspaces.rs b/backend/windmill-common/src/workspaces.rs index 94aca66eef..edd9c8abe5 100644 --- a/backend/windmill-common/src/workspaces.rs +++ b/backend/windmill-common/src/workspaces.rs @@ -1,4 +1,6 @@ use async_recursion::async_recursion; +#[cfg(feature = "cloud")] +use backon::{ConstantBuilder, Retryable}; use quick_cache::sync::Cache; use serde::{Deserialize, Serialize}; use serde_json::json; @@ -94,35 +96,56 @@ lazy_static::lazy_static! { } #[cfg(feature = "cloud")] -pub async fn get_team_plan_status(_db: &crate::DB, _w_id: &str) -> TeamPlanStatus { +pub async fn get_team_plan_status(_db: &crate::DB, _w_id: &str) -> Result { let cached = TEAM_PLAN_CACHE.get(_w_id); if let Some(cached) = cached { - return cached; + return Ok(cached); } - let team_plan_info = sqlx::query_as!( - TeamPlanStatus, - r#" - SELECT - w.premium, - COALESCE(cw.is_past_due, false) as "is_past_due!", - cw.max_tolerated_executions - FROM - workspace w - LEFT JOIN cloud_workspace_settings cw ON cw.workspace_id = w.id - WHERE - w.id = $1 - "#, - _w_id + + let team_plan_info = (|| async { + sqlx::query_as!( + TeamPlanStatus, + r#" + SELECT + w.premium, + COALESCE(cw.is_past_due, false) as "is_past_due!", + cw.max_tolerated_executions + FROM + workspace w + LEFT JOIN cloud_workspace_settings cw ON cw.workspace_id = w.id + WHERE + w.id = $1 + "#, + _w_id + ) + .fetch_optional(_db) + .await + }) + .retry( + ConstantBuilder::default() + .with_delay(std::time::Duration::from_secs(5)) + .with_max_times(10), ) - .fetch_one(_db) + .notify(|err, dur| { + tracing::error!( + "Failed to get team plan status for workspace {_w_id} (will retry in {dur:?}): {err:#}" + ); + }) .await - .unwrap_or_else(|_| TeamPlanStatus { + .map_err(|err| { + Error::internal_err(format!( + "Failed to get team plan status for workspace {_w_id} after 10 retries: {err:#}" + )) + })? + .unwrap_or_else(|| TeamPlanStatus { premium: false, is_past_due: false, max_tolerated_executions: None, }); + TEAM_PLAN_CACHE.insert(_w_id.to_string(), team_plan_info.clone()); - team_plan_info + + Ok(team_plan_info) } #[derive(Deserialize, Serialize, Debug)] diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 2ee22ae088..acf0b5df76 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -1435,9 +1435,6 @@ fn apply_completed_job_cloud_usage( let email2 = email.clone(); tokio::task::spawn(async move { let additional_usage = _duration / 1000; - let premium_workspace = windmill_common::workspaces::get_team_plan_status(&db, &w_id) - .await - .premium; let result = tokio::time::timeout(std::time::Duration::from_secs(10), async move { // Update workspace usage let workspace_result = sqlx::query!( @@ -1454,22 +1451,30 @@ fn apply_completed_job_cloud_usage( tracing::error!("Failed to update workspace usage for {}: {:#}", w_id, e); } - // Update user usage for non-premium workspaces - if !premium_workspace { - let user_result = sqlx::query!( - "INSERT INTO usage (id, is_workspace, month_, usage) - VALUES ($1, FALSE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), $2) - ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + EXCLUDED.usage", - &email, - additional_usage as i32 - ) - .execute(&db) - .await; + match windmill_common::workspaces::get_team_plan_status(&db, &w_id).await { + Ok(team_plan_status) => { + // Update user usage for non-premium workspaces + if !team_plan_status.premium { + let user_result = sqlx::query!( + "INSERT INTO usage (id, is_workspace, month_, usage) + VALUES ($1, FALSE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), $2) + ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + EXCLUDED.usage", + &email, + additional_usage as i32 + ) + .execute(&db) + .await; - if let Err(e) = user_result { - tracing::error!("Failed to update user usage for {}: {:#}", email, e); + if let Err(e) = user_result { + tracing::error!("Failed to update user usage for {}: {:#}", email, e); + } + } + }, + Err(err) => { + tracing::error!("Failed to get team plan status to update usage for workspace {w_id}: {err:#}"); } - } + }; + }).await; if let Err(_) = result { @@ -3273,10 +3278,7 @@ lazy_static::lazy_static! { const SUPERADMIN_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(60); #[cfg(feature = "cloud")] -async fn is_superadmin_cached( - db: &Pool, - email: &str, -) -> Result { +async fn is_superadmin_cached(db: &Pool, email: &str) -> Result { let now = std::time::Instant::now(); // Try to get from cache first @@ -3290,18 +3292,19 @@ async fn is_superadmin_cached( } // Cache miss or expired, fetch from database - let is_super_admin = sqlx::query_scalar!( - "SELECT super_admin FROM password WHERE email = $1", - email - ) - .fetch_optional(db) - .await? - .unwrap_or(false); + let is_super_admin = + sqlx::query_scalar!("SELECT super_admin FROM password WHERE email = $1", email) + .fetch_optional(db) + .await? + .unwrap_or(false); // Update cache { let mut cache = SUPERADMIN_CACHE.write().await; - cache.insert(email.to_string(), (is_super_admin, now + SUPERADMIN_CACHE_TTL)); + cache.insert( + email.to_string(), + (is_super_admin, now + SUPERADMIN_CACHE_TTL), + ); } Ok(is_super_admin) @@ -3347,11 +3350,7 @@ async fn check_usage_limits( } #[cfg(feature = "cloud")] -fn increment_usage_async( - db: Pool, - workspace_id: String, - email: Option, -) { +fn increment_usage_async(db: Pool, workspace_id: String, email: Option) { tokio::task::spawn(async move { let result = tokio::time::timeout(std::time::Duration::from_secs(10), async { // Update workspace usage @@ -3430,7 +3429,8 @@ pub async fn push<'c, 'd>( ) -> Result<(Uuid, Transaction<'c, Postgres>), Error> { #[cfg(feature = "cloud")] if *CLOUD_HOSTED { - let team_plan_status = windmill_common::workspaces::get_team_plan_status(_db, workspace_id).await; + let team_plan_status = + windmill_common::workspaces::get_team_plan_status(_db, workspace_id).await?; // we track only non flow steps let (workspace_usage, user_usage) = if !matches!( job_payload, @@ -3445,7 +3445,11 @@ pub async fn push<'c, 'd>( increment_usage_async( _db.clone(), workspace_id.to_string(), - if !team_plan_status.premium { Some(email.to_string()) } else { None }, + if !team_plan_status.premium { + Some(email.to_string()) + } else { + None + }, ); // Return the current usage + 1 to account for this job diff --git a/backend/windmill-worker/src/common.rs b/backend/windmill-worker/src/common.rs index b64b9558bf..bcc529258d 100644 --- a/backend/windmill-worker/src/common.rs +++ b/backend/windmill-worker/src/common.rs @@ -669,13 +669,16 @@ pub async fn resolve_job_timeout( ) -> (Duration, Option, bool) { let mut warn_msg: Option = None; #[cfg(feature = "cloud")] - let cloud_premium_workspace = *CLOUD_HOSTED + let cloud_premium_workspace = + *CLOUD_HOSTED && windmill_common::workspaces::get_team_plan_status( _conn.as_sql().expect("cloud cannot use http connection"), _w_id, ) .await - .premium; + .inspect_err(|err| tracing::error!("Failed to get team plan status to resolve job timeout for workspace {_w_id}: {err:#}")) + .map(|s| s.premium) + .unwrap_or(true); #[cfg(not(feature = "cloud"))] let cloud_premium_workspace = false;