Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Adding QueryType to AccessControlContext #23663

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
+187 −111
Open

Adding QueryType to AccessControlContext #23663

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .../src/test/java/com/facebook/presto/hive/security/ranger/TestRangerBasedAccessControl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@
public class TestRangerBasedAccessControl
{
public static final ConnectorTransactionHandle TRANSACTION_HANDLE = new ConnectorTransactionHandle() {};
public static final AccessControlContext CONTEXT = new AccessControlContext(new QueryId("query_id"), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats());
public static final AccessControlContext CONTEXT = new AccessControlContext(new QueryId("query_id"), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty());

@Test
public void testTablePriviledgesRolesNotAllowed()
Expand Down
157 changes: 97 additions & 60 deletions presto-main/src/main/java/com/facebook/presto/Session.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

import com.facebook.presto.common.RuntimeStats;
import com.facebook.presto.common.function.SqlFunctionProperties;
import com.facebook.presto.common.resourceGroups.QueryType;
import com.facebook.presto.common.transaction.TransactionId;
import com.facebook.presto.common.type.TimeZoneKey;
import com.facebook.presto.cost.PlanCostEstimate;
Expand Down Expand Up @@ -99,6 +100,7 @@ public final class Session
private final Optional<Tracer> tracer;
private final WarningCollector warningCollector;
private final RuntimeStats runtimeStats;
private final Optional<QueryType> queryType;

private final OptimizerInformationCollector optimizerInformationCollector = new OptimizerInformationCollector();
private final OptimizerResultCollector optimizerResultCollector = new OptimizerResultCollector();
Expand Down Expand Up @@ -131,7 +133,8 @@ public Session(
Map<SqlFunctionId, SqlInvokedFunction> sessionFunctions,
Optional<Tracer> tracer,
WarningCollector warningCollector,
RuntimeStats runtimeStats)
RuntimeStats runtimeStats,
Optional<QueryType> queryType)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Author

@RohanSidhu RohanSidhu Sep 20, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

@NikhilCollooru NikhilCollooru Sep 23, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

  1. Session is instantiated and passed to createAnalyzerOptions()
  2. Same session is passed to getAnalyzerType()
  3. 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.

{
this.queryId = requireNonNull(queryId, "queryId is null");
this.transactionId = requireNonNull(transactionId, "transactionId is null");
Expand Down Expand Up @@ -172,7 +175,8 @@ public Session(
this.tracer = requireNonNull(tracer, "tracer is null");
this.warningCollector = requireNonNull(warningCollector, "warningCollector is null");
this.runtimeStats = requireNonNull(runtimeStats, "runtimeStats is null");
this.context = new AccessControlContext(queryId, clientInfo, clientTags, source, warningCollector, runtimeStats);
this.queryType = requireNonNull(queryType, "queryType is null");
this.context = new AccessControlContext(queryId, clientInfo, clientTags, source, warningCollector, runtimeStats, queryType);
}

public QueryId getQueryId()
Expand Down Expand Up @@ -353,6 +357,11 @@ public Map<PlanNodeId, PlanCostEstimate> getPlanNodeCostMap()
return planNodeCostMap;
}

public Optional<QueryType> getQueryType()
{
return queryType;
}

public Session beginTransactionId(TransactionId transactionId, TransactionManager transactionManager, AccessControl accessControl)
{
requireNonNull(transactionId, "transactionId is null");
Expand Down Expand Up @@ -447,63 +456,8 @@ public Session beginTransactionId(TransactionId transactionId, TransactionManage
sessionFunctions,
tracer,
warningCollector,
runtimeStats);
}

