Send queries to proxy from the CLI

Author: cpg314

Cargo.toml

@@ -1,5 +1,5 @@
 [workspace.package]
-version = "0.1.2"
+version = "0.1.3"
 edition = "2021"
 
 [workspace.dependencies]
@@ -13,9 +13,11 @@ tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 glob = "0.3.1"
 serde_json = "1.0.117"
 lazy_static = "1.5.0"
+tokio = { version = "1.38.0", features = ["full"] }
+reqwest = { version = "0.12.4", features = ["rustls-tls", "json", "gzip", "stream"], default-features = false }
 
 [workspace]
-members = ["sqlsonnet", "clickhouse-proxy", "sqlsonnet-cli", "bindings/wasm"]
+members = ["sqlsonnet", "clickhouse-proxy", "sqlsonnet-cli", "clickhouse-client", "bindings/wasm"]
 resolver = "2"
 
 [profile.release]

Dockerfile

@@ -3,10 +3,11 @@ FROM debian:bookworm-slim
 LABEL org.opencontainers.image.source=https://github.com/cpg314/sqlsonnet
 LABEL org.opencontainers.image.licenses=MIT
 
-RUN apt update && apt install -y curl && apt clean
-
-RUN curl -L https://github.com/google/go-jsonnet/releases/download/v0.20.0/jsonnetfmt-go_0.20.0_linux_amd64.deb --output jsonnetfmt.deb && \
-    dpkg -i jsonnetfmt.deb
+RUN apt update && \
+    apt install -y curl && \
+    curl -L https://github.com/google/go-jsonnet/releases/download/v0.20.0/jsonnetfmt-go_0.20.0_linux_amd64.deb --output jsonnetfmt.deb && \
+    dpkg -i jsonnetfmt.deb && \
+    apt clean
 
 COPY target-cross/x86_64-unknown-linux-gnu/release/sqlsonnet /usr/bin/sqlsonnet
 COPY target-cross/x86_64-unknown-linux-gnu/release/sqlsonnet_clickhouse_proxy /usr/bin/sqlsonnet_clickhouse_proxy

README.md

@@ -76,6 +76,10 @@ Options:
           With --from-sql: Convert back to SQL and print the differences with the original, if any
       --display-format <DISPLAY_FORMAT>
           [possible values: sql, jsonnet, json]
+      --proxy-url <PROXY_URL>
+          sqlsonnet proxy URL [env: SQLSONNET_PROXY=]
+  -e, --execute
+          Send query to Clickhouse proxy (--proxy-url) for execution
   -h, --help
           Print help
   -V, --version

clickhouse-client/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "clickhouse-client"
+version.workspace = true
+edition = "2021"
+
+
+[dependencies]
+reqwest.workspace = true
+serde.workspace = true

clickhouse-client/src/lib.rs

@@ -0,0 +1,38 @@
+use std::collections::BTreeMap;
+
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone)]
+pub struct HttpClient {
+    client: reqwest::Client,
+    url: reqwest::Url,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize, Hash)]
+pub struct ClickhouseQuery {
+    pub query: String,
+    pub params: BTreeMap<String, String>,
+}
+
+impl HttpClient {
+    pub fn new(url: reqwest::Url) -> Self {
+        Self {
+            url,
+            client: reqwest::Client::new(),
+        }
+    }
+    pub fn prepare_request(&self, query: &ClickhouseQuery) -> reqwest::RequestBuilder {
+        self.client
+            .post(self.url.clone())
+            .body(query.query.clone())
+            .query(&query.params)
+            .header(reqwest::header::TRANSFER_ENCODING, "chunked")
+    }
+    pub async fn send_query(
+        &self,
+        query: &ClickhouseQuery,
+    ) -> Result<reqwest::Response, reqwest::Error> {
+        let request = self.prepare_request(query);
+        request.send().await?.error_for_status()
+    }
+}

clickhouse-proxy/Cargo.toml

@@ -8,6 +8,8 @@ name = "sqlsonnet_clickhouse_proxy"
 path = "src/clickhouse_proxy.rs"
 
 [dependencies]
+clickhouse-client = { path = "../clickhouse-client" }
+
 anyhow.workspace = true
 axum = { version = "0.7.5", features = ["ws"] }
 bincode = "1.3.3"
@@ -16,11 +18,11 @@ chrono = { version = "0.4.38", features = ["serde"] }
 clap.workspace = true
 futures = "0.3.30"
 itertools.workspace = true
