iota_rest_kv/
main.rs

1// Copyright (c) 2025 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{fs, 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::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}
49
50#[tokio::main]
51async fn main() -> Result<()> {
52    let cli = Cli::parse();
53
54    init_tracing(cli.log_level);
55
56    let raw_config = fs::read_to_string(cli.config).expect("failed to read config file");
57    let config = serde_yaml::from_str::<RestApiConfig>(&raw_config)?;
58
59    let token = CancellationToken::new();
60
61    shutdown_signal_listener(token.clone());
62
63    let server = Server::new(config, token).await?;
64    server.serve().await
65}
66
67/// Initialize the tracing with custom subscribers.
68fn init_tracing(log_level: Level) {
69    let subscriber = FmtSubscriber::builder().with_max_level(log_level).finish();
70    tracing::subscriber::set_global_default(subscriber).expect("setting default subscriber failed");
71}
72
73/// Set up a `CTRL+C` & `SIGTERM` handler for graceful shutdown and spawn a
74/// tokio task.
75fn shutdown_signal_listener(token: CancellationToken) {
76    tokio::spawn(async move {
77        #[cfg(unix)]
78        let terminate = async {
79            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
80                .expect("Cannot listen to SIGTERM signal")
81                .recv()
82                .await;
83        };
84
85        #[cfg(not(unix))]
86        let terminate = std::future::pending::<()>();
87
88        tokio::select! {
89            _ = tokio::signal::ctrl_c() => tracing::info!("CTRL+C signal received, shutting down"),
90            _ = terminate => tracing::info!("SIGTERM signal received, shutting down")
91        };
92
93        token.cancel();
94    });
95}
iota_rest_kv/main.rs

iota_rest_kv/
main.rs
