i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-07 01:47:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(admin-users): add plan entitlement revocation flow

Browse Source
This commit is contained in:
ZheFox
2026-09-02 11:23:24 +08:00
parent 6540c1e46a
commit 144a28f544
19 changed files with 726 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/adapters/sqlite/src/billing.rs
+167 -10
View File
@@ -920,6 +920,40 @@ ORDER BY expires_at ASC, created_at ASC
))
}
async fn revoke_user_plan_entitlement(
&self,
user_id: &str,
entitlement_id: &str,
) -> Result<AdminBillingMutationOutcome<()>, DataLayerError> {
let now = current_unix_secs_i64();
let result = sqlx::query(
r#"
UPDATE user_plan_entitlements
SET status = 'revoked',
expires_at = CASE WHEN expires_at > ? THEN ? ELSE expires_at END,
updated_at = ?
WHERE id = ?
AND user_id = ?
AND status = 'active'
AND expires_at > ?
"#,
)
.bind(now)
.bind(now)
.bind(now)
.bind(entitlement_id)
.bind(user_id)
.bind(now)
.execute(&self.pool)
.await
.map_sql_err()?;
if result.rows_affected() == 0 {
Ok(AdminBillingMutationOutcome::NotFound)
} else {
Ok(AdminBillingMutationOutcome::Applied(()))
}
}
async fn find_user_daily_quota_availability(
&self,
user_id: &str,
@@ -927,13 +961,19 @@ ORDER BY expires_at ASC, created_at ASC
let now_unix_secs = current_unix_secs_i64();
let rows = sqlx::query(
r#"
SELECT id, entitlements_snapshot
SELECT
user_plan_entitlements.id,
user_plan_entitlements.entitlements_snapshot,
billing_plans.entitlements_json AS plan_entitlements_json
FROM user_plan_entitlements
WHERE user_id = ?
AND status = 'active'
AND starts_at <= ?
AND expires_at > ?
ORDER BY expires_at ASC, created_at ASC, id ASC
JOIN billing_plans ON billing_plans.id = user_plan_entitlements.plan_id
WHERE user_plan_entitlements.user_id = ?
AND user_plan_entitlements.status = 'active'
AND user_plan_entitlements.starts_at <= ?
AND user_plan_entitlements.expires_at > ?
ORDER BY user_plan_entitlements.expires_at ASC,
user_plan_entitlements.created_at ASC,
user_plan_entitlements.id ASC
"#,
)
.bind(user_id)
@@ -948,9 +988,13 @@ ORDER BY expires_at ASC, created_at ASC, id ASC
let entitlement_id: String = row.try_get("id").map_sql_err()?;
let entitlements = parse_json(row.try_get("entitlements_snapshot").ok().flatten())?
.unwrap_or_else(|| serde_json::json!([]));
let plan_entitlements =
parse_json(row.try_get("plan_entitlements_json").ok().flatten())?
.unwrap_or_else(|| serde_json::json!([]));
grants.extend(daily_quota_grants_from_entitlement(
&entitlement_id,
&entitlements,
daily_quota_wallet_overage_policy(&plan_entitlements),
now,
)?);
}
@@ -1174,6 +1218,7 @@ fn daily_quota_usage_date(
fn daily_quota_grants_from_entitlement(
entitlement_id: &str,
entitlements: &serde_json::Value,
current_allow_wallet_overage: Option<bool>,
now: chrono::DateTime<chrono::Utc>,
) -> Result<Vec<DailyQuotaGrant>, DataLayerError> {
let mut grants = Vec::new();
@@ -1199,15 +1244,27 @@ fn daily_quota_grants_from_entitlement(
.and_then(serde_json::Value::as_str),
now,
)?,
allow_wallet_overage: item
.get("allow_wallet_overage")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false),
allow_wallet_overage: current_allow_wallet_overage.unwrap_or_else(|| {
item.get("allow_wallet_overage")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}),
});
}
Ok(grants)
}
fn daily_quota_wallet_overage_policy(entitlements: &serde_json::Value) -> Option<bool> {
entitlements.as_array()?.iter().find_map(|item| {
(item.get("type").and_then(serde_json::Value::as_str) == Some("daily_quota"))
.then(|| {
item.get("allow_wallet_overage")
.and_then(serde_json::Value::as_bool)
})
.flatten()
})
}
fn read_count_sqlite(row: &SqliteRow) -> Result<u64, DataLayerError> {
Ok(row.try_get::<i64, _>("total").map_sql_err()?.max(0) as u64)
}
@@ -1558,6 +1615,106 @@ mod tests {
assert_eq!(preset.errors, Vec::<String>::new());
}
#[tokio::test]
async fn sqlite_repository_revokes_active_user_plan_entitlement() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool should connect");
run_migrations(&pool)
.await
.expect("sqlite migrations should run");
let now = super::current_unix_secs_i64();
sqlx::query(
r#"
INSERT INTO users (
id, username, email, role, auth_source, password_hash, is_active,
is_deleted, created_at, updated_at
) VALUES (
'user-revoke', 'revoke-user', '[email protected]', 'user', 'local',
'hash', 1, 0, 1, 1
);
INSERT INTO wallets (
id, user_id, balance, gift_balance, limit_mode, created_at, updated_at
) VALUES (
'wallet-revoke', 'user-revoke', 5.0, 0.0, 'finite', 1, 1
);
INSERT INTO billing_plans (
id, title, price_amount, price_currency, duration_unit,
duration_value, entitlements_json, created_at, updated_at
) VALUES (
'plan-revoke', 'Revocable Plan', 0.0, 'USD', 'month', 1,
'[{"type":"daily_quota","daily_quota_usd":10.0,"allow_wallet_overage":true}]',
1, 1
);
INSERT INTO payment_orders (
id, order_no, wallet_id, user_id, amount_usd, refunded_amount_usd,
refundable_amount_usd, payment_method, gateway_response, status, created_at
) VALUES (
'order-revoke', 'order-revoke', 'wallet-revoke', 'user-revoke', 0.0, 0.0,
0.0, 'admin_manual', '{}', 'credited', 1
);
INSERT INTO user_plan_entitlements (
id, user_id, plan_id, payment_order_id, status, starts_at, expires_at,
entitlements_snapshot, created_at, updated_at
) VALUES (
'entitlement-revoke', 'user-revoke', 'plan-revoke', 'order-revoke',
'active', ?, ?,
'[{"type":"daily_quota","daily_quota_usd":10.0,"allow_wallet_overage":false}]',
?, ?
);
"#,
)
.bind(now - 60)
.bind(now + 3600)
.bind(now - 60)
.bind(now - 60)
.execute(&pool)
.await
.expect("revocable entitlement should seed");
let repository = SqliteBillingReadRepository::new(pool.clone());
let quota = repository
.find_user_daily_quota_availability("user-revoke")
.await
.expect("quota should load")
.expect("quota should be available");
assert!(quota.has_active_daily_quota);
assert!(quota.allow_wallet_overage);
let wrong_user = repository
.revoke_user_plan_entitlement("other-user", "entitlement-revoke")
.await
.expect("ownership check should run");
assert_eq!(wrong_user, AdminBillingMutationOutcome::NotFound);
let outcome = repository
.revoke_user_plan_entitlement("user-revoke", "entitlement-revoke")
.await
.expect("entitlement revoke should run");
assert_eq!(outcome, AdminBillingMutationOutcome::Applied(()));
let active = repository
.list_user_plan_entitlements("user-revoke")
.await
.expect("entitlements should load")
.expect("entitlements should be available");
assert!(active.is_empty());
let quota = repository
.find_user_daily_quota_availability("user-revoke")
.await
.expect("quota should load")
.expect("quota should be available");
assert!(!quota.has_active_daily_quota);
let status: String = sqlx::query_scalar(
"SELECT status FROM user_plan_entitlements WHERE id = 'entitlement-revoke'",
)
.fetch_one(&pool)
.await
.expect("entitlement status should load");
assert_eq!(status, "revoked");
}
#[tokio::test]
async fn sqlite_repository_deletes_unused_billing_plans_only() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()
crates/aether-data/adapters/sqlite/src/settlement.rs
+89 -10
View File
@@ -219,6 +219,7 @@ fn daily_quota_usage_date(
fn daily_quota_grants_from_entitlement(
entitlement_id: &str,
entitlements: &serde_json::Value,
current_allow_wallet_overage: Option<bool>,
now: chrono::DateTime<chrono::Utc>,
) -> Result<Vec<DailyQuotaGrant>, DataLayerError> {
let mut grants = Vec::new();
@@ -244,15 +245,27 @@ fn daily_quota_grants_from_entitlement(
.and_then(serde_json::Value::as_str),
now,
)?,
allow_wallet_overage: item
.get("allow_wallet_overage")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false),
allow_wallet_overage: current_allow_wallet_overage.unwrap_or_else(|| {
item.get("allow_wallet_overage")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}),
});
}
Ok(grants)
}
fn daily_quota_wallet_overage_policy(entitlements: &serde_json::Value) -> Option<bool> {
entitlements.as_array()?.iter().find_map(|item| {
(item.get("type").and_then(serde_json::Value::as_str) == Some("daily_quota"))
.then(|| {
item.get("allow_wallet_overage")
.and_then(serde_json::Value::as_bool)
})
.flatten()
})
}
async fn consume_daily_quota_sqlite(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
user_id: &str,
@@ -267,13 +280,19 @@ async fn consume_daily_quota_sqlite(
}
let rows = sqlx::query(
r#"
SELECT id, entitlements_snapshot
SELECT
user_plan_entitlements.id,
user_plan_entitlements.entitlements_snapshot,
billing_plans.entitlements_json AS plan_entitlements_json
FROM user_plan_entitlements
WHERE user_id = ?
AND status = 'active'
AND starts_at <= ?
AND expires_at > ?
ORDER BY expires_at ASC, created_at ASC, id ASC
JOIN billing_plans ON billing_plans.id = user_plan_entitlements.plan_id
WHERE user_plan_entitlements.user_id = ?
AND user_plan_entitlements.status = 'active'
AND user_plan_entitlements.starts_at <= ?
AND user_plan_entitlements.expires_at > ?
ORDER BY user_plan_entitlements.expires_at ASC,
user_plan_entitlements.created_at ASC,
user_plan_entitlements.id ASC
"#,
)
.bind(user_id)
@@ -293,9 +312,17 @@ ORDER BY expires_at ASC, created_at ASC, id ASC
"user_plan_entitlements.entitlements_snapshot invalid json: {err}"
))
})?;
let plan_entitlements_raw: String = row.try_get("plan_entitlements_json").map_sql_err()?;
let plan_entitlements = serde_json::from_str::<serde_json::Value>(&plan_entitlements_raw)
.map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"billing_plans.entitlements_json invalid json: {err}"
))
})?;
grants.extend(daily_quota_grants_from_entitlement(
&entitlement_id,
&entitlements,
daily_quota_wallet_overage_policy(&plan_entitlements),
now,
)?);
}
@@ -1012,6 +1039,58 @@ WHERE request_id = 'request-1'
assert_eq!(quota_used, 10.0);
}
#[tokio::test]
async fn sqlite_repository_uses_current_plan_wallet_overage_policy() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool should connect");
run_migrations(&pool)
.await
.expect("sqlite migrations should run");
seed_quota_covered_settlement_rows(&pool).await;
sqlx::query(
r#"
UPDATE wallets SET balance = 5.0 WHERE id = 'wallet-quota';
UPDATE billing_plans
SET entitlements_json = '[{"type":"daily_quota","daily_quota_usd":10.0,"reset_timezone":"Asia/Shanghai","allow_wallet_overage":true}]'
WHERE id = 'plan-quota';
"#,
)
.execute(&pool)
.await
.expect("plan overage policy should update");
let repository = SqliteSettlementRepository::new(pool.clone());
let settlement = repository
.settle_usage(UsageSettlementInput {
request_id: "request-quota-overrun".to_string(),
user_id: Some("user-quota".to_string()),
api_key_id: Some("key-quota".to_string()),
api_key_is_standalone: false,
provider_id: None,
status: "completed".to_string(),
billing_status: "pending".to_string(),
total_cost_usd: 12.0,
actual_total_cost_usd: 12.0,
finalized_at_unix_secs: Some(1_261),
})
.await
.expect("settlement should run")
.expect("usage should exist");
assert_eq!(settlement.billing_status, "settled");
assert_eq!(settlement.wallet_balance_after, Some(3.0));
let quota_used: f64 = sqlx::query_scalar(
"SELECT CAST(COALESCE(SUM(amount_usd), 0) AS REAL) FROM entitlement_usage_ledgers WHERE request_id = 'request-quota-overrun'",
)
.fetch_one(&pool)
.await
.expect("quota ledger should load");
assert_eq!(quota_used, 10.0);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn sqlite_repository_exhausts_strict_quota_across_concurrent_requests() {
let database_path = std::env::temp_dir().join(format!(