feat(clp-json): Add support for searching archives that are stored on S3. #674
Conversation
WalkthroughThe pull request introduces modifications across multiple components of the job orchestration and package utility systems. The changes primarily focus on enhancing command generation, error handling, and support for different storage types and engines. Key modifications include adding new functions for command generation, updating function signatures to support environment variables, and improving storage type validation logic. The changes enhance the error handling for specific combinations of storage types and engines, providing more flexible and robust handling of search and extraction tasks across various storage configurations. Changes
Possibly related PRs
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (4)
⏰ Context from checks skipped due to timeout of 90000ms (2)
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 (
|
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
components/clp-package-utils/clp_package_utils/scripts/search.py (1)
Line range hint
87-142: Consider enhancing error handling with specific error codes.The command generation and execution logic looks good, but consider improving error handling by:
- Using specific error codes instead of just -1 for different failure scenarios
- Adding error handling for the subprocess.run() call
- Cleaning up the generated config file in case of failures
Here's a suggested improvement:
if StorageType.S3 == storage_type and StorageEngine.CLP == storage_engine: logger.error( f"Search is not supported for archive storage type: {storage_type} while using the {storage_engine} storage engine." ) - return -1 + return 2 # Specific error code for unsupported storage configuration # ... existing code ... - subprocess.run(cmd, check=True) + try: + subprocess.run(cmd, check=True) + except subprocess.CalledProcessError as e: + logger.error(f"Search command failed with exit code {e.returncode}") + return 3 # Specific error code for search execution failure + finally: + # Clean up generated files even if an error occurs + if generated_config_path_on_host.exists(): + generated_config_path_on_host.unlink() - # Remove generated files - generated_config_path_on_host.unlink()
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
components/clp-package-utils/clp_package_utils/scripts/search.py(2 hunks)components/job-orchestration/job_orchestration/executor/query/extract_stream_task.py(1 hunks)components/job-orchestration/job_orchestration/executor/query/fs_search_task.py(5 hunks)components/job-orchestration/job_orchestration/executor/query/utils.py(3 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (2)
- GitHub Check: lint-check (ubuntu-latest)
- GitHub Check: lint-check (macos-latest)
🔇 Additional comments (8)
components/job-orchestration/job_orchestration/executor/query/fs_search_task.py (4)
4-4: LGTM! The new imports support S3 functionality.The addition of
Tuplefor type hints andgenerate_s3_virtual_hosted_style_urlfor S3 URL generation are appropriate for the new S3 search feature.Also applies to: 10-10
24-41: LGTM! Well-structured command generation for CLP storage engine.The function properly handles:
- Early validation of S3 storage type compatibility
- Clear error messaging
- Consistent command construction with optional path filtering
Line range hint
88-158: LGTM! Well-structured refactoring of command generation.The separation of concerns into specialized functions improves maintainability while preserving the existing functionality.
Line range hint
197-218: LGTM! Proper handling of environment variables in search task.The changes correctly integrate the new environment variables support while maintaining the existing functionality.
components/job-orchestration/job_orchestration/executor/query/utils.py (1)
46-46: LGTM! Proper environment variables support in task execution.The addition of the
env_varsparameter and its integration withsubprocess.Popenis well-implemented.Also applies to: 67-67
components/job-orchestration/job_orchestration/executor/query/extract_stream_task.py (1)
144-144: LGTM! Maintains backward compatibility.The addition of
env_vars=Noneproperly supports the new parameter while maintaining existing behaviour.components/clp-package-utils/clp_package_utils/scripts/search.py (2)
10-10: LGTM! Import change aligns with the new storage engine validation.The addition of
StorageEngineimport is necessary and follows Python's import style guidelines.
82-86:⚠️ Potential issueReview storage validation logic - possible contradiction with PR objectives.
The current implementation blocks S3 storage with CLP engine, but according to the PR objectives, we should be adding S3 search capability. This appears to contradict the PR's goal of "incorporating the capability to search clp-json archives stored in Amazon S3."
Consider the following:
- If S3 search should be supported, this validation block might need to be removed or modified.
- If certain S3 configurations aren't supported, the error message should clarify which configurations are supported.
- Based on the AI summary, there are new functions in fs_search_task.py for handling S3 storage - this code might need to leverage those instead of returning an error.
Let's verify the intended behaviour:
| env_vars = { | ||
| **os.environ, | ||
| "AWS_ACCESS_KEY_ID": s3_config.access_key_id, | ||
| "AWS_SECRET_ACCESS_KEY": s3_config.secret_access_key, | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Consider security implications of environment inheritance.
The current implementation inherits all existing environment variables (os.environ) before adding AWS credentials. This could potentially expose unnecessary environment variables to the subprocess.
Consider passing only the required environment variables:
- env_vars = {
- **os.environ,
- "AWS_ACCESS_KEY_ID": s3_config.access_key_id,
- "AWS_SECRET_ACCESS_KEY": s3_config.secret_access_key,
- }
+ env_vars = {
+ "AWS_ACCESS_KEY_ID": s3_config.access_key_id,
+ "AWS_SECRET_ACCESS_KEY": s3_config.secret_access_key,
+ "PATH": os.environ.get("PATH", ""), # Required for executing binaries
+ }📝 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.
| env_vars = { | |
| **os.environ, | |
| "AWS_ACCESS_KEY_ID": s3_config.access_key_id, | |
| "AWS_SECRET_ACCESS_KEY": s3_config.secret_access_key, | |
| } | |
| env_vars = { | |
| "AWS_ACCESS_KEY_ID": s3_config.access_key_id, | |
| "AWS_SECRET_ACCESS_KEY": s3_config.secret_access_key, | |
| "PATH": os.environ.get("PATH", ""), # Required for executing binaries | |
| } |
components/job-orchestration/job_orchestration/executor/query/fs_search_task.py
Outdated
Show resolved
Hide resolved
components/job-orchestration/job_orchestration/executor/query/fs_search_task.py
Outdated
Show resolved
Hide resolved
components/clp-package-utils/clp_package_utils/scripts/search.py
Outdated
Show resolved
Hide resolved
Co-authored-by: haiqi96 <14502009+haiqi96@users.noreply.github.com>
|
|
||
| from clp_py_utils.clp_config import CLPConfig, StorageEngine | ||
| from clp_py_utils.s3_utils import parse_aws_credentials_file | ||
|
|
||
| # from clp_py_utils.s3_utils import parse_aws_credentials_file |
There was a problem hiding this comment.
Yes sorry didn't mean to include this diff
2342742 to
3e50655
Compare
| } | ||
| else: | ||
| # fmt: off | ||
| command.extend(( |
There was a problem hiding this comment.
correct me if wrong, shouldn't the argument of extend be []?
There was a problem hiding this comment.
double checked, extend with a set also works.
I just get used to list more, but I guess set is also fine, unless our code guideline has a strict rule on this.
There was a problem hiding this comment.
It's consistent with how its used elsewhere in fs_search_task.py. Also its a tuple, not a set right?
| logger.error(f"Encountered error while generating s3 url: {ex}") | ||
| return None, None | ||
| # fmt: off | ||
| command.extend(( |
|
Nit: the private method (those make_blabla_commands) should start with "_". However I realized I forgot to add "" to the ones to compress_task.py in my previous PR and I guess you just copied them. So I will leave it to you to decide if you want to add "" in this PR. If not, we can open another PR to fix the method naming in both scripts |
I'll rename all the make commands used in this PR but won't touch the ones from compress_task.py |
There was a problem hiding this comment.
Tested manually and confirmed both S3 and FS search works for CLP-S.
Also confirmed FS search works for CLP.
will see if @kirkrodrigues wants to go through a final pass
|
Since #662 is merged, please also update the code that gets credential from the config |
Merged and pushed up the credentials changes. Github seems to be stuck processing that most recent push, though I can see when I check the branch that the CI passes without issue. If github is still stuck when I check tomorrow I might close this PR and open a new identical one. |
components/clp-package-utils/clp_package_utils/scripts/search.py
Outdated
Show resolved
Hide resolved
components/job-orchestration/job_orchestration/executor/query/fs_search_task.py
Outdated
Show resolved
Hide resolved
components/job-orchestration/job_orchestration/executor/query/fs_search_task.py
Outdated
Show resolved
Hide resolved
components/job-orchestration/job_orchestration/executor/query/fs_search_task.py
Outdated
Show resolved
Hide resolved
components/job-orchestration/job_orchestration/executor/query/fs_search_task.py
Outdated
Show resolved
Hide resolved
| "s3" | ||
| )) | ||
| # fmt: on | ||
| env_vars = { |
There was a problem hiding this comment.
When you create the new PR, can you move the error check for the credentials above this line?
There was a problem hiding this comment.
Made this change. Seems to have gotten un-stuck once I pushed up new changes, so looks like I won't have to put up a new PR.
Description
This PR adds support for searching clp-json archives from s3 in the package. This is achieved by simply pulling the storage configuration from WorkerConfig and adding appropriate arguments to the clp-s invocation.
Validation performed
Summary by CodeRabbit
New Features
Bug Fixes
Refactor