Class AWS2SQSPlugin

src/plugins/AWS2SQSPlugin.ts:32–174

import { getAWS, execute } from '../aws/SDK2';

class AWS2SQSPlugin implements SwPlugin {
  readonly module = 'aws-sdk';
  readonly versions = '2.*';

  install(installer: PluginInstaller): void {
    const AWS = getAWS(installer);
    const _SQS = AWS.SQS;

    function SQS(this: any) {
      const sqs = _SQS.apply(this, arguments);

      function instrumentSend(name: string, addTraceId: any): void {
        const _func = sqs[name];

        sqs[name] = function (params: any, callback: any) {
          const queueUrl = params.QueueUrl;
          const operation = `AWS/SQS/${name}/${queueUrl.slice(queueUrl.lastIndexOf('/') + 1)}`;
          const span = ContextManager.current.newExitSpan(operation, Component.AWS_SQS, Component.HTTP);

          span.component = Component.AWS_SQS;
          span.layer = SpanLayer.MQ;

          return execute(span, this, _func, addTraceId(params, span), callback, 'mqBroker');
        };
      }

      instrumentSend('sendMessage', (params: any, span: Span) => {
        params = Object.assign({}, params);
        params.MessageAttributes = params.MessageAttributes ? Object.assign({}, params.MessageAttributes) : {};
        params.MessageAttributes.__revdTraceId = {
          DataType: 'String',
          StringValue: `${span.inject().value}/${hostname()}`,
        };

        return params;
      });

      instrumentSend('sendMessageBatch', (params: any, span: Span) => {
        const traceId = { __revdTraceId: { DataType: 'String', StringValue: `${span.inject().value}/${hostname()}` } };
        params = Object.assign({}, params);
        params.Entries = params.Entries.map(
          (e: any) =>
            (e = Object.assign({}, e, {
              MessageAttributes: e.MessageAttributes ? Object.assign({}, e.MessageAttributes, traceId) : traceId,
            })),
        );

        return params;
      });

      const _receiveMessage = sqs.receiveMessage;

      sqs.receiveMessage = function (params: any, callback: any) {
        params = Object.assign({}, params);
        const _MessageAttributeNames = params.MessageAttributeNames;
        params.MessageAttributeNames = _MessageAttributeNames
          ? _MessageAttributeNames.concat(['__revdTraceId'])
          : ['__revdTraceId'];