Adding QueryType to AccessControlContext #23663

RohanSidhu · 2024-09-16T20:31:58Z

Description

Adding a QueryType field into the AccessControlContext object to allow for this information to be available when making access control related decisions.

Motivation and Context

The addition of QueryType allows for downstream function calls to access this context which can be used to alter decision making if needed.

Impact

Minor performance enhancement as rebuilding Session is no longer needed. If additional context needs to be added into Session, we can utilize SessionBuilder to add additional/edit material.

Test Plan

Tested code change by running sample queries and followed the call stack to ensure queryType was present.

Also verified the context was present in the AccessControlContext through unit tests under AbstractTestQueries.

Contributor checklist

Please make sure your submission complies with our development, formatting, commit message, and attribution guidelines.
PR description addresses the issue accurately and concisely. If the change is non-trivial, a GitHub Issue is referenced.
Documented new properties (with its default value), SQL syntax, functions, or other functionality.
If release notes are required, they follow the release notes guidelines.
Adequate tests were added if applicable.
CI passed.

Release Notes

== NO RELEASE NOTE ==

linux-foundation-easycla · 2024-09-16T20:32:02Z

The committers listed above are authorized under a signed CLA.

✅ login: RohanSidhu (33183e8)

ghelmling

Can you add test coverage for DispatchManager.createQuery(Internal) to validate that the session is getting populated correctly? I think you could plug in a stubbed out DispatchQueryFactory to validate the Session instance passed to createDispatchQuery().

presto-main/src/main/java/com/facebook/presto/Session.java

ghelmling · 2024-09-18T19:43:19Z

presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java

@@ -303,6 +306,11 @@ private <C> void createQueryInternal(QueryId queryId, String slug, int retryCoun
                    Optional.ofNullable(sessionContext.getSchema()),
                    sessionContext.getIdentity().getPrincipal().map(Principal::getName)));

+            session = sessionBuilder.build();
+            if (sessionContext.getTransactionId().isPresent()) {


There is a difference here in that previously, SessionSupplier.createSession() was calling beginTransaction up at line 285. I don't know if that makes a difference.

I reviewed the tasks that happen within beginTransactionId and I don't believe there is any dependency between what happens in there and the work done to create preparedQuery. This implies that the order does not matter on if we prepare the query or if we begin the transaction. From testing, I also found no impact in performance or functionality.

why are we starting a transaction here ?

The transaction was never started after I refactored the code to use SessionBuilder instead of Session. In the original construction of Session, the transaction was started once Session was built. Since we need queryType in the session, I moved the transaction piece out so we could utilize SessionBuilder to add the necessary field. Since the transaction hasn't started, I moved it out and placed it as early as I could while including the addition of queryType.

ghelmling

I noticed another bit in DispatchManager where we should consolidate the Session creation. Sorry for the add-on, but would be good to clean up some of the duplicate creation that existed.

presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java

NikhilCollooru · 2024-09-20T21:44:29Z

presto-main/src/main/java/com/facebook/presto/Session.java

@@ -131,7 +133,8 @@ public Session(
            Map<SqlFunctionId, SqlInvokedFunction> sessionFunctions,
            Optional<Tracer> tracer,
            WarningCollector warningCollector,
-            RuntimeStats runtimeStats)
+            RuntimeStats runtimeStats,
+            Optional<QueryType> queryType)


why is this optional ? can a session have this empty ?

I don't believe there is an instance where queryType can be empty. If my understanding is correct, all queries that make it to DispatchManager will have some enum of QueryType. However, as a precaution, I wrapped QueryType in an optional.

so then lets have it as NOT optional. And fix any issues that my arise.

After making the changes, I realized that there are some areas where I put an empty optional instead of null.
Example
Since the session ctor requires all fields to be non-null, I believe it may be the design appropriate choice to keep the queryType as an Optional.

Why are we not adding QueryType to SessionRepresentation as well ? Then we need NOT send Optional.empty() or null right ?

After making the change in SessionRepresentation and other places, there is another issue if we don't use an Optional. queryType is not defined when we build session. We pass in session in a few places before we even have queryType. So if we want to maintain the non-null rule in session ctor, I believe using an Optional is the best bet.

It is a little confusing but I am more than happy to discuss the thought process!

@NikhilCollooru the challenge here is the Session usage in DispatchManager#createQueryInternal:

Session is instantiated and passed to createAnalyzerOptions()

Same session is passed to getAnalyzerType()

The SQL string is only parsed on the following line in QueryPreparer#prepareQuery

So currently session is being used before we know what the QueryType actually is. Now the Session usage in (1) and (2) is only to extract system property values, so we could try to separate that concern from the Session representation itself, pulling out SystemPropertyManager and the Map<String, String> of system properties. But that looks like a big refactoring as the use of Session as an interface to system properties seems pretty deeply embedded.

Any thoughts on how we can better clean this up? The current use of Session in DispatchManager#createQueryInternal is a bit wonky to begin with as Session is first created from SessionSupplier (not the builder for some reason), then recreated by SessionPropertyDefaults#newSessionQithDefaultProperties. I don't really understand why the SessionPropertyDefaults part happens after the analyzer and query preparer setup.

NikhilCollooru · 2024-09-20T21:45:31Z

presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java

@@ -303,6 +306,11 @@ private <C> void createQueryInternal(QueryId queryId, String slug, int retryCoun
                    Optional.ofNullable(sessionContext.getSchema()),
                    sessionContext.getIdentity().getPrincipal().map(Principal::getName)));

