refactor: use rwlock for modifiable session data #4232
Conversation
WalkthroughThe recent changes involve refining the Changes
Sequence Diagram(s)Changes don't include a new feature or modifications to the control flow that necessitate a sequence diagram. Poem
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 (
|
sunng87
force-pushed
the
refactor/session-query-context-update
branch
from
8cec5ac to
6eba6f6
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
Files selected for processing (24)
- src/auth/src/permission.rs (1 hunks)
- src/auth/tests/mod.rs (2 hunks)
- src/common/datasource/Cargo.toml (1 hunks)
- src/common/function/src/system/database.rs (1 hunks)
- src/operator/src/expr_factory.rs (2 hunks)
- src/operator/src/statement.rs (2 hunks)
- src/pipeline/src/manager/table.rs (1 hunks)
- src/query/src/optimizer/type_conversion.rs (1 hunks)
- src/query/src/plan.rs (1 hunks)
- src/query/src/range_select/plan_rewrite.rs (1 hunks)
- src/query/src/sql.rs (1 hunks)
- src/script/src/table.rs (1 hunks)
- src/servers/src/export_metrics.rs (1 hunks)
- src/servers/src/grpc/authorize.rs (3 hunks)
- src/servers/src/grpc/greptime_handler.rs (3 hunks)
- src/servers/src/http/authorize.rs (3 hunks)
- src/servers/src/http/prometheus.rs (1 hunks)
- src/servers/src/mysql/handler.rs (1 hunks)
- src/servers/src/postgres/handler.rs (1 hunks)
- src/servers/tests/http/authorize.rs (3 hunks)
- src/servers/tests/http/http_handler_test.rs (4 hunks)
- src/servers/tests/py_script/mod.rs (1 hunks)
- src/session/src/context.rs (14 hunks)
- src/session/src/lib.rs (3 hunks)
Files skipped from review due to trivial changes (5)
- src/operator/src/statement.rs
- src/query/src/range_select/plan_rewrite.rs
- src/servers/src/grpc/authorize.rs
- src/servers/src/mysql/handler.rs
- src/servers/tests/http/http_handler_test.rs
Additional comments not posted (57)
src/common/datasource/Cargo.toml (1)
34-36: Addition of "async" feature toorc-rustdependency.The addition of the "async" feature to the
orc-rustdependency seems appropriate given the context of the PR. This change aligns with the need for asynchronous capabilities in the codebase.src/auth/src/permission.rs (2)
45-47: Simplified method signature by requiringUserInfoRef.The change from
Option<UserInfoRef>toUserInfoRefsimplifies the method signature and eliminates the need for handlingNonecases. This improves code readability and reduces potential errors.
53-55: Updated implementation to match the newPermissionCheckertrait signature.The implementation of
PermissionCheckerforOption<&PermissionCheckerRef>has been updated to match the new trait signature. This change ensures consistency and correctness.src/auth/tests/mod.rs (4)
30-32: Updated test implementation to match the newPermissionCheckertrait signature.The test implementation has been updated to match the new
PermissionCheckertrait signature. This ensures that the tests are consistent with the updated code.
48-50: Updated test cases to useUserInfoRefdirectly.The test cases have been updated to use
UserInfoRefdirectly instead ofOption<UserInfoRef>. This change ensures that the tests accurately reflect the new method signature and simplifies the test logic.
54-56: Updated SQL statement test case to useUserInfoRefdirectly.The SQL statement test case has been updated to use
UserInfoRefdirectly. This change ensures that the tests are consistent with the updated code and simplifies the test logic.
62-63: Updated error handling test case to useUserInfoRefdirectly.The error handling test case has been updated to use
UserInfoRefdirectly. This change ensures that the tests are consistent with the updated code and simplifies the test logic.src/common/function/src/system/database.rs (1)
81-81: Updatedcurrent_schemamethod call tocurrent_schema_raw.The method call has been updated to
current_schema_rawto align with the changes in theQueryContextstruct. This change ensures consistency and correctness in the codebase.src/servers/tests/py_script/mod.rs (1)
63-63: Updatedcurrent_schemamethod call tocurrent_schema_raw.The method call has been updated to
current_schema_rawto align with the changes in theQueryContextstruct. This change ensures consistency and correctness in the codebase.src/session/src/lib.rs (13)
20-20: Ensure correct usage ofRwLock.The
RwLockis used correctly here to wrap theStringtype, allowing for concurrent read and write access.
34-38: Ensure correct usage ofRwLockin struct fields.The
RwLockis used correctly to wrapString,UserInfoRef, andTimezonetypes, allowing for concurrent read and write access.
51-55: Ensure correct initialization ofRwLockfields.The
RwLockfields are correctly initialized withArc::new(RwLock::new(...))to ensure thread-safe access.
62-69: Ensure correct cloning ofRwLockfields.The
RwLockfields are correctly cloned for use in theQueryContextBuilder.
82-83: Ensure correct implementation oftimezonegetter.The
timezonemethod correctly reads theRwLockand clones theTimezonevalue.
87-87: Ensure correct implementation ofset_timezonemethod.The
set_timezonemethod correctly writes to theRwLock.
91-91: Ensure correct implementation ofuser_infogetter.The
user_infomethod correctly reads theRwLockand clones theUserInfoRefvalue.
95-95: Ensure correct implementation ofset_user_infomethod.The
set_user_infomethod correctly writes to theRwLock.
99-99: Ensure correct implementation ofset_catalogmethod.The
set_catalogmethod correctly writes to theRwLock.
103-103: Ensure correct implementation ofcataloggetter.The
catalogmethod correctly reads theRwLockand clones theStringvalue.
107-107: Ensure correct implementation ofschemagetter.The
schemamethod correctly reads theRwLockand clones theStringvalue.
111-111: Ensure correct implementation ofset_schemamethod.The
set_schemamethod correctly writes to theRwLock.
115-115: Ensure correct usage ofcatalogandschemamethods inget_db_string.The
get_db_stringmethod correctly uses thecatalogandschemamethods to build the database string.src/servers/tests/http/authorize.rs (3)
31-31: Ensure correct handling of user information in test.The
ctx.current_user()method is correctly used to retrieve the current user information.
42-42: Ensure correct handling of user information in test.The
ctx.current_user()method is correctly used to retrieve the current user information.
83-83: Ensure correct handling of user information in test.The
ctx.current_user()method is correctly used to retrieve the current user information.src/servers/src/grpc/greptime_handler.rs (3)
131-131: Ensure correct handling of user information inauthfunction.The
authfunction now directly returnsauth::userinfo_by_name(None)whenuser_providerisNone.
133-133: Ensure correct handling of user information inauthfunction.The
authfunction now directly returnsauth::userinfo_by_name(None)whenuser_providerisNone.
198-199: Ensure correct usage ofcurrent_schema_rawandtimezone_rawincreate_query_context.The
create_query_contextfunction correctly usescurrent_schema_rawandtimezone_rawto set the schema and timezone in the query context.src/query/src/plan.rs (1)
251-251: Ensure correct usage ofcurrent_schema_rawin test.The
current_schema_rawmethod is correctly used to set the schema in the query context.src/servers/src/export_metrics.rs (1)
250-252: Update usage ofcurrent_schema_rawThe change from
current_schematocurrent_schema_rawis consistent with the PR objectives. This ensures that the schema is wrapped inArc<RwLock>, allowing for modifiable session data.src/session/src/context.rs (17)
18-18: ImportRwLockThe import of
RwLockis necessary for the changes made to theQueryContextstruct.
40-43: ReplaceArcSwapwithArc<RwLock>inQueryContextThe change from
ArcSwaptoArc<RwLock>forcurrent_schema,current_user, andtimezonealigns with the PR objectives to allow modifiable session data. This is a crucial change for handling concurrent access safely.
63-73: Addcurrent_schema_rawandtimezone_rawmethodsThe addition of
current_schema_rawandtimezone_rawmethods in theQueryContextBuilderis consistent with the changes to useArc<RwLock>. These methods ensure that the schema and timezone are wrapped inArc<RwLock>.
79-81: CloneArc<RwLock>fields inQueryContextThe clone implementation for
QueryContextcorrectly clones theArc<RwLock>fields, ensuring that the cloned context maintains the same references.
95-96: Updatecurrent_schemaandtimezoneinFrom<RegionRequestHeader> for QueryContextThe updates to use
Arc<RwLock>forcurrent_schemaandtimezonein theFrom<RegionRequestHeader> for QueryContextimplementation are consistent with the overall changes.
107-108: Updatecurrent_schemaandtimezoneinFrom<api::v1::QueryContext> for QueryContextThe updates to use
Arc<RwLock>forcurrent_schemaandtimezonein theFrom<api::v1::QueryContext> for QueryContextimplementation are consistent with the overall changes.
126-127: ReadArc<RwLock>fields inFrom<QueryContext> for api::v1::QueryContextThe changes to read from
Arc<RwLock>forcurrent_schemaandtimezonein theFrom<QueryContext> for api::v1::QueryContextimplementation are necessary to obtain the values correctly.
141-141: UseArc<RwLock>forcurrent_schemainQueryContext::withThe change to use
Arc<RwLock>forcurrent_schemain theQueryContext::withmethod is consistent with the overall changes to theQueryContextstruct.
159-159: UseArc<RwLock>forcurrent_schemainQueryContext::with_db_nameThe change to use
Arc<RwLock>forcurrent_schemain theQueryContext::with_db_namemethod is consistent with the overall changes to theQueryContextstruct.
164-164: Readcurrent_schemafromArc<RwLock>The change to read
current_schemafromArc<RwLock>in thecurrent_schemamethod ofQueryContextis necessary to obtain the value correctly.
168-168: Writecurrent_schematoArc<RwLock>The change to write
current_schematoArc<RwLock>in theset_current_schemamethod ofQueryContextis necessary to update the value correctly.
185-186: ReadtimezonefromArc<RwLock>The change to read
timezonefromArc<RwLock>in thetimezonemethod ofQueryContextis necessary to obtain the value correctly.
189-190: WritetimezonetoArc<RwLock>The change to write
timezonetoArc<RwLock>in theset_timezonemethod ofQueryContextis necessary to update the value correctly.
193-194: Readcurrent_userfromArc<RwLock>The change to read
current_userfromArc<RwLock>in thecurrent_usermethod ofQueryContextis necessary to obtain the value correctly.
197-199: Writecurrent_usertoArc<RwLock>The change to write
current_usertoArc<RwLock>in theset_current_usermethod ofQueryContextis necessary to update the value correctly.
238-244: UseArc<RwLock>inQueryContextBuilder::buildThe change to use
Arc<RwLock>forcurrent_schema,current_user, andtimezonein theQueryContextBuilder::buildmethod is necessary to construct theQueryContextcorrectly.
263-265: CloneArc<RwLock>fields inQueryContextBuilder::from_existingThe change to clone
Arc<RwLock>fields in theQueryContextBuilder::from_existingmethod is consistent with the overall changes to theQueryContextstruct.src/script/src/table.rs (1)
340-340: Update usage ofcurrent_schema_rawThe change from
current_schematocurrent_schema_rawis consistent with the PR objectives. This ensures that the schema is wrapped inArc<RwLock>, allowing for modifiable session data.src/servers/src/http/authorize.rs (4)
65-65: Update usage ofcurrent_schema_rawThe change from
current_schematocurrent_schema_rawin theinner_authfunction is consistent with the PR objectives. This ensures that the schema is wrapped inArc<RwLock>, allowing for modifiable session data.
68-69: Update usage oftimezone_rawThe change from
timezonetotimezone_rawin theinner_authfunction is consistent with the PR objectives. This ensures that the timezone is wrapped inArc<RwLock>, allowing for modifiable session data.
78-78: Set current user directlyThe change to set the current user directly without wrapping in
Someis consistent with the updated function signature ofset_current_user.
106-106: Set authenticated userThe change to set the authenticated user directly in the
inner_authfunction is consistent with the updated function signature ofset_current_user.src/pipeline/src/manager/table.rs (1)
207-207: Update usage ofcurrent_schema_rawThe change from
current_schematocurrent_schema_rawis consistent with the PR objectives. This ensures that the schema is wrapped inArc<RwLock>, allowing for modifiable session data.src/query/src/optimizer/type_conversion.rs (1)
148-148: Verify the correctness of direct reference totimezone.The change from
self.query_ctx.timezone().as_ref()to&self.query_ctx.timezone()should be verified to ensure it correctly refers to thetimezonefield.src/operator/src/expr_factory.rs (1)
683-683: Verify the correctness of direct reference totimezone.The change from
QueryContextBuilder::default().timezone(Timezone::from_tz_string("+08:00").unwrap())toQueryContextBuilder::default().timezone_raw(Timezone::from_tz_string("+08:00").unwrap())should be verified to ensure it correctly refers to thetimezonefield.Verification successful
The
timezone_rawmethod inQueryContextBuilderis correctly implemented. The change to use.timezone_raw(Timezone::from_tz_string("+08:00").unwrap())is valid.
src/session/src/context.rs: Methodtimezone_rawis correctly defined and implemented.Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify the usage of `timezone` field in `QueryContextBuilder`. # Test: Search for the method definition and its usage. Expect: Correct references to the `timezone` field. rg --type rust $'QueryContextBuilder'Length of output: 4487
Script:
#!/bin/bash # Description: Verify the implementation of `timezone_raw` method in `QueryContextBuilder`. # Search for the method definition within the `QueryContextBuilder` implementation ast-grep --lang rust --pattern 'impl QueryContextBuilder { $$$ pub fn timezone_raw(&self, $_: Timezone) -> &Self { $$$ } $$$ }'Length of output: 141
Script:
#!/bin/bash # Description: Verify the existence of any method related to `timezone` in `QueryContextBuilder`. # Search for any method related to `timezone` in the `QueryContextBuilder` implementation rg --type rust 'impl QueryContextBuilder' -A 20 | rg 'timezone'Length of output: 529
src/servers/src/http/prometheus.rs (1)
625-625: Verify the correctness of direct reference toschema.The change from
current_schema(schema.to_string())tocurrent_schema_raw(schema.to_string())should be verified to ensure it correctly refers to theschemafield.Verification successful
Verified the correctness of direct reference to
schema.The change from
current_schema(schema.to_string())tocurrent_schema_raw(schema.to_string())correctly uses thecurrent_schema_rawmethod to set thecurrent_schemafield in theQueryContextBuilder.
src/servers/src/http/prometheus.rs, line 625Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify the usage of `schema` field in `QueryContextBuilder`. # Test: Search for the method definition and its usage. Expect: Correct references to the `schema` field. rg --type rust $'QueryContextBuilder'Length of output: 4487
Script:
#!/bin/bash # Description: Locate the definition of the `current_schema_raw` method in the `QueryContextBuilder` implementation and verify its usage. # Search for the `current_schema_raw` method definition in the `QueryContextBuilder` implementation. rg --type rust 'fn current_schema_raw' --context 5Length of output: 612
Script:
#!/bin/bash # Description: Verify the usage of the `current_schema` field within the `QueryContextBuilder` implementation. # Search for the `current_schema` field definition and its usage within the `QueryContextBuilder` implementation. rg --type rust 'current_schema' --context 5Length of output: 53321
src/query/src/sql.rs (1)
1024-1024: Ensuretimezone_rawmethod usage is appropriate.The method
timezone_rawis used here. Verify that it correctly replaces the previoustimezonemethod and that it aligns with the newRwLockstructure fortimezone.Verification successful
Verification successful for
timezone_rawmethod usage.The
timezone_rawmethod correctly replaces the previoustimezonemethod and aligns with theRwLockstructure fortimezone. The usage insrc/query/src/sql.rsis appropriate and consistent with the changes.
timezone_rawmethod implementation:src/session/src/context.rsRwLock<Timezone>structure:src/session/src/lib.rs,src/session/src/context.rsScripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify that `timezone_raw` method is correctly implemented and used in the codebase. # Test: Search for the `timezone_raw` method implementation. Expect: Find the correct implementation. rg --type rust $'fn timezone_raw' # Test: Search for all usages of `timezone_raw`. Expect: Verify all usages are appropriate. rg --type rust $'timezone_raw'Length of output: 792
Script:
#!/bin/bash # Description: Verify that the previous `timezone` method's implementation and usage. # Test: Search for the `timezone` method implementation. Expect: Find the correct implementation. rg --type rust $'fn timezone' # Test: Search for all usages of `timezone`. Expect: Verify all usages are appropriate. rg --type rust $'timezone'Length of output: 34797
Script:
#!/bin/bash # Description: Verify the `RwLock` structure for `timezone`. # Test: Search for `RwLock` usage with `timezone`. Expect: Verify the structure. rg --type rust 'RwLock<Timezone>'Length of output: 159
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #4232 +/- ##
==========================================
- Coverage 84.97% 84.65% -0.32%
==========================================
Files 1054 1057 +3
Lines 187154 187495 +341
==========================================
- Hits 159037 158732 -305
- Misses 28117 28763 +646 |
There was a problem hiding this comment.
Actionable comments posted: 0
Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Files selected for processing (7)
- src/operator/src/expr_factory.rs (2 hunks)
- src/operator/src/statement.rs (2 hunks)
- src/query/src/sql.rs (1 hunks)
- src/servers/src/grpc/greptime_handler.rs (3 hunks)
- src/servers/src/http/authorize.rs (3 hunks)
- src/session/src/context.rs (11 hunks)
- src/session/src/lib.rs (2 hunks)
Files skipped from review as they are similar to previous changes (5)
- src/operator/src/expr_factory.rs
- src/operator/src/statement.rs
- src/query/src/sql.rs
- src/servers/src/http/authorize.rs
- src/session/src/context.rs
Additional comments not posted (18)
src/session/src/lib.rs (15)
20-20: Ensure correct usage ofRwLock.The import of
RwLockis necessary for the changes. Ensure thatRwLockis used correctly throughout the codebase.
34-35: Consider the necessity ofArc.The
Sessionstruct usesArcforRwLockfields. Evaluate ifArcis necessary, given thatSessionis notClone.Would you like to explore alternatives for managing shared mutable state without
Arc?
42-48: Encapsulation of mutable fields.The introduction of
MutableInnerto encapsulate mutable fields is a good practice for better organization and potential future extensions.
50-58: Default implementation forMutableInner.The default implementation of
MutableInnerinitializes fields with default values. Ensure these defaults are appropriate for all use cases.
67-70: Initialization ofSessionwithRwLock.The
Sessionstruct is correctly initialized withRwLockforcatalogandmutable_inner. Ensure that all necessary initializations are covered.
76-79: Thread-safe access tocatalog.The
new_query_contextmethod provides thread-safe access tocatalogusingRwLock. Ensure similar patterns are followed for other mutable fields.
94-95: Thread-safe access totimezone.The
timezonemethod provides thread-safe read access totimezoneusingRwLock.
99-100: Thread-safe modification oftimezone.The
set_timezonemethod provides thread-safe write access totimezoneusingRwLock.
104-104: Thread-safe access touser_info.The
user_infomethod provides thread-safe read access touser_infousingRwLock.
108-108: Thread-safe modification ofuser_info.The
set_user_infomethod provides thread-safe write access touser_infousingRwLock.
112-112: Thread-safe modification ofcatalog.The
set_catalogmethod provides thread-safe write access tocatalogusingRwLock.
116-116: Thread-safe access tocatalog.The
catalogmethod provides thread-safe read access tocatalogusingRwLock.
120-120: Thread-safe access toschema.The
schemamethod provides thread-safe read access toschemausingRwLock.
124-124: Thread-safe modification ofschema.The
set_schemamethod provides thread-safe write access toschemausingRwLock.
128-128: Thread-safe access to database string.The
get_db_stringmethod provides thread-safe read access tocatalogandschemausingRwLock.src/servers/src/grpc/greptime_handler.rs (3)
131-131: Ensure correct handling ofUserInfoRef.The
authfunction now directly handlesUserInfoRef. Ensure that this change is consistent across the codebase.
133-133: Simplified handling ofNonecase foruser_provider.The early return for the
Nonecase ofuser_providersimplifies the logic.
199-199: Thread-safe access totimezone.The
create_query_contextfunction correctly sets thetimezoneusing the newRwLockbased implementation.
There was a problem hiding this comment.
Actionable comments posted: 0
Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Files selected for processing (1)
- src/session/src/lib.rs (2 hunks)
Files skipped from review as they are similar to previous changes (1)
- src/session/src/lib.rs
I hereby agree to the terms of the GreptimeDB CLA.
Refer to a related PR or issue link (optional)
What's changed and what's your intention?
Follow up for #4218 . Need a rebase after #4218 merged.
Some of the properties, like schema and timezone, are modifiable within a session.
Our original design is to generate a read-only
QueryContext, then pass it with all the execution process path. This works well before we have statements that modify some of the session data, likeSET time_zoneandUSE.After that, we introduce a new method called
update_sessionto give back modified properties to session, after query execution. Seegreptimedb/src/servers/src/mysql/handler.rs
Line 107 in b739c9f
update_sessionhas to check each modifiable field and update the session. This is becauseArcSwapdata structure is not clonable. So we cannot share reference between the parentSessionand the childQueryContext.This main point of this patch is to change those modifiable fields from
ArcSwaptoArc<RwLock>so that we can share it between parent and child. Note that becauseSessionis associated with a signle connection, so actually there is no race condition with reading and updating those modifiable fields.This patch also include a feature flag fix for orc-rust to exclude it datafusion dependency.
Checklist
Summary by CodeRabbit
Refactor
ArcSwapwithArc<RwLock>for fields inSessionstruct, enhancing concurrent access management.Bug Fixes
Chores
asyncfeature fororc-rust, enhancing asynchronous capabilities.