[graphql/rpc] split out graphiql server

crates/sui-graphql-rpc/src/main.rs

@@ -9,8 +9,9 @@ use sui_graphql_rpc::commands::Command;
 use sui_graphql_rpc::config::Ide;
 use sui_graphql_rpc::config::{ConnectionConfig, ServerConfig, ServiceConfig};
 use sui_graphql_rpc::schema_sdl_export;
-use sui_graphql_rpc::server::builder::Server;
-use sui_graphql_rpc::server::simple_server::start_example_server;
+use sui_graphql_rpc::server::graphiql_server::{
+    start_graphiql_server, start_graphiql_server_from_cfg_path,
+};
 use tracing::error;
 
 #[tokio::main]
@@ -74,20 +75,16 @@ async fn main() {
                 ..ServerConfig::default()
             };
 
-            start_example_server(&server_config).await.unwrap();
+            start_graphiql_server(&server_config).await.unwrap();
         }
         Command::FromConfig { path } => {
-            let server = Server::from_yaml_config(path.to_str().unwrap());
             println!("Starting server...");
-            server
+            start_graphiql_server_from_cfg_path(path.to_str().unwrap())
                 .await
                 .map_err(|x| {
                     error!("Error: {:?}", x);
                     x
                 })
-                .unwrap()
-                .run()
-                .await
                 .unwrap();
         }
     }

crates/sui-graphql-rpc/src/server/builder.rs

@@ -20,6 +20,7 @@ use async_graphql::{extensions::ExtensionFactory, Schema, SchemaBuilder};
 use async_graphql::{EmptyMutation, EmptySubscription};
 use async_graphql_axum::{GraphQLRequest, GraphQLResponse};
 use axum::http::HeaderMap;
