From 3476021bc8c513b223accffd7775c9b774fd33fa Mon Sep 17 00:00:00 2001
From: Johann Wagner <johann.wagner@wobcom.de>
Date: Tue, 14 Nov 2023 16:40:42 +0100
Subject: [PATCH] Implemented Rocket config loading for worker

---
 Cargo.lock                       |  1 +
 vicky-worker/Cargo.toml          |  1 +
 vicky-worker/Rocket.example.toml |  4 ++++
 vicky-worker/src/config.rs       |  7 -------
 vicky-worker/src/main.rs         | 36 +++++++++++++++++++++-----------
 5 files changed, 30 insertions(+), 19 deletions(-)
 create mode 100644 vicky-worker/Rocket.example.toml
 delete mode 100644 vicky-worker/src/config.rs

diff --git a/Cargo.lock b/Cargo.lock
index 89399db..30b7b7f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3018,6 +3018,7 @@ dependencies = [
  "futures-util",
  "hyper",
  "log",
+ "rocket",
  "serde",
  "serde_json",
  "serde_yaml",
diff --git a/vicky-worker/Cargo.toml b/vicky-worker/Cargo.toml
index 777d101..ba2b3d0 100644
--- a/vicky-worker/Cargo.toml
+++ b/vicky-worker/Cargo.toml
@@ -17,3 +17,4 @@ serde_yaml = "0.9.25"
 tokio = { version = "1.32.0", features = ["rt", "macros", "process"] }
 tokio-util = { version = "0.7.9", features = ["codec"] }
 uuid = { version = "1.4.1", features = ["serde"] }
+rocket = { version="=0.5.0-rc.3", features = ["json", "secrets"] } 
diff --git a/vicky-worker/Rocket.example.toml b/vicky-worker/Rocket.example.toml
new file mode 100644
index 0000000..d872024
--- /dev/null
+++ b/vicky-worker/Rocket.example.toml
@@ -0,0 +1,4 @@
+[default]
+
+vicky_url = ""
+machine_token = ""
\ No newline at end of file
diff --git a/vicky-worker/src/config.rs b/vicky-worker/src/config.rs
deleted file mode 100644
index 40d5dab..0000000
--- a/vicky-worker/src/config.rs
+++ /dev/null
@@ -1,7 +0,0 @@
-use serde::Deserialize;
-
-#[derive(Deserialize)]
-pub(crate) struct Config {
-    pub(crate) vicky_url: String,
-    pub(crate) machine_token: String,
-}
diff --git a/vicky-worker/src/main.rs b/vicky-worker/src/main.rs
index 4d129eb..3124c89 100644
--- a/vicky-worker/src/main.rs
+++ b/vicky-worker/src/main.rs
@@ -1,7 +1,3 @@
-mod config;
-
-use config::Config;
-
 use anyhow::anyhow;
 use futures_util::{Sink, StreamExt, TryStreamExt};
 use hyper::{Body, Client, Method, Request};
@@ -13,14 +9,30 @@ use tokio::process::Command;
 use tokio_util::codec::{FramedRead, LinesCodec};
 use uuid::Uuid;
 
+
+use rocket::figment::{Figment, Profile};
+use rocket::figment::providers::{Toml, Env, Format};
+#[derive(Deserialize)]
+pub(crate) struct AppConfig {
+    pub(crate) vicky_url: String,
+    pub(crate) machine_token: String,
+}
+
 fn main() -> anyhow::Result<()> {
     env_logger::init();
-    let cfg: Config = serde_yaml::from_slice(&std::fs::read("config.yaml")?)?;
-    run(cfg)
+
+    // Took from rocket source code and added .split("__") to be able to add keys in nested structures.
+    let rocket_config_figment = Figment::from(rocket::Config::default())
+        .merge(Toml::file(Env::var_or("ROCKET_CONFIG", "Rocket.toml")).nested())
+        .merge(Env::prefixed("ROCKET_").ignore(&["PROFILE"]).split("__").global())
+        .select(Profile::from_env_or("ROCKET_PROFILE", rocket::Config::DEFAULT_PROFILE));
+
+    let app_config = rocket_config_figment.extract::<AppConfig>()?; 
+    run(app_config)
 }
 
 async fn api<Q: Serialize, R: DeserializeOwned>(
-    cfg: &Config,
+    cfg: &AppConfig,
     method: Method,
     endpoint: &str,
     q: &Q,
@@ -74,7 +86,7 @@ pub struct Task {
     pub flake_ref: FlakeRef,
 }
 
-fn log_sink(cfg: Arc<Config>, task_id: Uuid) -> impl Sink<Vec<String>, Error = anyhow::Error> + Send {
+fn log_sink(cfg: Arc<AppConfig>, task_id: Uuid) -> impl Sink<Vec<String>, Error = anyhow::Error> + Send {
     futures_util::sink::unfold((), move |_, lines: Vec<String>| {
         println!("{}", lines.len());
         let cfg = cfg.clone();
@@ -90,7 +102,7 @@ fn log_sink(cfg: Arc<Config>, task_id: Uuid) -> impl Sink<Vec<String>, Error = a
     })
 }
 
-async fn try_run_task(cfg: Arc<Config>, task: &Task) -> anyhow::Result<()> {
+async fn try_run_task(cfg: Arc<AppConfig>, task: &Task) -> anyhow::Result<()> {
     let mut args = vec![
         "run".into(),
         "-v".into(),
@@ -129,7 +141,7 @@ async fn try_run_task(cfg: Arc<Config>, task: &Task) -> anyhow::Result<()> {
     }
 }
 
-async fn run_task(cfg: Arc<Config>, task: Task) {
+async fn run_task(cfg: Arc<AppConfig>, task: Task) {
     let result = match try_run_task(cfg.clone(), &task).await {
         Err(e) => {
             log::info!("task failed: {} {} {:?}", task.id, task.display_name, e);
@@ -147,7 +159,7 @@ async fn run_task(cfg: Arc<Config>, task: Task) {
     .await;
 }
 
-async fn try_claim(cfg: Arc<Config>) -> anyhow::Result<()> {
+async fn try_claim(cfg: Arc<AppConfig>) -> anyhow::Result<()> {
     if let Some(task) =
         api::<_, Option<Task>>(&cfg, Method::POST, "api/v1/tasks/claim", &None::<u32>).await?
     {
@@ -161,7 +173,7 @@ async fn try_claim(cfg: Arc<Config>) -> anyhow::Result<()> {
 }
 
 #[tokio::main(flavor = "current_thread")]
-async fn run(cfg: Config) -> anyhow::Result<()> {
+async fn run(cfg: AppConfig) -> anyhow::Result<()> {
     println!("Hello, world!");
 
     let cfg = Arc::new(cfg);