iota_rest_kv/
server.rs

1// Copyright (c) 2025 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4//! This module includes helper wrappers for building and starting a REST API
5//! server.
6use std::{net::SocketAddr, sync::Arc};
7
8use anyhow::Result;
9use axum::{
10    Router,
11    response::IntoResponse,
12    routing::{get, post},
13};
14use tokio_util::sync::CancellationToken;
15
16use crate::{
17    RestApiConfig,
18    bigtable::KvStoreClient,
19    errors::ApiError,
20    routes::{health, kv_store},
21    types::RestServerAppState,
22};
23
24/// A wrapper which builds the components needed for the REST API server and
25/// provides a simple way to start it.
26pub struct Server {
27    router: Router,
28    server_address: SocketAddr,
29    token: CancellationToken,
30}
31
32impl Server {
33    /// Create a new Server instance.
34    ///
35    /// Based on the config, it instantiates the [`KvStoreClient`] and
36    /// constructs the [`Router`].
37    pub async fn new(config: RestApiConfig, token: CancellationToken) -> Result<Self> {
38        let kv_store_client = KvStoreClient::new(config.kv_store_config).await?;
39
40        let shared_state = Arc::new(RestServerAppState {
41            kv_store_client: Arc::new(kv_store_client),
42            multiget_max_items: config.multiget_max_items,
43        });
44
45        let router = Router::new()
46            .route("/health", get(health::health))
47            .route("/{item_type}", post(kv_store::multi_get_data))
48            .route("/{item_type}/{key}", get(kv_store::data_as_bytes))
49            .with_state(shared_state)
50            .fallback(fallback);
51
52        Ok(Self {
53            router,
54            token,
55            server_address: config.server_address,
56        })
57    }
58
59    /// Start the server, this method is blocking.
60    pub async fn serve(self) -> Result<()> {
61        let listener = tokio::net::TcpListener::bind(self.server_address)
62            .await
63            .expect("failed to bind to socket");
64
65        tracing::info!("listening on: {}", self.server_address);
66
67        axum::serve(listener, self.router)
68            .with_graceful_shutdown(async move {
69                self.token.cancelled().await;
70                tracing::info!("shutdown signal received.");
71            })
72            .await
73            .inspect_err(|e| tracing::error!("server encountered an error: {e}"))
74            .map_err(Into::into)
75    }
76}
77
78/// Handles requests to routes that are not defined in the API.
79///
80/// This fallback handler is called when the requested URL path does not match
81/// any of the defined routes. It returns a `404 Not Found` error, indicating
82/// that the requested resource could not be found. This can happen if the user
83/// enters an incorrect URL or if the requested resource (identified by a
84/// [`Key`](iota_storage::http_key_value_store::Key)) cannot be extracted from
85/// the request.
86async fn fallback() -> impl IntoResponse {
87    ApiError::NotFound
88}
iota_rest_kv/server.rs

iota_rest_kv/
server.rs
