From d21459802c5e67295d0cf713adb35264133cf9cd Mon Sep 17 00:00:00 2001 From: GW_MC <72297530+GWMCwing@users.noreply.github.com> Date: Thu, 1 Jan 2026 10:40:44 +0800 Subject: [PATCH] Add total upstream count retrieval to UpstreamService --- .../restricted/nginx/upstream/get_upstream.rs | 50 ++++++++++++++----- apps/api/src/services/nginx/upstream.rs | 31 +++++++++++- 2 files changed, 67 insertions(+), 14 deletions(-) diff --git a/apps/api/src/routes/api/restricted/nginx/upstream/get_upstream.rs b/apps/api/src/routes/api/restricted/nginx/upstream/get_upstream.rs index 948f056..827f282 100644 --- a/apps/api/src/routes/api/restricted/nginx/upstream/get_upstream.rs +++ b/apps/api/src/routes/api/restricted/nginx/upstream/get_upstream.rs @@ -54,23 +54,32 @@ pub async fn get_upstream_list( State(state): State>, ) -> AxumResult, ServiceError> { let upstream_service = &state.service.nginx.get_upstream_service(); - let upstreams = upstream_service - .get_upstreams( + + let (upstreams_res, upstream_count_res) = tokio::join!( + upstream_service.get_upstreams( Some(pagination.clone().into()), Some(GetUpstreamOptions { include_targets: true, filter_by_enabled: false, }), None, - ) - .await?; + ), + upstream_service.get_total_upstreams(None, None), + ); + + let upstreams = upstreams_res?; + let upstream_count = upstream_count_res?; // Ok(Json(UpstreamListResponse { items: upstreams.into_iter().map(|u| u.into()).collect(), pagination: PaginationInfo { - total_items: 0, - total_pages: 0, + total_items: upstream_count, + total_pages: if upstream_count == 0 { + 0 + } else { + (upstream_count as f32 / pagination.per_page as f32).ceil() as u32 + }, current_page: pagination.page, per_page: pagination.per_page, }, @@ -118,19 +127,21 @@ pub async fn get_upstream( #[cfg(test)] mod tests { use super::*; - use std::sync::Arc; + use std::{collections::BTreeMap, sync::Arc}; use axum::http::StatusCode; use axum_test::TestServer; - use sea_orm::{DatabaseBackend, DatabaseConnection, MockDatabase}; + use sea_orm::{DatabaseBackend, DatabaseConnection, MockDatabase, Value}; use database::generated::entities::{upstream, upstream_target}; - use crate::configs::{FromConfig, ProgramSettings}; - - use crate::routes::api::restricted::nginx::upstream::get_upstream_router; - use crate::routes::api::restricted::nginx::upstream::info::response::UpstreamInfoResponse; - use crate::services::get_app_service; + use crate::{ + configs::{FromConfig, ProgramSettings}, + routes::api::restricted::nginx::upstream::{ + get_upstream_router, info::response::UpstreamInfoResponse, + }, + services::get_app_service, + }; fn get_router_with_state(db: DatabaseConnection) -> axum::Router { let program_settings = ProgramSettings::mock(); @@ -171,6 +182,10 @@ mod tests { (u1.clone(), None::), (u2.clone(), None::), ]]) + .append_query_results(vec![vec![BTreeMap::from([( + "count".to_string(), + Value::BigInt(Some(2)), + )])]]) .into_connection(); let router = get_router_with_state(db.clone()); @@ -181,6 +196,7 @@ mod tests { let body = res.json::(); assert_eq!(body.items.len(), 2); assert_eq!(body.pagination.current_page, 1u32); + assert_eq!(body.pagination.total_pages, 1u32); } #[tokio::test] @@ -233,6 +249,10 @@ mod tests { async fn extractor_pagination_validation_rejects_bad_values() { let db = MockDatabase::new(DatabaseBackend::Sqlite) .append_query_results(vec![Vec::::new()]) + .append_query_results(vec![vec![BTreeMap::from([( + "count".to_string(), + Value::BigInt(Some(0)), + )])]]) .into_connection(); let router = get_router_with_state(db.clone()); let server = TestServer::new(router).expect("failed to create test server"); @@ -304,6 +324,10 @@ mod tests { async fn handler_get_upstream_list_empty_returns_empty_items() { let db = MockDatabase::new(DatabaseBackend::Sqlite) .append_query_results(vec![Vec::::new()]) + .append_query_results(vec![vec![BTreeMap::from([( + "count".to_string(), + Value::BigInt(Some(0)), + )])]]) .into_connection(); let router = get_router_with_state(db.clone()); diff --git a/apps/api/src/services/nginx/upstream.rs b/apps/api/src/services/nginx/upstream.rs index db7f160..9678989 100644 --- a/apps/api/src/services/nginx/upstream.rs +++ b/apps/api/src/services/nginx/upstream.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use sea_orm::{ ActiveModelTrait, ColumnTrait, DatabaseConnection, DatabaseTransaction, EntityTrait, ExprTrait, - ModelTrait, QueryFilter, QuerySelect, QueryTrait, TransactionTrait, + FromQueryResult, ModelTrait, QueryFilter, QuerySelect, QueryTrait, TransactionTrait, }; use database::generated::entities::{upstream, upstream_target}; @@ -29,6 +29,11 @@ pub trait UpstreamService: Send + Sync { create_info: UpstreamCreateInfo, tx: Option<&mut DatabaseTransaction>, ) -> Result; + async fn get_total_upstreams( + &self, + options: Option, + tx: Option<&mut DatabaseTransaction>, + ) -> Result; async fn get_upstream( &self, upstream_id: uuid::Uuid, @@ -97,6 +102,9 @@ pub struct GetUpstreamOptions { pub filter_by_enabled: bool, } +#[allow(dead_code)] +pub struct UpstreamTotalCountOptions {} + #[derive(Default)] pub struct GetUpstreamTargetOptions { pub include_upstream: bool, @@ -152,6 +160,27 @@ impl UpstreamService for UpstreamServiceImpl { Ok(r.into()) } + async fn get_total_upstreams( + &self, + _options: Option, + tx: Option<&mut DatabaseTransaction>, + ) -> Result { + #[derive(Debug, FromQueryResult)] + struct CountResult { + // The field name must match the column alias in the query + count: i64, + } + let count_info = with_conn!(&*self.connection, tx, conn, { + upstream::Entity::find() + .select_only() + .column_as(upstream::Column::Id, "count") + .into_model::() + .one(*conn) + .await? + }); + Ok(count_info.map_or(0, |c| c.count) as u64) + } + async fn get_upstream( &self, upstream_id: uuid::Uuid,