dd-trace 5.62.0 → 5.63.0

package/packages/datadog-instrumentations/src/pg.js

@@ -2,8 +2,7 @@
 
 const {
   channel,
-  addHook,
-  AsyncResource
+  addHook
 } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
@@ -33,8 +32,6 @@ function wrapQuery (query) {
       return query.apply(this, arguments)
     }
 
-    const callbackResource = new AsyncResource('bound-anonymous-fn')
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
     const processId = this.processID
 
     const pgQuery = arguments[0] !== null && typeof arguments[0] === 'object'
@@ -46,8 +43,9 @@ function wrapQuery (query) {
     const stream = typeof textPropObj.read === 'function'
 
     // Only alter `text` property if safe to do so. Initially, it's a property, not a getter.
+    let originalText
     if (!textProp || textProp.configurable) {
-      const originalText = textPropObj.text
+      originalText = textPropObj.text
 
       Object.defineProperty(textPropObj, 'text', {
         get () {
@@ -55,25 +53,24 @@ function wrapQuery (query) {
         }
       })
     }
-
-    return asyncResource.runInAsyncScope(() => {
-      const abortController = new AbortController()
-
-      startCh.publish({
-        params: this.connectionParameters,
-        query: textPropObj,
-        processId,
-        abortController,
-        stream
-      })
-
-      const finish = asyncResource.bind(function (error, res) {
-        if (error) {
-          errorCh.publish(error)
-        }
-        finishCh.publish({ result: res?.rows })
-      })
-
+    const abortController = new AbortController()
+    const ctx = {
+      params: this.connectionParameters,
+      query: textPropObj,
+      originalText,
+      processId,
+      abortController,
+      stream
+    }
+    const finish = (error, res) => {
+      if (error) {
+        ctx.error = error
+        errorCh.publish(ctx)
+      }
+      ctx.result = res?.rows
+      return finishCh.publish(ctx)
+    }
+    return startCh.runStores(ctx, () => {
       if (abortController.signal.aborted) {
         const error = abortController.signal.reason || new Error('Aborted')
 
@@ -121,10 +118,10 @@ function wrapQuery (query) {
       }
 
       if (newQuery.callback) {
-        const originalCallback = callbackResource.bind(newQuery.callback)
-        newQuery.callback = function (err, res) {
-          finish(err, res)
-          return originalCallback.apply(this, arguments)
+        const originalCallback = newQuery.callback
+        newQuery.callback = function (err, ...args) {
+          finish(err, ...args)
+          return finishCh.runStores(ctx, originalCallback, this, err, ...args)
         }
       } else if (newQuery.once) {
         newQuery
@@ -139,40 +136,33 @@ function wrapQuery (query) {
 
       try {
         return retval
-      } catch (err) {
-        errorCh.publish(err)
+      } catch (error) {
+        ctx.error = error
+        errorCh.publish(ctx)
       }
     })
   }
 }
-
+const finish = (ctx) => {
+  finishPoolQueryCh.publish(ctx)
+}
 function wrapPoolQuery (query) {
   return function () {
     if (!startPoolQueryCh.hasSubscribers) {
       return query.apply(this, arguments)
     }
 
-    const asyncResource = new AsyncResource('bound-anonymous-fn')
-
     const pgQuery = arguments[0] !== null && typeof arguments[0] === 'object' ? arguments[0] : { text: arguments[0] }
+    const abortController = new AbortController()
 
-    return asyncResource.runInAsyncScope(() => {
-      const abortController = new AbortController()
-
-      startPoolQueryCh.publish({
-        query: pgQuery,
-        abortController
-      })
-
-      const finish = asyncResource.bind(function () {
-        finishPoolQueryCh.publish()
-      })
+    const ctx = { query: pgQuery, abortController }
 
+    return startPoolQueryCh.runStores(ctx, () => {
       const cb = arguments[arguments.length - 1]
 
       if (abortController.signal.aborted) {
         const error = abortController.signal.reason || new Error('Aborted')
-        finish()
+        finish(ctx)
 
         if (typeof cb === 'function') {
           cb(error)
@@ -184,7 +174,7 @@ function wrapPoolQuery (query) {
 
       if (typeof cb === 'function') {
         arguments[arguments.length - 1] = shimmer.wrapFunction(cb, cb => function () {
-          finish()
+          finish(ctx)
           return cb.apply(this, arguments)
         })
       }
@@ -193,9 +183,9 @@ function wrapPoolQuery (query) {
 
       if (retval?.then) {
         retval.then(() => {
-          finish()
+          finish(ctx)
         }).catch(() => {
-          finish()
+          finish(ctx)
         })
       }
 

package/packages/datadog-plugin-aerospike/src/index.js

@@ -64,8 +64,12 @@ class AerospikePlugin extends DatabasePlugin {
 function getMeta (resourceName, commandArgs) {
   let meta = {}
   if (resourceName.includes('Index')) {
-    const [ns, set, bin, index] = commandArgs
-    meta = getMetaForIndex(ns, set, bin, index)
+    const [ns, set, bin, exp, index] = commandArgs
+
+    // The `ext` argument was added to IndexCreate in 6.3.0
+    meta = commandArgs.length > 8
+      ? getMetaForIndex(ns, set, bin, index)
+      : getMetaForIndex(ns, set, bin, exp)
   } else if (resourceName === 'Query') {
     const { ns, set } = commandArgs[2]
     meta = getMetaForQuery({ ns, set })

package/packages/datadog-plugin-ai/src/index.js

@@ -0,0 +1,17 @@
+'use strict'
+
+const CompositePlugin = require('../../dd-trace/src/plugins/composite')
+const VercelAILLMObsPlugin = require('../../dd-trace/src/llmobs/plugins/ai')
+const VercelAITracingPlugin = require('./tracing')
+
+class VercelAIPlugin extends CompositePlugin {
+  static get id () { return 'ai' }
+  static get plugins () {
+    return {
+      llmobs: VercelAILLMObsPlugin,
+      tracing: VercelAITracingPlugin
+    }
+  }
+}
+
+module.exports = VercelAIPlugin

package/packages/datadog-plugin-ai/src/tracing.js

@@ -0,0 +1,33 @@
+'use strict'
+
+const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
+const { getModelProvider } = require('./utils')
+
+class VercelAITracingPlugin extends TracingPlugin {
+  static id = 'ai'
+  static prefix = 'tracing:dd-trace:vercel-ai'
+
+  bindStart (ctx) {
+    const attributes = ctx.attributes
+
+    const model = attributes['ai.model.id']
+    const modelProvider = getModelProvider(attributes)
+
+    this.startSpan(ctx.name, {
+      meta: {
+        'resource.name': ctx.name,
+        'ai.request.model': model,
+        'ai.request.model_provider': modelProvider
+      }
+    }, ctx)
+
+    return ctx.currentStore
+  }
+
+  asyncEnd (ctx) {
+    const span = ctx.currentStore?.span
+    span?.finish()
+  }
+}
+
+module.exports = VercelAITracingPlugin

package/packages/datadog-plugin-ai/src/utils.js

@@ -0,0 +1,28 @@
+'use strict'
+
+const { parseModelId } = require('../../datadog-plugin-aws-sdk/src/services/bedrockruntime/utils')
+
+/**
+ * Get the model provider from the span tags or attributes.
+ * This is normalized to LLM Observability model provider standards.
+ *
+ * @param {Record<string, string>} tags
+ * @returns {string}
+ */
+function getModelProvider (tags) {
+  const modelProviderTag = tags['ai.model.provider']
+  const providerParts = modelProviderTag?.split('.')
+  const provider = providerParts?.[0]
+
+  if (provider === 'amazon-bedrock') {
+    const modelId = tags['ai.model.id']
+    const model = modelId && parseModelId(modelId)
+    return model?.modelProvider ?? provider
+  }
+
+  return provider
+}
+
+module.exports = {
+  getModelProvider
+}

package/packages/datadog-plugin-couchbase/src/index.js

@@ -7,13 +7,15 @@ class CouchBasePlugin extends StoragePlugin {
   static id = 'couchbase'
   static peerServicePrecursors = ['db.couchbase.seed.nodes']
 
-  addSubs (func, start) {
-    this.addSub(`apm:couchbase:${func}:start`, start)
-    this.addSub(`apm:couchbase:${func}:error`, error => this.addError(error))
-    this.addSub(`apm:couchbase:${func}:finish`, message => this.finish(message))
+  addBinds (func, start) {
+    this.addBind(`apm:couchbase:${func}:start`, start)
+    this.addSub(`apm:couchbase:${func}:error`, ({ error }) => this.addError(error))
+    this.addSub(`apm:couchbase:${func}:finish`, ctx => this.finish(ctx))
+    this.addBind(`apm:couchbase:${func}:callback:start`, callbackStart)
+    this.addBind(`apm:couchbase:${func}:callback:finish`, callbackFinish)
   }
 
-  startSpan (operation, customTags, store, { bucket, collection, seedNodes }) {
+  startSpan (operation, customTags, { bucket, collection, seedNodes }, ctx) {
     const tags = {
       'db.type': 'couchbase',
       component: 'couchbase',
@@ -34,26 +36,34 @@ class CouchBasePlugin extends StoragePlugin {
       {
         service: this.serviceName({ pluginConfig: this.config }),
         meta: tags
-      }
+      },
+      ctx
     )
   }
 
   constructor (...args) {
     super(...args)
 
-    this.addSubs('query', ({ resource, bucket, seedNodes }) => {
-      const store = storage('legacy').getStore()
-      const span = this.startSpan(
-        'query', {
+    this.addBinds('query', (ctx) => {
+      const { resource, bucket, seedNodes } = ctx
+
+      this.startSpan(
+        'query',
+        {
           'span.type': 'sql',
           'resource.name': resource,
           'span.kind': this.constructor.kind
         },
-        store,
-        { bucket, seedNodes }
+        { bucket, seedNodes },
+        ctx
       )
-      this.enter(span, store)
+
+      return ctx.currentStore
     })
+    this.addBind('apm:couchbase:bucket:maybeInvoke:callback:start', callbackStart)
+    this.addBind('apm:couchbase:bucket:maybeInvoke:callback:finish', callbackFinish)
+    this.addBind('apm:couchbase:cluster:maybeInvoke:callback:start', callbackStart)
+    this.addBind('apm:couchbase:cluster:maybeInvoke:callback:finish', callbackFinish)
 
     this._addCommandSubs('upsert')
     this._addCommandSubs('insert')
@@ -63,12 +73,22 @@ class CouchBasePlugin extends StoragePlugin {
   }
 
   _addCommandSubs (name) {
-    this.addSubs(name, ({ bucket, collection, seedNodes }) => {
-      const store = storage('legacy').getStore()
-      const span = this.startSpan(name, {}, store, { bucket, collection, seedNodes })
-      this.enter(span, store)
+    this.addBinds(name, (ctx) => {
+      const { bucket, collection, seedNodes } = ctx
+
+      this.startSpan(name, {}, { bucket, collection, seedNodes }, ctx)
+      return ctx.currentStore
     })
   }
 }
 
+function callbackStart (ctx) {
+  ctx.parentStore = storage('legacy').getStore()
+  return ctx.parentStore
+}
+
+function callbackFinish (ctx) {
+  return ctx.parentStore
+}
+
 module.exports = CouchBasePlugin

package/packages/datadog-plugin-pg/src/index.js

@@ -8,7 +8,8 @@ class PGPlugin extends DatabasePlugin {
   static operation = 'query'
   static system = 'postgres'
 
-  start ({ params = {}, query, processId, stream }) {
+  bindStart (ctx) {
+    const { params = {}, query, processId, stream } = ctx
     const service = this.serviceName({ pluginConfig: this.config, params })
     const originalStatement = this.maybeTruncate(query.text)
 
@@ -25,13 +26,15 @@ class PGPlugin extends DatabasePlugin {
         'out.host': params.host,
         [CLIENT_PORT_KEY]: params.port
       }
-    })
+    }, ctx)
 
     if (stream) {
       span.setTag('db.stream', 1)
     }
 
     query.__ddInjectableQuery = this.injectDbmQuery(span, query.text, service, !!query.name)
+
+    return ctx.currentStore
   }
 }
 

package/packages/dd-trace/src/appsec/iast/analyzers/sql-injection-analyzer.js

@@ -16,7 +16,10 @@ class SqlInjectionAnalyzer extends StoredInjectionAnalyzer {
   onConfigure () {
     this.addSub('apm:mysql:query:start', ({ sql }) => this.analyze(sql, undefined, 'MYSQL'))
     this.addSub('datadog:mysql2:outerquery:start', ({ sql }) => this.analyze(sql, undefined, 'MYSQL'))
-    this.addSub('apm:pg:query:start', ({ query }) => this.analyze(query.text, undefined, 'POSTGRES'))
+    this.addSub(
+      'apm:pg:query:start',
+      ({ originalText, query }) => this.analyze(originalText || query.text, undefined, 'POSTGRES')
+    )
 
     this.addBind(
       'datadog:sequelize:query:start',
@@ -24,17 +27,22 @@ class SqlInjectionAnalyzer extends StoredInjectionAnalyzer {
     )
     this.addSub('datadog:sequelize:query:finish', () => this.returnToParentStore())
 
-    this.addSub('datadog:pg:pool:query:start', ({ query }) => this.setStoreAndAnalyze(query.text, 'POSTGRES'))
+    this.addBind('datadog:pg:pool:query:start', ({ query }) => this.getStoreAndAnalyze(query.text, 'POSTGRES'))
     this.addSub('datadog:pg:pool:query:finish', () => this.returnToParentStore())
 
     this.addSub('datadog:mysql:pool:query:start', ({ sql }) => this.setStoreAndAnalyze(sql, 'MYSQL'))
     this.addSub('datadog:mysql:pool:query:finish', () => this.returnToParentStore())
 
-    this.addSub('datadog:knex:raw:start', ({ sql, dialect: knexDialect }) => {
+    this.addBind('datadog:knex:raw:start', (context) => {
+      const { sql, dialect: knexDialect } = context
       const dialect = this.normalizeKnexDialect(knexDialect)
-      this.setStoreAndAnalyze(sql, dialect)
+      const currentStore = this.getStoreAndAnalyze(sql, dialect)
+      context.currentStore = currentStore
+      return currentStore
     })
-    this.addSub('datadog:knex:raw:finish', () => this.returnToParentStore())
+
+    this.addBind('datadog:knex:raw:subscribes', ({ currentStore }) => currentStore)
+    this.addBind('datadog:knex:raw:finish', ({ currentStore }) => currentStore?.sqlParentStore)
   }
 
   setStoreAndAnalyze (query, dialect) {
@@ -54,8 +62,7 @@ class SqlInjectionAnalyzer extends StoredInjectionAnalyzer {
     }
   }
 
-  returnToParentStore () {
-    const store = storage('legacy').getStore()
+  returnToParentStore (store = storage('legacy').getStore()) {
     if (store && store.sqlParentStore) {
       storage('legacy').enterWith(store.sqlParentStore)
     }

package/packages/dd-trace/src/appsec/iast/taint-tracking/plugin.js

@@ -139,12 +139,12 @@ class TaintTrackingPlugin extends SourceIastPlugin {
   addDatabaseSubscriptions () {
     this.addSub(
       { channelName: 'datadog:sequelize:query:finish', tag: SQL_ROW_VALUE },
-      ({ result }) => this._taintDatabaseResult(result, 'sequelize')
+      ({ result }) => this._taintDatabaseResult(result, 'sequelize', getIastContext(storage('legacy').getStore()))
     )
 
     this.addSub(
       { channelName: 'apm:pg:query:finish', tag: SQL_ROW_VALUE },
-      ({ result }) => this._taintDatabaseResult(result, 'pg')
+      ({ result, currentStore }) => this._taintDatabaseResult(result, 'pg', getIastContext(currentStore))
     )
   }
 
@@ -263,7 +263,7 @@ class TaintTrackingPlugin extends SourceIastPlugin {
     this.taintUrl(req, iastContext)
   }
 
-  _taintDatabaseResult (result, dbOrigin, iastContext = getIastContext(storage('legacy').getStore()), name) {
+  _taintDatabaseResult (result, dbOrigin, iastContext, name) {
     if (!iastContext) return result
 
     if (this._rowsToTaint === 0) return result