Add total upstream count retrieval to UpstreamService

apps/api/src/routes/api/restricted/nginx/upstream/get_upstream.rs

@@ -54,23 +54,32 @@ pub async fn get_upstream_list(
     State(state): State<Arc<AppState>>,
 ) -> AxumResult<Json<UpstreamListResponse>, 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();
@@ -167,7 +178,14 @@ mod tests {
         };
 
         let db = MockDatabase::new(DatabaseBackend::Sqlite)
-            .append_query_results(vec![vec![u1.clone(), u2.clone()]])
+            .append_query_results(vec![vec![
+                (u1.clone(), None::<upstream_target::Model>),
+                (u2.clone(), None::<upstream_target::Model>),
+            ]])
+            .append_query_results(vec![vec![BTreeMap::from([(
+                "count".to_string(),
+                Value::BigInt(Some(2)),
+            )])]])
             .into_connection();
 
         let router = get_router_with_state(db.clone());
@@ -178,6 +196,7 @@ mod tests {
         let body = res.json::<UpstreamListResponse>();
         assert_eq!(body.items.len(), 2);
         assert_eq!(body.pagination.current_page, 1u32);
+        assert_eq!(body.pagination.total_pages, 1u32);
     }
 
     #[tokio::test]
@@ -230,6 +249,10 @@ mod tests {
     async fn extractor_pagination_validation_rejects_bad_values() {
         let db = MockDatabase::new(DatabaseBackend::Sqlite)
             .append_query_results(vec![Vec::<sea_orm::MockRow>::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");
@@ -301,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::<sea_orm::MockRow>::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());

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<UpstreamInfo, ServiceError>;
+    async fn get_total_upstreams(
+        &self,
+        options: Option<UpstreamTotalCountOptions>,
+        tx: Option<&mut DatabaseTransaction>,
+    ) -> Result<u64, ServiceError>;
     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<UpstreamTotalCountOptions>,
+        tx: Option<&mut DatabaseTransaction>,
+    ) -> Result<u64, ServiceError> {
+        #[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::<CountResult>()
+                .one(*conn)
+                .await?
+        });
+        Ok(count_info.map_or(0, |c| c.count) as u64)
+    }
+
     async fn get_upstream(
         &self,
         upstream_id: uuid::Uuid,
@@ -580,7 +609,10 @@ mod tests {
         };
 
         let db = MockDatabase::new(DatabaseBackend::Sqlite)
-            .append_query_results(vec![vec![u1.clone(), u2.clone()]])
+            .append_query_results(vec![vec![
+                (u1.clone(), None::<upstream_target::Model>),
+                (u2.clone(), None::<upstream_target::Model>),
+            ]])
             .into_connection();
 
         let svc = UpstreamServiceImpl::new(Arc::new(db));