feat: Add capacity information when applicable to dynamodb spans (#1365)

Co-authored-by: Haddas Bronfman <[email protected]> Co-authored-by: Marc Pichler <[email protected]> Co-authored-by: Amir Blum <[email protected]>
open-telemetry · Aug 12, 2023 · ad94c5c · ad94c5c
1 parent 704f76f
commit ad94c5c
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 13 deletions.
diff --git a/plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/dynamodb.ts b/plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/dynamodb.ts
@@ -194,20 +194,16 @@ export class DynamodbServiceExtension implements ServiceExtension {
   responseHook(
     response: NormalizedResponse,
     span: Span,
-    tracer: Tracer,
-    config: AwsSdkInstrumentationConfig
+    _tracer: Tracer,
+    _config: AwsSdkInstrumentationConfig
   ) {
-    const operation = response.request.commandName;
-
-    if (operation === 'BatchGetItem') {
-      if (Array.isArray(response.data?.ConsumedCapacity)) {
-        span.setAttribute(
-          SemanticAttributes.AWS_DYNAMODB_CONSUMED_CAPACITY,
-          response.data.ConsumedCapacity.map(
-            (x: { [DictionaryKey: string]: any }) => JSON.stringify(x)
-          )
-        );
-      }
+    if (response.data?.ConsumedCapacity) {
+      span.setAttribute(
+        SemanticAttributes.AWS_DYNAMODB_CONSUMED_CAPACITY,
+        toArray(response.data.ConsumedCapacity).map(
+          (x: { [DictionaryKey: string]: any }) => JSON.stringify(x)
+        )
+      );
     }
 
     if (response.data?.ItemCollectionMetrics) {
@@ -241,3 +237,7 @@ export class DynamodbServiceExtension implements ServiceExtension {
     }
   }
 }
+
+function toArray<T>(values: T | T[]): T[] {
+  return Array.isArray(values) ? values : [values];
+}
diff --git a/plugins/node/opentelemetry-instrumentation-aws-sdk/test/dynamodb.test.ts b/plugins/node/opentelemetry-instrumentation-aws-sdk/test/dynamodb.test.ts
@@ -633,4 +633,78 @@ describe('DynamoDB', () => {
       );
     });
   });
+
+  describe('ConsumedCapacity', () => {
+    it('should populate ConsumedCapacity attributes when they exist', done => {
+      mockV2AwsSend(responseMockSuccess, {
+        ConsumedCapacity: {
+          TableName: 'test-table',
+          CapacityUnits: 0.5,
+          Table: { CapacityUnits: 0.5 },
+        },
+      } as AWS.DynamoDB.Types.PutItemOutput);
+
+      const dynamodb = new AWS.DynamoDB.DocumentClient();
+      dynamodb.put(
+        {
+          TableName: 'test-table',
+          Item: { key1: 'val1' },
+          ReturnConsumedCapacity: 'INDEXES',
+        },
+        (err: AWSError, data: AWS.DynamoDB.DocumentClient.PutItemOutput) => {
+          const spans = getTestSpans();
+          expect(spans.length).toStrictEqual(1);
+          const attrs = spans[0].attributes;
+          expect(attrs[SemanticAttributes.DB_SYSTEM]).toStrictEqual(
+            DbSystemValues.DYNAMODB
+          );
+          expect(attrs[SemanticAttributes.DB_OPERATION]).toStrictEqual(
+            'PutItem'
+          );
+          expect(
+            attrs[SemanticAttributes.AWS_DYNAMODB_CONSUMED_CAPACITY]
+          ).toStrictEqual([
+            JSON.stringify({
+              TableName: 'test-table',
+              CapacityUnits: 0.5,
+              Table: { CapacityUnits: 0.5 },
+            }),
+          ]);
+          expect(err).toBeFalsy();
+          done();
+        }
+      );
+    });
+
+    it('should not populate ConsumedCapacity attributes when it is not returned', done => {
+      mockV2AwsSend(responseMockSuccess, {
+        ConsumedCapacity: undefined,
+      } as AWS.DynamoDB.Types.PutItemOutput);
+
+      const dynamodb = new AWS.DynamoDB.DocumentClient();
+      dynamodb.put(
+        {
+          TableName: 'test-table',
+          Item: { key1: 'val1' },
+          ReturnConsumedCapacity: 'NONE',
+        },
+        (err: AWSError, data: AWS.DynamoDB.DocumentClient.PutItemOutput) => {
+          const spans = getTestSpans();
+          expect(spans.length).toStrictEqual(1);
+          const attrs = spans[0].attributes;
+          expect(attrs[SemanticAttributes.DB_SYSTEM]).toStrictEqual(
+            DbSystemValues.DYNAMODB
+          );
+          expect(attrs[SemanticAttributes.DB_OPERATION]).toStrictEqual(
+            'PutItem'
+          );
+          expect(attrs).not.toHaveProperty(
+            SemanticAttributes.AWS_DYNAMODB_CONSUMED_CAPACITY
+          );
+          expect(err).toBeFalsy();
+          done();
+        }
+      );
+    });
+  });
 });