async def initialize( cls, *, topic_arn: Optional[str] = None, region_name: Optional[str] = None, endpoint_url: Optional[str] = None, ): """ Class method to initialize the SNS Client and confirm the topic exists. We needed to do this because of asyncio functionality :param topic_arn: If you would like to initialize this client with a different topic. Defaults to :code:`INIESTA_SNS_PRODUCER_GLOBAL_TOPIC_ARN` if not passed. :param region_name: Takes priority or defaults to :code:`BotoSession.aws.default_region` settings. :param endpoint_url: Takes priority or defaults to :code:`INIESTA_SNS_ENDPOINT_URL` settings. :return: An initialized instance of :code:`cls` (:code:`SQSClient`). :rtype: :code:`SNSClient` """ topic_arn = topic_arn or settings.INIESTA_SNS_PRODUCER_GLOBAL_TOPIC_ARN try: await cls._confirm_topic(topic_arn, region_name=region_name, endpoint_url=endpoint_url) except botocore.exceptions.ClientError as e: error_message = f"[{e.response['Error']['Code']}]: {e.response['Error']['Message']} {topic_arn}" error_logger.critical(error_message) raise return cls(topic_arn, region_name=region_name, endpoint_url=endpoint_url)
async def publish(self) -> dict: """ Serializes this message and publishes this message to SNS. :return: The response of the publish request to SNS. """ session = BotoSession.get_session() try: async with session.create_client( "sns", region_name=BotoSession.aws_default_region, endpoint_url=self.client.endpoint_url, aws_access_key_id=BotoSession.aws_access_key_id, aws_secret_access_key=BotoSession.aws_secret_access_key, ) as client: message = await client.publish(TopicArn=self.client.topic_arn, **self) logger.debug(f"[INIESTA] Published ({self.event}) with " f"the following attributes: {self}") return message except botocore.exceptions.ClientError as e: error_logger.critical( f"[{e.response['Error']['Code']}]: {e.response['Error']['Message']}" ) raise except Exception: error_logger.exception("Publishing SNS Message Failed!") raise
async def initialize( cls, *, queue_name: Optional[str] = None, endpoint_url: Optional[str] = None, region_name: Optional[str] = None, ): """ The initialization classmethod that should be first run before any subsequent SQSClient initializations. :param queue_name: queue_name if want to initialize client with a different queue :rtype: :code:`SQSClient` """ session = BotoSession.get_session() endpoint_url = endpoint_url or getattr( settings, "INIESTA_SQS_ENDPOINT_URL", None ) if queue_name is None: queue_name = cls.default_queue_name() # check if queue exists if queue_name not in cls.queue_urls: try: async with session.create_client( "sqs", region_name=region_name or BotoSession.aws_default_region, endpoint_url=endpoint_url, aws_access_key_id=BotoSession.aws_access_key_id, aws_secret_access_key=BotoSession.aws_secret_access_key, ) as client: response = await client.get_queue_url(QueueName=queue_name) except botocore.exceptions.ClientError as e: error_message = f"[{e.response['Error']['Code']}]: {e.response['Error']['Message']} {queue_name}" error_logger.critical(error_message) raise else: queue_url = response["QueueUrl"] cls.queue_urls.update({queue_name: queue_url}) sqs_client = cls(queue_name=queue_name) # check if subscription exists # await cls._confirm_subscription(sqs_client, topic_arn, endpoint_url) return sqs_client
async def _list_subscriptions_by_topic(self, next_token=None): session = BotoSession.get_session() query_args = {"TopicArn": self.topic_arn} if next_token is not None: query_args.update({"NextToken": next_token}) try: async with session.create_client( "sns", region_name=BotoSession.aws_default_region, endpoint_url=self.endpoint_url, aws_access_key_id=BotoSession.aws_access_key_id, aws_secret_access_key=BotoSession.aws_secret_access_key, ) as client: return await client.list_subscriptions_by_topic(**query_args) except botocore.exceptions.ClientError as e: error_message = f"[{e.response['Error']['Code']}]: {e.response['Error']['Message']} {self.topic_arn}" error_logger.critical(error_message) raise
def handle_error(self, exc: Exception) -> None: """ If an exception occured while handling the message, log the error. """ message = exc.message handler = getattr(exc, "handler", None) extra = { "iniesta_pass": message.event, "sqs_message_id": message.message_id, "sqs_receipt_handle": message.receipt_handle, "sqs_md5_of_body": message.md5_of_body, "sqs_message_body": message.raw_body, "sqs_attributes": json.dumps(message.attributes), "handler_name": handler.__qualname__ if handler else None, } error_logger.critical( f"[INIESTA] Error while handling message: {str(exc)}", exc_info=exc, extra=extra, )
async def _poll(self) -> str: """ The long running method that consistently polls the SQS queue for messages. :return: """ session = BotoSession.get_session() async with session.create_client( "sqs", region_name=self.region_name, endpoint_url=self.endpoint_url, aws_access_key_id=BotoSession.aws_access_key_id, aws_secret_access_key=BotoSession.aws_secret_access_key, ) as client: try: while self._loop.is_running() and self._receive_messages: try: response = await client.receive_message( QueueUrl=self.queue_url, MaxNumberOfMessages=settings.INIESTA_SQS_RECEIVE_MESSAGE_MAX_NUMBER_OF_MESSAGES, WaitTimeSeconds=settings.INIESTA_SQS_RECEIVE_MESSAGE_WAIT_TIME_SECONDS, AttributeNames=["All"], MessageAttributeNames=["All"], ) except botocore.exceptions.ClientError as e: error_logger.critical( f"[INIESTA] [{e.response['Error']['Code']}]: {e.response['Error']['Message']}" ) else: event_tasks = [ asyncio.ensure_future( self.handle_message( SQSMessage.from_sqs(client, message) ) ) for message in response.get("Messages", []) ] for fut in asyncio.as_completed(event_tasks): # NOTE: must catch CancelledError and raise try: message_obj, result = await fut except asyncio.CancelledError: raise except Exception as e: # if error log failure and pass so sqs message persists and message becomes visible again self.handle_error(e) else: await self.handle_success(client, message_obj) await self.hook_post_receive_message_handler() except asyncio.CancelledError: logger.info("[INIESTA] POLLING TASK CANCELLED") return "Cancelled" except StopPolling: # mainly used for tests logger.info("[INIESTA] STOP POLLING") return "Stopped" except Exception: if self._receive_messages and self._loop.is_running(): error_logger.critical("[INIESTA] POLLING TASK RESTARTING") self._polling_task = asyncio.ensure_future(self._poll()) error_logger.exception("[INIESTA] POLLING EXCEPTION CAUGHT") finally: await client.close() return "Shutdown" # pragma: no cover