feature: add tracing to remote query (#9128)

packages/medusa/src/instrumentation/index.ts

@@ -2,7 +2,7 @@ import { snakeCase } from "lodash"
 import { NodeSDK } from "@opentelemetry/sdk-node"
 import { Resource } from "@opentelemetry/resources"
 import { SpanStatusCode } from "@opentelemetry/api"
-import { RoutesLoader, Tracer } from "@medusajs/framework"
+import { RoutesLoader, Tracer, Query } from "@medusajs/framework"
 import {
   type SpanExporter,
   SimpleSpanProcessor,
@@ -21,7 +21,7 @@ function shouldExcludeResource(resource: string) {
 }
 
 /**
- * Instrumenting the first touch point of the Http layer to report traces to
+ * Instrumenting the first touch point of the HTTP layer to report traces to
  * OpenTelemetry
  */
 export function instrumentHttpLayer() {
@@ -43,6 +43,12 @@ export function instrumentHttpLayer() {
       try {
         await requestHandler()
       } finally {
+        if (res.statusCode >= 500) {
+          span.setStatus({
+            code: SpanStatusCode.ERROR,
+            message: `Failed with ${res.statusMessage}`,
+          })
+        }
         span.setAttributes({ "http.statusCode": res.statusCode })
         span.end()
       }
@@ -71,6 +77,7 @@ export function instrumentHttpLayer() {
             code: SpanStatusCode.ERROR,
             message: error.message || "Failed",
           })
+          throw error
         } finally {
           span.end()
         }
@@ -117,6 +124,85 @@ export function instrumentHttpLayer() {
   })
 }
 
+/**
+ * Instrumenting the queries made using the remote query
+ */
+export function instrumentRemoteQuery() {
+  const QueryTracer = new Tracer("@medusajs/query", "2.0.0")
+
+  Query.instrument.graphQuery(async function (queryFn, queryOptions) {
+    return await QueryTracer.trace(
+      `query.graph: ${queryOptions.entryPoint}`,
+      async (span) => {
+        span.setAttributes({
+          "query.fields": queryOptions.fields,
+        })
+        return await queryFn()
+          .catch((error) => {
+            span.setStatus({
+              code: SpanStatusCode.ERROR,
+              message: error.message,
+            })
+          })
+          .finally(() => span.end())
+      }
+    )
+  })
+
+  Query.instrument.remoteQuery(async function (queryFn, queryOptions) {
+    const traceIdentifier =
+      "entryPoint" in queryOptions
+        ? queryOptions.entryPoint
+        : "service" in queryOptions
+        ? queryOptions.service
+        : "__value" in queryOptions
+        ? Object.keys(queryOptions.__value)[0]
+        : "unknown source"
+
+    return await QueryTracer.trace(
+      `remoteQuery: ${traceIdentifier}`,
+      async (span) => {
+        span.setAttributes({
+          "query.fields": "fields" in queryOptions ? queryOptions.fields : [],
+        })
+        return await queryFn()
+          .catch((error) => {
+            span.setStatus({
+              code: SpanStatusCode.ERROR,
+              message: error.message,
+            })
+          })
+          .finally(() => span.end())
+      }
+    )
+  })
+
+  Query.instrument.remoteDataFetch(async function (
+    fetchFn,
+    serviceName,
+    method,
+    options
+  ) {
+    return await QueryTracer.trace(
+      `${snakeCase(serviceName)}.${snakeCase(method)}`,
+      async (span) => {
+        span.setAttributes({
+          "fetch.select": options.select,
+          "fetch.relations": options.relations,
+        })
+        return await fetchFn()
+          .catch((error) => {
+            span.setStatus({
+              code: SpanStatusCode.ERROR,
+              message: error.message,
+            })
+          })
+          .finally(() => span.end())
+      }
+    )
+  })
+}
+
 /**
  * A helper function to configure the OpenTelemetry SDK with some defaults.
  * For better/more control, please configure the SDK manually.