@depup/dd-trace 5.90.0-depup.0

package/packages/datadog-plugin-avsc/src/schema_iterator.js

@@ -0,0 +1,176 @@
+'use strict'
+
+const AVRO = 'avro'
+const {
+  SCHEMA_DEFINITION,
+  SCHEMA_ID,
+  SCHEMA_NAME,
+  SCHEMA_OPERATION,
+  SCHEMA_WEIGHT,
+  SCHEMA_TYPE,
+} = require('../../dd-trace/src/constants')
+const log = require('../../dd-trace/src/log')
+const {
+  SchemaBuilder,
+} = require('../../dd-trace/src/datastreams')
+
+class SchemaExtractor {
+  constructor (schema) {
+    this.schema = schema
+  }
+
+  static getType (type) {
+    const typeMapping = {
+      string: 'string',
+      int: 'integer',
+      long: 'integer',
+      float: 'number',
+      double: 'number',
+      boolean: 'boolean',
+      bytes: 'string',
+      record: 'object',
+      enum: 'string',
+      array: 'array',
+      map: 'object',
+      fixed: 'string',
+    }
+    const typeName = type.typeName ?? type.name ?? type
+    return typeName === 'null' ? typeName : typeMapping[typeName] || 'string'
+  }
+
+  static extractProperty (field, schemaName, fieldName, builder, depth) {
+    let array = false
+    let type
+    let format
+    let enumValues
+    let description
+    let ref
+
+    const fieldType = field.type?.types ?? field.type?.typeName ?? field.type
+
+    if (Array.isArray(fieldType)) {
+      // Union Type
+      type = 'union[' + fieldType.map(t => SchemaExtractor.getType(t.type || t)).join(',') + ']'
+    } else if (fieldType === 'array') {
+      // Array Type
+      array = true
+      const nestedType = field.type.itemsType.typeName
+      type = SchemaExtractor.getType(nestedType)
+    } else if (fieldType === 'record') {
+      // Nested Record Type
+      type = 'object'
+      ref = `#/components/schemas/${field.type.name}`
+      if (!SchemaExtractor.extractSchema(field.type, builder, depth + 1, this)) {
+        return false
+      }
+    } else if (fieldType === 'enum') {
+      enumValues = []
+      let i = 0
+      type = 'string'
+      while (field.type.symbols[i]) {
+        enumValues.push(field.type.symbols[i])
+        i += 1
+      }
+    } else {
+      // Primitive type
+      type = SchemaExtractor.getType(fieldType.type || fieldType)
+      if (fieldType === 'bytes') {
+        format = 'byte'
+      } else if (fieldType === 'int') {
+        format = 'int32'
+      } else if (fieldType === 'long') {
+        format = 'int64'
+      } else if (fieldType === 'float') {
+        format = 'float'
+      } else if (fieldType === 'double') {
+        format = 'double'
+      }
+    }
+
+    return builder.addProperty(schemaName, fieldName, array, type, description, ref, format, enumValues)
+  }
+
+  static extractSchema (schema, builder, depth, extractor) {
+    depth += 1
+    const schemaName = schema.name
+    if (extractor) {
+      // if we already have a defined extractor, this is a nested schema. create a new extractor for the nested
+      // schema, ensure it is added to our schema builder's cache, and replace the builders iterator with our
+      // nested schema iterator / extractor. Once complete, add the new schema to our builder's schemas.
+      const nestedSchemaExtractor = new SchemaExtractor(schema)
+      builder.iterator = nestedSchemaExtractor
+      const nestedSchema = SchemaBuilder.getSchema(schemaName, nestedSchemaExtractor, builder)
+      for (const nestedSubSchemaName in nestedSchema.components.schemas) {
+        if (nestedSchema.components.schemas.hasOwnProperty(nestedSubSchemaName)) {
+          builder.schema.components.schemas[nestedSubSchemaName] = nestedSchema.components.schemas[nestedSubSchemaName]
+        }
+      }
+      return true
+    }
+    if (!builder.shouldExtractSchema(schemaName, depth)) {
+      return false
+    }
+    if (schema.fields?.[Symbol.iterator]) {
+      for (const field of schema.fields) {
+        if (!this.extractProperty(field, schemaName, field.name, builder, depth)) {
+          log.warn('DSM: Unable to extract field with name: %s from Avro schema with name: %s', field.name,
+            schemaName)
+        }
+      }
+    } else {
+      log.warn('DSM: schema.fields is not iterable from Avro schema with name: %s', schemaName)
+    }
+
+    return true
+  }
+
+  static extractSchemas (descriptor, dataStreamsProcessor) {
+    return dataStreamsProcessor.getSchema(descriptor.name, new SchemaExtractor(descriptor))
+  }
+
+  iterateOverSchema (builder) {
+    this.constructor.extractSchema(this.schema, builder, 0)
+  }
+
+  static attachSchemaOnSpan (args, span, operation, tracer) {
+    const { messageClass } = args
+    const descriptor = messageClass?.constructor?.type ?? messageClass
+
+    if (!descriptor || !span) {
+      return
+    }
+
+    if (span.context()._tags[SCHEMA_TYPE] && operation === 'serialization') {
+      // we have already added a schema to this span, this call is an encode of nested schema types
+      return
+    }
+
+    span.setTag(SCHEMA_TYPE, AVRO)
+    span.setTag(SCHEMA_NAME, descriptor.name)
+    span.setTag(SCHEMA_OPERATION, operation)
+
+    if (!tracer._dataStreamsProcessor.canSampleSchema(operation)) {
+      return
+    }
+
+    // if the span is unsampled, do not sample the schema
+    if (!tracer._prioritySampler.isSampled(span)) {
+      return
+    }
+
+    const weight = tracer._dataStreamsProcessor.trySampleSchema(operation)
+    if (weight === 0) {
+      return
+    }
+
+    const schemaData = SchemaBuilder.getSchemaDefinition(
+      this.extractSchemas(descriptor, tracer._dataStreamsProcessor)
+    )
+
+    span.setTag(SCHEMA_DEFINITION, schemaData.definition)
+    span.setTag(SCHEMA_WEIGHT, weight)
+    span.setTag(SCHEMA_ID, schemaData.id)
+  }
+}
+
+module.exports = SchemaExtractor

