mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-08-23 08:00:45 +00:00
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:
+15
@@ -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"
|
||||
}
|
||||
+2
-2
@@ -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"
|
||||
}
|
||||
-15
@@ -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"
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user