fix(cloud): better errors when failing to get team plan status (#6908)

* fix(cloud): better errors when failing to get team plan status

* better errors

* fix build

* fix build
This commit is contained in:
hugocasa
2025-10-23 22:19:10 +02:00
committed by GitHub
parent 242a57f840
commit 3208f908ca
7 changed files with 105 additions and 75 deletions
@@ -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"
}
@@ -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"
}
@@ -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"
}
+1 -1
View File
@@ -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;
+42 -19
View File
@@ -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<TeamPlanStatus> {
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)]
+40 -36
View File
@@ -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<Postgres>,
email: &str,
) -> Result<bool, Error> {
async fn is_superadmin_cached(db: &Pool<Postgres>, email: &str) -> Result<bool, Error> {
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<Postgres>,
workspace_id: String,
email: Option<String>,
) {
fn increment_usage_async(db: Pool<Postgres>, workspace_id: String, email: Option<String>) {
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
+5 -2
View File
@@ -669,13 +669,16 @@ pub async fn resolve_job_timeout(
) -> (Duration, Option<String>, bool) {
let mut warn_msg: Option<String> = 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;