[serverless] Add DynamoDB Span Pointers

packages/datadog-plugin-aws-sdk/src/services/dynamodb.js

@@ -1,6 +1,9 @@
 'use strict'
 
 const BaseAwsSdkPlugin = require('../base')
+const log = require('../../../dd-trace/src/log')
+const { DYNAMODB_PTR_KIND, SPAN_POINTER_DIRECTION } = require('../../../dd-trace/src/constants')
+const { extractPrimaryKeys, generatePointerHash } = require('../../../dd-trace/src/util')
 
 class DynamoDb extends BaseAwsSdkPlugin {
   static get id () { return 'dynamodb' }
@@ -48,6 +51,47 @@ class DynamoDb extends BaseAwsSdkPlugin {
 
     return tags
   }
+
+  addSpanPointers (span, response) {
+    const request = response?.request
+    const operationName = request?.operation
+
+    const hashes = []
+    switch (operationName) {
+      case 'updateItem':
+      case 'deleteItem': {
+        const hash = DynamoDb.calculateHashWithKnownKeys(request?.params?.TableName, request?.params?.Key)
+        if (hash) hashes.push(hash)
+        break
+      }
+    }
+
+    for (const hash of hashes) {
+      span.addSpanPointer(DYNAMODB_PTR_KIND, SPAN_POINTER_DIRECTION.DOWNSTREAM, hash)
+    }
+  }
+
+  /**
+   * Calculates a hash for DynamoDB operations that have keys provided (UpdateItem, DeleteItem).
+   *
+   * @param {string} tableName - Name of the DynamoDB table.
+   * @param {Object} keys - Object containing primary key/value attributes in DynamoDB format.
+   *                       (e.g., { userId: { S: "123" }, sortKey: { N: "456" } })
+   * @returns {string|undefined} Hash value combining table name and primary key/value pairs, or undefined if unable.
+   *
+   * @example
+   * calculateKeyBasedOperationsHash('UserTable', { userId: { S: "user123" }, timestamp: { N: "1234567" } })
+   */
+  static calculateHashWithKnownKeys (tableName, keys) {
+    if (!tableName || !keys) {
+      log.debug('Unable to calculate hash because missing parameters')
+      return
+    }
+    const keyValues = extractPrimaryKeys(keys, keys)
+    if (keyValues) {
+      return generatePointerHash([tableName, ...keyValues])
+    }
+  }
 }
 
 module.exports = DynamoDb

packages/dd-trace/src/constants.js

@@ -47,6 +47,7 @@ module.exports = {
   SCHEMA_NAME: 'schema.name',
   GRPC_CLIENT_ERROR_STATUSES: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
   GRPC_SERVER_ERROR_STATUSES: [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
+  DYNAMODB_PTR_KIND: 'aws.dynamodb.item',
   S3_PTR_KIND: 'aws.s3.object',
   SPAN_POINTER_DIRECTION: Object.freeze({
     UPSTREAM: 'u',

packages/dd-trace/src/util.js

@@ -87,12 +87,83 @@ function generatePointerHash (components) {
   return hash.substring(0, 32)
 }
 
+/**
+ * Encodes a DynamoDB attribute value to Buffer for span pointer hashing.
+ * @param {Object} valueObject - DynamoDB value in AWS format ({ S: string } or { N: string } or { B: Buffer })
+ * @returns {Buffer} Encoded value as Buffer, or empty Buffer if invalid input.
+ *
+ * @example
+ * encodeValue({ S: "user123" }) -> Buffer("user123")
+ * encodeValue({ N: "42" }) -> Buffer("42")
+ * encodeValue({ B: Buffer([1, 2, 3]) }) -> Buffer([1, 2, 3])
+ */
+function encodeValue (valueObject) {
+  if (!valueObject) {
+    return Buffer.from('')
+  }
+
+  try {
+    const type = Object.keys(valueObject)[0]
+    const value = valueObject[type]
+
+    switch (type) {
+      case 'S':
+        return Buffer.from(value)
+      case 'N':
+        return Buffer.from(value.toString())
+      case 'B':
+        return Buffer.isBuffer(value) ? value : Buffer.from(value)
+      default:
+        return Buffer.from('')
+    }
+  } catch (err) {
+    return Buffer.from('')
+  }
+}
+
+/**
+ * Extracts and encodes primary key values from a DynamoDB item.
+ * Handles tables with single-key and two-key scenarios.
+ *
+ * @param {Set<string>|Object} keySet - Set of key names or object of key names/value pairs.
+ * @param {Object} keyValuePairs - Object containing key/value pairs.
+ * @returns {Array|undefined} [key1Name, key1Value, key2Name, key2Value], or undefined if invalid input.
+ *                            key2 entries are empty strings in the single-key case.
+ * @example
+ * extractPrimaryKeys(new Set(['userId']), {userId: {S: "user123"}})
+ * // Returns ["userId", Buffer("user123"), "", ""]
+ * extractPrimaryKeys(new Set(['userId', 'timestamp']), {userId: {S: "user123"}, timestamp: {N: "1234}})
+ * // Returns ["timestamp", Buffer.from("1234"), "userId", Buffer.from("user123")]
+ */
+const extractPrimaryKeys = (keySet, keyValuePairs) => {
+  const keyNames = keySet instanceof Set
+    ? Array.from(keySet)
+    : Object.keys(keySet)
+  if (keyNames.length === 0) {
+    return
+  }
+
+  if (keyNames.length === 1) {
+    return [keyNames[0], encodeValue(keyValuePairs[keyNames[0]]), '', '']
+  } else {
+    const [key1, key2] = keyNames.sort()
+    return [
+      key1,
+      encodeValue(keyValuePairs[key1]),
+      key2,
+      encodeValue(keyValuePairs[key2])
+    ]
+  }
+}
+
 module.exports = {
   isTrue,
   isFalse,
   isError,
   globMatch,
   calculateDDBasePath,
   hasOwn,
-  generatePointerHash
+  generatePointerHash,
+  encodeValue,
+  extractPrimaryKeys
 }