public Session withDefaultProperties(
SystemSessionPropertyConfiguration systemPropertyConfiguration,
Map<String, Map<String, String>> catalogPropertyDefaults)
{
requireNonNull(systemPropertyConfiguration, "systemPropertyConfiguration is null");
requireNonNull(catalogPropertyDefaults, "catalogPropertyDefaults is null");

// to remove this check properties must be authenticated and validated as in beginTransactionId
checkState(
!this.transactionId.isPresent() && this.connectorProperties.isEmpty(),
"Session properties cannot be overridden once a transaction is active");

Map<String, String> systemProperties = new HashMap<>();
systemProperties.putAll(systemPropertyConfiguration.systemPropertyDefaults);
systemProperties.putAll(this.systemProperties);
systemProperties.putAll(systemPropertyConfiguration.systemPropertyOverrides);

Map<String, Map<String, String>> connectorProperties = catalogPropertyDefaults.entrySet().stream()
.map(entry -> Maps.immutableEntry(entry.getKey(), new HashMap<>(entry.getValue())))
.collect(Collectors.toMap(Entry::getKey, Entry::getValue));
for (Entry<String, Map<String, String>> catalogProperties : this.unprocessedCatalogProperties.entrySet()) {
String catalog = catalogProperties.getKey();
for (Entry<String, String> entry : catalogProperties.getValue().entrySet()) {
connectorProperties.computeIfAbsent(catalog, id -> new HashMap<>())
.put(entry.getKey(), entry.getValue());
}
}

return new Session(
queryId,
transactionId,
clientTransactionSupport,
identity,
source,
catalog,
schema,
traceToken,
timeZoneKey,
locale,
remoteUserAddress,
userAgent,
clientInfo,
clientTags,
resourceEstimates,
startTime,
systemProperties,
ImmutableMap.of(),
connectorProperties,
sessionPropertyManager,
preparedStatements,
sessionFunctions,
tracer,
warningCollector,
runtimeStats);
runtimeStats,
queryType);
}

public ConnectorSession toConnectorSession()
Expand Down Expand Up @@ -630,6 +584,7 @@ public static class SessionBuilder
private final SessionPropertyManager sessionPropertyManager;
private final Map<String, String> preparedStatements = new HashMap<>();
private final Map<SqlFunctionId, SqlInvokedFunction> sessionFunctions = new HashMap<>();
private Optional<QueryType> queryType = Optional.empty();
private WarningCollector warningCollector = WarningCollector.NOOP;
private RuntimeStats runtimeStats = new RuntimeStats();

Expand Down Expand Up @@ -665,6 +620,7 @@ private SessionBuilder(Session session)
this.tracer = requireNonNull(session.tracer, "tracer is null");
this.warningCollector = requireNonNull(session.warningCollector, "warningCollector is null");
this.runtimeStats = requireNonNull(session.runtimeStats, "runtimeStats is null");
this.queryType = requireNonNull(session.queryType, "queryType is null");
}

public SessionBuilder setQueryId(QueryId queryId)
Expand Down Expand Up @@ -821,11 +777,57 @@ public SessionBuilder setRuntimeStats(RuntimeStats runtimeStats)
return this;
}

public SessionBuilder setQueryType(Optional<QueryType> queryType)
{
this.queryType = requireNonNull(queryType, "queryType is null");
return this;
}

public <T> T getSystemProperty(String name, Class<T> type)
{
return sessionPropertyManager.decodeSystemPropertyValue(name, systemProperties.get(name), type);
}

public WarningCollector getWarningCollector()
{
return this.warningCollector;
}

public Map<String, String> getPreparedStatements()
{
return this.preparedStatements;
}

public Identity getIdentity()
{
return this.identity;
}

public Optional<String> getSource()
{
return Optional.ofNullable(this.source);
}

public Set<String> getClientTags()
{
return this.clientTags;
}

public Optional<String> getClientInfo()
{
return Optional.ofNullable(this.clientInfo);
}

public Map<String, String> getSystemProperties()
{
return this.systemProperties;
}

public Map<String, Map<String, String>> getUnprocessedCatalogProperties()
{
return this.catalogSessionProperties;
}

public Session build()
{
return new Session(
Expand Down Expand Up @@ -853,7 +855,42 @@ public Session build()
sessionFunctions,
tracer,
warningCollector,
runtimeStats);
runtimeStats,
queryType);
}

