reflactor(handler): move db handlers to state

Author: CeerDecy

mc-db/src/lib.rs

@@ -1,7 +1,6 @@
 mod mcp_handler;
 pub mod model;
 mod settings_handler;
-mod token_handler;
 
 pub use mcp_handler::*;
 pub use settings_handler::*;

mc-db/src/model/mcp_servers.rs

@@ -1,6 +1,6 @@
+use chrono::NaiveDateTime;
 use serde::{Deserialize, Serialize};
 use sqlx::FromRow;
-use chrono::NaiveDateTime;
 use std::fmt::Display;
 use std::str::FromStr;
 use uuid::Uuid;

mc-db/src/token_handler.rs

@@ -6,7 +6,7 @@ use cache::mcp_servers::Cache;
 use hyper_rustls::HttpsConnector;
 use hyper_util::client::legacy::Client;
 use hyper_util::client::legacy::connect::HttpConnector;
-use mc_db::DBClient;
+use mc_db::{DBClient, McpDBHandler, SystemSettingsDBHandler};
 pub use mcp_server::*;
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
@@ -21,6 +21,54 @@ pub struct AppState {
     pub event_sender: Sender<Event>,
     pub https_client: Arc<Client<HttpsConnector<HttpConnector>, Body>>,
     pub mcp_cache: Arc<Cache>,
+    handler_manager: HandlerManager,
+}
+impl AppState {
+    pub fn new(
+        db: Arc<DBClient>,
+        event_sender: Sender<Event>,
+        https_client: Arc<Client<HttpsConnector<HttpConnector>, Body>>,
+        mcp_cache: Arc<Cache>,
+        handler_manager: HandlerManager,
+    ) -> Self {
+        Self {
+            db,
+            event_sender,
+            https_client,
+            mcp_cache,
+            handler_manager,
+        }
+    }
+    pub fn handlers(&self) -> &HandlerManager {
+        &self.handler_manager
+    }
+}
+
+#[derive(Clone)]
+pub struct HandlerManager {
+    pub mcp_handler: Option<Arc<McpDBHandler>>,
+    pub system_settings_handler: Option<Arc<SystemSettingsDBHandler>>,
+    db: Arc<DBClient>,
+}
+
+impl HandlerManager {
+    pub fn new(db: Arc<DBClient>) -> Self {
+        HandlerManager {
+            db,
+            mcp_handler: None,
+            system_settings_handler: None,
+        }
+    }
+
+    pub fn with_mcp_handler(mut self) -> Self {
+        self.mcp_handler = Some(Arc::new(McpDBHandler::new(self.db.clone())));
+        self
+    }
+    pub fn with_system_settings_handler(mut self) -> Self {
+        self.system_settings_handler =
+            Some(Arc::new(SystemSettingsDBHandler::new(self.db.clone())));
+        self
+    }
 }
 
 #[derive(Deserialize, Serialize, Clone, Debug)]

mc-registry/src/lib.rs

@@ -4,7 +4,6 @@ use axum::Json;
 use axum::extract::{Query, State};
 use axum::http::StatusCode;
 use mc_db::model::{CreateFrom, McpServers, SettingKey};
-use mc_db::{McpDBHandler, SystemSettingsDBHandler};
 use serde::{Deserialize, Serialize};
 use uuid::Uuid;
 
@@ -17,8 +16,25 @@ pub async fn list_all(
     State(state): State<AppState>,
     Query(params): Query<ListAllRequest>,
 ) -> Result<Json<Response>, (StatusCode, String)> {
-    let mcp_handler = McpDBHandler::new(state.db.clone());
-    let settings_handler = SystemSettingsDBHandler::new(state.db.clone());
+    let mcp_handler = match &state.handlers().mcp_handler {
+        None => {
+            return Err((
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Can't get MCP handler not found".to_string(),
+            ));
+        }
+        Some(handler) => handler,
+    };
+
+    let settings_handler = match &state.handlers().system_settings_handler {
+        None => {
+            return Err((
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Can't get MCP handler not found".to_string(),
+            ));
+        }
+        Some(handler) => handler,
+    };
 
     let self_address = settings_handler
         .get_system_settings(SettingKey::SelfAddress)
@@ -67,8 +83,17 @@ pub async fn register_mcp_server(
     State(state): State<AppState>,
     Json(server): Json<McpRegisterRequest>,
 ) -> Result<Json<Response>, (StatusCode, String)> {
-    let db_client = state.db.clone();
-    let res = McpDBHandler::new(db_client)
+    let mcp_handler = match &state.handlers().mcp_handler {
+        None => {
+            return Err((
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Can't get MCP handler not found".to_string(),
+            ));
+        }
+        Some(handler) => handler,
+    };
+
+    let res = mcp_handler
         .create(&McpServers {
             id: Uuid::new_v4(),
             name: server.name.clone(),

mc-registry/src/mcp_server.rs

@@ -8,9 +8,9 @@ use hyper_util::client::legacy::Client;
 use hyper_util::rt::TokioExecutor;
 use mc_booter::app::application::Application;
 use mc_db::DBClient;
-use mc_registry::AppState;
 use mc_registry::cache::mcp_servers::Cache;
 use mc_registry::event::Event;
+use mc_registry::{AppState, HandlerManager};
 use std::error::Error;
 use std::sync::Arc;
 use tokio::runtime::Runtime;
@@ -161,12 +161,17 @@ impl Application for McpCenterServer {
             100,
         ));
 
-        let state = AppState {
-            db: db_client.clone(),
-            event_sender: tx.clone(),
-            https_client: client.clone(),
-            mcp_cache: cache.clone(),
-        };
+        let manager = HandlerManager::new(db_client.clone())
+            .with_mcp_handler()
+            .with_system_settings_handler();
+
+        let state = AppState::new(
+            db_client.clone(),
+            tx.clone(),
+            client.clone(),
+            cache.clone(),
+            manager,
+        );
 
         self.state = Some(state);