From 29b8cd57a380603d81cef9e8c0cfa63c6083dea2 Mon Sep 17 00:00:00 2001
From: Kirill Fomichev <fanatid@ya.ru>
Date: Mon, 9 Oct 2023 17:41:17 +0400
Subject: [PATCH] kafka: support YAML config (#197)

---
 CHANGELOG.md                         |  1 +
 Cargo.lock                           | 83 ++++++++++++++++++++++++++++
 yellowstone-grpc-kafka/Cargo.toml    |  2 +
 yellowstone-grpc-kafka/src/config.rs | 10 +++-
 4 files changed, 94 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 28526405..71b7ee63 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,7 @@ The minor version will be incremented upon a breaking change and the patch versi
 
 - client: add `GeyserGrpcClient::subscribe_once2` ([#195](https://github.com/rpcpool/yellowstone-grpc/pull/195)).
 - kafka: add metrics (stats, sent, recv) ([#196](https://github.com/rpcpool/yellowstone-grpc/pull/196)).
+- kafka: support YAML config ([#197](https://github.com/rpcpool/yellowstone-grpc/pull/197)).
 
 ### Fixes
 
diff --git a/Cargo.lock b/Cargo.lock
index cbbf30b7..b7d70f31 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1791,6 +1791,17 @@ dependencies = [
  "wasm-bindgen",
 ]
 
+[[package]]
+name = "json5"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "96b0db21af676c1ce64250b5f40f3ce2cf27e4e47cb91ed91eb6fe9350b430c1"
+dependencies = [
+ "pest",
+ "pest_derive",
+ "serde",
+]
+
 [[package]]
 name = "keccak"
 version = "0.1.4"
@@ -2302,6 +2313,51 @@ dependencies = [
  "num",
 ]
 
+[[package]]
+name = "pest"
+version = "2.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c022f1e7b65d6a24c0dbbd5fb344c66881bc01f3e5ae74a1c8100f2f985d98a4"
+dependencies = [
+ "memchr",
+ "thiserror",
+ "ucd-trie",
+]
+
+[[package]]
+name = "pest_derive"
+version = "2.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "35513f630d46400a977c4cb58f78e1bfbe01434316e60c37d27b9ad6139c66d8"
+dependencies = [
+ "pest",
+ "pest_generator",
+]
+
+[[package]]
+name = "pest_generator"
+version = "2.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bc9fc1b9e7057baba189b5c626e2d6f40681ae5b6eb064dc7c7834101ec8123a"
+dependencies = [
+ "pest",
+ "pest_meta",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.37",
+]
+
+[[package]]
+name = "pest_meta"
+version = "2.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1df74e9e7ec4053ceb980e7c0c8bd3594e977fde1af91daba9c928e8e8c6708d"
+dependencies = [
+ "once_cell",
+ "pest",
+ "sha2 0.10.8",
+]
+
 [[package]]
 name = "petgraph"
 version = "0.6.4"
@@ -3023,6 +3079,19 @@ dependencies = [
  "syn 2.0.37",
 ]
 
+[[package]]
+name = "serde_yaml"
+version = "0.9.25"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1a49e178e4452f45cb61d0cd8cebc1b0fafd3e41929e996cef79aa3aca91f574"
+dependencies = [
+ "indexmap 2.0.2",
+ "itoa",
+ "ryu",
+ "serde",
+ "unsafe-libyaml",
+]
+
 [[package]]
 name = "sha2"
 version = "0.9.9"
@@ -4220,6 +4289,12 @@ version = "1.17.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825"
 
+[[package]]
+name = "ucd-trie"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ed646292ffc8188ef8ea4d1e0e0150fb15a5c2e12ad9b8fc191ae7a8a7f3c4b9"
+
 [[package]]
 name = "unicode-bidi"
 version = "0.3.13"
@@ -4260,6 +4335,12 @@ dependencies = [
  "void",
 ]
 
+[[package]]
+name = "unsafe-libyaml"
+version = "0.2.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f28467d3e1d3c6586d8f25fa243f544f5800fec42d97032474e17222c2b75cfa"
+
 [[package]]
 name = "untrusted"
 version = "0.7.1"
@@ -4650,11 +4731,13 @@ dependencies = [
  "futures",
  "git-version",
  "hyper",
+ "json5",
  "lazy_static",
  "prometheus",
  "rdkafka",
  "serde",
  "serde_json",
+ "serde_yaml",
  "sha2 0.10.8",
  "tokio",
  "tokio-stream",
diff --git a/yellowstone-grpc-kafka/Cargo.toml b/yellowstone-grpc-kafka/Cargo.toml
index c697a32b..958c9005 100644
--- a/yellowstone-grpc-kafka/Cargo.toml
+++ b/yellowstone-grpc-kafka/Cargo.toml
@@ -14,11 +14,13 @@ clap = { version = "4.3.0", features = ["cargo", "derive"] }
 const-hex = "1.6.2"
 futures = "0.3.24"
 hyper = { version = "0.14.27", features = ["server"] }
+json5 = "0.4.1"
 lazy_static = "1.4.0"
 prometheus = "0.13.2"
 rdkafka = { version = "0.33.2", features = ["ssl", "sasl"] }
 serde = { version = "1.0.145", features = ["derive"] }
 serde_json = "1.0.86"
+serde_yaml = "0.9.25"
 sha2 = "0.10.7"
 tokio = { version = "1.21.2", features = ["rt-multi-thread", "macros", "time", "fs", "signal"] }
 tokio-stream = "0.1.11"
diff --git a/yellowstone-grpc-kafka/src/config.rs b/yellowstone-grpc-kafka/src/config.rs
index bbd96247..9a8f2814 100644
--- a/yellowstone-grpc-kafka/src/config.rs
+++ b/yellowstone-grpc-kafka/src/config.rs
@@ -35,12 +35,18 @@ pub struct Config {
 }
 
 impl Config {
-    pub async fn load(path: impl AsRef<Path>) -> anyhow::Result<Self> {
+    pub async fn load(path: impl AsRef<Path> + Copy) -> anyhow::Result<Self> {
         let text = fs::read_to_string(path)
             .await
             .context("failed to read config from file")?;
 
-        serde_json::from_str(&text).context("failed to parse config from file")
+        match path.as_ref().extension().and_then(|e| e.to_str()) {
+            Some("yaml") | Some("yml") => {
+                serde_yaml::from_str(&text).context("failed to parse config from file")
+            }
+            Some("json") => serde_yaml::from_str(&text).context("failed to parse config from file"),
+            value => anyhow::bail!("unknown config extension: {value:?}"),
+        }
     }
 }