updating receiveFromSQS.py

slack-bot/traffic-scraper/prd/receiveFromSQS.py

@@ -11,6 +11,7 @@ class receiveFromSQS(object):
         self.session = session
         self.sqs = session.client('sqs')
         self.queueURL = queueURL
+        self.messages = []
 
     @classmethod
     def createSession(cls, profileName, queueURL):
@@ -35,7 +36,7 @@ class receiveFromSQS(object):
     ):
         self.response = []
         loops = int(math.ceil(totalNumberOfMessages / maxNumberOfMessages))
-        for i in range(0, loops):
+        for i in range(0, loops * 2):
             self.response.append(
                 self.sqs.receive_message(
                     QueueUrl=self.queueURL,
@@ -44,55 +45,58 @@ class receiveFromSQS(object):
             )
         return self
 
-    def _extractMessageFromSQS(self, totalNumberOfMessages):
-        self.messages = []
+    def _extractMessageFromSQS(self, totalNumberOfMessages, deleteOnReceipt):
+        self.extractedMessages = []
+        self.receiptHandles = []
         try:
             loops = len(self.response)
             for i in range(0, loops):
                 if 'Messages' in self.response[i]:
                     _message = self.response[i]['Messages']
                     for rawMessage in _message:
-                        self.messages.append(rawMessage['Body'])
+                        self.extractedMessages.append(rawMessage['Body'])
+                        # self.receiptHandles.append(rawMessage['ReceiptHandle'])
+                        if deleteOnReceipt:
+                            self.sqs.delete_message(
+                                QueueUrl=self.queueURL,
+                                ReceiptHandle=rawMessage['ReceiptHandle'],
+                            )
                 else:
-                    print("No messages in the queue")
+                    print('No messages in the queue')
         except KeyError:
-            print("No messages in the queue")
+            print('No messages in the queue')
         return self
 
-    # def receiveMessage(
-    #     self, b64=True, MaxNumberOfMessages=10, delete=False, endString='END'
-    # ):
-    #     while True:
-    #         response = self.sqs.receive_message(
-    #             QueueUrl=self.queueURL,
-    # MaxNumberOfMessages=MaxNumberOfMessages
-    #         )
-    #         if 'Messages' in response:
-    #             for message in response['Messages']:
-    #                 if b64:
-    #                     decoded = base64.b64decode(message['Body'])
-    #                     msg = decoded.decode()
-    #                     if msg == endString:
-    #                         return self
-    #                 else:
-    #                     msg = print(message['Body'])
-    #                     if msg == endString:
-    #                         return self
-    #         else:
-    #             print('Queue is empty')
-    #             break
+    def _decodeMessages(self):
+        if len(self.extractedMessages) == 0:
+            print('No messages to process')
+        else:
+            for message in self.extractedMessages:
+                decoded = base64.b64decode(message).decode()
+                self.messages.append(decoded)
+        return self
+
+    def _deleteMessage(self, ReceiptHandle):
+        self.sqs.delete_message(
+            QueueUrl=self.queueURL, ReceiptHandle=ReceiptHandle
+        )
 
     def receiveAllMessages(self, b64=True, deleteOnReceipt=False):
         totalNumberOfMessages = self.getQueueLength()
         self._receiveSQSMessage(
-            totalNumberOfMessages,
-            deleteOnReceipt=deleteOnReceipt
+            totalNumberOfMessages, deleteOnReceipt=deleteOnReceipt
         )
         self._extractMessageFromSQS(
-            totalNumberOfMessages=totalNumberOfMessages
+            totalNumberOfMessages=totalNumberOfMessages,
+            deleteOnReceipt=deleteOnReceipt
         )
         if b64:
-            pass  # WRITE THIS IN!
+            self._decodeMessages()
+        else:
+            self.messages = self.extractedMessages
+        # if deleteOnReceipt:
+        #     for receipt in self.receiptHandles:
+        #         self._deleteMessage(receipt)
         return self
 
 
@@ -102,7 +106,7 @@ inst = receiveFromSQS.createSession(
     '.com/745437999005/slack-bot.fifo',
 )
 
-inst.receiveAllMessages()
+inst.receiveAllMessages(deleteOnReceipt=False)
 for item in inst.messages:
     print(item)