-reqwest = { version = "0.12.4", features = ["rustls-tls", "json", "gzip", "stream"], default-features = false }
+reqwest.workspace = true
 serde.workspace = true
 sqlsonnet = { path = "../sqlsonnet" }
 thiserror.workspace = true
-tokio = { version = "1.38.0", features = ["full"] }
+tokio.workspace = true
 tokio-util = { version = "0.7.11", features = ["io", "io-util"] }
 tower-http = { version = "0.5.2", features = ["trace"] }
 tracing.workspace = true

clickhouse-proxy/src/lib.rs

@@ -14,9 +14,9 @@ use axum::response::IntoResponse;
 use clap::Parser;
 use itertools::Itertools;
 use metrics_exporter_prometheus::PrometheusBuilder;
-use serde::{Deserialize, Serialize};
 use tracing::*;
 
+use clickhouse_client::ClickhouseQuery;
 use sqlsonnet::{FsResolver, Queries, Query};
 
 lazy_static::lazy_static! {
@@ -30,14 +30,9 @@ lazy_static::lazy_static! {
 #[derive(Clone, Parser)]
 #[clap(version)]
 pub struct Flags {
-    // URL to the Clickhouse HTTP endpoint
+    // URL to the Clickhouse HTTP endpoint, with username and password if necessary
     #[clap(long, env = "CLICKHOUSE_URL")]
     pub url: reqwest::Url,
-    /// Clickhouse username
-    #[clap(long, env = "CLICKHOUSE_USERNAME")]
-    pub username: String,
-    #[clap(long, env = "CLICKHOUSE_PASSWORD")]
-    pub password: Option<String>,
     #[clap(long)]
     pub cache: Option<PathBuf>,
     /// Folder with Jsonnet library files
@@ -102,19 +97,6 @@ async fn handle_query(
         .await
 }
 
-#[derive(Clone, Debug, Serialize, Deserialize, Hash)]
-pub struct ClickhouseQuery {
-    query: String,
-    params: BTreeMap<String, String>,
-}
-impl ClickhouseQuery {
-    fn heartbeat() -> Self {
-        Self {
-            query: "SELECT 1+1".into(),
-            params: Default::default(),
-        }
-    }
-}
 pub struct PreparedRequest {
     id: u64,
     query: Option<ClickhouseQuery>,
@@ -138,7 +120,7 @@ impl PreparedRequest {
 
 #[derive(Clone)]
 struct State {
-    client: reqwest::Client,
+    client: clickhouse_client::HttpClient,
     args: Arc<Flags>,
     resolver: Arc<FsResolver>,
     cache: Option<Arc<cache::Cache>>,
@@ -147,7 +129,7 @@ struct State {
 impl State {
     fn new(args: &Flags) -> Result<Self, Error> {
         Ok(Self {
-            client: reqwest::Client::new(),
+            client: clickhouse_client::HttpClient::new(args.url.clone()),
             resolver: sqlsonnet::FsResolver::new(
                 args.library.clone().map(|p| vec![p]).unwrap_or_default(),
             )
@@ -177,24 +159,15 @@ impl State {
         // Hash query
         let mut hasher = DefaultHasher::default();
         query.hash(&mut hasher);
-        self.args.username.hash(&mut hasher);
-        self.args.password.hash(&mut hasher);
-        let inner = self
-            .client
-            .post(self.args.url.clone())
-            .body(query.query.clone())
-            .query(&query.params)
-            .header(reqwest::header::TRANSFER_ENCODING, "chunked")
-            .basic_auth(&self.args.username, self.args.password.clone());
+        let builder = self.client.prepare_request(&query);
         PreparedRequest {
             id: hasher.finish(),
             query: Some(query),
-            builder: inner,
+            builder,
         }
     }
     async fn send_query(&self, query: ClickhouseQuery) -> Result<axum::response::Response, Error> {
         let request = self.prepare_request(query);
-
         if let Some(cache) = &self.cache {
             Ok(cache.process(request).await?)
         } else {
@@ -208,7 +181,10 @@ impl State {
     }
     async fn test_clickhouse(&self) -> Result<(), ClickhouseError> {
         let resp = self
-            .prepare_request(ClickhouseQuery::heartbeat())
+            .prepare_request(ClickhouseQuery {
+                query: "SELECT 1+1".into(),
+                params: Default::default(),
+            })
             .send()
             .await?;
         let headers = resp.headers().clone();

clickhouse-proxy/tests/main.rs

@@ -27,9 +27,7 @@ async fn main() -> anyhow::Result<()> {
     drop(listener);
     let cache = tempfile::tempdir()?;
     let _server = tokio::spawn(clickhouse_proxy::main_impl(clickhouse_proxy::Flags {
-        url: reqwest::Url::parse(&format!("http://{}", fake_chaddr))?,
-        username: "default".into(),
-        password: None,
+        url: reqwest::Url::parse(&format!("http://default@{}", fake_chaddr))?,
         cache: Some(cache.path().into()),
         library: None,
         prelude: None,

sqlsonnet-cli/Cargo.toml

@@ -9,6 +9,7 @@ path = "src/sqlsonnet.rs"
 
 [dependencies]
 sqlsonnet = { path = "../sqlsonnet", features = ["fancy"] }
+clickhouse-client = { path = "../clickhouse-client" }
 
 bat = { version = "0.24.0", default-features = false }
 clap-stdin = "0.4.0"
@@ -17,7 +18,9 @@ miette = { version = "7.2.0", features = ["fancy", "syntect-highlighter"] }
 lazy_static.workspace = true
 pretty_assertions = "1.4.0"
 thiserror.workspace = true
+tokio.workspace = true
 tracing.workspace = true
+reqwest.workspace = true
 
 [dev-dependencies]
 anyhow.workspace = true

sqlsonnet-cli/src/sqlsonnet.rs

@@ -1,3 +1,4 @@
+use std::collections::BTreeMap;
 use std::io::IsTerminal;
 use std::str::FromStr;
 
@@ -22,6 +23,8 @@ enum Error {
     Bat(#[from] bat::error::Error),
     #[error(transparent)]
     Miette(#[from] miette::InstallError),
+    #[error("Failed to execute query")]
+    Clickhouse(#[from] reqwest::Error),
 }
 
 #[derive(Parser)]
@@ -44,6 +47,12 @@ struct Flags {
     diff: bool,
     #[clap(long, value_delimiter = ',')]
     display_format: Option<Vec<Language>>,
+    /// sqlsonnet proxy URL
+    #[clap(long, env = "SQLSONNET_PROXY")]
+    proxy_url: Option<reqwest::Url>,
+    /// Send query to Clickhouse proxy (--proxy-url) for execution
+    #[clap(long, short, conflicts_with = "from_sql", requires = "proxy_url")]
+    execute: bool,
 }
 
 #[derive(Clone)]
@@ -115,14 +124,18 @@ fn highlight<T: std::fmt::Display>(
     Ok(())
 }
 
-fn main() -> miette::Result<()> {
-    Ok(main_impl()?)
+#[tokio::main]
+async fn main() -> miette::Result<()> {
+    Ok(main_impl().await?)
 }
 
-fn main_impl() -> Result<(), Error> {
+async fn main_impl() -> Result<(), Error> {
     let start = std::time::Instant::now();
     sqlsonnet::setup_logging();
-    let args = Flags::parse();
+    let mut args = Flags::parse();
+    if !args.execute {
+        args.proxy_url = None;
+    }
 
     let assets = bat::assets::HighlightingAssets::from_binary();
     let theme = assets.get_theme(
@@ -151,6 +164,12 @@ fn main_impl() -> Result<(), Error> {
             Language::Sql
         }]
     });
+
+    let client = args
+        .proxy_url
+        .clone()
+        .map(clickhouse_client::HttpClient::new);
+
     let filename = args.input.filename();
     let input = args.input.contents()?;
     if args.from_sql {
@@ -184,6 +203,23 @@ fn main_impl() -> Result<(), Error> {
         if has(Language::Sql) {
             highlight(queries.to_sql(args.compact), Language::Sql, &args)?;
         }
+        if let Some(client) = client {
+            info!("Executing query on Clickhouse");
+            for query in queries {
+                let resp = client
+                    .send_query(&clickhouse_client::ClickhouseQuery {
+                        query: query.to_sql(false),
+                        params: BTreeMap::from([(
+                            "default_format".into(),
+                            "PrettyMonoBlock".into(),
+                        )]),
+                    })
+                    .await?
+                    .text()
+                    .await?;
+                println!("{}", resp);
+            }
+        }
     }
 
     info!(elapsed=?start.elapsed(), "Done");