feat(plugin-aws-sdk): add custom attributes hook

Author: Amir Blum

packages/plugin-aws-sdk/README.md

@@ -25,3 +25,11 @@ const provider = new NodeTracerProvider({
   },
 });
 ```
+
+### aws-sdk Plugin Options
+
+aws-sdk plugin has few options available to choose from. You can set the following:
+
+| Options        | Type                                   | Description                                                                                     |
+| -------------- | -------------------------------------- | ----------------------------------------------------------------------------------------------- |
+| `preRequestHook` | `AwsSdkRequestCustomAttributeFunction` | Hook called before request send, which allow to add custom attributes to span. |

packages/plugin-aws-sdk/src/aws-sdk.ts

@@ -9,19 +9,21 @@
     callback    |       1           |       2   
  */
 import { BasePlugin } from "@opentelemetry/core";
-import { Span } from "@opentelemetry/api";
+import { Span, CanonicalCode } from "@opentelemetry/api";
 import * as shimmer from "shimmer";
 import AWS from "aws-sdk";
 import { AttributeNames } from "./enums";
 import {
   getRequestServiceAttributes,
   getResponseServiceAttributes,
 } from "./services";
+import { AwsSdkPluginConfig } from "./types";
 
 const VERSION = "0.0.3";
 
 class AwsPlugin extends BasePlugin<typeof AWS> {
   readonly component: string;
+  protected _config: AwsSdkPluginConfig;
   private activeRequests: Set<AWS.Request<any, any>> = new Set();
 
   constructor(readonly moduleName: string) {
@@ -60,13 +62,12 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
 
   private _getPatchedRequestMethod = (original: Function) => {
     const thisPlugin = this;
-
     return function () {
       let span: Span | null = null;
       /* 
-               if the span was already started, we don't want to start a new one 
-               when Request.promise() is called
-            */
+        if the span was already started, we don't want to start a new one 
+        when Request.promise() is called
+      */
 
       if (
         this._asm.currentState !== "complete" &&
@@ -78,7 +79,6 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
           attributes: {
             [AttributeNames.COMPONENT]: thisPlugin.moduleName,
             [AttributeNames.AWS_OPERATION]: this.operation,
-            [AttributeNames.AWS_PARAMS]: JSON.stringify(this.params),
             [AttributeNames.AWS_SIGNATURE_VERSION]: this.service?.config
               ?.signatureVersion,
             [AttributeNames.AWS_REGION]: this.service?.config?.region,
@@ -90,6 +90,14 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
           },
         });
 
+        if (thisPlugin._config?.preRequestHook) {
+          thisPlugin._safeExecute(
+            span,
+            () => thisPlugin._config.preRequestHook(span, this),
+            false
+          );
+        }
+
         (this as AWS.Request<any, any>).on("complete", (response) => {
           if (thisPlugin.activeRequests.has(this)) {
             thisPlugin.activeRequests.delete(this);
@@ -121,6 +129,33 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
       request.operation
     }`;
   };
+
+  private _safeExecute<
+    T extends (...args: unknown[]) => ReturnType<T>,
+    K extends boolean
+  >(
+    span: Span,
+    execute: T,
+    rethrow: K
+  ): K extends true ? ReturnType<T> : ReturnType<T> | void;
+  private _safeExecute<T extends (...args: unknown[]) => ReturnType<T>>(
+    span: Span,
+    execute: T,
+    rethrow: boolean
+  ): ReturnType<T> | void {
+    try {
+      return execute();
+    } catch (error) {
+      if (rethrow) {
+        span.setStatus({
+          code: CanonicalCode.UNKNOWN,
+        });
+        span.end();
+        throw error;
+      }
+      this._logger.error("caught error ", error);
+    }
+  }
 }
 
 export const plugin = new AwsPlugin("aws-sdk");

packages/plugin-aws-sdk/src/enums.ts

@@ -8,6 +8,5 @@ export enum AttributeNames {
   AWS_SERVICE_NAME = "aws.service.name",
   AWS_SERVICE_IDENTIFIER = "aws.service.identifier",
   AWS_REQUEST_ID = "aws.request.id",
-  AWS_PARAMS = "aws.params",
   AWS_SIGNATURE_VERSION = "aws.signature.version",
 }

packages/plugin-aws-sdk/src/types.ts

@@ -0,0 +1,11 @@
+import { PluginConfig, Span } from "@opentelemetry/api";
+import AWS from "aws-sdk";
+
+export interface AwsSdkRequestCustomAttributeFunction {
+  (span: Span, request: AWS.Request<any, any>): void;
+}
+
+export interface AwsSdkPluginConfig extends PluginConfig {
+  /** hook for adding custom attributes before producer message is sent */
+  preRequestHook?: AwsSdkRequestCustomAttributeFunction;
+}