+use axum::routing::{post, MethodRouter};
 use axum::{
     extract::{connect_info::IntoMakeServiceWithConnectInfo, ConnectInfo},
     middleware,
@@ -42,69 +43,6 @@ impl Server {
             .await
             .map_err(|e| Error::Internal(format!("Server run failed: {}", e)))
     }
-
-    pub async fn from_yaml_config(path: &str) -> Result<Self, Error> {
-        let config = ServerConfig::from_yaml(path)?;
-        Self::from_config(&config).await
-    }
-
-    pub async fn from_config(config: &ServerConfig) -> Result<Self, Error> {
-        let mut builder =
-            ServerBuilder::new(config.connection.port, config.connection.host.clone());
-
-        let name_service_config = config.name_service.clone();
-        let reader = PgManager::reader_with_config(
-            config.connection.db_url.clone(),
-            config.connection.db_pool_size,
-        )
-        .map_err(|e| Error::Internal(format!("Failed to create pg connection pool: {}", e)))?;
-        let pg_conn_pool = PgManager::new(reader.clone(), config.service.limits);
-        let package_store = DbPackageStore(reader);
-        let package_cache = PackageStoreWithLruCache::new(package_store);
-
-        let prom_addr: SocketAddr = format!(
-            "{}:{}",
-            config.connection.prom_url, config.connection.prom_port
-        )
-        .parse()
-        .map_err(|_| {
-            Error::Internal(format!(
-                "Failed to parse url {}, port {} into socket address",
-                config.connection.prom_url, config.connection.prom_port
-            ))
-        })?;
-        let registry_service = mysten_metrics::start_prometheus_server(prom_addr);
-        println!("Starting Prometheus HTTP endpoint at {}", prom_addr);
-        let registry = registry_service.default_registry();
-
-        let metrics = RequestMetrics::new(&registry);
-
-        builder = builder
-            .max_query_depth(config.service.limits.max_query_depth)
-            .max_query_nodes(config.service.limits.max_query_nodes)
-            .context_data(config.service.clone())
-            .context_data(pg_conn_pool)
-            .context_data(Resolver::new(package_cache))
-            .context_data(name_service_config)
-            .ide_title(config.ide.ide_title.clone())
-            .context_data(Arc::new(metrics))
-            .context_data(config.clone());
-
-        if config.internal_features.feature_gate {
-            builder = builder.extension(FeatureGate);
-        }
-        if config.internal_features.logger {
-            builder = builder.extension(Logger::default());
-        }
-        if config.internal_features.query_limits_checker {
-            builder = builder.extension(QueryLimitsChecker::default());
-        }
-        if config.internal_features.query_timeout {
-            builder = builder.extension(Timeout);
-        }
-
-        builder.build()
-    }
 }
 
 pub(crate) struct ServerBuilder {
@@ -113,6 +51,8 @@ pub(crate) struct ServerBuilder {
 
     schema: SchemaBuilder<Query, EmptyMutation, EmptySubscription>,
     ide_title: Option<String>,
+
+    router: Option<Router>,
 }
 
 impl ServerBuilder {
@@ -122,6 +62,7 @@ impl ServerBuilder {
             host,
             schema: async_graphql::Schema::build(Query, EmptyMutation, EmptySubscription),
             ide_title: None,
+            router: None,
         }
     }
 
@@ -158,19 +99,58 @@ impl ServerBuilder {
         self.schema.finish()
     }
 
-    pub fn build(self) -> Result<Server, Error> {
+    fn build_components(
+        self,
+    ) -> (
+        String,
+        Option<String>,
+        Schema<Query, EmptyMutation, EmptySubscription>,
+        Router,
+    ) {
         let address = self.address();
-        let ide_title = self.ide_title.clone();
-        let schema = self.build_schema();
+        let ServerBuilder {
+            schema,
+            // TODO: remove this once we have expose layer in builder.
+            // This should be set in the builder.
+            ide_title,
+            router,
+            ..
+        } = self;
+        (
+            address,
+            ide_title,
+            schema.finish(),
+            router.expect("Router not initialized"),
+        )
+    }
 
-        let app = axum::Router::new()
-            .route("/", axum::routing::get(graphiql).post(graphql_handler))
-            .route("/schema", axum::routing::get(get_schema))
-            .route("/health", axum::routing::get(health_checks))
+    fn init_router(&mut self) {
+        if self.router.is_none() {
+            let router: Router = Router::new()
+                .route("/", post(graphql_handler))
+                .route("/health", axum::routing::get(health_checks))
+                .route("/schema", axum::routing::get(get_schema))
+                .layer(middleware::from_fn(check_version_middleware))
+                .layer(middleware::from_fn(set_version_middleware));
+            self.router = Some(router);
+        }
+    }
+
+    pub fn route(mut self, path: &str, method_handler: MethodRouter) -> Self {
+        self.init_router();
+        self.router = self.router.map(|router| router.route(path, method_handler));
+        self
+    }
+
+    pub fn build(self) -> Result<Server, Error> {
+        let (address, ide_title, schema, router) = self.build_components();
+
+        let app = router
             .layer(axum::extract::Extension(schema))
-            .layer(axum::extract::Extension(ide_title))
-            .layer(middleware::from_fn(check_version_middleware))
-            .layer(middleware::from_fn(set_version_middleware));
+            // TODO: remove this once we have expose layer in builder.
+            // This should be set in the builder.
+            .layer(axum::extract::Extension(ide_title));
+
         Ok(Server {
             server: axum::Server::bind(
                 &address
@@ -180,6 +160,69 @@ impl ServerBuilder {
             .serve(app.into_make_service_with_connect_info::<SocketAddr>()),
         })
     }
+
+    pub async fn from_yaml_config(path: &str) -> Result<Self, Error> {
+        let config = ServerConfig::from_yaml(path)?;
+        Self::from_config(&config).await
+    }
+
+    pub async fn from_config(config: &ServerConfig) -> Result<Self, Error> {
+        let mut builder =
+            ServerBuilder::new(config.connection.port, config.connection.host.clone());
+
+        let name_service_config = config.name_service.clone();
+        let reader = PgManager::reader_with_config(
+            config.connection.db_url.clone(),
+            config.connection.db_pool_size,
+        )
+        .map_err(|e| Error::Internal(format!("Failed to create pg connection pool: {}", e)))?;
+        let pg_conn_pool = PgManager::new(reader.clone(), config.service.limits);
+        let package_store = DbPackageStore(reader);
+        let package_cache = PackageStoreWithLruCache::new(package_store);
+
+        let prom_addr: SocketAddr = format!(
+            "{}:{}",
+            config.connection.prom_url, config.connection.prom_port
+        )
+        .parse()
+        .map_err(|_| {
+            Error::Internal(format!(
+                "Failed to parse url {}, port {} into socket address",
+                config.connection.prom_url, config.connection.prom_port
+            ))
+        })?;
+        let registry_service = mysten_metrics::start_prometheus_server(prom_addr);
+        println!("Starting Prometheus HTTP endpoint at {}", prom_addr);
+        let registry = registry_service.default_registry();
+
+        let metrics = RequestMetrics::new(&registry);
+
+        builder = builder
+            .max_query_depth(config.service.limits.max_query_depth)
+            .max_query_nodes(config.service.limits.max_query_nodes)
+            .context_data(config.service.clone())
+            .context_data(pg_conn_pool)
+            .context_data(Resolver::new(package_cache))
+            .context_data(name_service_config)
+            .ide_title(config.ide.ide_title.clone())
+            .context_data(Arc::new(metrics))
+            .context_data(config.clone());
+
+        if config.internal_features.feature_gate {
+            builder = builder.extension(FeatureGate);
+        }
+        if config.internal_features.logger {
+            builder = builder.extension(Logger::default());
+        }
+        if config.internal_features.query_limits_checker {
+            builder = builder.extension(QueryLimitsChecker::default());
+        }
+        if config.internal_features.query_timeout {
+            builder = builder.extension(Timeout);
+        }
+
+        Ok(builder)
+    }
 }
 
 async fn get_schema() -> impl axum::response::IntoResponse {
@@ -211,15 +254,6 @@ async fn graphql_handler(
     schema.execute(req).await.into()
 }
 
-async fn graphiql(ide_title: axum::Extension<Option<String>>) -> impl axum::response::IntoResponse {
-    let gq = async_graphql::http::GraphiQLSource::build().endpoint("/");
-    if let axum::Extension(Some(title)) = ide_title {
-        axum::response::Html(gq.title(&title).finish())
-    } else {
-        axum::response::Html(gq.finish())
-    }
-}
-
 async fn health_checks(
     schema: axum::Extension<SuiGraphQLSchema>,
 ) -> impl axum::response::IntoResponse {

crates/sui-graphql-rpc/src/server/graphiql_server.rs

@@ -0,0 +1,39 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use tracing::info;
+
+use crate::config::ServerConfig;
+use crate::error::Error;
+use crate::server::builder::ServerBuilder;
+
+async fn graphiql(ide_title: axum::Extension<Option<String>>) -> impl axum::response::IntoResponse {
+    let gq = async_graphql::http::GraphiQLSource::build().endpoint("/");
+    if let axum::Extension(Some(title)) = ide_title {
+        axum::response::Html(gq.title(&title).finish())
+    } else {
+        axum::response::Html(gq.finish())
+    }
+}
+
+pub async fn start_graphiql_server(server_config: &ServerConfig) -> Result<(), Error> {
+    info!("Starting server with config: {:?}", server_config);
+    start_graphiql_server_impl(ServerBuilder::from_config(server_config).await?).await
+}
+
+pub async fn start_graphiql_server_from_cfg_path(server_config_path: &str) -> Result<(), Error> {
+    start_graphiql_server_impl(ServerBuilder::from_yaml_config(server_config_path).await?).await
+}
+
+async fn start_graphiql_server_impl(server_builder: ServerBuilder) -> Result<(), Error> {
+    let address = server_builder.address();
+
+    // Add GraphiQL IDE handler on GET request to `/`` endpoint
+    let server = server_builder
+        .route("/", axum::routing::get(graphiql))
+        .build()?;
+
+    info!("Launch GraphiQL IDE at: http://{}", address);
+
+    server.run().await
+}

crates/sui-graphql-rpc/src/server/mod.rs

@@ -1,7 +1,7 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-pub mod simple_server;
+pub mod graphiql_server;
 
 pub mod builder;
 pub mod version;

crates/sui-graphql-rpc/src/test_infra/cluster.rs

@@ -4,7 +4,7 @@
 use crate::client::simple_client::SimpleClient;
 use crate::config::ConnectionConfig;
 use crate::config::ServerConfig;
-use crate::server::simple_server::start_example_server;
+use crate::server::graphiql_server::start_graphiql_server;
 use mysten_metrics::init_metrics;
 use std::env;
 use std::net::SocketAddr;
@@ -135,7 +135,7 @@ pub async fn start_graphql_server(graphql_connection_config: ConnectionConfig) -
 
     // Starts graphql server
     tokio::spawn(async move {
-        start_example_server(&server_config).await.unwrap();
+        start_graphiql_server(&server_config).await.unwrap();
     })
 }