From 353f4fcfdbe6ba819dc9b1c2955b3ca5a7ea5b32 Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@rubenfiszel.com>
Date: Fri, 27 Oct 2023 11:51:26 +0200
Subject: [PATCH] add loki as optional feature flag

---
 backend/Cargo.lock                          | 70 +++++++++++++++++++++
 backend/Cargo.toml                          |  1 +
 backend/windmill-common/Cargo.toml          |  3 +-
 backend/windmill-common/src/tracing_init.rs |  9 +++
 4 files changed, 82 insertions(+), 1 deletion(-)

diff --git a/backend/Cargo.lock b/backend/Cargo.lock
index f04a13e56b..7d47abe90f 100644
--- a/backend/Cargo.lock
+++ b/backend/Cargo.lock
@@ -3196,6 +3196,16 @@ dependencies = [
  "value-bag",
 ]
 
+[[package]]
+name = "loki-api"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f56d36f573486ba7f462b62cbae597fef7d5d93665e7047956b457531b8a1ced"
+dependencies = [
+ "prost",
+ "prost-types",
+]
+
 [[package]]
 name = "lru"
 version = "0.10.1"
@@ -4421,6 +4431,38 @@ dependencies = [
  "thiserror",
 ]
 
+[[package]]
+name = "prost"
+version = "0.11.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b82eaa1d779e9a4bc1c3217db8ffbeabaae1dca241bf70183242128d48681cd"
+dependencies = [
+ "bytes",
+ "prost-derive",
+]
+
+[[package]]
+name = "prost-derive"
+version = "0.11.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e5d2d8d10f3c6ded6da8b05b5fb3b8a5082514344d56c9f871412d29b4e075b4"
+dependencies = [
+ "anyhow",
+ "itertools 0.10.5",
+ "proc-macro2",
+ "quote",
+ "syn 1.0.109",
+]
+
+[[package]]
+name = "prost-types"
+version = "0.11.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "213622a1460818959ac1181aaeb2dc9c7f63df720db7d788b3e24eacd1983e13"
+dependencies = [
+ "prost",
+]
+
 [[package]]
 name = "psm"
 version = "0.1.21"
@@ -5670,6 +5712,12 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d4b756ac662e92a0e5b360349bea5f0b0784d4be4541eff2972049dfdfd7f862"
 
+[[package]]
+name = "snap"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5e9f0ab6ef7eb7353d9119c170a436d1bf248eea575ac42d19d12f4e34130831"
+
 [[package]]
 name = "socket2"
 version = "0.4.10"
@@ -7031,6 +7079,27 @@ dependencies = [
  "tracing-core",
 ]
 
+[[package]]
+name = "tracing-loki"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "49bbc87d08020d7c2a9f4bb0b7d10da5381d3867f8ae57fcc54621b34567e963"
+dependencies = [
+ "loki-api",
+ "reqwest",
+ "serde",
+ "serde_json",
+ "snap",
+ "tokio",
+ "tokio-stream",
+ "tracing",
+ "tracing-core",
+ "tracing-log",
+ "tracing-serde",
+ "tracing-subscriber",
+ "url",
+]
+
 [[package]]
 name = "tracing-serde"
 version = "0.1.3"
@@ -7833,6 +7902,7 @@ dependencies = [
  "tokio",
  "tracing",
  "tracing-flame",
+ "tracing-loki",
  "tracing-subscriber",
  "uuid 1.5.0",
 ]
diff --git a/backend/Cargo.toml b/backend/Cargo.toml
index 9d55165dcb..a349dbc931 100644
--- a/backend/Cargo.toml
+++ b/backend/Cargo.toml
@@ -37,6 +37,7 @@ incremental = true
 enterprise = ["windmill-worker/enterprise", "windmill-queue/enterprise", "windmill-api/enterprise"]
 benchmark = ["windmill-api/benchmark", "windmill-worker/benchmark", "windmill-queue/benchmark"]
 flamegraph = ["windmill-common/flamegraph", "windmill-worker/flamegraph"]
+loki = ["windmill-common/loki"]
 
 [dependencies]
 anyhow.workspace = true
diff --git a/backend/windmill-common/Cargo.toml b/backend/windmill-common/Cargo.toml
index a90a353171..5bff350550 100644
--- a/backend/windmill-common/Cargo.toml
+++ b/backend/windmill-common/Cargo.toml
@@ -17,7 +17,7 @@ tracing_init = [
     "dep:tracing-subscriber",
 ]
 flamegraph = ["dep:tracing-flame"]
-
+loki = ["dep:tracing-loki"]
 
 [lib]
 name = "windmill_common"
@@ -48,3 +48,4 @@ itertools.workspace = true
 regex.workspace = true
 git-version.workspace = true
 cron.workspace = true
+tracing-loki = { version = "^0", optional = true }
\ No newline at end of file
diff --git a/backend/windmill-common/src/tracing_init.rs b/backend/windmill-common/src/tracing_init.rs
index 67ad871e19..0f31072ba5 100644
--- a/backend/windmill-common/src/tracing_init.rs
+++ b/backend/windmill-common/src/tracing_init.rs
@@ -41,6 +41,15 @@ pub fn initialize_tracing() {
 
     let ts_base = tracing_subscriber::registry().with(env_filter);
 
+    #[cfg(feature = "loki")]
+    let ts_base = {
+        let (layer, task) = tracing_loki::builder()
+            .build_url(reqwest::Url::parse("http://127.0.0.1:3100").unwrap())
+            .expect("build loki url");
+        tokio::spawn(task);
+        ts_base.with(layer)
+    };
+
     match json_fmt {
         true => ts_base.with(json_layer().flatten_event(true)).init(),
         false => ts_base