chore(maintenance): add powertools to user-agent in SDK clients (#1567)

* add custom user agent middleware

* remove singleton for ddb client, add useragent in constructor

* remove conditional, because SDK will always resolve user-agent header

* simplify test

* remove retry, no longer needed

* remove ua from idempotency, will be done in separate PR

* review changes

* revert client lazy loading

* Revert "remove ua from idempotency, will be done in separate PR"

This reverts commit bdda143.

* revert the revert, misunderstanding

* add explicit ts-ignore instead of the entire file

* parameterized tests for useragent

* in case of failure, warn and don't block user code

* add test to check if middleware absorbs an error and continue

* add more client to useragent test

* fix imports

Author: Alexander Schueren

package-lock.json

@@ -14,6 +14,7 @@
     "test:unit": "jest --group=unit --detectOpenHandles --coverage --verbose",
     "test:e2e": "echo 'Not Applicable'",
     "watch": "jest --watch",
+    "generateVersionFile": "echo \"// this file is auto generated, do not modify\nexport const PT_VERSION = '$(jq -r '.version' package.json)';\" > src/version.ts",
     "build": "tsc",
     "lint": "eslint --ext .ts,.js --no-error-on-unmatched-pattern .",
     "lint-fix": "eslint --fix --ext .ts,.js --no-error-on-unmatched-pattern .",
@@ -47,6 +48,10 @@
   ],
   "devDependencies": {
     "@aws-sdk/client-lambda": "^3.360.0",
+    "@aws-sdk/client-dynamodb": "^3.360.0",
+    "@aws-sdk/client-appconfigdata": "^3.360.0",
+    "@aws-sdk/client-secrets-manager": "^3.360.0",
+    "@aws-sdk/client-ssm": "^3.360.0",
     "@aws-sdk/util-utf8-node": "^3.259.0"
   }
 }

packages/commons/package.json

@@ -5,3 +5,4 @@ export * as ContextExamples from './samples/resources/contexts';
 export * as Events from './samples/resources/events';
 export * from './types/middy';
 export * from './types/utils';
+export * from './userAgentMiddleware';

packages/commons/src/index.ts

@@ -0,0 +1,45 @@
+import { PT_VERSION } from './version';
+
+/**
+ * @internal
+ */
+const EXEC_ENV = process.env.AWS_EXECUTION_ENV || 'NA';
+const middlewareOptions = {
+  relation: 'after',
+  toMiddleware: 'getUserAgentMiddleware',
+  name: 'addPowertoolsToUserAgent',
+  tags: ['POWERTOOLS', 'USER_AGENT'],
+};
+
+/**
+ * @internal
+ * returns a middleware function for the MiddlewareStack, that can be used for the SDK clients
+ * @param feature
+ */
+const customUserAgentMiddleware = (feature: string) => {
+  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+  // @ts-ignore
+  return (next, _context) => async (args) => {
+    const powertoolsUserAgent = `PT/${feature}/${PT_VERSION} PTEnv/${EXEC_ENV}`;
+    args.request.headers[
+      'user-agent'
+    ] = `${args.request.headers['user-agent']} ${powertoolsUserAgent}`;
+
+    return await next(args);
+  };
+};
+
+// eslint-disable-next-line @typescript-eslint/ban-ts-comment
+// @ts-ignore
+const addUserAgentMiddleware = (client, feature: string): void => {
+  try {
+    client.middlewareStack.addRelativeTo(
+      customUserAgentMiddleware(feature),
+      middlewareOptions
+    );
+  } catch (e) {
+    console.warn('Failed to add user agent middleware', e);
+  }
+};
+
+export { addUserAgentMiddleware };

packages/commons/src/userAgentMiddleware.ts

@@ -0,0 +1,2 @@
+// this file is auto generated, do not modify
+export const PT_VERSION = '1.11.0';

packages/commons/src/version.ts

