Instrument query with spans

crates/ndc-sqlserver/Cargo.toml

@@ -17,8 +17,9 @@ path = "bin/main.rs"
 ndc-sdk = { git = "https://github.com/hasura/ndc-hub.git", rev = "f2a2a75", package = "ndc-sdk" }
 
 query-engine-execution = { path = "../query-engine/execution" }
-query-engine-translation = { path = "../query-engine/translation" }
 query-engine-metadata = {path = "../query-engine/metadata"}
+query-engine-sql = { path = "../query-engine/sql" }
+query-engine-translation = { path = "../query-engine/translation" }
 
 tiberius = { version = "0.12.2", default-features = false, features = ["rustls"] }
 bb8 = "0.8.1"

crates/ndc-sqlserver/src/query.rs

@@ -2,13 +2,14 @@
 //! See the Hasura
 //! [Native Data Connector Specification](https://hasura.github.io/ndc-spec/specification/queries/index.html)
 //! for further details.
-
 use super::configuration;
 use ndc_sdk::connector;
 use ndc_sdk::json_response::JsonResponse;
 use ndc_sdk::models;
 use query_engine_execution::execution;
+use query_engine_sql::sql;
 use query_engine_translation::translation;
+use tracing::{info_span, Instrument};
 
 /// Execute a query
 ///
@@ -22,23 +23,57 @@ pub async fn query(
     tracing::info!("{}", serde_json::to_string(&query_request).unwrap());
     tracing::info!("{:?}", query_request);
 
-    // Compile the query.
-    let plan = match translation::query::translate(&configuration.config.metadata, query_request) {
-        Ok(plan) => Ok(plan),
-        Err(err) => {
+    let timer = state.metrics.time_query_total();
+    let result = async move {
+        // Plan the query
+        let plan = async { plan_query(configuration, state, query_request) }
+            .instrument(info_span!("Plan query"))
+            .await?;
+
+        // Execute the query.
+        let result = execute_query(state, plan)
+            .instrument(info_span!("Execute query"))
+            .await?;
+
+        // assuming query succeeded, increment counter
+        state.metrics.record_successful_query();
+
+        // TODO: return raw JSON
+        Ok(result)
+    }
+    .instrument(info_span!("Execute query"))
+    .await;
+
+    timer.complete_with(result)
+}
+
+// Compile the query.
+fn plan_query(
+    configuration: &configuration::Configuration,
+    state: &configuration::State,
+    query_request: models::QueryRequest,
+) -> Result<sql::execution_plan::ExecutionPlan, connector::QueryError> {
+    let timer = state.metrics.time_query_plan();
+    let result = translation::query::translate(&configuration.config.metadata, query_request)
+        .map_err(|err| {
             tracing::error!("{}", err);
             match err {
                 translation::query::error::Error::NotSupported(_) => {
-                    Err(connector::QueryError::UnsupportedOperation(err.to_string()))
+                    connector::QueryError::UnsupportedOperation(err.to_string())
                 }
-                _ => Err(connector::QueryError::InvalidRequest(err.to_string())),
+                _ => connector::QueryError::InvalidRequest(err.to_string()),
             }
-        }
-    }?;
+        });
+    timer.complete_with(result)
+}
 
-    // Execute the query.
-    let result = execution::mssql_execute(&state.mssql_pool, &state.metrics, plan)
+async fn execute_query(
+    state: &configuration::State,
+    plan: sql::execution_plan::ExecutionPlan,
+) -> Result<JsonResponse<models::QueryResponse>, connector::QueryError> {
+    execution::mssql_execute(&state.mssql_pool, &state.metrics, plan)
         .await
+        .map(JsonResponse::Value)
         .map_err(|err| match err {
             execution::Error::Query(err) => {
                 tracing::error!("{}", err);
@@ -48,11 +83,5 @@ pub async fn query(
                 tracing::error!("{}", err);
                 connector::QueryError::Other(err.into())
             }
-        })?;
-
-    // assuming query succeeded, increment counter
-    state.metrics.record_successful_query();
-
-    // TODO: return raw JSON
-    Ok(JsonResponse::Value(result))
+        })
 }

crates/query-engine/execution/src/execution.rs

@@ -7,6 +7,7 @@ use serde_json;
 use std::collections::BTreeMap;
 use tiberius::QueryItem;
 use tokio_stream::StreamExt;
+use tracing::{info_span, Instrument};
 
 /// Execute a query against sqlserver.
 pub async fn mssql_execute(
@@ -24,11 +25,17 @@ pub async fn mssql_execute(
     );
 
     let acquisition_timer = metrics.time_connection_acquisition_wait();
-    let connection_result = mssql_pool.get().await.map_err(Error::ConnectionPool);
+    let connection_result = mssql_pool
+        .get()
+        .instrument(info_span!("Acquire connection"))
+        .await
+        .map_err(Error::ConnectionPool);
     let mut connection = acquisition_timer.complete_with(connection_result)?;
 
     let query_timer = metrics.time_query_execution();
-    let rows_result = execute_query(&mut connection, plan).await;
+    let rows_result = execute_query(&mut connection, plan)
+        .instrument(info_span!("Database request"))
+        .await;
     let rows = query_timer.complete_with(rows_result)?;
 
     tracing::info!("Database rows result: {:?}", rows);