mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 01:47:47 +08:00
feat: add payment gateway and billing plans
This commit is contained in:
Entropy.Xu
@@ -1,7 +1,10 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use super::{SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput};
|
||||
use super::{
|
||||
finite_wallet_available_usd, plan_finite_wallet_debit, SettlementWriteRepository,
|
||||
StoredUsageSettlement, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
|
||||
};
|
||||
use crate::driver::postgres::PostgresTransactionRunner;
|
||||
use crate::error::SqlxResultExt;
|
||||
use crate::DataLayerError;
|
||||
@@ -190,6 +193,192 @@ where
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct DailyQuotaDebitResult {
|
||||
debited_usd: f64,
|
||||
insufficient: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct DailyQuotaGrant {
|
||||
entitlement_id: String,
|
||||
daily_quota_usd: f64,
|
||||
usage_date: String,
|
||||
allow_wallet_overage: bool,
|
||||
}
|
||||
|
||||
fn daily_quota_usage_date(
|
||||
reset_timezone: Option<&str>,
|
||||
now: chrono::DateTime<chrono::Utc>,
|
||||
) -> Result<String, DataLayerError> {
|
||||
let timezone = reset_timezone
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("Asia/Shanghai")
|
||||
.parse::<chrono_tz::Tz>()
|
||||
.map_err(|err| DataLayerError::InvalidInput(format!("invalid reset_timezone: {err}")))?;
|
||||
Ok(now.with_timezone(&timezone).date_naive().to_string())
|
||||
}
|
||||
|
||||
fn daily_quota_grants_from_entitlement(
|
||||
entitlement_id: &str,
|
||||
entitlements: &serde_json::Value,
|
||||
now: chrono::DateTime<chrono::Utc>,
|
||||
) -> Result<Vec<DailyQuotaGrant>, DataLayerError> {
|
||||
let mut grants = Vec::new();
|
||||
let Some(items) = entitlements.as_array() else {
|
||||
return Ok(grants);
|
||||
};
|
||||
for item in items {
|
||||
if item.get("type").and_then(serde_json::Value::as_str) != Some("daily_quota") {
|
||||
continue;
|
||||
}
|
||||
let daily_quota_usd = item
|
||||
.get("daily_quota_usd")
|
||||
.and_then(serde_json::Value::as_f64)
|
||||
.unwrap_or(0.0);
|
||||
if !daily_quota_usd.is_finite() || daily_quota_usd <= 0.0 {
|
||||
continue;
|
||||
}
|
||||
let usage_date = daily_quota_usage_date(
|
||||
item.get("reset_timezone")
|
||||
.and_then(serde_json::Value::as_str),
|
||||
now,
|
||||
)?;
|
||||
grants.push(DailyQuotaGrant {
|
||||
entitlement_id: entitlement_id.to_string(),
|
||||
daily_quota_usd,
|
||||
usage_date,
|
||||
allow_wallet_overage: item
|
||||
.get("allow_wallet_overage")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
});
|
||||
}
|
||||
Ok(grants)
|
||||
}
|
||||
|
||||
async fn consume_daily_quota_postgres(
|
||||
tx: &mut crate::driver::postgres::PostgresTransaction,
|
||||
user_id: &str,
|
||||
request_id: &str,
|
||||
total_cost_usd: f64,
|
||||
wallet_available_usd: Option<f64>,
|
||||
) -> Result<DailyQuotaDebitResult, DataLayerError> {
|
||||
if total_cost_usd <= 0.0 {
|
||||
return Ok(DailyQuotaDebitResult::default());
|
||||
}
|
||||
let now = chrono::Utc::now();
|
||||
let entitlement_rows = sqlx::query(
|
||||
r#"
|
||||
SELECT id, entitlements_snapshot
|
||||
FROM user_plan_entitlements
|
||||
WHERE user_id = $1
|
||||
AND status = 'active'
|
||||
AND starts_at <= NOW()
|
||||
AND expires_at > NOW()
|
||||
ORDER BY expires_at ASC, created_at ASC, id ASC
|
||||
FOR UPDATE
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_all(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
let mut grants = Vec::new();
|
||||
for row in entitlement_rows {
|
||||
let entitlement_id: String = row.try_get("id").map_postgres_err()?;
|
||||
let entitlements: serde_json::Value =
|
||||
row.try_get("entitlements_snapshot").map_postgres_err()?;
|
||||
grants.extend(daily_quota_grants_from_entitlement(
|
||||
&entitlement_id,
|
||||
&entitlements,
|
||||
now,
|
||||
)?);
|
||||
}
|
||||
if grants.is_empty() {
|
||||
return Ok(DailyQuotaDebitResult::default());
|
||||
}
|
||||
|
||||
let mut grants_with_remaining = Vec::new();
|
||||
let mut total_remaining = 0.0;
|
||||
let mut allow_wallet_overage = true;
|
||||
for grant in grants {
|
||||
allow_wallet_overage &= grant.allow_wallet_overage;
|
||||
let used = sqlx::query_scalar::<_, Option<f64>>(
|
||||
r#"
|
||||
SELECT CAST(COALESCE(SUM(amount_usd), 0) AS DOUBLE PRECISION)
|
||||
FROM entitlement_usage_ledgers
|
||||
WHERE user_entitlement_id = $1
|
||||
AND usage_date = $2
|
||||
"#,
|
||||
)
|
||||
.bind(&grant.entitlement_id)
|
||||
.bind(&grant.usage_date)
|
||||
.fetch_one(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.unwrap_or(0.0);
|
||||
let remaining = (grant.daily_quota_usd - used).max(0.0);
|
||||
total_remaining += remaining;
|
||||
grants_with_remaining.push((grant, remaining));
|
||||
}
|
||||
|
||||
if !allow_wallet_overage && total_remaining + 0.000_000_01 < total_cost_usd {
|
||||
return Ok(DailyQuotaDebitResult {
|
||||
debited_usd: 0.0,
|
||||
insufficient: true,
|
||||
});
|
||||
}
|
||||
if allow_wallet_overage
|
||||
&& wallet_available_usd.is_some_and(|available| {
|
||||
total_remaining + available + SETTLEMENT_EPSILON_USD < total_cost_usd
|
||||
})
|
||||
{
|
||||
return Ok(DailyQuotaDebitResult {
|
||||
debited_usd: 0.0,
|
||||
insufficient: true,
|
||||
});
|
||||
}
|
||||
|
||||
let mut remaining_cost = total_cost_usd;
|
||||
let mut debited = 0.0;
|
||||
for (grant, balance_before) in grants_with_remaining {
|
||||
if remaining_cost <= 0.000_000_01 || balance_before <= 0.0 {
|
||||
continue;
|
||||
}
|
||||
let amount = remaining_cost.min(balance_before);
|
||||
let balance_after = balance_before - amount;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO entitlement_usage_ledgers (
|
||||
id, user_entitlement_id, user_id, request_id, amount_usd,
|
||||
balance_before, balance_after, usage_date, created_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, NOW())
|
||||
ON CONFLICT (user_entitlement_id, request_id) DO NOTHING
|
||||
"#,
|
||||
)
|
||||
.bind(uuid::Uuid::new_v4().to_string())
|
||||
.bind(&grant.entitlement_id)
|
||||
.bind(user_id)
|
||||
.bind(request_id)
|
||||
.bind(amount)
|
||||
.bind(balance_before)
|
||||
.bind(balance_after)
|
||||
.bind(&grant.usage_date)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
remaining_cost -= amount;
|
||||
debited += amount;
|
||||
}
|
||||
Ok(DailyQuotaDebitResult {
|
||||
debited_usd: debited,
|
||||
insufficient: false,
|
||||
})
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SettlementWriteRepository for SqlxSettlementRepository {
|
||||
async fn settle_usage(
|
||||
@@ -212,14 +401,17 @@ impl SettlementWriteRepository for SqlxSettlementRepository {
|
||||
|
||||
let current_billing_status: String =
|
||||
usage_row.try_get("billing_status").map_postgres_err()?;
|
||||
if current_billing_status == "settled" || current_billing_status == "void" {
|
||||
if matches!(
|
||||
current_billing_status.as_str(),
|
||||
"settled" | "void" | "insufficient_quota"
|
||||
) {
|
||||
return settlement_from_row(&usage_row).map(Some);
|
||||
}
|
||||
|
||||
let final_billing_status = if input.status == "completed" {
|
||||
"settled"
|
||||
let mut final_billing_status = if input.status == "completed" {
|
||||
"settled".to_string()
|
||||
} else {
|
||||
"void"
|
||||
"void".to_string()
|
||||
};
|
||||
let finalized_at =
|
||||
i64::try_from(input.finalized_at_unix_secs.unwrap_or_else(|| {
|
||||
@@ -323,25 +515,92 @@ LIMIT 1
|
||||
None
|
||||
};
|
||||
|
||||
if let Some(wallet_row) = wallet_row {
|
||||
let wallet_id: String = wallet_row.try_get("id").map_postgres_err()?;
|
||||
let before_recharge: f64 =
|
||||
wallet_row.try_get("balance").map_postgres_err()?;
|
||||
let before_gift: f64 =
|
||||
wallet_row.try_get("gift_balance").map_postgres_err()?;
|
||||
let limit_mode: String =
|
||||
wallet_row.try_get("limit_mode").map_postgres_err()?;
|
||||
let before_total = before_recharge + before_gift;
|
||||
let mut after_recharge = before_recharge;
|
||||
let mut after_gift = before_gift;
|
||||
if !limit_mode.eq_ignore_ascii_case("unlimited") {
|
||||
let gift_deduction = before_gift.max(0.0).min(input.total_cost_usd);
|
||||
let recharge_deduction = input.total_cost_usd - gift_deduction;
|
||||
after_gift = before_gift - gift_deduction;
|
||||
after_recharge = before_recharge - recharge_deduction;
|
||||
let wallet_available_usd = match wallet_row.as_ref() {
|
||||
Some(row) => {
|
||||
let limit_mode: String =
|
||||
row.try_get("limit_mode").map_postgres_err()?;
|
||||
if limit_mode.eq_ignore_ascii_case("unlimited") {
|
||||
None
|
||||
} else {
|
||||
Some(finite_wallet_available_usd(
|
||||
row.try_get("balance").map_postgres_err()?,
|
||||
row.try_get("gift_balance").map_postgres_err()?,
|
||||
))
|
||||
}
|
||||
}
|
||||
sqlx::query(
|
||||
r#"
|
||||
None => Some(0.0),
|
||||
};
|
||||
|
||||
let wallet_debit_cost_usd = if !api_key_is_standalone {
|
||||
if let Some(user_id) =
|
||||
input.user_id.as_deref().filter(|value| !value.is_empty())
|
||||
{
|
||||
let quota = consume_daily_quota_postgres(
|
||||
tx,
|
||||
user_id,
|
||||
&input.request_id,
|
||||
input.total_cost_usd,
|
||||
wallet_available_usd,
|
||||
)
|
||||
.await?;
|
||||
if quota.insufficient {
|
||||
final_billing_status = "insufficient_quota".to_string();
|
||||
settlement.billing_status = final_billing_status.clone();
|
||||
0.0
|
||||
} else {
|
||||
(input.total_cost_usd - quota.debited_usd).max(0.0)
|
||||
}
|
||||
} else {
|
||||
input.total_cost_usd
|
||||
}
|
||||
} else {
|
||||
input.total_cost_usd
|
||||
};
|
||||
if final_billing_status != "settled" {
|
||||
sync_usage_settlement_snapshot(&mut **tx, &settlement).await?;
|
||||
sqlx::query(FINALIZE_USAGE_BILLING_SQL)
|
||||
.bind(&input.request_id)
|
||||
.bind(&final_billing_status)
|
||||
.bind(finalized_at)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
return Ok(Some(settlement));
|
||||
}
|
||||
|
||||
if wallet_debit_cost_usd > SETTLEMENT_EPSILON_USD {
|
||||
if let Some(wallet_row) = wallet_row {
|
||||
let wallet_id: String =
|
||||
wallet_row.try_get("id").map_postgres_err()?;
|
||||
let before_recharge: f64 =
|
||||
wallet_row.try_get("balance").map_postgres_err()?;
|
||||
let before_gift: f64 =
|
||||
wallet_row.try_get("gift_balance").map_postgres_err()?;
|
||||
let limit_mode: String =
|
||||
wallet_row.try_get("limit_mode").map_postgres_err()?;
|
||||
let before_total = before_recharge + before_gift;
|
||||
let mut after_recharge = before_recharge;
|
||||
let mut after_gift = before_gift;
|
||||
if !limit_mode.eq_ignore_ascii_case("unlimited") {
|
||||
let debit_plan = plan_finite_wallet_debit(
|
||||
before_recharge,
|
||||
before_gift,
|
||||
wallet_debit_cost_usd,
|
||||
);
|
||||
if debit_plan.covered_usd() + SETTLEMENT_EPSILON_USD
|
||||
< wallet_debit_cost_usd
|
||||
{
|
||||
final_billing_status = "insufficient_quota".to_string();
|
||||
settlement.billing_status = final_billing_status.clone();
|
||||
} else {
|
||||
after_recharge =
|
||||
before_recharge - debit_plan.recharge_deduction;
|
||||
after_gift = before_gift - debit_plan.gift_deduction;
|
||||
}
|
||||
}
|
||||
if final_billing_status == "settled" {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE wallets
|
||||
SET
|
||||
balance = $2,
|
||||
@@ -350,22 +609,39 @@ SET
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(&wallet_id)
|
||||
.bind(after_recharge)
|
||||
.bind(after_gift)
|
||||
.bind(input.total_cost_usd)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
)
|
||||
.bind(&wallet_id)
|
||||
.bind(after_recharge)
|
||||
.bind(after_gift)
|
||||
.bind(wallet_debit_cost_usd)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
}
|
||||
|
||||
settlement.wallet_id = Some(wallet_id.clone());
|
||||
settlement.wallet_balance_before = Some(before_total);
|
||||
settlement.wallet_balance_after = Some(after_recharge + after_gift);
|
||||
settlement.wallet_recharge_balance_before = Some(before_recharge);
|
||||
settlement.wallet_recharge_balance_after = Some(after_recharge);
|
||||
settlement.wallet_gift_balance_before = Some(before_gift);
|
||||
settlement.wallet_gift_balance_after = Some(after_gift);
|
||||
settlement.wallet_id = Some(wallet_id.clone());
|
||||
settlement.wallet_balance_before = Some(before_total);
|
||||
settlement.wallet_balance_after = Some(after_recharge + after_gift);
|
||||
settlement.wallet_recharge_balance_before = Some(before_recharge);
|
||||
settlement.wallet_recharge_balance_after = Some(after_recharge);
|
||||
settlement.wallet_gift_balance_before = Some(before_gift);
|
||||
settlement.wallet_gift_balance_after = Some(after_gift);
|
||||
} else {
|
||||
final_billing_status = "insufficient_quota".to_string();
|
||||
settlement.billing_status = final_billing_status.clone();
|
||||
}
|
||||
}
|
||||
|
||||
if final_billing_status != "settled" {
|
||||
sync_usage_settlement_snapshot(&mut **tx, &settlement).await?;
|
||||
sqlx::query(FINALIZE_USAGE_BILLING_SQL)
|
||||
.bind(&input.request_id)
|
||||
.bind(&final_billing_status)
|
||||
.bind(finalized_at)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
return Ok(Some(settlement));
|
||||
}
|
||||
|
||||
if let Some(provider_id) = input
|
||||
@@ -396,7 +672,7 @@ RETURNING CAST(monthly_used_usd AS DOUBLE PRECISION) AS monthly_used_usd
|
||||
sync_usage_settlement_snapshot(&mut **tx, &settlement).await?;
|
||||
sqlx::query(FINALIZE_USAGE_BILLING_SQL)
|
||||
.bind(&input.request_id)
|
||||
.bind(final_billing_status)
|
||||
.bind(&final_billing_status)
|
||||
.bind(finalized_at)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
|
||||
Reference in New Issue
Block a user