Merge pull request #17 from hasura:stream-response

add trace spans, don't parse db response before forwarding it

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.2.6]
+
+- Add additional trace spans for SQL generation and query execution
+- Do not parse db response as JSON, instead send it back directly
+
 ## [0.2.5]
 
 - Implement validate cli command

Cargo.toml

@@ -6,5 +6,5 @@ members = [
 ]
 resolver = "2"
 
-package.version = "0.2.5"
+package.version = "0.2.6"
 package.edition = "2021"

crates/common/Cargo.toml

@@ -4,11 +4,13 @@ version.workspace = true
 edition.workspace = true
 
 [dependencies]
+bytes = "1.6.0"
+peg = "0.8.2"
 reqwest = { version = "0.12.3", features = [
     "json",
     "rustls-tls",
 ], default-features = false }
+schemars = "0.8.16"
 serde = { version = "1.0.197", features = ["derive"] }
 serde_json = "1.0.114"
-peg = "0.8.2"
-schemars = "0.8.16"
+tracing = "0.1.40"

crates/common/src/client.rs

@@ -1,6 +1,8 @@
 use std::error::Error;
 
+use bytes::Bytes;
 use serde::{de::DeserializeOwned, Deserialize};
+use tracing::Instrument;
 
 use crate::config::ConnectionConfig;
 
@@ -12,7 +14,35 @@ pub fn get_http_client(
     Ok(client)
 }
 
-pub async fn execute_query<T: DeserializeOwned>(
+pub async fn execute_query(
+    client: &reqwest::Client,
+    connection_config: &ConnectionConfig,
+    statement: &str,
+    parameters: &Vec<(String, String)>,
+) -> Result<Bytes, Box<dyn Error>> {
+    let response = client
+        .post(&connection_config.url)
+        .header("X-ClickHouse-User", &connection_config.username)
+        .header("X-ClickHouse-Key", &connection_config.password)
+        .query(parameters)
+        .body(statement.to_owned())
+        .send()
+        .instrument(tracing::info_span!("Execute HTTP request"))
+        .await?;
+
+    if response.error_for_status_ref().is_err() {
+        return Err(response.text().await?.into());
+    }
+
+    let response = response
+        .bytes()
+        .instrument(tracing::info_span!("Read HTTP response"))
+        .await?;
+
+    Ok(response)
+}
+
+pub async fn execute_json_query<T: DeserializeOwned>(
     client: &reqwest::Client,
     connection_config: &ConnectionConfig,
     statement: &str,
@@ -25,13 +55,17 @@ pub async fn execute_query<T: DeserializeOwned>(
         .query(parameters)
         .body(statement.to_owned())
         .send()
+        .instrument(tracing::info_span!("Execute HTTP request"))
         .await?;
 
     if response.error_for_status_ref().is_err() {
         return Err(response.text().await?.into());
     }
 
-    let payload: ClickHouseResponse<T> = response.json().await?;
+    let payload: ClickHouseResponse<T> = response
+        .json()
+        .instrument(tracing::info_span!("Parse HTTP response"))
+        .await?;
 
     Ok(payload.data)
 }

crates/ndc-clickhouse-cli/src/database_introspection.rs

@@ -3,7 +3,7 @@ use std::error::Error;
 use serde::Deserialize;
 
 use common::{
-    client::{execute_query, get_http_client},
+    client::{execute_json_query, get_http_client},
     config::ConnectionConfig,
 };
 
@@ -44,5 +44,5 @@ pub async fn introspect_database(
 ) -> Result<Vec<TableInfo>, Box<dyn Error>> {
     let introspection_sql = include_str!("./database_introspection.sql");
     let client = get_http_client(connection_config)?;
-    execute_query::<TableInfo>(&client, connection_config, introspection_sql, &vec![]).await
+    execute_json_query::<TableInfo>(&client, connection_config, introspection_sql, &vec![]).await
 }

crates/ndc-clickhouse/Cargo.toml

@@ -5,9 +5,10 @@ edition.workspace = true
 
 [dependencies]
 async-trait = "0.1.78"
+bytes = "1.6.0"
 common = { path = "../common" }
 indexmap = "2.1.0"
-ndc-sdk = { git = "https://github.com/hasura/ndc-sdk-rs", rev = "7b56fac", package = "ndc-sdk" }
+ndc-sdk = { git = "https://github.com/hasura/ndc-sdk-rs", rev = "972dba6", package = "ndc-sdk" }
 prometheus = "0.13.3"
 reqwest = { version = "0.12.3", features = [
   "json",
@@ -18,3 +19,4 @@ serde_json = "1.0.114"
 sqlformat = "0.2.3"
 strum = { version = "0.26.2", features = ["derive"] }
 tokio = "1.36.0"
+tracing = "0.1.40"