@@ -0,0 +1,109 @@
+import { addUserAgentMiddleware } from '../../src/userAgentMiddleware';
+import { InvokeCommand, LambdaClient } from '@aws-sdk/client-lambda';
+import { DynamoDBClient } from '@aws-sdk/client-dynamodb';
+import { ScanCommand } from '@aws-sdk/lib-dynamodb';
+import { GetParameterCommand, SSMClient } from '@aws-sdk/client-ssm';
+import { version as PT_VERSION } from '../../package.json';
+import { AppConfigDataClient } from '@aws-sdk/client-appconfigdata';
+import {
+  GetSecretValueCommand,
+  SecretsManagerClient,
+} from '@aws-sdk/client-secrets-manager';
+
+const options = {
+  region: 'us-east-1',
+  endpoint: 'http://localhost:9001',
+  credentials: {
+    accessKeyId: 'test',
+    secretAccessKey: 'test',
+    sessionToken: 'test',
+  },
+};
+
+describe('Given a client of instance: ', () => {
+  it.each([
+    {
+      name: 'LambdaClient',
+      client: new LambdaClient(options),
+      command: new InvokeCommand({ FunctionName: 'test', Payload: '' }),
+    },
+    {
+      name: 'DynamoDBClient',
+      client: new DynamoDBClient(options),
+      command: new ScanCommand({ TableName: 'test' }),
+    },
+    {
+      name: 'SSMClient',
+      client: new SSMClient(options),
+      command: new GetParameterCommand({ Name: 'test' }),
+    },
+    {
+      name: 'AppConfigDataClient',
+      client: new AppConfigDataClient(options),
+      command: new GetParameterCommand({ Name: 'test' }),
+    },
+    {
+      name: 'SecretsManagerClient',
+      client: new SecretsManagerClient(options),
+      command: new GetSecretValueCommand({ SecretId: 'test' }),
+    },
+  ])(
+    `using $name, add powertools user agent to request header at the end`,
+    async ({ client, command }) => {
+      addUserAgentMiddleware(client, 'my-feature');
+
+      expect(client.middlewareStack.identify()).toContain(
+        'addPowertoolsToUserAgent: POWERTOOLS,USER_AGENT'
+      );
+
+      client.middlewareStack.addRelativeTo(
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        (next) => (args) => {
+          const userAgent = args?.request?.headers['user-agent'];
+          expect(userAgent).toContain(`PT/my-feature/${PT_VERSION} PTEnv/NA`);
+          // make sure it's at the end of the user agent
+          expect(
+            // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+            // @ts-ignore
+            userAgent
+              ?.split(' ')
+              .slice(userAgent?.split(' ').length - 2) // take the last to entries of the user-agent header
+              .join(' ')
+          ).toEqual(`PT/my-feature/${PT_VERSION} PTEnv/NA`);
+
+          return next(args);
+        },
+        {
+          relation: 'after',
+          toMiddleware: 'addPowertoolsToUserAgent',
+          name: 'testUserAgentHeader',
+          tags: ['TEST'],
+        }
+      );
+
+      try {
+        // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+        // @ts-ignore
+        await client.send(command);
+      } catch (e) {
+        if (e instanceof Error && e.name === 'JestAssertionError') {
+          throw e;
+        }
+      }
+    }
+  );
+
+  it('should not throw erro, when client fails to add middleware', () => {
+    // create mock client that throws error when adding middleware
+    const client = {
+      middlewareStack: {
+        addRelativeTo: () => {
+          throw new Error('test');
+        },
+      },
+    };
+
+    expect(() => addUserAgentMiddleware(client, 'my-feature')).not.toThrow();
+  });
+});

packages/commons/tests/unit/userAgentMiddleware.test.ts

@@ -8,11 +8,11 @@ import {
 } from '../errors';
 import { IdempotencyRecord } from '../persistence';
 import { MAX_RETRIES } from '../constants';
