refactor(plugin-aws-sdk): make service extension a class

Author: Amir Blum

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

@@ -13,10 +13,7 @@ import { Span, CanonicalCode, Attributes } from "@opentelemetry/api";
 import * as shimmer from "shimmer";
 import AWS from "aws-sdk";
 import { AttributeNames } from "./enums";
-import {
-  getRequestServiceAttributes,
-  getResponseServiceAttributes,
-} from "./services";
+import { ServicesExtensions } from "./services";
 import { AwsSdkPluginConfig } from "./types";
 
 const VERSION = "0.0.4";
@@ -25,12 +22,15 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
   readonly component: string;
   protected _config: AwsSdkPluginConfig;
   private REQUEST_SPAN_KEY = Symbol("opentelemetry.plugin.aws-sdk.span");
+  private servicesExtensions: ServicesExtensions;
 
   constructor(readonly moduleName: string) {
     super(`opentelemetry-plugin-aws-sdk`, VERSION);
   }
 
   protected patch() {
+    this.servicesExtensions = new ServicesExtensions(this._tracer);
+
     this._logger.debug(
       "applying patch to %s@%s",
       this.moduleName,
@@ -112,11 +112,7 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
         span.setAttribute(AttributeNames.AWS_ERROR, response.error);
       }
 
-      const responseAttributes = getResponseServiceAttributes(
-        response,
-        span,
-        thisPlugin._tracer
-      );
+      this.servicesExtensions.responseHook(response, span);
 
       span.setAttributes({
         [AttributeNames.AWS_REQUEST_ID]: response.requestId,
@@ -143,7 +139,7 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
       }
 
       const span = thisPlugin._startAwsSpan(awsRequest);
-      const requestMetadata = getRequestServiceAttributes(awsRequest, span);
+      thisPlugin.servicesExtensions.requestHook(awsRequest, span);
       thisPlugin._callPreRequestHooks(span, awsRequest);
       thisPlugin._registerCompletedEvent(span, awsRequest);
 
@@ -170,7 +166,10 @@ class AwsPlugin extends BasePlugin<typeof AWS> {
       }
 
       const span = thisPlugin._startAwsSpan(awsRequest);
-      const requestMetadata = getRequestServiceAttributes(awsRequest, span);
+      const requestMetadata = thisPlugin.servicesExtensions.requestHook(
+        awsRequest,
+        span
+      );
       thisPlugin._callPreRequestHooks(span, awsRequest);
       thisPlugin._registerCompletedEvent(span, awsRequest);
 

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

@@ -0,0 +1,13 @@
+import { Span } from "@opentelemetry/api";
+
+/*
+isIncoming - if true, then the operation callback / promise should be bind with the operation's span 
+*/
+export interface RequestMetadata {
+  isIncoming: boolean;
+}
+
+export interface ServiceExtension {
+  requestHook: (request: AWS.Request<any, any>, span: Span) => RequestMetadata;
+  responseHook: (response: AWS.Response<any, any>, span: Span) => void;
+}

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

@@ -0,0 +1,26 @@
+import { Tracer, Span } from "@opentelemetry/api";
+import { ServiceExtension, RequestMetadata } from "./ServiceExtension";
+import { SqsServiceExtension } from "./sqs";
+import * as AWS from "aws-sdk";
+
+export class ServicesExtensions implements ServiceExtension {
+  services: Map<string, ServiceExtension> = new Map();
+
+  constructor(tracer: Tracer) {
+    this.services.set("sqs", new SqsServiceExtension(tracer));
+  }
+
+  requestHook(request: AWS.Request<any, any>, span: Span): RequestMetadata {
+    const serviceId = (request as any)?.service?.serviceIdentifier;
+    const serviceExtension = this.services.get(serviceId);
+    if (!serviceExtension) return { isIncoming: false };
+    return serviceExtension.requestHook(request, span);
+  }
+
+  responseHook(response: AWS.Response<any, any>, span: Span) {
+    const serviceId = (response as any)?.request?.service?.serviceIdentifier;
+    const serviceExtension = this.services.get(serviceId);
+    if (!serviceExtension) return;
+    serviceExtension.responseHook(response, span);
+  }
+}

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

@@ -1,4 +1 @@
-export {
-  getRequestServiceAttributes,
-  getResponseServiceAttributes,
-} from "./service-attributes";
+export { ServicesExtensions } from "./ServicesExtensions";