feat(azure/cosmos): add Azure CosmosDB integration

packages/datadog-instrumentations/src/azure-cosmos.js

@@ -0,0 +1,7 @@
+'use strict'
+
+const { addHook, getHooks } = require('./helpers/instrument')
+
+for (const hook of getHooks('@azure/cosmos')) {
+  addHook(hook, exports => exports)
+}

packages/datadog-instrumentations/src/helpers/hooks.js

@@ -7,6 +7,7 @@ module.exports = {
   '@langchain/langgraph': { esmFirst: true, fn: () => require('../langgraph') },
   'apollo-server-core': () => require('../apollo-server-core'),
   '@aws-sdk/smithy-client': () => require('../aws-sdk'),
+  '@azure/cosmos': () => require('../azure-cosmos'),
   '@azure/event-hubs': () => require('../azure-event-hubs'),
   '@azure/functions': () => require('../azure-functions'),
   'durable-functions': () => require('../azure-durable-functions'),

packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/azure-cosmos.js

@@ -0,0 +1,48 @@
+module.exports = [{
+  module: {
+    name: '@azure/cosmos',
+    versionRange: '>=4.4.0',
+    filePath: 'dist/browser/plugins/Plugin.js'  // file containing the target method
+  },
+  functionQuery: {
+    functionName: 'executePlugins',
+    kind: 'Async'  // Async | Callback | Sync
+  },
+  channelName: 'executePlugins'
+},
+{
+  module: {
+    name: '@azure/cosmos',
+    versionRange: '>=4.4.0',
+    filePath: 'dist/commonjs/plugins/Plugin.js'  // file containing the target method
+  },
+  functionQuery: {
+    functionName: 'executePlugins',
+    kind: 'Async'  // Async | Callback | Sync
+  },
+  channelName: 'executePlugins'
+},
+{
+  module: {
+    name: '@azure/cosmos',
+    versionRange: '>=4.4.0',
+    filePath: 'dist/esm/plugins/Plugin.js'  // file containing the target method
+  },
+  functionQuery: {
+    functionName: 'executePlugins',
+    kind: 'Async'  // Async | Callback | Sync
+  },
+  channelName: 'executePlugins'
+},
+{
+  module: {
+    name: '@azure/cosmos',
+    versionRange: '>=4.4.0',
+    filePath: 'dist/react-native/plugins/Plugin.js'  // file containing the target method
+  },
+  functionQuery: {
+    functionName: 'executePlugins',
+    kind: 'Async'  // Async | Callback | Sync
+  },
+  channelName: 'executePlugins'
+}]

packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/index.js

@@ -2,6 +2,7 @@
 
 module.exports = [
   ...require('./ai'),
+  ...require('./azure-cosmos'),
   ...require('./bullmq'),
   ...require('./langchain'),
   ...require('./langgraph'),

packages/datadog-plugin-azure-cosmos/src/index.js

@@ -0,0 +1,158 @@
+'use strict'
+
+const { storage } = require('../../datadog-core')
+const log = require('../../dd-trace/src/log')
+
+const DatabasePlugin = require('../../dd-trace/src/plugins/database')
+
+class AzureCosmosPlugin extends DatabasePlugin {
+  static id = 'azure-cosmos'
+  // Channel prefix determines how the plugin subscribes to instrumentation events.
+  // Three patterns exist — set `static prefix` explicitly based on instrumentation type:
+  //
+  // Orchestrion:              static prefix = 'tracing:orchestrion:<npm-package>:<channelName>'
+  // Shimmer + tracingChannel: static prefix = 'tracing:apm:<name>:<operation>'
+  // Shimmer + manual channels: omit prefix — defaults to `apm:${id}:${operation}`
+  static prefix = 'tracing:orchestrion:@azure/cosmos:executePlugins'
+  static peerServicePrecursors = ['db.name']
+
+  operationName() {
+    return 'cosmosdb.query'
+  }
+
+  asyncEnd(ctx) {
+    var span = ctx.currentStore?.span;
+    if (span != null) {
+      const result = ctx.result;
+      if (result != null) {
+        if (result.code != null) {
+          span.setTag("http.status_code", result.code);
+        }
+        if (result.substatus != undefined) {
+          span.setTag("http.status_subcode", result.substatus);
+        }
+      }
+      span.finish();
+    }
+  }
+
+  error(ctx) {
+    var span = ctx.currentStore?.span;
+    this.addError(ctx.error, span);
+    if (span != null) {
+      const error = ctx.error;
+      if (error != null) {
+        if (error.code != null) {
+          span.setTag("http.status_code", error.code);
+        }
+        if (error.substatus != undefined) {
+          span.setTag("http.status_subcode", error.substatus);
+        }
+      }
+    }
+  }
+
+  bindStart(ctx) {
+    // executePlugins(diagnosticNode, requestContext, next, on) — `on` is PluginOn.request | PluginOn.operation
+    const pluginOn = ctx.arguments?.[3];
+
+
+    const requestContext = ctx.arguments?.[1];
+    const resource = this.getResource(requestContext);
+    const { dbName, containerName } = this.getDbInfo(requestContext);
+    const connectionMode = this.getConnectionMode(requestContext);
+    const { outHost, userAgent } = this.getHttpInfo(requestContext);
+    // getting really specific here but otherwise we get doubled up read spans
+    if (pluginOn === "request" && ((!resource.includes("read") && !resource.includes("query")) || (resource.includes("read") && requestContext.resourceType != "docs"))) {
+      ctx.currentStore = { ...(storage('legacy').getStore() || {}) };
+      return ctx.currentStore;
+    }
+
+    this.startSpan(this.operationName(), {
+      resource: resource,
+      type: 'cosmosdb',
+      kind: 'client',
+      meta: {
+        component: 'azure_cosmos',
+        'db.system': 'cosmosdb',
+        'db.name': dbName,
+        'cosmosdb.container': containerName,
+        'cosmosdb.connection.mode': connectionMode,
+        'http.useragent': userAgent,
+        'out.host': outHost,
+      }
+    }, ctx);
+
+    return ctx.currentStore;
+  }
+
+  getResource(requestContext) {
+    if (requestContext != null) {
+      const operationType = requestContext.operationType;
+      const resourceLink = requestContext.path;
+      return operationType + " " + resourceLink;
+    }
+    return null;
+  }
+
+
+  getDbInfo(requestContext) {
+    var dbName = null;
+    var containerName = null;
+    if (requestContext != null) {
+      if (requestContext.operationType === "create" && requestContext.resourceType === "dbs" && requestContext.body != null) {
+        if (requestContext.body.id != null) {
+          dbName = requestContext.body.id;
+        }
+      }
+
+      var resourceLink = requestContext.path;
+      if (resourceLink != null) {
+        if (resourceLink.startsWith("/") && resourceLink.length > 1) {
+          resourceLink = resourceLink.slice(1);
+        }
+        const parts = resourceLink.split("/");
+        if (parts.length > 0 && parts[0].toLowerCase() === "dbs" && parts.length >= 2) {
+          dbName = parts[1];
+          if (parts.length >= 4) {
+            if (parts[2].toLowerCase() === "colls" && parts[3].toLowerCase() !== "") {
+              containerName = parts[3];
+            }
+          }
+        }
+      }
+    }
+
+    return { dbName, containerName };
+  }
+
+  getConnectionMode(requestContext) {
+    if (requestContext != null) {
+      const mode = requestContext.client?.connectionPolicy?.connectionMode;
+      if (mode == 0) {
+        return "gateway";
+      } else if (mode == 1) {
+        return "direct";
+      } else {
+        return "other";
+      }
+    }
+    return null;
+  }
+
+  getHttpInfo(requestContext) {
+    var outHost = null;
+    var userAgent = null;
+    if (requestContext != null) {
+      outHost = requestContext.client?.cosmosClientOptions?.endpoint;
+      const headers = requestContext.headers;
+      if (headers != null) {
+        userAgent = headers['User-Agent'];
+      }
+    }
+    return { outHost, userAgent };
+  }
+
+}
+
+module.exports = AzureCosmosPlugin

packages/datadog-plugin-http/src/index.js

@@ -14,7 +14,7 @@ const HttpServerPlugin = require('./server')
  */
 class HttpPlugin extends CompositePlugin {
   static id = 'http'
-  static get plugins () {
+  static get plugins() {
     const plugins = {}
 
     // Load push subscription plugin first (if enabled) for GCP Cloud Run

packages/dd-trace/src/config/supported-configurations.json

@@ -2206,6 +2206,13 @@
         "default": "true"
       }
     ],
+    "DD_TRACE_AZURE_COSMOS_ENABLED": [
+      {
+        "implementation": "A",
+        "type": "boolean",
+        "default": "true"
+      }
+    ],
     "DD_TRACE_AZURE_DURABLE_FUNCTIONS_ENABLED": [
       {
         "implementation": "B",