+            session = sessionBuilder.build();
+            if (sessionContext.getTransactionId().isPresent()) {


why are we starting a transaction here ?

NikhilCollooru · 2024-09-23T22:14:02Z

presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java

            query = preparedQuery.getFormattedQuery().orElse(query);

            // select resource group
            Optional<QueryType> queryType = preparedQuery.getQueryType();
+            sessionBuilder.setQueryType(queryType);


why make all these changes instead when you can just do
session = Session.builder(session).setQueryType(queryType).build();

This is because of this line. My change makes it so we still have access to the SessionBuilder object before we begin the transaction. After we start the transaction, we can't use the ctor that passes in Session. I could use the other ctor but would need to refill in all the Session fields.

NikhilCollooru · 2024-09-25T22:10:02Z

I took a closer look and the changes look okay to me.

please squash all commits into one commit
Please click "Please click here to be authorized" to get CLA authorized
Please build a custom presto version and run verifier tests to confirm the changes run fine for a bigger set of queries

RohanSidhu · 2024-09-27T07:31:08Z

I have completed the verifier tests and did not find any significant failures. I ran over 10000 queries and only 12 have failed. Reasons for their failures are environment errors (i.e. com.facebook.presto.ExceededMemoryLimitException). If there is anything else that I needed or proof for the test, please let me know!

RohanSidhu requested review from agrawalreetika, a team, jaystarshot, feilong-liu, elharo and ClarenceThreepwood as code owners September 16, 2024 20:32

RohanSidhu requested a review from presto-oss September 16, 2024 20:32

ghelmling reviewed Sep 18, 2024

View reviewed changes

ghelmling reviewed Sep 20, 2024

View reviewed changes

presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java Outdated Show resolved Hide resolved

presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java Outdated Show resolved Hide resolved

NikhilCollooru reviewed Sep 20, 2024

View reviewed changes

RohanSidhu requested a review from shrinidhijoshi as a code owner September 20, 2024 23:53

RohanSidhu requested review from NikhilCollooru and ghelmling September 23, 2024 18:47

NikhilCollooru reviewed Sep 23, 2024

View reviewed changes

RohanSidhu requested a review from NikhilCollooru September 24, 2024 19:47

RohanSidhu force-pushed the query-reset branch from 8140210 to 7bf7299 Compare September 26, 2024 04:38

RohanSidhu requested review from steveburnett, hantangwangd, ZacBlanco and a team as code owners September 26, 2024 04:38

RohanSidhu force-pushed the query-reset branch 4 times, most recently from f66077a to af469aa Compare September 26, 2024 05:57

steveburnett removed request for ghelmling and steveburnett September 26, 2024 13:53

RohanSidhu force-pushed the query-reset branch from eacfc11 to 437f9b5 Compare September 26, 2024 22:08

Adding QueryType to AccessControlContext

33183e8

RohanSidhu force-pushed the query-reset branch from 437f9b5 to 33183e8 Compare September 27, 2024 07:32

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Adding QueryType to AccessControlContext #23663

Adding QueryType to AccessControlContext #23663

RohanSidhu commented Sep 16, 2024 •

edited

Loading

linux-foundation-easycla bot commented Sep 16, 2024 •

edited

Loading

ghelmling left a comment

ghelmling Sep 18, 2024

RohanSidhu Sep 18, 2024 •

edited

Loading

NikhilCollooru Sep 20, 2024

RohanSidhu Sep 20, 2024 •

edited

Loading

ghelmling left a comment

NikhilCollooru Sep 20, 2024

RohanSidhu Sep 20, 2024

NikhilCollooru Sep 20, 2024

RohanSidhu Sep 20, 2024 •

edited

Loading

NikhilCollooru Sep 23, 2024 •

edited

Loading

RohanSidhu Sep 24, 2024

RohanSidhu Sep 24, 2024

ghelmling Sep 24, 2024

NikhilCollooru Sep 20, 2024

NikhilCollooru Sep 23, 2024

RohanSidhu Sep 24, 2024

NikhilCollooru commented Sep 25, 2024

RohanSidhu commented Sep 27, 2024

Adding QueryType to AccessControlContext #23663

Are you sure you want to change the base?

Adding QueryType to AccessControlContext #23663

Conversation

RohanSidhu commented Sep 16, 2024 • edited Loading

Description

Motivation and Context

Impact

Test Plan

Contributor checklist

Release Notes

linux-foundation-easycla bot commented Sep 16, 2024 • edited Loading

ghelmling left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

RohanSidhu Sep 18, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

RohanSidhu Sep 20, 2024 • edited Loading

Choose a reason for hiding this comment

ghelmling left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

RohanSidhu Sep 20, 2024 • edited Loading

Choose a reason for hiding this comment

NikhilCollooru Sep 23, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

NikhilCollooru commented Sep 25, 2024

RohanSidhu commented Sep 27, 2024

RohanSidhu commented Sep 16, 2024 •

edited

Loading

linux-foundation-easycla bot commented Sep 16, 2024 •

edited

Loading

RohanSidhu Sep 18, 2024 •

edited

Loading

RohanSidhu Sep 20, 2024 •

edited

Loading

RohanSidhu Sep 20, 2024 •

edited

Loading

NikhilCollooru Sep 23, 2024 •

edited

Loading