mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 18:07:47 +08:00
refactor(workspace): enforce layered crate boundaries
This commit is contained in:
elky
@@ -0,0 +1,782 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
|
||||
use aether_data_contracts::repository::settlement::{
|
||||
finite_wallet_available_usd, plan_finite_wallet_debit, settlement_billable_cost_usd,
|
||||
settlement_billing_status_for_usage_status, SettlementWriteRepository, StoredUsageSettlement,
|
||||
UsageSettlementInput, SETTLEMENT_EPSILON_USD,
|
||||
};
|
||||
use aether_data_contracts::DataLayerError;
|
||||
|
||||
use crate::error::SqlxResultExt;
|
||||
use crate::PostgresTransactionRunner;
|
||||
|
||||
const FIND_USAGE_FOR_SETTLEMENT_SQL: &str = r#"
|
||||
SELECT
|
||||
usage_record.request_id,
|
||||
COALESCE(usage_settlement_snapshots.wallet_id, usage_record.wallet_id) AS wallet_id,
|
||||
COALESCE(usage_settlement_snapshots.billing_status, usage_record.billing_status) AS billing_status,
|
||||
COALESCE(
|
||||
CAST(usage_settlement_snapshots.wallet_balance_before AS DOUBLE PRECISION),
|
||||
CAST(usage_record.wallet_balance_before AS DOUBLE PRECISION)
|
||||
) AS wallet_balance_before,
|
||||
COALESCE(
|
||||
CAST(usage_settlement_snapshots.wallet_balance_after AS DOUBLE PRECISION),
|
||||
CAST(usage_record.wallet_balance_after AS DOUBLE PRECISION)
|
||||
) AS wallet_balance_after,
|
||||
COALESCE(
|
||||
CAST(usage_settlement_snapshots.wallet_recharge_balance_before AS DOUBLE PRECISION),
|
||||
CAST(usage_record.wallet_recharge_balance_before AS DOUBLE PRECISION)
|
||||
) AS wallet_recharge_balance_before,
|
||||
COALESCE(
|
||||
CAST(usage_settlement_snapshots.wallet_recharge_balance_after AS DOUBLE PRECISION),
|
||||
CAST(usage_record.wallet_recharge_balance_after AS DOUBLE PRECISION)
|
||||
) AS wallet_recharge_balance_after,
|
||||
COALESCE(
|
||||
CAST(usage_settlement_snapshots.wallet_gift_balance_before AS DOUBLE PRECISION),
|
||||
CAST(usage_record.wallet_gift_balance_before AS DOUBLE PRECISION)
|
||||
) AS wallet_gift_balance_before,
|
||||
COALESCE(
|
||||
CAST(usage_settlement_snapshots.wallet_gift_balance_after AS DOUBLE PRECISION),
|
||||
CAST(usage_record.wallet_gift_balance_after AS DOUBLE PRECISION)
|
||||
) AS wallet_gift_balance_after,
|
||||
CAST(usage_settlement_snapshots.provider_monthly_used_usd AS DOUBLE PRECISION) AS provider_monthly_used_usd,
|
||||
usage_record.provider_id,
|
||||
CAST(
|
||||
EXTRACT(
|
||||
EPOCH FROM COALESCE(usage_settlement_snapshots.finalized_at, usage_record.finalized_at)
|
||||
) AS BIGINT
|
||||
) AS finalized_at_unix_secs
|
||||
FROM "usage" AS usage_record
|
||||
LEFT JOIN usage_settlement_snapshots
|
||||
ON usage_settlement_snapshots.request_id = usage_record.request_id
|
||||
WHERE usage_record.request_id = $1
|
||||
FOR UPDATE OF usage_record
|
||||
"#;
|
||||
|
||||
const FINALIZE_USAGE_BILLING_SQL: &str = r#"
|
||||
UPDATE "usage"
|
||||
SET
|
||||
billing_status = $2,
|
||||
finalized_at = COALESCE(finalized_at, to_timestamp($3))
|
||||
WHERE request_id = $1
|
||||
"#;
|
||||
|
||||
const UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL: &str = r#"
|
||||
INSERT INTO usage_settlement_snapshots (
|
||||
request_id,
|
||||
billing_status,
|
||||
wallet_id,
|
||||
wallet_balance_before,
|
||||
wallet_balance_after,
|
||||
wallet_recharge_balance_before,
|
||||
wallet_recharge_balance_after,
|
||||
wallet_gift_balance_before,
|
||||
wallet_gift_balance_after,
|
||||
provider_monthly_used_usd,
|
||||
finalized_at
|
||||
) VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4,
|
||||
$5,
|
||||
$6,
|
||||
$7,
|
||||
$8,
|
||||
$9,
|
||||
$10,
|
||||
CASE
|
||||
WHEN $11 IS NULL THEN NULL
|
||||
ELSE TO_TIMESTAMP($11::double precision)
|
||||
END
|
||||
)
|
||||
ON CONFLICT (request_id)
|
||||
DO UPDATE SET
|
||||
billing_status = EXCLUDED.billing_status,
|
||||
wallet_id = COALESCE(EXCLUDED.wallet_id, usage_settlement_snapshots.wallet_id),
|
||||
wallet_balance_before = COALESCE(
|
||||
EXCLUDED.wallet_balance_before,
|
||||
usage_settlement_snapshots.wallet_balance_before
|
||||
),
|
||||
wallet_balance_after = COALESCE(
|
||||
EXCLUDED.wallet_balance_after,
|
||||
usage_settlement_snapshots.wallet_balance_after
|
||||
),
|
||||
wallet_recharge_balance_before = COALESCE(
|
||||
EXCLUDED.wallet_recharge_balance_before,
|
||||
usage_settlement_snapshots.wallet_recharge_balance_before
|
||||
),
|
||||
wallet_recharge_balance_after = COALESCE(
|
||||
EXCLUDED.wallet_recharge_balance_after,
|
||||
usage_settlement_snapshots.wallet_recharge_balance_after
|
||||
),
|
||||
wallet_gift_balance_before = COALESCE(
|
||||
EXCLUDED.wallet_gift_balance_before,
|
||||
usage_settlement_snapshots.wallet_gift_balance_before
|
||||
),
|
||||
wallet_gift_balance_after = COALESCE(
|
||||
EXCLUDED.wallet_gift_balance_after,
|
||||
usage_settlement_snapshots.wallet_gift_balance_after
|
||||
),
|
||||
provider_monthly_used_usd = COALESCE(
|
||||
EXCLUDED.provider_monthly_used_usd,
|
||||
usage_settlement_snapshots.provider_monthly_used_usd
|
||||
),
|
||||
finalized_at = COALESCE(EXCLUDED.finalized_at, usage_settlement_snapshots.finalized_at),
|
||||
updated_at = NOW()
|
||||
"#;
|
||||
|
||||
const ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL: &str = r#"
|
||||
INSERT INTO usage_counter_deltas (
|
||||
id,
|
||||
request_id,
|
||||
kind,
|
||||
target_id,
|
||||
total_cost_usd_delta
|
||||
) VALUES (
|
||||
$1,
|
||||
$2,
|
||||
'provider_monthly',
|
||||
$3,
|
||||
$4
|
||||
)
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxSettlementRepository {
|
||||
tx_runner: PostgresTransactionRunner,
|
||||
}
|
||||
|
||||
impl SqlxSettlementRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
let tx_runner = PostgresTransactionRunner::new(pool);
|
||||
Self { tx_runner }
|
||||
}
|
||||
}
|
||||
|
||||
fn settlement_from_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> Result<StoredUsageSettlement, DataLayerError> {
|
||||
Ok(StoredUsageSettlement {
|
||||
request_id: row.try_get("request_id").map_postgres_err()?,
|
||||
wallet_id: row.try_get("wallet_id").map_postgres_err()?,
|
||||
billing_status: row.try_get("billing_status").map_postgres_err()?,
|
||||
wallet_balance_before: row.try_get("wallet_balance_before").map_postgres_err()?,
|
||||
wallet_balance_after: row.try_get("wallet_balance_after").map_postgres_err()?,
|
||||
wallet_recharge_balance_before: row
|
||||
.try_get("wallet_recharge_balance_before")
|
||||
.map_postgres_err()?,
|
||||
wallet_recharge_balance_after: row
|
||||
.try_get("wallet_recharge_balance_after")
|
||||
.map_postgres_err()?,
|
||||
wallet_gift_balance_before: row
|
||||
.try_get("wallet_gift_balance_before")
|
||||
.map_postgres_err()?,
|
||||
wallet_gift_balance_after: row
|
||||
.try_get("wallet_gift_balance_after")
|
||||
.map_postgres_err()?,
|
||||
provider_monthly_used_usd: row
|
||||
.try_get("provider_monthly_used_usd")
|
||||
.map_postgres_err()?,
|
||||
finalized_at_unix_secs: row
|
||||
.try_get::<Option<i64>, _>("finalized_at_unix_secs")
|
||||
.map_postgres_err()?
|
||||
.map(|value| value as u64),
|
||||
})
|
||||
}
|
||||
|
||||
async fn sync_usage_settlement_snapshot<'e, E>(
|
||||
executor: E,
|
||||
settlement: &StoredUsageSettlement,
|
||||
) -> Result<(), DataLayerError>
|
||||
where
|
||||
E: sqlx::Executor<'e, Database = sqlx::Postgres>,
|
||||
{
|
||||
sqlx::query(UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL)
|
||||
.bind(&settlement.request_id)
|
||||
.bind(&settlement.billing_status)
|
||||
.bind(settlement.wallet_id.as_deref())
|
||||
.bind(settlement.wallet_balance_before)
|
||||
.bind(settlement.wallet_balance_after)
|
||||
.bind(settlement.wallet_recharge_balance_before)
|
||||
.bind(settlement.wallet_recharge_balance_after)
|
||||
.bind(settlement.wallet_gift_balance_before)
|
||||
.bind(settlement.wallet_gift_balance_after)
|
||||
.bind(settlement.provider_monthly_used_usd)
|
||||
.bind(settlement.finalized_at_unix_secs.map(|value| value as f64))
|
||||
.execute(executor)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn enqueue_provider_monthly_usage_delta<'e, E>(
|
||||
executor: E,
|
||||
request_id: &str,
|
||||
provider_id: &str,
|
||||
total_cost_usd_delta: f64,
|
||||
) -> Result<(), DataLayerError>
|
||||
where
|
||||
E: sqlx::Executor<'e, Database = sqlx::Postgres>,
|
||||
{
|
||||
let request_id = request_id.trim();
|
||||
let provider_id = provider_id.trim();
|
||||
if request_id.is_empty() || provider_id.is_empty() || total_cost_usd_delta == 0.0 {
|
||||
return Ok(());
|
||||
}
|
||||
if !total_cost_usd_delta.is_finite() {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"provider monthly usage delta is not finite for {provider_id}"
|
||||
)));
|
||||
}
|
||||
|
||||
sqlx::query(ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL)
|
||||
.bind(uuid::Uuid::new_v4().to_string())
|
||||
.bind(request_id)
|
||||
.bind(provider_id)
|
||||
.bind(total_cost_usd_delta)
|
||||
.execute(executor)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
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::PostgresTransaction,
|
||||
user_id: &str,
|
||||
request_id: &str,
|
||||
total_cost_usd: f64,
|
||||
wallet_available_usd: Option<f64>,
|
||||
wallet_can_overdraft: bool,
|
||||
) -> 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_can_overdraft
|
||||
&& 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(
|
||||
&self,
|
||||
input: UsageSettlementInput,
|
||||
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
|
||||
input.validate()?;
|
||||
self.tx_runner
|
||||
.run_read_write(|tx| {
|
||||
Box::pin(async move {
|
||||
let row = sqlx::query(FIND_USAGE_FOR_SETTLEMENT_SQL)
|
||||
.bind(&input.request_id)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
|
||||
let Some(usage_row) = row else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let current_billing_status: String =
|
||||
usage_row.try_get("billing_status").map_postgres_err()?;
|
||||
if matches!(
|
||||
current_billing_status.as_str(),
|
||||
"settled" | "void" | "insufficient_quota"
|
||||
) {
|
||||
return settlement_from_row(&usage_row).map(Some);
|
||||
}
|
||||
|
||||
let mut final_billing_status =
|
||||
settlement_billing_status_for_usage_status(&input.status).to_string();
|
||||
let finalized_at =
|
||||
i64::try_from(input.finalized_at_unix_secs.unwrap_or_else(|| {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}))
|
||||
.map_err(|_| {
|
||||
DataLayerError::InvalidInput("finalized_at overflow".to_string())
|
||||
})?;
|
||||
|
||||
let mut settlement = StoredUsageSettlement {
|
||||
request_id: input.request_id.clone(),
|
||||
wallet_id: None,
|
||||
billing_status: final_billing_status.to_string(),
|
||||
wallet_balance_before: None,
|
||||
wallet_balance_after: None,
|
||||
wallet_recharge_balance_before: None,
|
||||
wallet_recharge_balance_after: None,
|
||||
wallet_gift_balance_before: None,
|
||||
wallet_gift_balance_after: None,
|
||||
provider_monthly_used_usd: None,
|
||||
finalized_at_unix_secs: Some(finalized_at as u64),
|
||||
};
|
||||
|
||||
if final_billing_status == "settled" {
|
||||
let api_key_id = input
|
||||
.api_key_id
|
||||
.as_deref()
|
||||
.filter(|value| !value.is_empty());
|
||||
let api_key_is_standalone = if input.api_key_is_standalone {
|
||||
true
|
||||
} else if let Some(api_key_id) = api_key_id {
|
||||
sqlx::query_scalar::<_, bool>(
|
||||
r#"
|
||||
SELECT is_standalone
|
||||
FROM api_keys
|
||||
WHERE id = $1
|
||||
LIMIT 1
|
||||
"#,
|
||||
)
|
||||
.bind(api_key_id)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.unwrap_or(false)
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
let wallet_row = if let Some(api_key_id) = api_key_id {
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
CAST(balance AS DOUBLE PRECISION) AS balance,
|
||||
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
|
||||
limit_mode
|
||||
FROM wallets
|
||||
WHERE api_key_id = $1
|
||||
FOR UPDATE
|
||||
LIMIT 1
|
||||
"#,
|
||||
)
|
||||
.bind(api_key_id)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let wallet_row = if wallet_row.is_some() {
|
||||
wallet_row
|
||||
} else if !api_key_is_standalone {
|
||||
if let Some(user_id) =
|
||||
input.user_id.as_deref().filter(|value| !value.is_empty())
|
||||
{
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
CAST(balance AS DOUBLE PRECISION) AS balance,
|
||||
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
|
||||
limit_mode
|
||||
FROM wallets
|
||||
WHERE user_id = $1
|
||||
FOR UPDATE
|
||||
LIMIT 1
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
} else {
|
||||
None
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let wallet_can_overdraft = wallet_row.is_some();
|
||||
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()?,
|
||||
))
|
||||
}
|
||||
}
|
||||
None => Some(0.0),
|
||||
};
|
||||
if let Some(row) = wallet_row.as_ref() {
|
||||
let wallet_id: String = row.try_get("id").map_postgres_err()?;
|
||||
let before_recharge: f64 = row.try_get("balance").map_postgres_err()?;
|
||||
let before_gift: f64 =
|
||||
row.try_get("gift_balance").map_postgres_err()?;
|
||||
let before_total = before_recharge + before_gift;
|
||||
settlement.wallet_id = Some(wallet_id);
|
||||
settlement.wallet_balance_before = Some(before_total);
|
||||
settlement.wallet_balance_after = Some(before_total);
|
||||
settlement.wallet_recharge_balance_before = Some(before_recharge);
|
||||
settlement.wallet_recharge_balance_after = Some(before_recharge);
|
||||
settlement.wallet_gift_balance_before = Some(before_gift);
|
||||
settlement.wallet_gift_balance_after = Some(before_gift);
|
||||
}
|
||||
|
||||
let billable_cost_usd = settlement_billable_cost_usd(&input);
|
||||
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,
|
||||
billable_cost_usd,
|
||||
wallet_available_usd,
|
||||
wallet_can_overdraft,
|
||||
)
|
||||
.await?;
|
||||
if quota.insufficient {
|
||||
final_billing_status = "insufficient_quota".to_string();
|
||||
settlement.billing_status = final_billing_status.clone();
|
||||
0.0
|
||||
} else {
|
||||
(billable_cost_usd - quota.debited_usd).max(0.0)
|
||||
}
|
||||
} else {
|
||||
billable_cost_usd
|
||||
}
|
||||
} else {
|
||||
billable_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,
|
||||
);
|
||||
(after_recharge, after_gift) =
|
||||
debit_plan.after_balances(before_recharge, before_gift);
|
||||
}
|
||||
if final_billing_status == "settled" {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE wallets
|
||||
SET
|
||||
balance = $2,
|
||||
gift_balance = $3,
|
||||
total_consumed = CAST(total_consumed AS DOUBLE PRECISION) + $4,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.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);
|
||||
} 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
|
||||
.provider_id
|
||||
.as_deref()
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
enqueue_provider_monthly_usage_delta(
|
||||
&mut **tx,
|
||||
&input.request_id,
|
||||
provider_id,
|
||||
input.actual_total_cost_usd,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
|
||||
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()?;
|
||||
|
||||
Ok(Some(settlement))
|
||||
})
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[test]
|
||||
fn finalize_usage_billing_sql_does_not_require_usage_updated_at_column() {
|
||||
assert!(!super::FINALIZE_USAGE_BILLING_SQL.contains("updated_at"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settlement_sql_reads_settlement_snapshots_before_legacy_usage_columns() {
|
||||
assert!(
|
||||
super::FIND_USAGE_FOR_SETTLEMENT_SQL.contains("LEFT JOIN usage_settlement_snapshots")
|
||||
);
|
||||
assert!(super::FIND_USAGE_FOR_SETTLEMENT_SQL.contains(
|
||||
"COALESCE(usage_settlement_snapshots.billing_status, usage_record.billing_status)"
|
||||
));
|
||||
assert!(super::FIND_USAGE_FOR_SETTLEMENT_SQL.contains("FOR UPDATE OF usage_record"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settlement_sql_dual_writes_usage_settlement_snapshots() {
|
||||
assert!(super::UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL
|
||||
.contains("INSERT INTO usage_settlement_snapshots"));
|
||||
assert!(super::UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL.contains("provider_monthly_used_usd"));
|
||||
assert!(super::UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL
|
||||
.contains("TO_TIMESTAMP($11::double precision)"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settlement_sql_no_longer_dual_writes_wallet_snapshots_to_usage_rows() {
|
||||
let source = include_str!("settlement.rs");
|
||||
assert!(!source.contains("UPDATE \"usage\"\nSET\n wallet_id = $2"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settlement_sql_enqueues_provider_monthly_usage_delta() {
|
||||
let source = include_str!("settlement.rs");
|
||||
assert!(super::ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL.contains("usage_counter_deltas"));
|
||||
assert!(super::ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL.contains("'provider_monthly'"));
|
||||
assert!(!source.contains("UPDATE providers\nSET\n monthly_used_usd"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settlement_sql_blocks_standalone_key_owner_wallet_fallback() {
|
||||
let source = include_str!("settlement.rs");
|
||||
assert!(source.contains("SELECT is_standalone"));
|
||||
assert!(source.contains("} else if !api_key_is_standalone {"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user