JSON input, envvars, themes

- Allow passing dashboard as JSON - List available themes - Pass the dashboard URL via an envvar - Improve print output
cpg314 · Jun 11, 2024 · df47d05 · df47d05
1 parent 3e2507f
commit df47d05
Show file tree

Hide file tree

Showing 3 changed files with 94 additions and 36 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,19 +1,21 @@
 [package]
 name = "ch-grafana-cache"
-version = "0.1.1"
+version = "0.1.2"
 edition = "2021"
 
 [dependencies]
 anyhow = "1.0.86"
 bat = "0.24.0"
 clap = { version = "4.5.4", features = ["derive", "env"] }
+colored = "2.1.0"
 itertools = "0.13.0"
 lazy_static = "1.4.0"
 regex = "1.10.4"
 reqwest = { version = "0.12.4", features = ["rustls-tls", "json", "gzip"], default_features = false }
 reqwest-middleware = "0.3.1"
 reqwest-retry = "0.5.0"
 serde = { version = "1.0.203", features = ["derive"] }
+serde_json = "1.0.117"
 thiserror = "1.0.61"
 tokio = { version = "1.38.0", features = ["full"] }
 tracing = "0.1.40"

diff --git a/README.md b/README.md
@@ -9,42 +9,57 @@ Variables are supported, even those depending on others. The tool runs over all
 ## Usage
 
 ```console
-$ ch-grafana-cache --help
-Execute Clickhouse SQL queries from a Grafana dashboard
+Execute Clickhouse SQL queries from a Grafana dashboard.
 
-Usage: ch-grafana-cache [OPTIONS] --grafana <GRAFANA> --dashboard <DASHBOARD> <COMMAND>
+Call with either --grafana-url and --dashboard, or with --json
+
+Usage: ch-grafana-cache [OPTIONS] <COMMAND>
 
 Commands:
   print    Print SQL statements, with syntax highlighting
   execute  Execute the queries
   help     Print this message or the help of the given subcommand(s)
 
 Options:
-      --grafana <GRAFANA>      Base Grafana URL
-      --dashboard <DASHBOARD>  Grafana dashboard id
-      --theme <THEME>          Synctect for syntax highlighting [env: CH_GRAFANA_CACHE_THEME=]
-  -h, --help                   Print helptext
+      --grafana-url <GRAFANA_URL>
+          Base Grafana URL
 
-$ ch-grafana-cache execute --help
-Execute the queries
+          [env: GRAFANA_URL=https://grafana.corp.com/]
 
-Usage: ch-grafana-cache --grafana <GRAFANA> --dashboard <DASHBOARD> execute --clickhouse <CLICKHOUSE> --username <USERNAME>
+      --dashboard <DASHBOARD>
+          Grafana dashboard id
 
-Options:
-      --clickhouse <CLICKHOUSE>  URL to the Clickhouse HTTP endpoint
-      --username <USERNAME>      Clickhouse username
-  -h, --help                     Print help
+      --json <JSON>
+          Dashboard JSON file
+
+      --theme <THEME>
+          Synctect for syntax highlighting. Pass any invalid value to see the list of available themes
+
+          [env: CH_GRAFANA_CACHE_THEME=Nord]
+
+  -h, --help
+          Print help (see a summary with '-h')
 ```
 
-Sample output:
+Examples
+
+```console
+$ # Printing the SQL queries in the dashboard
+$ ch-grafana-cache --grafana https://grafana.corp.com --dashboard mydashboard print
+Variables:
+
+...
 
-```text
+Panels:
+...
+
+$ # Executing the SQL queries in the dashboard across all combinations
+$ ch-grafana-cache --grafana https://grafana.corp.com --dashboard mydashboard execute
 INFO ch_grafana_cache: Retrieving dashboard
 INFO ch_grafana_cache: Retrieved dashboard 'mydashboard'
 INFO ch_grafana_cache: 166 variables combinations found. Executing queries...
 INFO ch_grafana_cache: Executing combination i=0 n_combinations=166
 INFO ch_grafana_cache: Executed combination duration=178.932498ms size_mb=0.107275
