sns_publish_batch.py

import uuid
from dataclasses import dataclass
from itertools import islice
import json
import boto3
# pyrefly: ignore [missing-import]
from tenacity import Retrying, stop_after_attempt, wait_exponential
from typing import Iterable, Iterator, List, Dict, Any, Tuple


@dataclass(frozen=True)
class MessagePayload:
    patient_id: str
    value: int

    def to_json_string(self) -> str:
        """Converte o payload em uma string JSON."""
        return json.dumps(self.__dict__)


class SNSBatchPublisher:
    """
    Class responsible for managing batch sending to AWS SNS.
    """
    def __init__(self, topic_arn: str, region_name: str = 'us-east-1') -> None:
        self.topic_arn = topic_arn
        self.sns_client = boto3.client('sns', region_name=region_name)
        # Store pending items for the current block being processed
        self.pending_items: List[Any] = []
        # Store all definitively failed items across multiple blocks
        self.failed_items: List[Any] = []

    @staticmethod
    def _chunk_messages(iterable: Iterable[Any], size: int) -> Iterator[List[Any]]:
        """Divides an iterable into chunks of fixed size."""
        iterator = iter(iterable)
        while chunk := list(islice(iterator, size)):
            yield chunk

    def process_block_with_strict_retry(self, initial_block: List[Any]) -> List[str]:
        """
        Controls the lifecycle of a block of up to 10 messages.
        Ensures that ONLY the items that fail enter the next attempt.
        Returns the list of success_ids.
        Unprocessed items will be left in self.pending_items.
        """
        self.pending_items = list(initial_block)
        total_successes: List[str] = []

        # Manually configures the Tenacity engine to manage the 3 attempts and wait time
        config_retry = Retrying(
            stop=stop_after_attempt(3),
            wait=wait_exponential(multiplier=1, min=2, max=10),
            reraise=True # If the 3 attempts are exhausted, raises the exception for the block
        )

        try:
            # The loop runs controlling the state of what is left
            for attempt in config_retry:
                with attempt:
                    if not self.pending_items:
                        return total_successes

                    print(f" -> Attempting to send {len(self.pending_items)} pending items...")

                    # Prepares the batch structure with unique IDs for SNS
                    request_entries: List[Dict[str, str]] = []
                    map_id_to_msg: Dict[str, Any] = {}
                    
                    for msg in self.pending_items:
                        msg_id = str(uuid.uuid4())[:8]
                        map_id_to_msg[msg_id] = msg
                        
                        # Use to_json_string() if available, else convert to string
                        if hasattr(msg, 'to_json_string'):
                            content = msg.to_json_string()
                        else:
                            content = str(msg)
                            
                        request_entries.append({"Id": msg_id, "Message": content})

                    try:
                        # API call
                        response = self.sns_client.publish_batch(
                            TopicArn=self.topic_arn,
                            PublishBatchRequestEntries=request_entries
                        )

                        # 1. Store what succeeded
                        if 'Successful' in response:
                            total_successes.extend([s['MessageId'] for s in response['Successful']])

                        # 2. Discover what partially failed for the next attempt
                        if 'Failed' in response and response['Failed']:
                            # Replace the pending items ONLY with the messages that failed
                            self.pending_items = [map_id_to_msg[f['Id']] for f in response['Failed']]
                            
                            # Raise a generic exception to FORCE the Retrying object to trigger the next loop and backoff
                            raise RuntimeWarning("Partial failures detected in the SNS batch.")
                        else:
                            # If there was no failure in the JSON, we clear the block queue
                            self.pending_items = []

                    except Exception as e:
                        # If it falls here, it could be our RuntimeWarning OR a raw network error (boto3)
                        if isinstance(e, RuntimeWarning):
                            raise # Just pass it on so Tenacity applies the partial failure delay
                        
                        # If it is a network/AWS error, pending_items remains the same, and it will try the whole block again
                        print(f"⚠️ Network/AWS error on attempt: {e}")
                        raise

            return total_successes

        except Exception as e:
            # Catches the final exception raised by tenacity when attempts are exhausted
            print(f"⚠️ Block exhausted attempts! Error: {e}")
            return total_successes

    def publish_messages(self, messages: Iterable[Any], block_size: int = 10) -> List[str]:
        """
        Publishes a list of messages by dividing them into chunks.
        Returns a list of all success_ids. 
        Definitively failed items can be accessed via self.failed_items.
        """
        success_ids: List[str] = []
        self.failed_items = []
        
        for block in self._chunk_messages(messages, size=block_size):
            successes = self.process_block_with_strict_retry(block)
            success_ids.extend(successes)
            
            if self.pending_items:
                self.failed_items.extend(self.pending_items)
                print(f"❌ DEFINITIVE FAILURE: {len(self.pending_items)} items could not be sent after 3 attempts.")
            else:
                print(f"✅ Block successfully completed!")
                
        return success_ids

# --- MAIN FLOW ---
if __name__ == "__main__":
    MY_TOPIC_ARN: str = "arn:aws:sns:us-east-1:123456789012:MyTopic"
    
    # 1. Instantiate the class
    publisher = SNSBatchPublisher(topic_arn=MY_TOPIC_ARN)
    
    # 2. Prepare the messages
    messages_to_send: List[MessagePayload] = [MessagePayload(patient_id=f"patient_{i}", value=i) for i in range(1, 25)]
    
    # 3. Do the sending
    successes = publisher.publish_messages(messages_to_send)
    
    # 4. Results
    print(f"\nFinal Result:")
    print(f"Total successes: {len(successes)}")
    print(f"Total failures: {len(publisher.failed_items)}")
    if publisher.failed_items:
        print("Failed items:")
        for item in publisher.failed_items:
            print(f" - {item}")