Add axum-tracing-opentelemetry to add parent traces etc

hasura · Sep 27, 2023 · a6155bf · a6155bf
1 parent d55e05e
commit a6155bf
Show file tree

Hide file tree

Showing 3 changed files with 153 additions and 23 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/rust-connector-sdk/Cargo.toml b/rust-connector-sdk/Cargo.toml
@@ -15,6 +15,7 @@ path = "bin/main.rs"
 async-trait = "^0.1.68"
 axum = "^0.6.18"
 axum-macros = "^0.3.7"
+axum-tracing-opentelemetry = { version = "0.10.0", features = [] }
 clap = { version = "^4.3.9", features = ["derive", "env"] }
 ndc-client = { git = "http://github.com/hasura/ndc-spec.git", tag = "v0.1.0-rc.6" }
 ndc-test = { git = "http://github.com/hasura/ndc-spec.git", tag = "v0.1.0-rc.6" }

diff --git a/rust-connector-sdk/src/default_main.rs b/rust-connector-sdk/src/default_main.rs
@@ -5,7 +5,6 @@ use crate::{
     connector::{Connector, InvalidRange, SchemaError, UpdateConfigurationError},
     routes,
 };
-
 use async_trait::async_trait;
 use axum::{
     body::Body,
@@ -15,6 +14,7 @@ use axum::{
     routing::{get, post},
     Json, Router,
 };
+use axum_tracing_opentelemetry::{opentelemetry_tracing_layer, response_with_trace_layer};
 use tower_http::validate_request::ValidateRequestHeaderLayer;
 
 use clap::{Parser, Subcommand};
@@ -322,7 +322,11 @@ where
         .route("/query", post(post_query::<C>))
         .route("/explain", post(post_explain::<C>))
         .route("/mutation", post(post_mutation::<C>))
-        .with_state(state);
+        .with_state(state)
+        // include trace context as header in responses to handlers above this line
+        .layer(response_with_trace_layer())
+        // set up `TraceLayer` from tower-http
+        .layer(opentelemetry_tracing_layer());
 
     let expected_auth_header: Option<HeaderValue> =
         service_token_secret.and_then(|service_token_secret| {
@@ -331,10 +335,7 @@ where
         });
 
     router
-        .layer(
-            TraceLayer::new_for_http()
-                .make_span_with(DefaultMakeSpan::default().level(Level::INFO)),
-        )
+        .layer(TraceLayer::new_for_http())
         .layer(ValidateRequestHeaderLayer::custom(
             move |request: &mut Request<Body>| {
                 // Validate the request