mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat: paginate provider keys from backend
This commit is contained in:
@@ -12,6 +12,38 @@ use axum::{
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
fn parse_provider_keys_page_param(raw: Option<String>) -> Result<usize, String> {
|
||||
match raw {
|
||||
None => Ok(1),
|
||||
Some(value) => {
|
||||
let parsed = value
|
||||
.parse::<usize>()
|
||||
.map_err(|_| "page must be an integer between 1 and 10000".to_string())?;
|
||||
if (1..=10_000).contains(&parsed) {
|
||||
Ok(parsed)
|
||||
} else {
|
||||
Err("page must be an integer between 1 and 10000".to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_provider_keys_page_size_param(raw: Option<String>) -> Result<usize, String> {
|
||||
match raw {
|
||||
None => Ok(20),
|
||||
Some(value) => {
|
||||
let parsed = value
|
||||
.parse::<usize>()
|
||||
.map_err(|_| "page_size must be an integer between 1 and 1000".to_string())?;
|
||||
if (1..=1000).contains(&parsed) {
|
||||
Ok(parsed)
|
||||
} else {
|
||||
Err("page_size must be an integer between 1 and 1000".to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn maybe_handle(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
@@ -141,6 +173,47 @@ pub(super) async fn maybe_handle(
|
||||
.into_response(),
|
||||
));
|
||||
};
|
||||
let page_param = query_param_value(request_context.query_string(), "page");
|
||||
let page_size_param = query_param_value(request_context.query_string(), "page_size");
|
||||
if page_param.is_some() || page_size_param.is_some() {
|
||||
let page = match parse_provider_keys_page_param(page_param) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => {
|
||||
return Ok(Some(
|
||||
(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response(),
|
||||
));
|
||||
}
|
||||
};
|
||||
let page_size = match parse_provider_keys_page_size_param(page_size_param) {
|
||||
Ok(value) => value,
|
||||
Err(detail) => {
|
||||
return Ok(Some(
|
||||
(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail })),
|
||||
)
|
||||
.into_response(),
|
||||
));
|
||||
}
|
||||
};
|
||||
return Ok(Some(
|
||||
match state
|
||||
.build_admin_provider_keys_page_payload(&provider_id, page, page_size)
|
||||
.await
|
||||
{
|
||||
Some(payload) => Json(payload).into_response(),
|
||||
None => (
|
||||
http::StatusCode::NOT_FOUND,
|
||||
Json(json!({ "detail": format!("Provider {provider_id} 不存在") })),
|
||||
)
|
||||
.into_response(),
|
||||
},
|
||||
));
|
||||
}
|
||||
let skip = query_param_value(request_context.query_string(), "skip")
|
||||
.and_then(|value| value.parse::<usize>().ok())
|
||||
.unwrap_or(0);
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
pub(crate) use self::{
|
||||
create::build_admin_create_provider_key_record, payload::build_admin_provider_keys_payload,
|
||||
create::build_admin_create_provider_key_record,
|
||||
payload::{build_admin_provider_keys_page_payload, build_admin_provider_keys_payload},
|
||||
update::build_admin_update_provider_key_record,
|
||||
};
|
||||
|
||||
|
||||
@@ -3,14 +3,15 @@ use crate::provider_key_auth::provider_key_effective_api_formats;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
pub(crate) async fn build_admin_provider_keys_payload(
|
||||
async fn build_admin_provider_key_items_payload(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
skip: usize,
|
||||
limit: usize,
|
||||
) -> Option<serde_json::Value> {
|
||||
) -> Option<(Vec<Value>, usize)> {
|
||||
if !state.has_provider_catalog_data_reader() {
|
||||
return None;
|
||||
}
|
||||
@@ -40,20 +41,47 @@ pub(crate) async fn build_admin_provider_keys_payload(
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0);
|
||||
Some(serde_json::Value::Array(
|
||||
key_page
|
||||
.items
|
||||
.into_iter()
|
||||
.map(|key| {
|
||||
let api_formats =
|
||||
provider_key_effective_api_formats(&key, &provider.provider_type, &endpoints);
|
||||
state.build_admin_provider_key_response(
|
||||
&key,
|
||||
&provider.provider_type,
|
||||
&api_formats,
|
||||
now_unix_secs,
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
))
|
||||
let items = key_page
|
||||
.items
|
||||
.into_iter()
|
||||
.map(|key| {
|
||||
let api_formats =
|
||||
provider_key_effective_api_formats(&key, &provider.provider_type, &endpoints);
|
||||
state.build_admin_provider_key_response(
|
||||
&key,
|
||||
&provider.provider_type,
|
||||
&api_formats,
|
||||
now_unix_secs,
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
Some((items, key_page.total))
|
||||
}
|
||||
|
||||
pub(crate) async fn build_admin_provider_keys_payload(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
skip: usize,
|
||||
limit: usize,
|
||||
) -> Option<Value> {
|
||||
let (items, _) =
|
||||
build_admin_provider_key_items_payload(state, provider_id, skip, limit).await?;
|
||||
Some(Value::Array(items))
|
||||
}
|
||||
|
||||
pub(crate) async fn build_admin_provider_keys_page_payload(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
page: usize,
|
||||
page_size: usize,
|
||||
) -> Option<Value> {
|
||||
let skip = page.saturating_sub(1).saturating_mul(page_size);
|
||||
let (items, total) =
|
||||
build_admin_provider_key_items_payload(state, provider_id, skip, page_size).await?;
|
||||
Some(json!({
|
||||
"total": total,
|
||||
"page": page,
|
||||
"page_size": page_size,
|
||||
"keys": items,
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -70,6 +70,21 @@ impl<'a> AdminAppState<'a> {
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn build_admin_provider_keys_page_payload(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
page: usize,
|
||||
page_size: usize,
|
||||
) -> Option<serde_json::Value> {
|
||||
crate::handlers::admin::provider::write::keys::build_admin_provider_keys_page_payload(
|
||||
self,
|
||||
provider_id,
|
||||
page,
|
||||
page_size,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) fn build_admin_reveal_key_payload(
|
||||
&self,
|
||||
key: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey,
|
||||
|
||||
@@ -118,6 +118,90 @@ async fn gateway_handles_admin_provider_keys_locally_with_trusted_admin_principa
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_keys_page_locally_with_total() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/providers/provider-openai/keys",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut key_a = sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
);
|
||||
key_a.internal_priority = 10;
|
||||
key_a.created_at_unix_ms = Some(1_711_000_000);
|
||||
|
||||
let mut key_b = sample_key(
|
||||
"key-openai-b",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-b",
|
||||
);
|
||||
key_b.internal_priority = 20;
|
||||
key_b.created_at_unix_ms = Some(1_711_100_000);
|
||||
|
||||
let mut key_c = sample_key(
|
||||
"key-openai-c",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-c",
|
||||
);
|
||||
key_c.internal_priority = 30;
|
||||
key_c.created_at_unix_ms = Some(1_711_200_000);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![key_a, key_b, key_c],
|
||||
));
|
||||
|
||||
let (_upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-openai/keys?page=2&page_size=1"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total"], 3);
|
||||
assert_eq!(payload["page"], 2);
|
||||
assert_eq!(payload["page_size"], 1);
|
||||
let items = payload["keys"].as_array().expect("keys should be an array");
|
||||
assert_eq!(items.len(), 1);
|
||||
assert_eq!(items[0]["id"], "key-openai-b");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_admin_provider_keys_prefers_upstream_plan_type_over_auth_config() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
|
||||
Reference in New Issue
Block a user