public void applyDefaultProperties(SystemSessionPropertyConfiguration systemPropertyConfiguration, Map<String, Map<String, String>> catalogPropertyDefaults)
{
requireNonNull(systemPropertyConfiguration, "systemPropertyConfiguration is null");
requireNonNull(catalogPropertyDefaults, "catalogPropertyDefaults is null");

// to remove this check properties must be authenticated and validated as in beginTransactionId
checkState(
this.transactionId == null && this.connectorProperties.isEmpty(),
"Session properties cannot be overridden once a transaction is active");

Map<String, String> systemProperties = new HashMap<>();
systemProperties.putAll(systemPropertyConfiguration.systemPropertyDefaults);
systemProperties.putAll(this.systemProperties);
systemProperties.putAll(systemPropertyConfiguration.systemPropertyOverrides);
this.systemProperties.putAll(systemProperties);

Map<String, Map<String, String>> connectorProperties = catalogPropertyDefaults.entrySet().stream()
.map(entry -> Maps.immutableEntry(entry.getKey(), new HashMap<>(entry.getValue())))
.collect(Collectors.toMap(Entry::getKey, Entry::getValue));
for (Entry<String, Map<String, String>> catalogProperties : this.catalogSessionProperties.entrySet()) {
String catalog = catalogProperties.getKey();
for (Entry<String, String> entry : catalogProperties.getValue().entrySet()) {
connectorProperties.computeIfAbsent(catalog, id -> new HashMap<>()).put(entry.getKey(), entry.getValue());
}
}

for (Entry<String, Map<String, String>> catalogProperties : connectorProperties.entrySet()) {
String catalog = catalogProperties.getKey();
for (Entry<String, String> entry : catalogProperties.getValue().entrySet()) {
setCatalogSessionProperty(catalog, entry.getKey(), entry.getValue());
}
}
}
}

Expand Down
3 changes: 2 additions & 1 deletion presto-main/src/main/java/com/facebook/presto/SessionRepresentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -338,6 +338,7 @@ public Session toSession(SessionPropertyManager sessionPropertyManager, Map<Stri
Optional.empty(),
// we use NOOP to create a session from the representation as worker does not require warning collectors
WarningCollector.NOOP,
new RuntimeStats());
new RuntimeStats(),
Optional.empty());
}
}
17 changes: 13 additions & 4 deletions presto-main/src/main/java/com/facebook/presto/dispatcher/DispatchManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@
import java.util.Optional;
import java.util.concurrent.Executor;

