feat: Add capacity information when applicable to dynamodb spans

open-telemetry · May 5, 2023 · 2d2d098 · 2d2d098
1 parent 69d1f17
commit 2d2d098
Show file tree

Hide file tree

Showing 2 changed files with 74 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
@@ -58,20 +58,20 @@ 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)
+        )
+      );
     }
   }
 }
+
+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
@@ -242,4 +242,65 @@ 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();
+      });
+    });
+  });
 });