iota_rest_kv/
main.rs

1// Copyright (c) 2025 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{fs, num::NonZeroUsize, path::PathBuf};
5
6use anyhow::Result;
7use clap::Parser;
8use serde::{Deserialize, Serialize};
9use server::Server;
10use tokio_util::sync::CancellationToken;
11use tracing::Level;
12use tracing_subscriber::{EnvFilter, FmtSubscriber};
13
14/// This module contains the DynamoDb and S3 implementation of the KV store
15/// client.
16#[allow(dead_code)]
17mod aws;
18/// This module contains the Bigtable implementation of the KV store client.
19mod bigtable;
20mod errors;
21mod extractors;
22mod routes;
23mod server;
24mod types;
25
26use bigtable::KvStoreConfig;
27
28/// The main CLI application.
29#[derive(Parser, Clone, Debug)]
30#[clap(
31    name = "KV Store REST API",
32    about = "A HTTP server exposing key-value data of the IOTA network through a REST API."
33)]
34struct Cli {
35    #[clap(long, default_value = "INFO", env = "LOG_LEVEL")]
36    log_level: Level,
37    /// The yaml config file path.
38    #[clap(short, long)]
39    config: PathBuf,
40}
41
42#[derive(Serialize, Deserialize, Clone, Debug)]
43#[serde(rename_all = "kebab-case")]
44pub struct RestApiConfig {
45    #[serde(flatten)]
46    pub kv_store_config: KvStoreConfig,
47    pub server_address: std::net::SocketAddr,
48    #[serde(default = "default_multiget_max_items")]
49    pub multiget_max_items: NonZeroUsize,
50}
51
52fn default_multiget_max_items() -> NonZeroUsize {
53    NonZeroUsize::new(100).expect("value should be greater than 0")
54}
55
56#[tokio::main]
57async fn main() -> Result<()> {
58    rustls::crypto::ring::default_provider()
59        .install_default()
60        .expect("failed to install rustls crypto provider");
61
62    let cli = Cli::parse();
63
64    init_tracing(cli.log_level);
65
66    let raw_config = fs::read_to_string(cli.config).expect("failed to read config file");
67    let config = serde_yaml::from_str::<RestApiConfig>(&raw_config)?;
68
69    let token = CancellationToken::new();
70
71    shutdown_signal_listener(token.clone());
72
73    let server = Server::new(config, token).await?;
74    server.serve().await
75}
76
77/// Initialize the tracing with custom subscribers.
78///
79/// The default level can be overridden with per-module directives (e.g.
80/// `RUST_LOG=iota_rest_kv=debug,info`).
81fn init_tracing(log_level: Level) {
82    let env_filter =
83        EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(log_level.to_string()));
84    let subscriber = FmtSubscriber::builder()
85        .with_env_filter(env_filter)
86        .finish();
87
88    tracing::subscriber::set_global_default(subscriber).expect("setting default subscriber failed");
89}
90
91/// Set up a `CTRL+C` & `SIGTERM` handler for graceful shutdown and spawn a
92/// tokio task.
93fn shutdown_signal_listener(token: CancellationToken) {
94    tokio::spawn(async move {
95        #[cfg(unix)]
96        let terminate = async {
97            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
98                .expect("cannot listen to SIGTERM signal")
99                .recv()
100                .await;
101        };
102
103        #[cfg(not(unix))]
104        let terminate = std::future::pending::<()>();
105
106        tokio::select! {
107            _ = tokio::signal::ctrl_c() => tracing::info!("shutting down, CTRL+C signal received"),
108            _ = terminate => tracing::info!("shutting down, SIGTERM signal received")
109        };
110
111        token.cancel();
112    });
113}
iota_rest_kv/main.rs

iota_rest_kv/
main.rs