package/packages/datadog-plugin-aws-sdk/src/base.js

@@ -0,0 +1,335 @@
+'use strict'
+
+const analyticsSampler = require('../../dd-trace/src/analytics_sampler')
+const ClientPlugin = require('../../dd-trace/src/plugins/client')
+const { storage } = require('../../datadog-core')
+const { isTrue } = require('../../dd-trace/src/util')
+const { tagsFromRequest, tagsFromResponse } = require('../../dd-trace/src/payload-tagging')
+const { getValueFromEnvSources } = require('../../dd-trace/src/config/helper')
+const { IS_SERVERLESS } = require('../../dd-trace/src/serverless')
+
+class BaseAwsSdkPlugin extends ClientPlugin {
+  static id = 'aws'
+  static isPayloadReporter = false
+
+  get serviceIdentifier () {
+    const id = this.constructor.id.toLowerCase()
+    Object.defineProperty(this, 'serviceIdentifier', {
+      configurable: true,
+      writable: true,
+      enumerable: true,
+      value: id,
+    })
+    return id
+  }
+
+  get cloudTaggingConfig () {
+    return this._tracerConfig.cloudPayloadTagging
+  }
+
+  get payloadTaggingRules () {
+    return this.cloudTaggingConfig.rules.aws?.[this.constructor.id]
+  }
+
+  constructor (...args) {
+    super(...args)
+
+    this._parentMap = new WeakMap()
+
+    this.addBind(`apm:aws:request:start:${this.serviceIdentifier}`, (ctx) => {
+      const {
+        request,
+        operation,
+        awsRegion,
+        awsService,
+      } = ctx
+
+      const parentStore = ctx.parentStore = storage('legacy').getStore()
+      const childOf = parentStore?.span
+
+      this._parentMap.set(request, parentStore)
+
+      if (!this.isEnabled(request)) {
+        return parentStore
+      }
+
+      const meta = {
+        'span.kind': 'client',
+        'service.name': this.serviceName(),
+        'aws.operation': operation,
+        'aws.region': awsRegion,
+        region: awsRegion,
+        'aws.partition': getPartition(awsRegion),
+        aws_service: awsService,
+        'aws.service': awsService,
+        component: 'aws-sdk',
+      }
+      if (this.requestTags) this.requestTags.set(request, meta)
+
+      const span = this.startSpan(this.operationFromRequest(request), {
+        childOf,
+        meta,
+        integrationName: 'aws-sdk',
+      }, ctx)
+
+      analyticsSampler.sample(span, this.config.measured)
+
+      storage('legacy').run(ctx.currentStore, () => {
+        this.requestInject(span, request)
+      })
+
+      if (this.constructor.isPayloadReporter && this.cloudTaggingConfig.requestsEnabled) {
+        const maxDepth = this.cloudTaggingConfig.maxDepth
+        const requestTags = tagsFromRequest(this.payloadTaggingRules, request.params, { maxDepth })
+        span.addTags(requestTags)
+      }
+
+      return ctx.currentStore
+    })
+
+    this.addSub(`apm:aws:request:start:${this.serviceIdentifier}`, (ctx) => {
+      if (!IS_SERVERLESS) return
+
+      const { awsRegion, awsService, currentStore, request } = ctx
+      const peerServerlessStorage = storage('peerServerless')
+
+      // Try to resolve the hostname immediately; if not possible, keep enough
+      // information so the region callback can resolve it later.
+      const hostname = getHostname({ awsParams: request.params, awsService }, awsRegion)
+      const peerServerlessStore = {}
+      peerServerlessStorage.enterWith(peerServerlessStore)
+
+      if (hostname) {
+        currentStore.span.setTag('peer.service', hostname)
+        peerServerlessStore.peerHostname = hostname
+      } else {
+        currentStore.awsParams = request.params
+        currentStore.awsService = awsService
+      }
+    })
+
+    this.addSub(`apm:aws:request:region:${this.serviceIdentifier}`, ({ region }) => {
+      const store = storage('legacy').getStore()
+      if (!store) return
+      const { span } = store
+      if (!span) return
+      span.setTag('aws.region', region)
+      span.setTag('region', region)
+
+      const partition = getPartition(region)
+      if (partition) {
+        span.setTag('aws.partition', partition)
+      }
+
+      if (!IS_SERVERLESS) return
+
+      const hostname = getHostname(store, region)
+      if (!hostname) return
+
+      span.setTag('peer.service', hostname)
+      const peerServerlessStore = storage('peerServerless').getStore()
+      if (peerServerlessStore) {
+        peerServerlessStore.peerHostname = hostname
+      }
+    })
+
+    this.addSub(`apm:aws:request:complete:${this.serviceIdentifier}`, ctx => {
+      const { response, cbExists = false, currentStore } = ctx
+      if (!currentStore) return
+      const { span } = currentStore
+      if (!span) return
+
+      storage('legacy').run(currentStore, () => {
+        // try to extract DSM context from response if no callback exists as extraction normally happens in CB
+        if (!cbExists && this.serviceIdentifier === 'sqs') {
+          const params = response.request.params
+          const operation = response.request.operation
+          this.responseExtractDSMContext(operation, params, response.data ?? response, span)
+        }
+        this.addResponseTags(span, response)
+
+        if (this._tracerConfig?.trace?.aws?.addSpanPointers) {
+          this.addSpanPointers(span, response)
+        }
+      })
+
+      this.finish(ctx)
+    })
+
+    this.addBind(`apm:aws:response:start:${this.serviceIdentifier}`, ctx => {
+      return this._parentMap.get(ctx.request)
+    })
+  }
+
+  requestInject (span, request) {
+    // implemented by subclasses, or not
+  }
+
+  addSpanPointers (span, response) {
+    // Optionally implemented by subclasses, for services where we're unable to inject trace context
+  }
+
+  operationFromRequest (request) {
+    // can be overriden by subclasses
+    return this.operationName({
+      id: 'aws',
+      type: 'web',
+      kind: 'client',
+      awsService: this.serviceIdentifier,
+    })
+  }
+
+  serviceName () {
+    return this.config.service ||
+      super.serviceName({
+        id: 'aws',
+        type: 'web',
+        kind: 'client',
+        awsService: this.serviceIdentifier,
+      })
+  }
+
+  isEnabled (request) {
+    const serviceId = this.serviceIdentifier.toUpperCase()
+    const envVarValue = getValueFromEnvSources(`DD_TRACE_AWS_SDK_${serviceId}_ENABLED`)
+    return envVarValue ? isTrue(envVarValue) : true
+  }
+
+  addResponseTags (span, response) {
+    if (!span || !response.request) return
+    const params = response.request.params
+    const operation = response.request.operation
+    const extraTags = this.generateTags(params, operation, response) || {}
+
+    const tags = {
+      'aws.response.request_id': response.requestId,
+      'resource.name': operation,
+      'span.kind': 'client',
+      ...extraTags,
+    }
+
+    span.addTags(tags)
+
+    if (this.constructor.isPayloadReporter && this.cloudTaggingConfig.responsesEnabled) {
+      const maxDepth = this.cloudTaggingConfig.maxDepth
+      const responseBody = this.extractResponseBody(response)
+      const responseTags = tagsFromResponse(this.payloadTaggingRules, responseBody, { maxDepth })
+      span.addTags(responseTags)
+    }
+  }
+
+  extractResponseBody (response) {
+    if (response.hasOwnProperty('data')) {
+      return response.data
+    }
+    return Object.fromEntries(
+      Object.entries(response).filter(([key]) => !['request', 'requestId', 'error', '$metadata'].includes(key))
+    )
+  }
+
+  generateTags () {
+    // implemented by subclasses, or not
+  }
+
+  finish (ctx) {
+    const { currentStore, response } = ctx
+    const { span } = currentStore
+    const error = response?.error || ctx.error
+
+    if (error) {
+      span.setTag('error', error)
+
+      const requestId = error.RequestId || error.requestId
+      if (requestId) {
+        span.addTags({ 'aws.response.request_id': requestId })
+      }
+    }
+
+    if (response) {
+      this.config.hooks.request(span, response)
+    }
+
+    super.finish(ctx)
+  }
+
+  configure (config) {
+    super.configure(normalizeConfig(config, this.serviceIdentifier))
+  }
+}
+
+function normalizeConfig (config, serviceIdentifier) {
+  const hooks = getHooks(config)
+
+  let specificConfig = config[serviceIdentifier]
+  switch (typeof specificConfig) {
+    case 'undefined':
+      specificConfig = {}
+      break
+    case 'boolean':
+      specificConfig = { enabled: specificConfig }
+      break
+  }
+
+  // check if AWS batch propagation or AWS_[SERVICE] batch propagation is enabled via env variable
+  const serviceId = serviceIdentifier.toUpperCase()
+  const batchPropagationEnabled = isTrue(
+    specificConfig.batchPropagationEnabled ??
+    getValueFromEnvSources(`DD_TRACE_AWS_SDK_${serviceId}_BATCH_PROPAGATION_ENABLED`) ??
+    config.batchPropagationEnabled ??
+    getValueFromEnvSources('DD_TRACE_AWS_SDK_BATCH_PROPAGATION_ENABLED')
+  )
+
+  // Merge the specific config back into the main config
+  return {
+    ...config,
+    ...specificConfig,
+    batchPropagationEnabled,
+    hooks,
+  }
+}
+
+const noop = () => {}
+
+function getHooks (config) {
+  const request = config.hooks?.request || noop
+
+  return { request }
+}
+
+function getHostname (store, region) {
+  if (!store) return
+  if (!region) return
+  const { awsParams, awsService } = store
+  switch (awsService) {
+    case 'EventBridge':
+      return `events.${region}.amazonaws.com`
+    case 'SQS':
+      return `sqs.${region}.amazonaws.com`
+    case 'SNS':
+      return `sns.${region}.amazonaws.com`
+    case 'Kinesis':
+      return `kinesis.${region}.amazonaws.com`
+    case 'DynamoDBDocument':
+    case 'DynamoDB':
+      return `dynamodb.${region}.amazonaws.com`
+    case 'S3':
+      return awsParams?.Bucket
+        ? `${awsParams.Bucket}.s3.${region}.amazonaws.com`
+        : `s3.${region}.amazonaws.com`
+  }
+}
+
+function getPartition (region) {
+  if (!region) return
+
+  let partition = 'aws'
+  if (region.startsWith('cn-')) {
+    partition = 'aws-cn'
+  } else if (region.startsWith('us-gov-')) {
+    partition = 'aws-us-gov'
+  }
+
+  return partition
+}
+
+module.exports = BaseAwsSdkPlugin

