Assume another iam role #32
Conversation
|
Hi @adiadd Thanks for your effort in this great project. Would you please help review this PR? |
adiadd
left a comment
adiadd
left a comment
There was a problem hiding this comment.
Really solid work overall! I love the mixin approach - it's clean and reusable. The automatic credential refresh with the 5-minute buffer is a nice touch too. This is going to be super useful for cross-account log analysis.
Most of my comments are about defensive programming and consistency. Once we clean up the error handling and imports, this should be good to go!
Would also love a quick working video showcasing the updated MCP server and functionality!
| sts_client = source_session.client("sts") | ||
| assume_role_params = { | ||
| "RoleArn": self.role_arn, | ||
| "RoleSessionName": f"CloudWatchLogs-{int(time.time())}", |
There was a problem hiding this comment.
Minor thing, but I noticed we're only using timestamp for the session name. In high-frequency scenarios (like multiple MCP servers running), we might get collisions. What do you think about making it a bit more unique?
Maybe something like:
import uuid
"RoleSessionName": f"CloudWatchLogs-{int(time.time())}-{uuid.uuid4().hex[:8]}",Not a blocker, but could save us some headaches down the road!
| try: | ||
| response = sts_client.assume_role(**assume_role_params) | ||
| except Exception as e: | ||
| raise Exception(f"Failed to assume role {self.role_arn}: {str(e)}") |
There was a problem hiding this comment.
Really like the overall approach here, but I'm a bit worried about the generic exception handling. This could potentially expose sensitive AWS error details to end users, and we might miss important specific error cases.
Would you mind updating this to catch specific AWS exceptions? Something like:
from botocore.exceptions import ClientError
try:
response = sts_client.assume_role(**assume_role_params)
except ClientError as e:
error_code = e.response['Error']['Code']
if error_code == 'AccessDenied':
raise Exception(f"Access denied when assuming role {self.role_arn}")
elif error_code == 'InvalidUserType':
raise Exception(f"Invalid user type for role assumption")
else:
raise Exception(f"Failed to assume role: {error_code}")This way we can give users more helpful error messages while keeping things secure.
| import boto3 | ||
| import time | ||
| from datetime import datetime, timedelta |
There was a problem hiding this comment.
From the other code review comment, we'll likely need to import ClientError. Could you add:
from botocore.exceptions import ClientErrorThanks!
| self._assumed_session = None | ||
| self._credentials_expiry = None |
There was a problem hiding this comment.
Since MCP servers can handle concurrent requests, I'm wondering if we should consider some synchronization here? The session state could potentially get mixed up if multiple threads are assuming roles simultaneously.
Might be worth thinking about a simple lock/check around the credential refresh logic. What's your take on this?
| import sys | ||
| import os | ||
|
|
||
| sys.path.append(os.path.dirname(os.path.dirname(__file__))) | ||
| from aws_credentials_mixin import AWSCredentialsMixin |
There was a problem hiding this comment.
This works, but modifying sys.path at runtime can make the code a bit fragile and harder to test. Since we're doing this in multiple files, what do you think about either:
- Moving the mixin to a proper package location, or
- Using relative imports like
from ..aws_credentials_mixin import AWSCredentialsMixin
It would make the imports cleaner and more maintainable. Same with the src/cw-mcp-server/tools/correlation_tools.py and src/cw-mcp-server/tools/search_tools.py as well
| if __name__ == "__main__": | ||
| # Run the MCP server | ||
| mcp.run() | ||
| mcp.run(transport="streamable-http") |
There was a problem hiding this comment.
Curious if this works as intended with the addition of this, as this is an issue we have open (#29) - just want to make sure this doesn't break existing setups. Maybe worth a quick note in the PR description + any update to docs
|
Thanks for your review. I happened to find these tickets. What do you think? boto/botocore#761 |
2 participants
Issue number:
#31
Summary
Changes
User experience
Retrieve cross-account cloudwatch logs
Checklist
If your change doesn't seem to apply, please leave them unchecked.
Is this a breaking change?
No **RFC issue number**:Checklist:
Acknowledgment
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of the project license.