File size: 6,555 Bytes
bbb1195 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 | use axum::{
extract::Path,
Json,
};
use serde::{Deserialize, Serialize};
use crate::models::{Account, TokenData, QuotaData};
use crate::modules;
use crate::error::AppError;
/// API response wrapper
#[derive(Serialize)]
pub struct ApiResponse<T> {
pub success: bool,
pub data: Option<T>,
pub error: Option<String>,
}
impl<T: Serialize> ApiResponse<T> {
pub fn success(data: T) -> Json<Self> {
Json(Self {
success: true,
data: Some(data),
error: None,
})
}
pub fn error(message: String) -> Json<Self> {
Json(Self {
success: false,
data: None,
error: Some(message),
})
}
}
/// List all accounts
pub async fn list_accounts() -> Result<Json<ApiResponse<Vec<Account>>>, AppError> {
let accounts = modules::list_accounts().map_err(AppError::Account)?;
Ok(ApiResponse::success(accounts))
}
/// Add account request
#[derive(Deserialize)]
pub struct AddAccountRequest {
pub email: String,
pub refresh_token: String,
}
/// Add new account
pub async fn add_account(
Json(req): Json<AddAccountRequest>,
) -> Result<Json<ApiResponse<Account>>, AppError> {
modules::logger::log_info(&format!("Adding account: {}", req.email));
// 1. Use refresh_token to get access_token
let token_res = modules::oauth::refresh_access_token(&req.refresh_token)
.await
.map_err(AppError::OAuth)?;
// 2. Get user info
let user_info = modules::oauth::get_user_info(&token_res.access_token)
.await
.map_err(AppError::OAuth)?;
// 3. Construct TokenData
let token = TokenData::new(
token_res.access_token,
req.refresh_token,
token_res.expires_in,
Some(user_info.email.clone()),
None,
None,
);
// 4. Add or update account using real email
let mut account = modules::upsert_account(
user_info.email.clone(),
user_info.get_display_name(),
token,
).map_err(AppError::Account)?;
modules::logger::log_info(&format!("Account added successfully: {}", account.email));
// 5. Auto refresh quota
let _ = internal_refresh_account_quota(&mut account).await;
Ok(ApiResponse::success(account))
}
/// Delete account
pub async fn delete_account(
Path(account_id): Path<String>,
) -> Result<Json<ApiResponse<()>>, AppError> {
modules::logger::log_info(&format!("Deleting account: {}", account_id));
modules::delete_account(&account_id).map_err(AppError::Account)?;
modules::logger::log_info(&format!("Account deleted: {}", account_id));
Ok(ApiResponse::success(()))
}
/// Refresh account quota
pub async fn refresh_quota(
Path(account_id): Path<String>,
) -> Result<Json<ApiResponse<QuotaData>>, AppError> {
modules::logger::log_info(&format!("Refreshing quota for: {}", account_id));
let mut account = modules::load_account(&account_id).map_err(AppError::Account)?;
let quota = modules::account::fetch_quota_with_retry(&mut account).await?;
// Update account quota
modules::update_account_quota(&account_id, quota.clone()).map_err(AppError::Account)?;
Ok(ApiResponse::success(quota))
}
/// Refresh stats response
#[derive(Serialize)]
pub struct RefreshStats {
total: usize,
success: usize,
failed: usize,
details: Vec<String>,
}
/// Refresh all account quotas
pub async fn refresh_all_quotas() -> Result<Json<ApiResponse<RefreshStats>>, AppError> {
modules::logger::log_info("Starting batch quota refresh for all accounts");
let accounts = modules::list_accounts().map_err(AppError::Account)?;
let mut success = 0;
let mut failed = 0;
let mut details = Vec::new();
// Serial processing to ensure persistence safety
for mut account in accounts {
if let Some(ref q) = account.quota {
if q.is_forbidden {
modules::logger::log_info(&format!("Skipping {} (Forbidden)", account.email));
continue;
}
}
modules::logger::log_info(&format!("Processing {}", account.email));
match modules::account::fetch_quota_with_retry(&mut account).await {
Ok(quota) => {
if let Err(e) = modules::update_account_quota(&account.id, quota) {
failed += 1;
let msg = format!("Account {}: Save quota failed - {}", account.email, e);
details.push(msg.clone());
modules::logger::log_error(&msg);
} else {
success += 1;
modules::logger::log_info("Success");
}
},
Err(e) => {
failed += 1;
let msg = format!("Account {}: Fetch quota failed - {}", account.email, e);
details.push(msg.clone());
modules::logger::log_error(&msg);
}
}
}
modules::logger::log_info(&format!("Batch refresh completed: {} success, {} failed", success, failed));
Ok(ApiResponse::success(RefreshStats { total: success + failed, success, failed, details }))
}
/// Get current account
pub async fn get_current_account() -> Result<Json<ApiResponse<Option<Account>>>, AppError> {
let account_id = modules::get_current_account_id().map_err(AppError::Account)?;
if let Some(id) = account_id {
let account = modules::load_account(&id).map_err(AppError::Account)?;
Ok(ApiResponse::success(Some(account)))
} else {
Ok(ApiResponse::success(None))
}
}
/// Set current account
pub async fn set_current_account(
Path(account_id): Path<String>,
) -> Result<Json<ApiResponse<()>>, AppError> {
modules::logger::log_info(&format!("Setting current account: {}", account_id));
modules::set_current_account_id(&account_id).map_err(AppError::Account)?;
Ok(ApiResponse::success(()))
}
/// Internal helper: auto refresh quota after adding account
async fn internal_refresh_account_quota(account: &mut Account) -> Result<QuotaData, String> {
modules::logger::log_info(&format!("Auto refreshing quota: {}", account.email));
match modules::account::fetch_quota_with_retry(account).await {
Ok(quota) => {
let _ = modules::update_account_quota(&account.id, quota.clone());
Ok(quota)
},
Err(e) => {
modules::logger::log_warn(&format!("Auto refresh quota failed ({}): {}", account.email, e));
Err(e.to_string())
}
}
}
|