-...
 ```
 
 ## Verifying that `chproxy` caching works

diff --git a/src/main.rs b/src/main.rs
@@ -2,24 +2,67 @@ mod clickhouse;
 mod grafana;
 mod variables;
 
+use std::collections::HashSet;
+use std::path::PathBuf;
+
 use clap::Parser;
+use colored::Colorize;
 use tracing::*;
 
 /// Execute Clickhouse SQL queries from a Grafana dashboard.
+///
+/// Call with either --grafana-url and --dashboard, or with --json
 #[derive(clap::Parser)]
 struct Flags {
     /// Base Grafana URL
-    #[clap(long)]
-    grafana: reqwest::Url,
+    #[clap(long, env = "GRAFANA_URL")]
+    grafana_url: Option<reqwest::Url>,
     /// Grafana dashboard id
-    #[clap(long)]
-    dashboard: String,
-    /// Synctect for syntax highlighting
+    #[clap(long, requires = "grafana")]
+    dashboard: Option<String>,
+    /// Dashboard JSON file.
+    #[clap(long, conflicts_with = "dashboard")]
+    json: Option<PathBuf>,
+    /// Synctect for syntax highlighting. Pass any invalid value to see the list of available themes.
     #[clap(long, env = "CH_GRAFANA_CACHE_THEME")]
     theme: Option<String>,
     #[clap(subcommand)]
     command: Command,
 }
+#[derive(serde::Deserialize)]
+#[serde(untagged)]
+enum Json {
+    Resp(grafana::DashboardResponse),
+    Dashboard(grafana::Dashboard),
+}
+impl From<Json> for grafana::Dashboard {
+    fn from(js: Json) -> grafana::Dashboard {
+        match js {
+            Json::Resp(r) => r.dashboard,
+            Json::Dashboard(d) => d,
+        }
+    }
+}
+impl Flags {
+    async fn get_dashboard(&self) -> anyhow::Result<grafana::Dashboard> {
+        let resp: Json = match (&self.json, &self.grafana_url, &self.dashboard) {
+            (Some(json), _, _) => serde_json::from_str(&std::fs::read_to_string(json)?)?,
+            (None, Some(grafana), Some(dashboard)) => {
+                info!("Retrieving dashboard from {}", grafana);
+                reqwest::get(grafana.join("api/dashboards/uid/")?.join(dashboard)?)
+                    .await?
+                    .json::<Json>()
+                    .await?
+            }
+            _ => {
+                anyhow::bail!("Use --json, or --grafana and --dashboard")
+            }
+        };
+        let dashboard = grafana::Dashboard::from(resp);
+        Ok(dashboard)
+    }
+}
+
 #[derive(clap::Parser)]
 enum Command {
     /// Print SQL statements, with syntax highlighting
@@ -43,6 +86,13 @@ fn print_sql(sql: &str, theme: Option<&String>) -> anyhow::Result<()> {
     let mut printer = bat::PrettyPrinter::new();
     printer.input_from_reader(sql).language("sql");
     if let Some(theme) = theme {
+        let themes: HashSet<_> = printer.themes().collect();
+        anyhow::ensure!(
+            themes.contains(theme.as_str()),
+            "Theme {} not found. Available themes: {:?}",
+            theme,
+            themes
+        );
         printer.theme(theme);
     }
     printer.print()?;
@@ -53,25 +103,16 @@ async fn main_impl() -> anyhow::Result<()> {
     let args = Flags::parse();
     let start = std::time::Instant::now();
 
-    info!("Retrieving dashboard");
-    let resp: grafana::DashboardResponse = reqwest::get(
-        args.grafana
-            .join("api/dashboards/uid/")?
-            .join(&args.dashboard)?,
-    )
-    .await?
-    .json()
-    .await?;
-    let dashboard = &resp.dashboard;
+    let dashboard = args.get_dashboard().await?;
     info!("Retrieved dashboard '{}'", dashboard.title);
-
+    println!();
     match args.command {
         Command::Print => {
-            info!("Variables");
+            println!("{}", "Variables:\n".yellow());
             for sql in dashboard.variables_sql() {
                 print_sql(sql, args.theme.as_ref())?;
             }
-            info!("Panels");
+            println!("{}", "Panels:\n".yellow());
             for sql in dashboard.panels_sql() {
                 print_sql(sql, args.theme.as_ref())?;
             }