omni_orchestrator/
server.rs

1use rocket::{Rocket, Build};
2use std::sync::Arc;
3use tokio::sync::RwLock;
4use colored::Colorize;
5use libomni::types::db::auth::AuthConfig;
6
7use crate::cluster::ClusterManager;
8use crate::state::SharedState;
9use crate::db_manager::DatabaseManager;
10use crate::cors::CORS;
11use crate::endpoints::{health_check, cluster_status};
12use crate::cors::cors_preflight;
13use crate::schemas::v1::api;
14
15pub trait RocketExt {
16    fn mount_routes(self, routes: Vec<(&'static str, Vec<rocket::Route>)>) -> Self;
17}
18
19impl RocketExt for Rocket<Build> {
20    fn mount_routes(self, routes: Vec<(&'static str, Vec<rocket::Route>)>) -> Self {
21        let mut rocket = self;
22        for (path, routes) in routes {
23            log::info!("{}", format!("Mounting routes at {}", path).green());
24            rocket = rocket.mount(path, routes);
25        }
26        rocket
27    }
28}
29
30pub fn build_rocket(
31    port: u16,
32    db_manager: Arc<DatabaseManager>,
33    pool: sqlx::Pool<sqlx::MySql>,
34    cluster_manager: Arc<RwLock<ClusterManager>>,
35    clickhouse_client: clickhouse::Client,
36    shared_state: Arc<RwLock<SharedState>>,
37    auth_config: AuthConfig,
38) -> Rocket<Build> {
39    println!(
40        "{}",
41        "╔═══════════════════════════════════════════════════════════════╗".bright_cyan()
42    );
43    println!(
44        "{}",
45        "║                       SERVER STARTUP                          ║".bright_cyan()
46    );
47    println!(
48        "{}",
49        "╚═══════════════════════════════════════════════════════════════╝".bright_cyan()
50    );
51
52    log::info!("{}", "Defining API routes".cyan());
53    let routes = vec![
54        (
55            "/",
56            routes![
57                health_check,
58                api::index::routes_ui,
59                cluster_status,
60                cors_preflight
61            ],
62        ),
63        ("/api/v1", api::routes()),
64    ];
65
66    log::info!("{}", "Building Rocket instance".cyan());
67    let rocket_instance = rocket::build()
68        .configure(rocket::Config {
69            port,
70            address: std::net::IpAddr::V4(std::net::Ipv4Addr::new(0, 0, 0, 0)),
71            ..Default::default()
72        })
73        .manage(db_manager)
74        .manage(pool)
75        .manage(cluster_manager)
76        .manage(clickhouse_client)
77        .manage(shared_state)
78        .manage(auth_config)
79        .attach(CORS);
80
81    log::info!("{}", "Mounting API routes".cyan());
82    let rocket_with_routes = rocket_instance.mount_routes(routes);
83
84    api::index::collect_routes(&rocket_with_routes);
85
86    rocket_with_routes
87}