feat: add configurable verification and redemption codes
This commit is contained in:
237899745
@@ -1,8 +1,247 @@
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::services::billing::BillingContext;
|
||||
use crate::state::AppState;
|
||||
|
||||
use chrono::{Duration, Utc};
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use sqlx::{FromRow, Postgres, Transaction};
|
||||
use std::net::IpAddr;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct UserUsageBalance {
|
||||
pub used_units: i64,
|
||||
pub included_units: i64,
|
||||
pub bonus_units: i64,
|
||||
pub redeemed_units: i64,
|
||||
pub total_units: i64,
|
||||
pub remaining_units: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct UsagePeriodRow {
|
||||
used_units: i32,
|
||||
bonus_units: i32,
|
||||
grant_used_units: i32,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct AvailableGrantRow {
|
||||
id: Uuid,
|
||||
expires_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
pub async fn user_usage_balance(
|
||||
state: &AppState,
|
||||
billing: &BillingContext,
|
||||
) -> Result<UserUsageBalance, AppError> {
|
||||
let usage = sqlx::query_as::<_, UsagePeriodRow>(
|
||||
r#"
|
||||
SELECT used_units, bonus_units, grant_used_units
|
||||
FROM usage_periods
|
||||
WHERE user_id = $1 AND period_start = $2 AND period_end = $3
|
||||
"#,
|
||||
)
|
||||
.bind(billing.user_id)
|
||||
.bind(billing.period_start)
|
||||
.bind(billing.period_end)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用量失败").with_source(err))?
|
||||
.unwrap_or(UsagePeriodRow {
|
||||
used_units: 0,
|
||||
bonus_units: 0,
|
||||
grant_used_units: 0,
|
||||
});
|
||||
|
||||
let redeemed_units: i64 = sqlx::query_scalar(
|
||||
r#"
|
||||
SELECT COALESCE(SUM(remaining_units), 0)::bigint
|
||||
FROM unit_grants
|
||||
WHERE user_id = $1
|
||||
AND starts_at <= NOW()
|
||||
AND expires_at > NOW()
|
||||
AND remaining_units > 0
|
||||
"#,
|
||||
)
|
||||
.bind(billing.user_id)
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询兑换额度失败").with_source(err))?;
|
||||
|
||||
Ok(calculate_user_balance(
|
||||
billing.plan.included_units_per_period,
|
||||
usage.used_units,
|
||||
usage.bonus_units,
|
||||
usage.grant_used_units,
|
||||
redeemed_units,
|
||||
))
|
||||
}
|
||||
|
||||
pub async fn ensure_user_units(
|
||||
state: &AppState,
|
||||
billing: &BillingContext,
|
||||
needed_units: i32,
|
||||
) -> Result<(), AppError> {
|
||||
if needed_units <= 0 {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let balance = user_usage_balance(state, billing).await?;
|
||||
if balance.remaining_units < i64::from(needed_units) {
|
||||
return Err(AppError::new(ErrorCode::QuotaExceeded, "可用配额已用完"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn consume_user_unit(
|
||||
tx: &mut Transaction<'_, Postgres>,
|
||||
billing: &BillingContext,
|
||||
bytes_in: u64,
|
||||
bytes_out: u64,
|
||||
) -> Result<(), AppError> {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO usage_periods (user_id, subscription_id, period_start, period_end)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
ON CONFLICT (user_id, period_start, period_end) DO NOTHING
|
||||
"#,
|
||||
)
|
||||
.bind(billing.user_id)
|
||||
.bind(billing.subscription_id)
|
||||
.bind(billing.period_start)
|
||||
.bind(billing.period_end)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "初始化用量周期失败").with_source(err))?;
|
||||
|
||||
let usage = sqlx::query_as::<_, UsagePeriodRow>(
|
||||
r#"
|
||||
SELECT used_units, bonus_units, grant_used_units
|
||||
FROM usage_periods
|
||||
WHERE user_id = $1 AND period_start = $2 AND period_end = $3
|
||||
FOR UPDATE
|
||||
"#,
|
||||
)
|
||||
.bind(billing.user_id)
|
||||
.bind(billing.period_start)
|
||||
.bind(billing.period_end)
|
||||
.fetch_one(&mut **tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "锁定用量周期失败").with_source(err))?;
|
||||
|
||||
let plan_used = usage.used_units.saturating_sub(usage.grant_used_units);
|
||||
let plan_capacity = billing
|
||||
.plan
|
||||
.included_units_per_period
|
||||
.saturating_add(usage.bonus_units);
|
||||
let plan_available = plan_used < plan_capacity;
|
||||
let grant = sqlx::query_as::<_, AvailableGrantRow>(
|
||||
r#"
|
||||
SELECT id, expires_at
|
||||
FROM unit_grants
|
||||
WHERE user_id = $1
|
||||
AND starts_at <= NOW()
|
||||
AND expires_at > NOW()
|
||||
AND remaining_units > 0
|
||||
ORDER BY expires_at ASC, created_at ASC
|
||||
FOR UPDATE
|
||||
LIMIT 1
|
||||
"#,
|
||||
)
|
||||
.bind(billing.user_id)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "锁定兑换额度失败").with_source(err))?;
|
||||
|
||||
let use_grant = should_consume_grant(
|
||||
plan_available,
|
||||
grant.as_ref().map(|grant| grant.expires_at),
|
||||
billing.period_end,
|
||||
);
|
||||
|
||||
if !plan_available && !use_grant {
|
||||
return Err(AppError::new(ErrorCode::QuotaExceeded, "可用配额已用完"));
|
||||
}
|
||||
|
||||
if use_grant {
|
||||
let grant_id = grant
|
||||
.ok_or_else(|| AppError::new(ErrorCode::QuotaExceeded, "可用配额已用完"))?
|
||||
.id;
|
||||
|
||||
let updated = sqlx::query(
|
||||
r#"
|
||||
UPDATE unit_grants
|
||||
SET remaining_units = remaining_units - 1,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
AND remaining_units > 0
|
||||
AND expires_at > NOW()
|
||||
"#,
|
||||
)
|
||||
.bind(grant_id)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "扣减兑换额度失败").with_source(err))?;
|
||||
|
||||
if updated.rows_affected() != 1 {
|
||||
return Err(AppError::new(ErrorCode::QuotaExceeded, "可用配额已用完"));
|
||||
}
|
||||
}
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE usage_periods
|
||||
SET used_units = used_units + 1,
|
||||
grant_used_units = grant_used_units + $1,
|
||||
bytes_in = bytes_in + $2,
|
||||
bytes_out = bytes_out + $3,
|
||||
updated_at = NOW()
|
||||
WHERE user_id = $4 AND period_start = $5 AND period_end = $6
|
||||
"#,
|
||||
)
|
||||
.bind(use_grant as i32)
|
||||
.bind(bytes_in as i64)
|
||||
.bind(bytes_out as i64)
|
||||
.bind(billing.user_id)
|
||||
.bind(billing.period_start)
|
||||
.bind(billing.period_end)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "记录用量失败").with_source(err))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn calculate_user_balance(
|
||||
included_units: i32,
|
||||
used_units: i32,
|
||||
bonus_units: i32,
|
||||
grant_used_units: i32,
|
||||
redeemed_units: i64,
|
||||
) -> UserUsageBalance {
|
||||
let base_capacity = i64::from(included_units.saturating_add(bonus_units));
|
||||
let base_used = i64::from(used_units.saturating_sub(grant_used_units));
|
||||
let base_remaining = base_capacity.saturating_sub(base_used).max(0);
|
||||
let remaining_units = base_remaining.saturating_add(redeemed_units.max(0));
|
||||
let used_units = i64::from(used_units.max(0));
|
||||
|
||||
UserUsageBalance {
|
||||
used_units,
|
||||
included_units: i64::from(included_units),
|
||||
bonus_units: i64::from(bonus_units),
|
||||
redeemed_units: redeemed_units.max(0),
|
||||
total_units: used_units.saturating_add(remaining_units),
|
||||
remaining_units,
|
||||
}
|
||||
}
|
||||
|
||||
fn should_consume_grant(
|
||||
plan_available: bool,
|
||||
grant_expires_at: Option<DateTime<Utc>>,
|
||||
plan_expires_at: DateTime<Utc>,
|
||||
) -> bool {
|
||||
grant_expires_at.is_some_and(|expires_at| !plan_available || expires_at <= plan_expires_at)
|
||||
}
|
||||
|
||||
pub async fn consume_anonymous_units(
|
||||
state: &AppState,
|
||||
@@ -71,3 +310,42 @@ fn utc8_date() -> String {
|
||||
let now = Utc::now() + Duration::hours(8);
|
||||
now.format("%Y-%m-%d").to_string()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn balance_keeps_redeemed_units_separate_from_plan_usage() {
|
||||
let balance = calculate_user_balance(10, 13, 0, 3, 7);
|
||||
assert_eq!(balance.remaining_units, 7);
|
||||
assert_eq!(balance.total_units, 20);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn balance_uses_plan_capacity_before_redeemed_units() {
|
||||
let balance = calculate_user_balance(10, 4, 2, 0, 5);
|
||||
assert_eq!(balance.remaining_units, 13);
|
||||
assert_eq!(balance.total_units, 17);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn earlier_expiring_entitlement_is_consumed_first() {
|
||||
let now = Utc::now();
|
||||
assert!(should_consume_grant(
|
||||
true,
|
||||
Some(now + Duration::days(2)),
|
||||
now + Duration::days(20),
|
||||
));
|
||||
assert!(!should_consume_grant(
|
||||
true,
|
||||
Some(now + Duration::days(30)),
|
||||
now + Duration::days(20),
|
||||
));
|
||||
assert!(should_consume_grant(
|
||||
false,
|
||||
Some(now + Duration::days(30)),
|
||||
now + Duration::days(20),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user