fix(clp-package): Add support for loading credentials from boto3 session. #681
Conversation
WalkthroughThis pull request introduces enhancements to AWS credential management across multiple components of the job orchestration system. The changes focus on improving the handling of S3 credentials by adding a new function Changes
Sequence DiagramsequenceDiagram
participant Task as Execution Task
participant Utils as load_session_credentials()
participant S3Utils as s3_get_frozen_credentials()
Task->>Utils: Request AWS credentials
Utils->>S3Utils: Attempt to retrieve credentials
alt Credentials Found
S3Utils-->>Utils: Return Credentials
Utils-->>Task: Provide Credentials
else No Credentials
S3Utils-->>Utils: Return None
Utils-->>Task: Return None
end
Possibly Related PRs
Suggested Reviewers
✨ Finishing Touches
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
haiqi96
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (2)
components/job-orchestration/job_orchestration/executor/utils.py (2)
30-30: Improve error message clarity without exposing sensitive information.The current error messages could be more descriptive while maintaining security. Consider including troubleshooting guidance.
Apply this diff:
- logger.error("Failed to get s3 credentials from local session") + logger.error("Unable to load AWS credentials. Verify AWS configuration is properly set up.") - logger.error("Not supporting session token at the moment") + logger.error("Temporary session credentials are not supported. Use long-term credentials instead.")Also applies to: 33-33
27-35: Consider adding retry mechanism for transient failures.The function could benefit from retrying the credential fetch operation in case of temporary AWS connectivity issues.
Consider implementing an exponential backoff retry mechanism using the
tenacitylibrary. Would you like me to provide an example implementation?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
components/clp-py-utils/clp_py_utils/s3_utils.py(2 hunks)components/job-orchestration/job_orchestration/executor/compress/compression_task.py(4 hunks)components/job-orchestration/job_orchestration/executor/query/extract_stream_task.py(2 hunks)components/job-orchestration/job_orchestration/executor/query/fs_search_task.py(2 hunks)components/job-orchestration/job_orchestration/executor/utils.py(2 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: lint-check (ubuntu-latest)
🔇 Additional comments (2)
components/job-orchestration/job_orchestration/executor/compress/compression_task.py (1)
287-291: LGTM!The added error handling for the compression command is consistent with the function's error handling pattern.
components/job-orchestration/job_orchestration/executor/utils.py (1)
3-3: LGTM! Import statements are well-organized.The new imports are properly organized and necessary for the added functionality.
Also applies to: 7-7
| def s3_get_frozen_credentials() -> Optional[ReadOnlyCredentials]: | ||
| session = boto3.Session() | ||
| credentials = session.get_credentials() | ||
| if credentials is None: | ||
| return None | ||
| frozen_credentials = credentials.get_frozen_credentials() | ||
| return frozen_credentials |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Enhance error handling and logging.
The function should handle potential boto3 exceptions and log appropriate messages when credentials are not found.
Apply this diff to improve error handling:
def s3_get_frozen_credentials() -> Optional[ReadOnlyCredentials]:
+ try:
session = boto3.Session()
credentials = session.get_credentials()
if credentials is None:
+ logger.warning("No AWS credentials found in the session")
return None
frozen_credentials = credentials.get_frozen_credentials()
return frozen_credentials
+ except Exception as ex:
+ logger.error(f"Failed to get AWS credentials: {ex}")
+ return NoneCommittable suggestion skipped: line range outside the PR's diff.
| aws_access_key_id, aws_access_key_id = load_session_credentials(logger) | ||
| if aws_access_key_id is None or aws_secret_access_key is None: | ||
| return None, None | ||
|
|
There was a problem hiding this comment.
Fix the credential variable assignment bug.
There is a bug in the variable assignment where both variables are assigned the access key ID.
Apply this diff to fix the variable assignment:
- aws_access_key_id, aws_access_key_id = load_session_credentials(logger)
+ aws_access_key_id, aws_secret_access_key = load_session_credentials(logger)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| aws_access_key_id, aws_access_key_id = load_session_credentials(logger) | |
| if aws_access_key_id is None or aws_secret_access_key is None: | |
| return None, None | |
| aws_access_key_id, aws_secret_access_key = load_session_credentials(logger) | |
| if aws_access_key_id is None or aws_secret_access_key is None: | |
| return None, None | |
| aws_access_key_id, aws_access_key_id = load_session_credentials(logger) | ||
| if aws_access_key_id is None or aws_secret_access_key is None: | ||
| return None, None | ||
|
|
There was a problem hiding this comment.
Fix the credential variable assignment bug.
There is a bug in the variable assignment where both variables are assigned the access key ID.
Apply this diff to fix the variable assignment:
- aws_access_key_id, aws_access_key_id = load_session_credentials(logger)
+ aws_access_key_id, aws_secret_access_key = load_session_credentials(logger)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| aws_access_key_id, aws_access_key_id = load_session_credentials(logger) | |
| if aws_access_key_id is None or aws_secret_access_key is None: | |
| return None, None | |
| aws_access_key_id, aws_secret_access_key = load_session_credentials(logger) | |
| if aws_access_key_id is None or aws_secret_access_key is None: | |
| return None, None | |
| aws_access_key_id = clp_config.input.aws_access_key_id | ||
| aws_secret_access_key = clp_config.input.aws_secret_access_key | ||
| if aws_access_key_id is None or aws_secret_access_key is None: | ||
| aws_access_key_id, aws_access_key_id = load_session_credentials(logger) | ||
| if aws_access_key_id is None or aws_secret_access_key is None: | ||
| return None, None | ||
|
|
There was a problem hiding this comment.
Fix the credential variable assignment bug.
There is a bug in the variable assignment where both variables are assigned the access key ID.
Apply this diff to fix the variable assignment:
- aws_access_key_id, aws_access_key_id = load_session_credentials(logger)
+ aws_access_key_id, aws_secret_access_key = load_session_credentials(logger)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| aws_access_key_id = clp_config.input.aws_access_key_id | |
| aws_secret_access_key = clp_config.input.aws_secret_access_key | |
| if aws_access_key_id is None or aws_secret_access_key is None: | |
| aws_access_key_id, aws_access_key_id = load_session_credentials(logger) | |
| if aws_access_key_id is None or aws_secret_access_key is None: | |
| return None, None | |
| aws_access_key_id = clp_config.input.aws_access_key_id | |
| aws_secret_access_key = clp_config.input.aws_secret_access_key | |
| if aws_access_key_id is None or aws_secret_access_key is None: | |
| aws_access_key_id, aws_secret_access_key = load_session_credentials(logger) | |
| if aws_access_key_id is None or aws_secret_access_key is None: | |
| return None, None |
| def load_session_credentials(logger: Logger) -> Tuple[Optional[str], Optional[str]]: | ||
| s3_frozen_credentials = s3_get_frozen_credentials() | ||
| if s3_frozen_credentials is None: | ||
| logger.error("Failed to get s3 credentials from local session") | ||
| return None, None | ||
| if s3_frozen_credentials.token is not None: | ||
| logger.error("Not supporting session token at the moment") | ||
| return None, None | ||
| return s3_frozen_credentials.access_key, s3_frozen_credentials.secret_key |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Add comprehensive docstring with type hints.
The function lacks documentation explaining its purpose, parameters, return values, and possible error conditions.
Add a docstring following this pattern:
def load_session_credentials(logger: Logger) -> Tuple[Optional[str], Optional[str]]:
+ """
+ Load AWS credentials from the current session.
+
+ Args:
+ logger: Logger instance for error reporting
+
+ Returns:
+ Tuple[Optional[str], Optional[str]]: A tuple containing (access_key, secret_key).
+ Both values will be None if credentials cannot be loaded or contain unsupported features.
+ """📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def load_session_credentials(logger: Logger) -> Tuple[Optional[str], Optional[str]]: | |
| s3_frozen_credentials = s3_get_frozen_credentials() | |
| if s3_frozen_credentials is None: | |
| logger.error("Failed to get s3 credentials from local session") | |
| return None, None | |
| if s3_frozen_credentials.token is not None: | |
| logger.error("Not supporting session token at the moment") | |
| return None, None | |
| return s3_frozen_credentials.access_key, s3_frozen_credentials.secret_key | |
| def load_session_credentials(logger: Logger) -> Tuple[Optional[str], Optional[str]]: | |
| """ | |
| Load AWS credentials from the current session. | |
| Args: | |
| logger: Logger instance for error reporting | |
| Returns: | |
| Tuple[Optional[str], Optional[str]]: A tuple containing (access_key, secret_key). | |
| Both values will be None if credentials cannot be loaded or contain unsupported features. | |
| """ | |
| s3_frozen_credentials = s3_get_frozen_credentials() | |
| if s3_frozen_credentials is None: | |
| logger.error("Failed to get s3 credentials from local session") | |
| return None, None | |
| if s3_frozen_credentials.token is not None: | |
| logger.error("Not supporting session token at the moment") | |
| return None, None | |
| return s3_frozen_credentials.access_key, s3_frozen_credentials.secret_key |
Description
Validation performed
Summary by CodeRabbit
Release Notes
New Features
Improvements
Bug Fixes