-import type {
+import type { IdempotencyLambdaHandlerOptions } from '../types';
+import {
   MiddlewareLikeObj,
   MiddyLikeRequest,
 } from '@aws-lambda-powertools/commons';
-import type { IdempotencyLambdaHandlerOptions } from '../types';
 
 /**
  * A middy middleware to make your Lambda Handler idempotent.

packages/idempotency/src/middleware/makeHandlerIdempotent.ts

@@ -5,18 +5,19 @@ import {
 import { IdempotencyRecordStatus } from '../types';
 import type { DynamoDBPersistenceOptions } from '../types';
 import {
+  AttributeValue,
+  DeleteItemCommand,
   DynamoDBClient,
   DynamoDBClientConfig,
   DynamoDBServiceException,
-  DeleteItemCommand,
   GetItemCommand,
   PutItemCommand,
   UpdateItemCommand,
-  AttributeValue,
 } from '@aws-sdk/client-dynamodb';
 import { marshall, unmarshall } from '@aws-sdk/util-dynamodb';
 import { IdempotencyRecord } from './IdempotencyRecord';
 import { BasePersistenceLayer } from './BasePersistenceLayer';
+import { addUserAgentMiddleware } from '@aws-lambda-powertools/commons';
 
 /**
  * DynamoDB persistence layer for idempotency records. This class will use the AWS SDK V3 to write and read idempotency records from DynamoDB.
@@ -28,7 +29,7 @@ import { BasePersistenceLayer } from './BasePersistenceLayer';
  * @implements {BasePersistenceLayer}
  */
 class DynamoDBPersistenceLayer extends BasePersistenceLayer {
-  private client?: DynamoDBClient;
+  private client: DynamoDBClient;
   private clientConfig: DynamoDBClientConfig = {};
   private dataAttr: string;
   private expiryAttr: string;
@@ -64,18 +65,18 @@ class DynamoDBPersistenceLayer extends BasePersistenceLayer {
       if (config?.awsSdkV3Client instanceof DynamoDBClient) {
         this.client = config.awsSdkV3Client;
       } else {
-        console.warn(
-          'Invalid AWS SDK V3 client passed to DynamoDBPersistenceLayer. Using default client.'
-        );
+        throw Error('Not valid DynamoDBClient provided.');
       }
     } else {
       this.clientConfig = config?.clientConfig ?? {};
+      this.client = new DynamoDBClient(this.clientConfig);
     }
+
+    addUserAgentMiddleware(this.client, 'idempotency');
   }
 
   protected async _deleteRecord(record: IdempotencyRecord): Promise<void> {
-    const client = this.getClient();
-    await client.send(
+    await this.client.send(
       new DeleteItemCommand({
         TableName: this.tableName,
         Key: this.getKey(record.idempotencyKey),
@@ -86,8 +87,7 @@ class DynamoDBPersistenceLayer extends BasePersistenceLayer {
   protected async _getRecord(
     idempotencyKey: string
   ): Promise<IdempotencyRecord> {
-    const client = this.getClient();
-    const result = await client.send(
+    const result = await this.client.send(
       new GetItemCommand({
         TableName: this.tableName,
         Key: this.getKey(idempotencyKey),
@@ -111,8 +111,6 @@ class DynamoDBPersistenceLayer extends BasePersistenceLayer {
   }
 
   protected async _putRecord(record: IdempotencyRecord): Promise<void> {
-    const client = this.getClient();
-
     const item = {
       ...this.getKey(record.idempotencyKey),
       ...marshall({
@@ -163,7 +161,7 @@ class DynamoDBPersistenceLayer extends BasePersistenceLayer {
       ].join(' OR ');
 
       const now = Date.now();
-      await client.send(
+      await this.client.send(
         new PutItemCommand({
           TableName: this.tableName,
           Item: item,
@@ -195,8 +193,6 @@ class DynamoDBPersistenceLayer extends BasePersistenceLayer {
   }
 
   protected async _updateRecord(record: IdempotencyRecord): Promise<void> {
-    const client = this.getClient();
-
     const updateExpressionFields: string[] = [
       '#response_data = :response_data',
       '#expiry = :expiry',
@@ -219,7 +215,7 @@ class DynamoDBPersistenceLayer extends BasePersistenceLayer {
       expressionAttributeValues[':validation_key'] = record.payloadHash;
     }
 
-    await client.send(
+    await this.client.send(
       new UpdateItemCommand({
         TableName: this.tableName,
         Key: this.getKey(record.idempotencyKey),
@@ -230,14 +226,6 @@ class DynamoDBPersistenceLayer extends BasePersistenceLayer {
     );
   }
 
-  private getClient(): DynamoDBClient {
-    if (!this.client) {
-      this.client = new DynamoDBClient(this.clientConfig);
-    }
-
-    return this.client;
-  }
-
   /**
    * Build primary key attribute simple or composite based on params.
    *

packages/idempotency/src/persistence/DynamoDBPersistenceLayer.ts

@@ -163,20 +163,14 @@ describe('Class: DynamoDBPersistenceLayer', () => {
       );
     });
 
-    test('when passed an invalid AWS SDK client it logs a warning', () => {
-      // Prepare
-      const consoleWarnSpy = jest.spyOn(console, 'warn').mockImplementation();
-
-      // Act
-      new TestDynamoDBPersistenceLayer({
-        tableName: dummyTableName,
-        awsSdkV3Client: {} as DynamoDBClient,
-      });
-
-      // Assess
-      expect(consoleWarnSpy).toHaveBeenCalledWith(
-        'Invalid AWS SDK V3 client passed to DynamoDBPersistenceLayer. Using default client.'
-      );
+    test('when passed an invalid AWS SDK client, it throws an error', () => {
+      // Act & Assess
+      expect(() => {
+        new TestDynamoDBPersistenceLayer({
+          tableName: dummyTableName,
+          awsSdkV3Client: {} as DynamoDBClient,
+        });
+      }).toThrow();
     });
 
     test('when passed a client config it stores it for later use', () => {