async def test_boto_session_singleton(self):
session1 = BotoSession.get_session()
assert session1 is not None
session2 = BotoSession.get_session()
assert session1 is session2
async def confirm_permission(self) -> None:
"""
Confirms correct permissions are in place.
:raises ImproperlyConfigured: If the permissions were not found.
:raises AssertionError: If the permissions are not correctly configured on AWS.
"""
session = BotoSession.get_session()
async with session.create_client(
"sqs",
region_name=self.region_name or 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:
policy_attributes = await client.get_queue_attributes(
QueueUrl=self.queue_url, AttributeNames=["Policy"]
)
try:
policies = json.loads(policy_attributes["Attributes"]["Policy"])
statement = policies["Statement"][0]
except KeyError:
raise ImproperlyConfigured("Permissions not found.")
# need "Effect": "Allow", "Action": "SQS:SendMessage"
assert statement["Effect"] == "Allow"
assert "SQS:SendMessage" in statement["Action"]
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 get_subscription_attributes(self, subscription_arn: str) -> dict:
"""
Retrieves the attributes of the subscription from AWS.
Refer to https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sns.html#SNS.Client.get_subscription_attributes
for more information.
"""
async with BotoSession.get_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.get_subscription_attributes(
SubscriptionArn=subscription_arn)
async def _confirm_topic(
cls,
topic_arn: str,
*,
region_name: Optional[str] = None,
endpoint_url: Optional[str] = None,
) -> None:
"""
Confirm that the topic exists by request :code:`get_topic_attributes` to AWS.
"""
session = BotoSession.get_session()
async with session.create_client(
"sns",
region_name=region_name or BotoSession.aws_default_region,
endpoint_url=endpoint_url or settings.INIESTA_SNS_ENDPOINT_URL,
aws_access_key_id=BotoSession.aws_access_key_id,
aws_secret_access_key=BotoSession.aws_secret_access_key,
) as client:
await client.get_topic_attributes(TopicArn=topic_arn)
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 reset_session(self):
yield
BotoSession.reset_aws_credentials()
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
def reset_boto_session():
BotoSession.session = None
yield
BotoSession.session = None
BotoSession.reset_aws_credentials()