[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,155 @@ class DynamoDb extends BaseAwsSdkPlugin {
 
     return tags
   }
+
+  addSpanPointers (span, response) {
+    const request = response?.request
+    const operationName = request?.operation
+
+    const hashes = []
+    switch (operationName) {
+      case 'putItem': {
+        const hash = DynamoDb.calculatePutItemHash(
+          request?.params?.TableName,
+          request?.params?.Item,
+          DynamoDb.getPrimaryKeyConfig()
+        )
+        if (hash) hashes.push(hash)
+        break
+      }
+      case 'updateItem':
+      case 'deleteItem': {
+        const hash = DynamoDb.calculateHashWithKnownKeys(request?.params?.TableName, request?.params?.Key)
+        if (hash) hashes.push(hash)
+        break
+      }
+      case 'transactWriteItems': {
+        const transactItems = request?.params?.TransactItems || []
+        for (const item of transactItems) {
+          if (item.Put) {
+            const hash =
+              DynamoDb.calculatePutItemHash(item.Put.TableName, item.Put.Item, DynamoDb.getPrimaryKeyConfig())
+            if (hash) hashes.push(hash)
+          } else if (item.Update || item.Delete) {
+            const operation = item.Update ? item.Update : item.Delete
+            const hash = DynamoDb.calculateHashWithKnownKeys(operation.TableName, operation.Key)
+            if (hash) hashes.push(hash)
+          }
+        }
+        break
+      }
+      case 'batchWriteItem': {
+        const requestItems = request?.params.RequestItems || {}
+        for (const [tableName, operations] of Object.entries(requestItems)) {
+          if (!Array.isArray(operations)) continue
+          for (const operation of operations) {
+            if (operation?.PutRequest) {
+              const hash =
+                DynamoDb.calculatePutItemHash(tableName, operation.PutRequest.Item, DynamoDb.getPrimaryKeyConfig())
+              if (hash) hashes.push(hash)
+            } else if (operation?.DeleteRequest) {
+              const hash = DynamoDb.calculateHashWithKnownKeys(tableName, operation.DeleteRequest.Key)
+              if (hash) hashes.push(hash)
+            }
+          }
+        }
+        break
+      }
+    }
+
+    for (const hash of hashes) {
+      span.addSpanPointer(DYNAMODB_PTR_KIND, SPAN_POINTER_DIRECTION.DOWNSTREAM, hash)
+    }
+  }
+
+  /**
+   * Loads primary key config from the `DD_AWS_SDK_DYNAMODB_TABLE_PRIMARY_KEYS` env var.
+   * Only runs when needed, and warns when missing or invalid config.
+   * @returns {Object|null} Parsed config from env var or null if empty/missing/invalid config.
+   */
+  static getPrimaryKeyConfig () {
+    const config = DynamoDb.dynamoPrimaryKeyConfig || {}
+    // Return cached config if it exists
+    if (Object.keys(config).length > 0) {
+      return config
+    }
+
+    const primaryKeysEnvVar = process.env.DD_AWS_SDK_DYNAMODB_TABLE_PRIMARY_KEYS
+    if (!primaryKeysEnvVar) {
+      log.warn('Missing DD_AWS_SDK_DYNAMODB_TABLE_PRIMARY_KEYS env variable')
+      return null
+    }
+
+    try {
+      const parsedConfig = JSON.parse(primaryKeysEnvVar)
+      for (const [tableName, primaryKeys] of Object.entries(parsedConfig)) {
+        if (Array.isArray(primaryKeys) && primaryKeys.length > 0) {
+          config[tableName] = new Set(primaryKeys)
+        } else {
+          log.warn(`Invalid primary key configuration for table: ${tableName}`)
+        }
+      }
+
+      DynamoDb.dynamoPrimaryKeyConfig = config
+      return config
+    } catch (err) {
+      log.warn('Failed to parse DD_AWS_SDK_DYNAMODB_TABLE_PRIMARY_KEYS:', err)
+      return null
+    }
+  }
+
+  /**
+   * Calculates a hash for DynamoDB PutItem operations using table's configured primary keys.
+   * @param {string} tableName - Name of the DynamoDB table.
+   * @param {Object} item - Complete PutItem item parameter to be put.
+   * @param {Object.<string, Set<string>>} primaryKeyConfig - Mapping of table names to Sets of primary key names
+   *                                                         loaded from DD_AWS_SDK_DYNAMODB_TABLE_PRIMARY_KEYS.
+   * @returns {string|undefined} Hash combining table name and primary key/value pairs, or undefined if unable.
+   *
+   * @example
+   * // With env var DD_AWS_SDK_DYNAMODB_TABLE_PRIMARY_KEYS='{"UserTable":["userId","timestamp"]}'
+   * calculatePutItemHash(
+   *   'UserTable',
+   *   { userId: { S: "user123" }, timestamp: { N: "1234567" }, name: { S: "John" } },
+   *   { UserTable: new Set(['userId', 'timestamp']) }
+   * )
+   */
+  static calculatePutItemHash (tableName, item, primaryKeyConfig) {
+    if (!tableName || !item || !primaryKeyConfig) {
+      log.debug('Unable to calculate hash because missing required parameters')
+      return
+    }
+    const primaryKeySet = primaryKeyConfig[tableName]
+    if (!primaryKeySet || !(primaryKeySet instanceof Set) || primaryKeySet.size === 0 || primaryKeySet.size > 2) {
+      log.warn(`Invalid dynamo primary key config for table ${tableName}: ${JSON.stringify(primaryKeyConfig)}`)
+      return
+    }
+    const keyValues = extractPrimaryKeys(primaryKeySet, item)
+    if (keyValues) {
+      return generatePointerHash([tableName, ...keyValues])
+    }
+  }
+
+  /**
+   * 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