import static com.facebook.presto.Session.SessionBuilder;
import static com.facebook.presto.SystemSessionProperties.getAnalyzerType;
import static com.facebook.presto.spi.StandardErrorCode.QUERY_TEXT_TOO_LARGE;
import static com.facebook.presto.util.AnalyzerUtil.createAnalyzerOptions;
Expand Down Expand Up @@ -259,6 +260,7 @@ public ListenableFuture<?> createQuery(QueryId queryId, String slug, int retryCo
private <C> void createQueryInternal(QueryId queryId, String slug, int retryCount, SessionContext sessionContext, String query, ResourceGroupManager<C> resourceGroupManager)
{
Session session = null;
SessionBuilder sessionBuilder = null;
PreparedQuery preparedQuery;
try {
if (query.length() > maxQueryLength) {
Expand All @@ -268,16 +270,18 @@ private <C> void createQueryInternal(QueryId queryId, String slug, int retryCoun
}

// decode session
session = sessionSupplier.createSession(queryId, sessionContext, warningCollectorFactory);
sessionBuilder = sessionSupplier.createSessionBuilder(queryId, sessionContext, warningCollectorFactory);
session = sessionBuilder.build();

// prepare query
AnalyzerOptions analyzerOptions = createAnalyzerOptions(session, session.getWarningCollector());
AnalyzerOptions analyzerOptions = createAnalyzerOptions(session, sessionBuilder.getWarningCollector());
QueryPreparerProvider queryPreparerProvider = queryPreparerProviderManager.getQueryPreparerProvider(getAnalyzerType(session));
preparedQuery = queryPreparerProvider.getQueryPreparer().prepareQuery(analyzerOptions, query, session.getPreparedStatements(), session.getWarningCollector());
preparedQuery = queryPreparerProvider.getQueryPreparer().prepareQuery(analyzerOptions, query, sessionBuilder.getPreparedStatements(), sessionBuilder.getWarningCollector());
query = preparedQuery.getFormattedQuery().orElse(query);

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

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

SelectionContext<C> selectionContext = resourceGroupManager.selectGroup(new SelectionCriteria(
sessionContext.getIdentity().getPrincipal().isPresent(),
sessionContext.getIdentity().getUser(),
Expand All @@ -290,7 +294,12 @@ private <C> void createQueryInternal(QueryId queryId, String slug, int retryCoun
sessionContext.getIdentity().getPrincipal().map(Principal::getName)));

// apply system default session properties (does not override user set properties)
session = sessionPropertyDefaults.newSessionWithDefaultProperties(session, queryType.map(Enum::name), Optional.of(selectionContext.getResourceGroupId()));
sessionPropertyDefaults.applyDefaultProperties(sessionBuilder, queryType.map(Enum::name), Optional.of(selectionContext.getResourceGroupId()));

session = sessionBuilder.build();
if (sessionContext.getTransactionId().isPresent()) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Author

@RohanSidhu RohanSidhu Sep 18, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why are we starting a transaction here ?

Copy link
Author

@RohanSidhu RohanSidhu Sep 20, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

session = session.beginTransactionId(sessionContext.getTransactionId().get(), transactionManager, accessControl);
}

// mark existing transaction as active
transactionManager.activateTransaction(session, preparedQuery.isTransactionControlStatement(), accessControl);
Expand Down
6 changes: 4 additions & 2 deletions presto-main/src/main/java/com/facebook/presto/security/AccessControlUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,8 @@ public static void checkPermissions(AccessControl accessControl, SecurityConfig
sessionContext.getClientTags(),
Optional.ofNullable(sessionContext.getSource()),
WarningCollector.NOOP,
sessionContext.getRuntimeStats()),
sessionContext.getRuntimeStats(),
Optional.empty()),
identity.getPrincipal(),
identity.getUser());
}
Expand All @@ -71,7 +72,8 @@ public static Optional<AuthorizedIdentity> getAuthorizedIdentity(AccessControl a
sessionContext.getClientTags(),
Optional.ofNullable(sessionContext.getSource()),
WarningCollector.NOOP,
sessionContext.getRuntimeStats()),
sessionContext.getRuntimeStats(),
Optional.empty()),
identity.getUser(),
sessionContext.getCertificates());
return Optional.of(authorizedIdentity);
Expand Down
8 changes: 8 additions & 0 deletions presto-main/src/main/java/com/facebook/presto/server/NoOpSessionSupplier.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
import com.facebook.presto.execution.warnings.WarningCollectorFactory;
import com.facebook.presto.spi.QueryId;

import static com.facebook.presto.Session.SessionBuilder;

/**
* Used on workers.
*/
Expand All @@ -28,4 +30,10 @@ public Session createSession(QueryId queryId, SessionContext context, WarningCol
{
throw new UnsupportedOperationException();
}

@Override
public SessionBuilder createSessionBuilder(QueryId queryId, SessionContext context, WarningCollectorFactory warningCollectorFactory)
{
throw new UnsupportedOperationException();
}
}
16 changes: 11 additions & 5 deletions presto-main/src/main/java/com/facebook/presto/server/QuerySessionSupplier.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,16 @@ public QuerySessionSupplier(

@Override
public Session createSession(QueryId queryId, SessionContext context, WarningCollectorFactory warningCollectorFactory)
{
Session session = createSessionBuilder(queryId, context, warningCollectorFactory).build();
if (context.getTransactionId().isPresent()) {
session = session.beginTransactionId(context.getTransactionId().get(), transactionManager, accessControl);
}
return session;
}

@Override
public SessionBuilder createSessionBuilder(QueryId queryId, SessionContext context, WarningCollectorFactory warningCollectorFactory)
{
SessionBuilder sessionBuilder = Session.builder(sessionPropertyManager)
.setQueryId(queryId)
Expand Down Expand Up @@ -128,11 +138,7 @@ else if (context.getTimeZoneId() != null) {
WarningCollector warningCollector = warningCollectorFactory.create(sessionBuilder.getSystemProperty(WARNING_HANDLING, WarningHandlingLevel.class));
sessionBuilder.setWarningCollector(warningCollector);

Session session = sessionBuilder.build();
if (context.getTransactionId().isPresent()) {
session = session.beginTransactionId(context.getTransactionId().get(), transactionManager, accessControl);
}
return session;
return sessionBuilder;
}

private Identity authenticateIdentity(QueryId queryId, SessionContext context)
Expand Down
Loading
Loading