feat(core): add support for logging stats when consumed capacity is t…

…urned on
typedorm · May 20, 2021 · b179c48 · b179c48
1 parent d3dc72b
commit b179c48
Show file tree

Hide file tree

Showing 11 changed files with 218 additions and 18 deletions.
diff --git a/packages/common/src/enums.ts b/packages/common/src/enums.ts
@@ -32,3 +32,9 @@ export enum AUTO_GENERATE_ATTRIBUTE_STRATEGY {
 export enum INTERNAL_ENTITY_ATTRIBUTE {
   ENTITY_NAME = '__en',
 }
+
+export enum CONSUMED_CAPACITY_TYPE {
+  INDEXES = 'INDEXES',
+  TOTAL = 'TOTAL',
+  NONE = 'NONE',
+}
diff --git a/packages/common/src/logger/debug-logger.ts b/packages/common/src/logger/debug-logger.ts
@@ -169,16 +169,22 @@ export class DebugLogger {
 
   logStats({
     requestId,
+    requestSegment,
     statsType,
     consumedCapacityData,
   }: {
     requestId?: string;
+    requestSegment?: number; // if for any reason, request was segmented into multiple request, this indicates segment index
     statsType: STATS_TYPE;
     consumedCapacityData: any;
   }) {
     if (this.debugStatsLog.enabled) {
       this.debugStatsLog(
-        `${chalk.bold.bgCyanBright(requestId)} ${chalk.green(statsType)}:`,
+        `${
+          chalk.bold.bgCyanBright(requestId) + requestSegment
+            ? ':' + chalk.bold.bgCyanBright(requestId)
+            : ''
+        } ${chalk.green(statsType)}:`,
         chalk.white(this.ensurePrintable(consumedCapacityData))
       );
     }

diff --git a/packages/core/src/classes/manager/__test__/batch-manager.spec.ts b/packages/core/src/classes/manager/__test__/batch-manager.spec.ts
@@ -1,4 +1,4 @@
-import {Replace} from '@typedorm/common';
+import {CONSUMED_CAPACITY_TYPE, Replace} from '@typedorm/common';
 import {User} from '@typedorm/core/__mocks__/user';
 import {UserUniqueEmail} from '@typedorm/core/__mocks__/user-unique-email';
 import {createTestConnection, resetTestConnection} from '@typedorm/testing';
@@ -371,7 +371,13 @@ test('processes simple batch read request', async () => {
       },
     },
   ]);
-  const response = await manager.read(readTestBatch);
+  const response = await manager.read(
+    readTestBatch,
+    {},
+    {
+      returnConsumedCapacity: CONSUMED_CAPACITY_TYPE.TOTAL,
+    }
+  );
 
   expect(originalPromiseAll).toHaveBeenCalledTimes(1);
   expect(originalPromiseAll).toHaveBeenCalledWith([expect.any(Promise)]);

diff --git a/packages/core/src/classes/manager/__test__/entity-manager.spec.ts b/packages/core/src/classes/manager/__test__/entity-manager.spec.ts
@@ -12,6 +12,7 @@ import {
   UserAutoGenerateAttributes,
 } from '../../../../__mocks__/user-auto-generate-attributes';
 import {Connection} from '../../connection/connection';
+import {CONSUMED_CAPACITY_TYPE} from '@typedorm/common';
 
 let manager: EntityManager;
 let connection: Connection;
@@ -49,7 +50,9 @@ test('creates entity', async () => {
   user.name = 'Test User';
   user.status = 'active';
 
-  const userEntity = await manager.create(user);
+  const userEntity = await manager.create(user, undefined, {
+    returnConsumedCapacity: CONSUMED_CAPACITY_TYPE.TOTAL,
+  });
   expect(dcMock.put).toHaveBeenCalledTimes(1);
   expect(dcMock.put).toHaveBeenCalledWith({
     Item: {
@@ -63,6 +66,7 @@ test('creates entity', async () => {
       status: 'active',
     },
     TableName: 'test-table',
+    ReturnConsumedCapacity: 'TOTAL',
     ConditionExpression:
       '(attribute_not_exists(#CE_PK)) AND (attribute_not_exists(#CE_SK))',
     ExpressionAttributeNames: {
@@ -285,6 +289,9 @@ test('checks if given item exists', async () => {
     User,
     {
       id: '1',
+    },
+    {
+      returnConsumedCapacity: CONSUMED_CAPACITY_TYPE.INDEXES,
     }
   );
 
@@ -294,6 +301,7 @@ test('checks if given item exists', async () => {
       SK: 'USER#1',
     },
     TableName: 'test-table',
+    ReturnConsumedCapacity: 'INDEXES',
   });
   expect(userEntity).toEqual(true);
 });

diff --git a/packages/core/src/classes/manager/__test__/transaction-manager.spec.ts b/packages/core/src/classes/manager/__test__/transaction-manager.spec.ts
@@ -10,6 +10,7 @@ import {
 } from '@typedorm/core/__mocks__/user-unique-email';
 import {Organisation} from '@typedorm/core/__mocks__/organisation';
 import {ReadTransaction} from '../../transaction/read-transaction';
+import {CONSUMED_CAPACITY_TYPE} from '@typedorm/common';
 
 let manager: TransactionManager;
 const dcMock = {
@@ -79,10 +80,13 @@ test('performs write transactions for simple writes', async () => {
       },
     ]);
 
-  const response = await manager.write(transaction);
+  const response = await manager.write(transaction, {
+    returnConsumedCapacity: CONSUMED_CAPACITY_TYPE.TOTAL,
+  });
 
   expect(dcMock.transactWrite).toHaveBeenCalledTimes(1);
   expect(dcMock.transactWrite).toHaveBeenCalledWith({
+    ReturnConsumedCapacity: 'TOTAL',
     TransactItems: [
       {
         Put: {

diff --git a/packages/core/src/classes/manager/batch-manager.ts b/packages/core/src/classes/manager/batch-manager.ts
@@ -8,6 +8,7 @@ import {
   BATCH_WRITE_MAX_ALLOWED_ATTEMPTS,
   INTERNAL_ENTITY_ATTRIBUTE,
   MANAGER_NAME,
+  STATS_TYPE,
 } from '@typedorm/common';
 import {
   BatchGetResponseMap,
@@ -110,6 +111,7 @@ export class BatchManager {
           () =>
             this.connection.transactionManger.writeRaw(transformedInput, {
               requestId,
+              returnConsumedCapacity: metadataOptions?.returnConsumedCapacity,
             }),
           // return original item when failed to process
           rawInput,
@@ -130,6 +132,7 @@ export class BatchManager {
               undefined,
               {
                 requestId,
+                returnConsumedCapacity: metadataOptions?.returnConsumedCapacity,
               }
             );
 
@@ -151,6 +154,7 @@ export class BatchManager {
               deleteTransactionItemList,
               {
                 requestId,
+                returnConsumedCapacity: metadataOptions?.returnConsumedCapacity,
               }
             );
           },
@@ -169,6 +173,7 @@ export class BatchManager {
           this.connection.documentClient
             .batchWrite({
               RequestItems: {...batchRequestMap},
+              ReturnConsumedCapacity: metadataOptions?.returnConsumedCapacity,
             })
             .promise(),
         // for batch requests this returning item will be transformed to
@@ -189,6 +194,18 @@ export class BatchManager {
     // 2. wait for all promises to finish
     const responses = await Promise.all(allRequests);
 
+    // log stats
+    responses.forEach((response, index) => {
+      if (response.ConsumedCapacity) {
+        this.connection.logger.logStats({
+          requestId,
+          requestSegment: index,
+          statsType: STATS_TYPE.CONSUMED_CAPACITY,
+          consumedCapacityData: response.ConsumedCapacity,
+        });
+      }
+    });
+
     // 3. run retry attempts
     // process all unprocessed items recursively until all are either done
     // or reached the retry limit
@@ -198,6 +215,7 @@ export class BatchManager {
       options,
       {
         requestId,
+        returnConsumedCapacity: metadataOptions?.returnConsumedCapacity,
       }
     );
 
@@ -272,6 +290,7 @@ export class BatchManager {
           this.connection.documentClient
             .batchGet({
               RequestItems: {...batchRequestItems},
+              ReturnConsumedCapacity: metadataOptions?.returnConsumedCapacity,
             })
             .promise(),
         batchRequestItems,
@@ -282,6 +301,18 @@ export class BatchManager {
     // 2. wait for all promises to finish, either failed or hit the limit
     const initialResponses = await Promise.all(batchRequests);
 
+    // log stats
+    initialResponses.forEach((response, index) => {
+      if (response.ConsumedCapacity) {
+        this.connection.logger.logStats({
+          requestId,
+          requestSegment: index,
+          statsType: STATS_TYPE.CONSUMED_CAPACITY,
+          consumedCapacityData: response.ConsumedCapacity,
+        });
+      }
+    });
+
     // 3. run retries
     const {
       items,
@@ -293,6 +324,7 @@ export class BatchManager {
       [],
       {
         requestId,
+        returnConsumedCapacity: metadataOptions?.returnConsumedCapacity,
       }
     );
 
@@ -416,6 +448,8 @@ export class BatchManager {
           this.connection.documentClient
             .batchWrite({
               RequestItems: {...batchRequestMap},
+              ReturnItemCollectionMetrics:
+                metadataOptions?.returnConsumedCapacity,
             })
             .promise(),
         batchRequestMap,
@@ -426,6 +460,19 @@ export class BatchManager {
     const batchRequestsResponses = (await Promise.all(
       batchRequests
     )) as DocumentClient.BatchWriteItemOutput[];
+
+    // log stats
+    batchRequestsResponses.forEach((response, index) => {
+      if (response.ConsumedCapacity) {
+        this.connection.logger.logStats({
+          requestId: metadataOptions?.requestId,
+          requestSegment: index,
+          statsType: STATS_TYPE.CONSUMED_CAPACITY,
+          consumedCapacityData: response.ConsumedCapacity,
+        });
+      }
+    });
+
     return this.recursiveHandleBatchWriteItemsResponse(
       batchRequestsResponses,
       ++totalAttemptsSoFar,
@@ -529,6 +576,7 @@ export class BatchManager {
           this.connection.documentClient
             .batchGet({
               RequestItems: {...batchRequestMap},
+              ReturnConsumedCapacity: metadataOptions?.returnConsumedCapacity,
             })
             .promise(),
         batchRequestMap,
@@ -540,6 +588,18 @@ export class BatchManager {
       batchRequests
     )) as DocumentClient.BatchGetItemOutput[];
 
+    // log stats
+    batchRequestsResponses.forEach((response, index) => {
+      if (response.ConsumedCapacity) {
+        this.connection.logger.logStats({
+          requestId: metadataOptions?.requestId,
+          requestSegment: index,
+          statsType: STATS_TYPE.CONSUMED_CAPACITY,
+          consumedCapacityData: response.ConsumedCapacity,
+        });
+      }
+    });
+
     return this.recursiveHandleBatchReadItemsResponse(
       batchRequestsResponses,
       ++totalAttemptsSoFar,