From c35eb4c534c42e7db5007deb1296ca6e9a5eb8cf Mon Sep 17 00:00:00 2001
From: Nicholas Griffin <nicholas.griffin@bbc.co.uk>
Date: Tue, 1 Oct 2024 21:08:38 +0100
Subject: [PATCH] feat: adding MessageSystemAttributeNames (#518)

---
 src/consumer.ts             |  5 +++++
 src/types.ts                | 12 +++++++++++-
 test/tests/consumer.test.ts |  3 +++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/consumer.ts b/src/consumer.ts
index 4fce267..c70a5c9 100644
--- a/src/consumer.ts
+++ b/src/consumer.ts
@@ -15,6 +15,7 @@ import {
   ReceiveMessageCommandInput,
   ReceiveMessageCommandOutput,
   QueueAttributeName,
+  MessageSystemAttributeName,
 } from "@aws-sdk/client-sqs";
 
 import { ConsumerOptions, StopOptions, UpdatableOptions } from "./types.js";
@@ -45,6 +46,7 @@ export class Consumer extends TypedEventEmitter {
   private handleMessageTimeout: number;
   private attributeNames: QueueAttributeName[];
   private messageAttributeNames: string[];
+  private messageSystemAttributeNames: MessageSystemAttributeName[];
   private shouldDeleteMessages: boolean;
   private alwaysAcknowledge: boolean;
   private batchSize: number;
@@ -71,6 +73,8 @@ export class Consumer extends TypedEventEmitter {
     this.handleMessageTimeout = options.handleMessageTimeout;
     this.attributeNames = options.attributeNames || [];
     this.messageAttributeNames = options.messageAttributeNames || [];
+    this.messageSystemAttributeNames =
+      options.messageSystemAttributeNames || [];
     this.batchSize = options.batchSize || 1;
     this.visibilityTimeout = options.visibilityTimeout;
     this.terminateVisibilityTimeout =
@@ -242,6 +246,7 @@ export class Consumer extends TypedEventEmitter {
       QueueUrl: this.queueUrl,
       AttributeNames: this.attributeNames,
       MessageAttributeNames: this.messageAttributeNames,
+      MessageSystemAttributeNames: this.messageSystemAttributeNames,
       MaxNumberOfMessages: this.batchSize,
       WaitTimeSeconds: this.waitTimeSeconds,
       VisibilityTimeout: this.visibilityTimeout,
diff --git a/src/types.ts b/src/types.ts
index 075407e..aeec5c8 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -1,4 +1,9 @@
-import { SQSClient, Message, QueueAttributeName } from "@aws-sdk/client-sqs";
+import {
+  SQSClient,
+  Message,
+  QueueAttributeName,
+  MessageSystemAttributeName,
+} from "@aws-sdk/client-sqs";
 
 /**
  * The options for the consumer.
@@ -18,6 +23,11 @@ export interface ConsumerOptions {
    * @defaultvalue `[]`
    */
   messageAttributeNames?: string[];
+  /**
+   * A list of attributes that need to be returned along with each message.
+   * @defaultvalue `[]`
+   */
+  messageSystemAttributeNames?: MessageSystemAttributeName[];
   /** @hidden */
   stopped?: boolean;
   /**
diff --git a/test/tests/consumer.test.ts b/test/tests/consumer.test.ts
index d70219e..df81845 100644
--- a/test/tests/consumer.test.ts
+++ b/test/tests/consumer.test.ts
@@ -926,6 +926,7 @@ describe("Consumer", () => {
       consumer = new Consumer({
         queueUrl: QUEUE_URL,
         messageAttributeNames: ["attribute-1", "attribute-2"],
+        messageSystemAttributeNames: ["All"],
         region: REGION,
         handleMessage,
         batchSize: 3,
@@ -944,6 +945,7 @@ describe("Consumer", () => {
           QueueUrl: QUEUE_URL,
           AttributeNames: [],
           MessageAttributeNames: ["attribute-1", "attribute-2"],
+          MessageSystemAttributeNames: ["All"],
           MaxNumberOfMessages: 3,
           WaitTimeSeconds: AUTHENTICATION_ERROR_TIMEOUT,
           VisibilityTimeout: undefined,
@@ -988,6 +990,7 @@ describe("Consumer", () => {
           QueueUrl: QUEUE_URL,
           AttributeNames: ["ApproximateReceiveCount"],
           MessageAttributeNames: [],
+          MessageSystemAttributeNames: [],
           MaxNumberOfMessages: 1,
           WaitTimeSeconds: AUTHENTICATION_ERROR_TIMEOUT,
           VisibilityTimeout: undefined,