Add total upstream count retrieval to UpstreamService

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,