package/packages/datadog-plugin-aws-sdk/src/index.js

@@ -0,0 +1,26 @@
+'use strict'
+
+const Plugin = require('../../dd-trace/src/plugins/plugin')
+
+const services = require('./services')
+
+class AwsSdkPlugin extends Plugin {
+  static id = 'aws-sdk'
+
+  constructor (...args) {
+    super(...args)
+    this.services = {}
+    for (const name in services) {
+      const ServicePlugin = services[name]
+      this.services[name] = new ServicePlugin(...args)
+    }
+  }
+
+  configure (config) {
+    for (const service of Object.values(this.services)) {
+      service.configure(config)
+    }
+  }
+}
+
+module.exports = AwsSdkPlugin

package/packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/index.js

@@ -0,0 +1,16 @@
+'use strict'
+
+const CompositePlugin = require('../../../../dd-trace/src/plugins/composite')
+const BedrockRuntimeLLMObsPlugin = require('../../../../dd-trace/src/llmobs/plugins/bedrockruntime')
+const BedrockRuntimeTracing = require('./tracing')
+class BedrockRuntimePlugin extends CompositePlugin {
+  static id = 'bedrockruntime'
+
+  static get plugins () {
+    return {
+      llmobs: BedrockRuntimeLLMObsPlugin,
+      tracing: BedrockRuntimeTracing,
+    }
+  }
+}
+module.exports = BedrockRuntimePlugin

package/packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/tracing.js

@@ -0,0 +1,31 @@
+'use strict'
+
+const BaseAwsSdkPlugin = require('../../base')
+const { parseModelId } = require('./utils')
+
+const enabledOperations = new Set(['invokeModel', 'invokeModelWithResponseStream'])
+
+class BedrockRuntime extends BaseAwsSdkPlugin {
+  static id = 'bedrockruntime'
+
+  isEnabled (request) {
+    const operation = request.operation
+    if (!enabledOperations.has(operation)) {
+      return false
+    }
+
+    return super.isEnabled(request)
+  }
+
+  generateTags (params, operation) {
+    const { modelProvider, modelName } = parseModelId(params.modelId)
+
+    return {
+      'resource.name': operation,
+      'aws.bedrock.request.model': modelName,
+      'aws.bedrock.request.model_provider': modelProvider.toLowerCase(),
+    }
+  }
+}
+
+module.exports = BedrockRuntime