feat: support user level search work group #3790
Conversation
WalkthroughThe recent code revisions focus on improving token management, enhancing search functionalities, and refining concurrency controls in a multi-user context. Key changes include explicit token handling, support for query cancellation, addition of user identification in search requests, and removing potential panics with Changes
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 as PR comments)
Additionally, you can add CodeRabbit Configration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 9
Outside diff range and nitpick comments (2)
src/service/search/cluster/mod.rs (1)
Line range hint
655-694: Review lock release and final merge task handling.The changes properly handle the final merge task and ensure the lock is released. However, consider refactoring to avoid duplication in the lock release logic.
- if user_id.is_some() { - work_group.as_ref().unwrap().done(trace_id, user_id).await?; - } - work_group.as_ref().unwrap().done(trace_id, None).await?; + work_group.as_ref().unwrap().finalize_search(trace_id, user_id).await?;src/service/search/sql.rs (1)
457-457: Handling ofignore_caselogic for full text search.The handling of
ignore_caseis done based on the configuration ofinverted_index_enabledand whether anyfulltextterms are present. This logic seems correct, but adding a comment explaining why this decision is made would improve code readability.
Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Files selected for processing (13)
- src/common/utils/auth.rs (1 hunks)
- src/config/src/meta/search.rs (1 hunks)
- src/handler/http/auth/validator.rs (1 hunks)
- src/handler/http/request/search/job.rs (3 hunks)
- src/handler/http/request/search/mod.rs (1 hunks)
- src/handler/http/router/mod.rs (1 hunks)
- src/infra/src/schema/mod.rs (3 hunks)
- src/ingester/src/partition.rs (1 hunks)
- src/proto/proto/cluster/search.proto (1 hunks)
- src/service/compact/merge.rs (1 hunks)
- src/service/search/cluster/mod.rs (6 hunks)
- src/service/search/mod.rs (1 hunks)
- src/service/search/sql.rs (7 hunks)
Files skipped from review due to trivial changes (2)
- src/config/src/meta/search.rs
- src/handler/http/request/search/mod.rs
Additional context used
Path-based instructions (10)
src/handler/http/request/search/job.rs (1)
Pattern
**/*.rs: You are a smart rustlang pull request reviewer.
You are going to review all the rustlang files.
Be concise, and add a brief explanation to your suggestionsMake sure the responses are not too verbose and keep the suggestions to the point i.e. actionable.
src/ingester/src/partition.rs (1)
Pattern
**/*.rs: You are a smart rustlang pull request reviewer.
You are going to review all the rustlang files.
Be concise, and add a brief explanation to your suggestionsMake sure the responses are not too verbose and keep the suggestions to the point i.e. actionable.
src/handler/http/router/mod.rs (1)
Pattern
**/*.rs: You are a smart rustlang pull request reviewer.
You are going to review all the rustlang files.
Be concise, and add a brief explanation to your suggestionsMake sure the responses are not too verbose and keep the suggestions to the point i.e. actionable.
src/common/utils/auth.rs (1)
Pattern
**/*.rs: You are a smart rustlang pull request reviewer.
You are going to review all the rustlang files.
Be concise, and add a brief explanation to your suggestionsMake sure the responses are not too verbose and keep the suggestions to the point i.e. actionable.
src/infra/src/schema/mod.rs (1)
Pattern
**/*.rs: You are a smart rustlang pull request reviewer.
You are going to review all the rustlang files.
Be concise, and add a brief explanation to your suggestionsMake sure the responses are not too verbose and keep the suggestions to the point i.e. actionable.
src/handler/http/auth/validator.rs (1)
Pattern
**/*.rs: You are a smart rustlang pull request reviewer.
You are going to review all the rustlang files.
Be concise, and add a brief explanation to your suggestionsMake sure the responses are not too verbose and keep the suggestions to the point i.e. actionable.
src/service/search/mod.rs (1)
Pattern
**/*.rs: You are a smart rustlang pull request reviewer.
You are going to review all the rustlang files.
Be concise, and add a brief explanation to your suggestionsMake sure the responses are not too verbose and keep the suggestions to the point i.e. actionable.
src/service/search/cluster/mod.rs (1)
Pattern
**/*.rs: You are a smart rustlang pull request reviewer.
You are going to review all the rustlang files.
Be concise, and add a brief explanation to your suggestionsMake sure the responses are not too verbose and keep the suggestions to the point i.e. actionable.
src/service/compact/merge.rs (1)
Pattern
**/*.rs: You are a smart rustlang pull request reviewer.
You are going to review all the rustlang files.
Be concise, and add a brief explanation to your suggestionsMake sure the responses are not too verbose and keep the suggestions to the point i.e. actionable.
src/service/search/sql.rs (1)
Pattern
**/*.rs: You are a smart rustlang pull request reviewer.
You are going to review all the rustlang files.
Be concise, and add a brief explanation to your suggestionsMake sure the responses are not too verbose and keep the suggestions to the point i.e. actionable.
Gitleaks
src/common/utils/auth.rs
577-577: Identified a HashiCorp Terraform password field, risking unauthorized infrastructure configuration and security breaches. (hashicorp-tf-password)
Additional comments not posted (11)
src/proto/proto/cluster/search.proto (1)
47-47: Addition ofuser_idfield toSearchRequestis a positive change.This change supports user-level search work group limits, aligning with the PR's objectives to enhance user-specific functionalities.
src/handler/http/request/search/job.rs (3)
18-18: ImportingputHTTP method is appropriate.This is necessary for handling the new
putroute for canceling multiple queries.
39-48: Ensure proper handling and error messages incancel_multiple_query.The function correctly parses and handles multiple trace IDs for cancellation. However, consider adding more specific error messages for different types of failures to enhance debugging and user feedback.
[REFACTOR_SUGGESTion]- return Ok(MetaHttpResponse::bad_request(e)); + return Ok(MetaHttpResponse::bad_request(format!("Failed to parse trace IDs: {}", e)));
75-87: Check for empty trace ID list is crucial for preventing unnecessary operations.The logic to prevent operations on an empty list of trace IDs is a good safeguard. This prevents the system from proceeding with no actionable items, which is efficient.
src/ingester/src/partition.rs (1)
113-115: Direct use ofget_stream_setting_bloom_filter_fieldsandget_stream_setting_fts_fieldssimplifies the code.Removing the
unwrap()calls and directly returning the fields as vectors simplifies error handling and makes the code more robust.src/handler/http/router/mod.rs (1)
387-387: Inclusion of new search-related routes enhances functionality.Adding routes for canceling multiple queries and checking query statuses aligns with the PR's objectives to improve search handling and concurrency management.
src/handler/http/auth/validator.rs (1)
592-592: Refactor to handle potential errors gracefully in JSON parsing.The use of
unwrap_or_default()in JSON parsing could mask underlying issues with the authentication string format, leading to unexpected behavior.
[REFACTOR_SUGGESTion]- config::utils::json::from_str(auth_str).unwrap_or_default(); + match config::utils::json::from_str(auth_str) { + Ok(tokens) => tokens, + Err(e) => { + log::error!("Failed to parse auth tokens: {:?}", e); + return None; + } + }src/service/search/mod.rs (1)
126-126: Ensure proper handling ofOptiontype foruser_id.Assigning
user_iddirectly without checking if it'sNonecould lead to unexpected behavior or errors in downstream processing whereuser_idis expected to be a non-null string.
[REFACTOR_SUGGESTion]- req.user_id = user_id.clone(); + req.user_id = user_id.clone().unwrap_or_default();This change ensures that
user_idis always a string, even if it's empty, which may be safer depending on how it's used downstream.src/service/search/cluster/mod.rs (1)
73-75: Ensure default timeout is set correctly.The code correctly sets a default timeout if none is specified. This aligns with the PR objectives to handle timeouts effectively.
src/service/search/sql.rs (2)
410-410: Optimization of SQL query string construction.The construction of the
select_fieldsstring is done in a way that could be optimized. Usingjoindirectly on an array of strings is more efficient and cleaner than concatenating strings in a loop.
[REFACTOR_SUGGESTion]- let select_fields = "SELECT ".to_string() + &fields.join(","); + let select_fields = format!("SELECT {}", fields.join(","));
496-496: Dynamic selection of SQL function for case sensitivity.The dynamic selection of SQL function based on
ignore_caseis a clever use of runtime configuration to adjust query behavior. This is a good implementation but ensure that all possible edge cases are handled correctly.
| } else if access_token.starts_with("Basic") || access_token.starts_with("Bearer") { | ||
| access_token | ||
| } else if access_token.starts_with("session") { | ||
| let session_key = access_token.strip_prefix("session ").unwrap().to_string(); | ||
| match USER_SESSIONS.get(&session_key) { | ||
| Some(token) => { | ||
| format!("Bearer {}", *token) | ||
| } | ||
| } else { | ||
| format!("Bearer {}", access_token) | ||
| None => access_token, | ||
| } | ||
| } else { | ||
| format!("Bearer {}", access_token) |
There was a problem hiding this comment.
Refactor extract_auth_str for clarity and security.
The function extract_auth_str has been modified to handle different token types, including a session-based token. Consider using a match statement instead of multiple if-else for clarity and maintainability. Additionally, ensure that session tokens are validated securely.
pub fn extract_auth_str(req: &HttpRequest) -> String {
let auth_token = req.cookie("auth_tokens")
.map(|cookie| json::from_str::<AuthTokens>(cookie.value()).unwrap_or_default().access_token)
.unwrap_or_default();
match auth_token.as_str() {
"" => req.cookie("auth_ext").map_or_else(|| "".to_string(), |cookie| cookie.value().to_string()),
"Basic" | "Bearer" => auth_token,
token if token.starts_with("session ") => {
let session_key = token.strip_prefix("session ").unwrap().to_string();
USER_SESSIONs.get(&session_key).map_or(token, |t| format!("Bearer {}", t))
},
_ => format!("Bearer {}", auth_token),
}
}
| for trace_id in trace_ids { | ||
| match crate::service::search::cancel_query(trace_id).await { | ||
| Ok(status) => res.push(status), | ||
| Err(e) => return Ok(MetaHttpResponse::bad_request(e)), |
There was a problem hiding this comment.
maybe we can skip failed trace_id and process the rest and return the result, instead of returning right away?
There was a problem hiding this comment.
Yep. both are okay.
|
keep improving #2506 |
New ENV
Summary by CodeRabbit
New Features
Improvements
access_tokencases.Bug Fixes
.unwrap()calls in several functions.Refactor