Skip to content

Commit 8385163

Browse files
RohanSidhuNikhilCollooru
RohanSidhu
authored and
NikhilCollooru
committed
Add QueryType to AccessControlContext
1 parent 375dd5a commit 8385163

File tree

17 files changed

+187
-111
lines changed

17 files changed

+187
-111
lines changed

‎presto-hive/src/test/java/com/facebook/presto/hive/security/ranger/TestRangerBasedAccessControl.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@
4848
public class TestRangerBasedAccessControl
4949
{
5050
public static final ConnectorTransactionHandle TRANSACTION_HANDLE = new ConnectorTransactionHandle() {};
51-
public static final AccessControlContext CONTEXT = new AccessControlContext(new QueryId("query_id"), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats());
51+
public static final AccessControlContext CONTEXT = new AccessControlContext(new QueryId("query_id"), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty());
5252

5353
@Test
5454
public void testTablePriviledgesRolesNotAllowed()

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

Lines changed: 97 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515

1616
import com.facebook.presto.common.RuntimeStats;
1717
import com.facebook.presto.common.function.SqlFunctionProperties;
18+
import com.facebook.presto.common.resourceGroups.QueryType;
1819
import com.facebook.presto.common.transaction.TransactionId;
1920
import com.facebook.presto.common.type.TimeZoneKey;
2021
import com.facebook.presto.cost.PlanCostEstimate;
@@ -99,6 +100,7 @@ public final class Session
99100
private final Optional<Tracer> tracer;
100101
private final WarningCollector warningCollector;
101102
private final RuntimeStats runtimeStats;
103+
private final Optional<QueryType> queryType;
102104

103105
private final OptimizerInformationCollector optimizerInformationCollector = new OptimizerInformationCollector();
104106
private final OptimizerResultCollector optimizerResultCollector = new OptimizerResultCollector();
@@ -131,7 +133,8 @@ public Session(
131133
Map<SqlFunctionId, SqlInvokedFunction> sessionFunctions,
132134
Optional<Tracer> tracer,
133135
WarningCollector warningCollector,
134-
RuntimeStats runtimeStats)
136+
RuntimeStats runtimeStats,
137+
Optional<QueryType> queryType)
135138
{
136139
this.queryId = requireNonNull(queryId, "queryId is null");
137140
this.transactionId = requireNonNull(transactionId, "transactionId is null");
@@ -172,7 +175,8 @@ public Session(
172175
this.tracer = requireNonNull(tracer, "tracer is null");
173176
this.warningCollector = requireNonNull(warningCollector, "warningCollector is null");
174177
this.runtimeStats = requireNonNull(runtimeStats, "runtimeStats is null");
175-
this.context = new AccessControlContext(queryId, clientInfo, clientTags, source, warningCollector, runtimeStats);
178+
this.queryType = requireNonNull(queryType, "queryType is null");
179+
this.context = new AccessControlContext(queryId, clientInfo, clientTags, source, warningCollector, runtimeStats, queryType);
176180
}
177181

178182
public QueryId getQueryId()
@@ -353,6 +357,11 @@ public Map<PlanNodeId, PlanCostEstimate> getPlanNodeCostMap()
353357
return planNodeCostMap;
354358
}
355359

360+
public Optional<QueryType> getQueryType()
361+
{
362+
return queryType;
363+
}
364+
356365
public Session beginTransactionId(TransactionId transactionId, TransactionManager transactionManager, AccessControl accessControl)
357366
{
358367
requireNonNull(transactionId, "transactionId is null");
@@ -447,63 +456,8 @@ public Session beginTransactionId(TransactionId transactionId, TransactionManage
447456
sessionFunctions,
448457
tracer,
449458
warningCollector,
450-
runtimeStats);
451-
}
452-
453-
public Session withDefaultProperties(
454-
SystemSessionPropertyConfiguration systemPropertyConfiguration,
455-
Map<String, Map<String, String>> catalogPropertyDefaults)
456-
{
457-
requireNonNull(systemPropertyConfiguration, "systemPropertyConfiguration is null");
458-
requireNonNull(catalogPropertyDefaults, "catalogPropertyDefaults is null");
459-
460-
// to remove this check properties must be authenticated and validated as in beginTransactionId
461-
checkState(
462-
!this.transactionId.isPresent() && this.connectorProperties.isEmpty(),
463-
"Session properties cannot be overridden once a transaction is active");
464-
465-
Map<String, String> systemProperties = new HashMap<>();
466-
systemProperties.putAll(systemPropertyConfiguration.systemPropertyDefaults);
467-
systemProperties.putAll(this.systemProperties);
468-
systemProperties.putAll(systemPropertyConfiguration.systemPropertyOverrides);
469-
470-
Map<String, Map<String, String>> connectorProperties = catalogPropertyDefaults.entrySet().stream()
471-
.map(entry -> Maps.immutableEntry(entry.getKey(), new HashMap<>(entry.getValue())))
472-
.collect(Collectors.toMap(Entry::getKey, Entry::getValue));
473-
for (Entry<String, Map<String, String>> catalogProperties : this.unprocessedCatalogProperties.entrySet()) {
474-
String catalog = catalogProperties.getKey();
475-
for (Entry<String, String> entry : catalogProperties.getValue().entrySet()) {
476-
connectorProperties.computeIfAbsent(catalog, id -> new HashMap<>())
477-
.put(entry.getKey(), entry.getValue());
478-
}
479-
}
480-
481-
return new Session(
482-
queryId,
483-
transactionId,
484-
clientTransactionSupport,
485-
identity,
486-
source,
487-
catalog,
488-
schema,
489-
traceToken,
490-
timeZoneKey,
491-
locale,
492-
remoteUserAddress,
493-
userAgent,
494-
clientInfo,
495-
clientTags,
496-
resourceEstimates,
497-
startTime,
498-
systemProperties,
499-
ImmutableMap.of(),
500-
connectorProperties,
501-
sessionPropertyManager,
502-
preparedStatements,
503-
sessionFunctions,
504-
tracer,
505-
warningCollector,
506-
runtimeStats);
459+
runtimeStats,
460+
queryType);
507461
}
508462

509463
public ConnectorSession toConnectorSession()
@@ -630,6 +584,7 @@ public static class SessionBuilder
630584
private final SessionPropertyManager sessionPropertyManager;
631585
private final Map<String, String> preparedStatements = new HashMap<>();
632586
private final Map<SqlFunctionId, SqlInvokedFunction> sessionFunctions = new HashMap<>();
587+
private Optional<QueryType> queryType = Optional.empty();
633588
private WarningCollector warningCollector = WarningCollector.NOOP;
634589
private RuntimeStats runtimeStats = new RuntimeStats();
635590

@@ -665,6 +620,7 @@ private SessionBuilder(Session session)
665620
this.tracer = requireNonNull(session.tracer, "tracer is null");
666621
this.warningCollector = requireNonNull(session.warningCollector, "warningCollector is null");
667622
this.runtimeStats = requireNonNull(session.runtimeStats, "runtimeStats is null");
623+
this.queryType = requireNonNull(session.queryType, "queryType is null");
668624
}
669625

670626
public SessionBuilder setQueryId(QueryId queryId)
@@ -821,11 +777,57 @@ public SessionBuilder setRuntimeStats(RuntimeStats runtimeStats)
821777
return this;
822778
}
823779

780+
public SessionBuilder setQueryType(Optional<QueryType> queryType)
781+
{
782+
this.queryType = requireNonNull(queryType, "queryType is null");
783+
return this;
784+
}
785+
824786
public <T> T getSystemProperty(String name, Class<T> type)
825787
{
826788
return sessionPropertyManager.decodeSystemPropertyValue(name, systemProperties.get(name), type);
827789
}
828790

791+
public WarningCollector getWarningCollector()
792+
{
793+
return this.warningCollector;
794+
}
795+
796+
public Map<String, String> getPreparedStatements()
797+
{
798+
return this.preparedStatements;
799+
}
800+
801+
public Identity getIdentity()
802+
{
803+
return this.identity;
804+
}
805+
806+
public Optional<String> getSource()
807+
{
808+
return Optional.ofNullable(this.source);
809+
}
810+
811+
public Set<String> getClientTags()
812+
{
813+
return this.clientTags;
814+
}
815+
816+
public Optional<String> getClientInfo()
817+
{
818+
return Optional.ofNullable(this.clientInfo);
819+
}
820+
821+
public Map<String, String> getSystemProperties()
822+
{
823+
return this.systemProperties;
824+
}
825+
826+
public Map<String, Map<String, String>> getUnprocessedCatalogProperties()
827+
{
828+
return this.catalogSessionProperties;
829+
}
830+
829831
public Session build()
830832
{
831833
return new Session(
@@ -853,7 +855,42 @@ public Session build()
853855
sessionFunctions,
854856
tracer,
855857
warningCollector,
856-
runtimeStats);
858+
runtimeStats,
859+
queryType);
860+
}
861+
862+
public void applyDefaultProperties(SystemSessionPropertyConfiguration systemPropertyConfiguration, Map<String, Map<String, String>> catalogPropertyDefaults)
863+
{
864+
requireNonNull(systemPropertyConfiguration, "systemPropertyConfiguration is null");
865+
requireNonNull(catalogPropertyDefaults, "catalogPropertyDefaults is null");
866+
867+
// to remove this check properties must be authenticated and validated as in beginTransactionId
868+
checkState(
869+
this.transactionId == null && this.connectorProperties.isEmpty(),
870+
"Session properties cannot be overridden once a transaction is active");
871+
872+
Map<String, String> systemProperties = new HashMap<>();
873+
systemProperties.putAll(systemPropertyConfiguration.systemPropertyDefaults);
874+
systemProperties.putAll(this.systemProperties);
875+
systemProperties.putAll(systemPropertyConfiguration.systemPropertyOverrides);
876+
this.systemProperties.putAll(systemProperties);
877+
878+
Map<String, Map<String, String>> connectorProperties = catalogPropertyDefaults.entrySet().stream()
879+
.map(entry -> Maps.immutableEntry(entry.getKey(), new HashMap<>(entry.getValue())))
880+
.collect(Collectors.toMap(Entry::getKey, Entry::getValue));
881+
for (Entry<String, Map<String, String>> catalogProperties : this.catalogSessionProperties.entrySet()) {
882+
String catalog = catalogProperties.getKey();
883+
for (Entry<String, String> entry : catalogProperties.getValue().entrySet()) {
884+
connectorProperties.computeIfAbsent(catalog, id -> new HashMap<>()).put(entry.getKey(), entry.getValue());
885+
}
886+
}
887+
888+
for (Entry<String, Map<String, String>> catalogProperties : connectorProperties.entrySet()) {
889+
String catalog = catalogProperties.getKey();
890+
for (Entry<String, String> entry : catalogProperties.getValue().entrySet()) {
891+
setCatalogSessionProperty(catalog, entry.getKey(), entry.getValue());
892+
}
893+
}
857894
}
858895
}
859896

‎presto-main/src/main/java/com/facebook/presto/SessionRepresentation.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,7 @@ public Session toSession(SessionPropertyManager sessionPropertyManager, Map<Stri
338338
Optional.empty(),
339339
// we use NOOP to create a session from the representation as worker does not require warning collectors
340340
WarningCollector.NOOP,
341-
new RuntimeStats());
341+
new RuntimeStats(),
342+
Optional.empty());
342343
}
343344
}

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

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
import java.util.Optional;
5555
import java.util.concurrent.Executor;
5656

57+
import static com.facebook.presto.Session.SessionBuilder;
5758
import static com.facebook.presto.SystemSessionProperties.getAnalyzerType;
5859
import static com.facebook.presto.spi.StandardErrorCode.QUERY_TEXT_TOO_LARGE;
5960
import static com.facebook.presto.util.AnalyzerUtil.createAnalyzerOptions;
@@ -259,6 +260,7 @@ public ListenableFuture<?> createQuery(QueryId queryId, String slug, int retryCo
259260
private <C> void createQueryInternal(QueryId queryId, String slug, int retryCount, SessionContext sessionContext, String query, ResourceGroupManager<C> resourceGroupManager)
260261
{
261262
Session session = null;
263+
SessionBuilder sessionBuilder = null;
262264
PreparedQuery preparedQuery;
263265
try {
264266
if (query.length() > maxQueryLength) {
@@ -268,16 +270,18 @@ private <C> void createQueryInternal(QueryId queryId, String slug, int retryCoun
268270
}
269271

270272
// decode session
271-
session = sessionSupplier.createSession(queryId, sessionContext, warningCollectorFactory);
273+
sessionBuilder = sessionSupplier.createSessionBuilder(queryId, sessionContext, warningCollectorFactory);
274+
session = sessionBuilder.build();
272275

273276
// prepare query
274-
AnalyzerOptions analyzerOptions = createAnalyzerOptions(session, session.getWarningCollector());
277+
AnalyzerOptions analyzerOptions = createAnalyzerOptions(session, sessionBuilder.getWarningCollector());
275278
QueryPreparerProvider queryPreparerProvider = queryPreparerProviderManager.getQueryPreparerProvider(getAnalyzerType(session));
276-
preparedQuery = queryPreparerProvider.getQueryPreparer().prepareQuery(analyzerOptions, query, session.getPreparedStatements(), session.getWarningCollector());
279+
preparedQuery = queryPreparerProvider.getQueryPreparer().prepareQuery(analyzerOptions, query, sessionBuilder.getPreparedStatements(), sessionBuilder.getWarningCollector());
277280
query = preparedQuery.getFormattedQuery().orElse(query);
278281

279282
// select resource group
280283
Optional<QueryType> queryType = preparedQuery.getQueryType();
284+
sessionBuilder.setQueryType(queryType);
281285
SelectionContext<C> selectionContext = resourceGroupManager.selectGroup(new SelectionCriteria(
282286
sessionContext.getIdentity().getPrincipal().isPresent(),
283287
sessionContext.getIdentity().getUser(),
@@ -290,7 +294,12 @@ private <C> void createQueryInternal(QueryId queryId, String slug, int retryCoun
290294
sessionContext.getIdentity().getPrincipal().map(Principal::getName)));
291295

292296
// apply system default session properties (does not override user set properties)
293-
session = sessionPropertyDefaults.newSessionWithDefaultProperties(session, queryType.map(Enum::name), Optional.of(selectionContext.getResourceGroupId()));
297+
sessionPropertyDefaults.applyDefaultProperties(sessionBuilder, queryType.map(Enum::name), Optional.of(selectionContext.getResourceGroupId()));
298+
299+
session = sessionBuilder.build();
300+
if (sessionContext.getTransactionId().isPresent()) {
301+
session = session.beginTransactionId(sessionContext.getTransactionId().get(), transactionManager, accessControl);
302+
}
294303

295304
// mark existing transaction as active
296305
transactionManager.activateTransaction(session, preparedQuery.isTransactionControlStatement(), accessControl);

‎presto-main/src/main/java/com/facebook/presto/security/AccessControlUtils.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,8 @@ public static void checkPermissions(AccessControl accessControl, SecurityConfig
4646
sessionContext.getClientTags(),
4747
Optional.ofNullable(sessionContext.getSource()),
4848
WarningCollector.NOOP,
49-
sessionContext.getRuntimeStats()),
49+
sessionContext.getRuntimeStats(),
50+
Optional.empty()),
5051
identity.getPrincipal(),
5152
identity.getUser());
5253
}
@@ -71,7 +72,8 @@ public static Optional<AuthorizedIdentity> getAuthorizedIdentity(AccessControl a
7172
sessionContext.getClientTags(),
7273
Optional.ofNullable(sessionContext.getSource()),
7374
WarningCollector.NOOP,
74-
sessionContext.getRuntimeStats()),
75+
sessionContext.getRuntimeStats(),
76+
Optional.empty()),
7577
identity.getUser(),
7678
sessionContext.getCertificates());
7779
return Optional.of(authorizedIdentity);

‎presto-main/src/main/java/com/facebook/presto/server/NoOpSessionSupplier.java

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
import com.facebook.presto.execution.warnings.WarningCollectorFactory;
1818
import com.facebook.presto.spi.QueryId;
1919

20+
import static com.facebook.presto.Session.SessionBuilder;
21+
2022
/**
2123
* Used on workers.
2224
*/
@@ -28,4 +30,10 @@ public Session createSession(QueryId queryId, SessionContext context, WarningCol
2830
{
2931
throw new UnsupportedOperationException();
3032
}
33+
34+
@Override
35+
public SessionBuilder createSessionBuilder(QueryId queryId, SessionContext context, WarningCollectorFactory warningCollectorFactory)
36+
{
37+
throw new UnsupportedOperationException();
38+
}
3139
}

‎presto-main/src/main/java/com/facebook/presto/server/QuerySessionSupplier.java

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,16 @@ public QuerySessionSupplier(
7575

7676
@Override
7777
public Session createSession(QueryId queryId, SessionContext context, WarningCollectorFactory warningCollectorFactory)
78+
{
79+
Session session = createSessionBuilder(queryId, context, warningCollectorFactory).build();
80+
if (context.getTransactionId().isPresent()) {
81+
session = session.beginTransactionId(context.getTransactionId().get(), transactionManager, accessControl);
82+
}
83+
return session;
84+
}
85+
86+
@Override
87+
public SessionBuilder createSessionBuilder(QueryId queryId, SessionContext context, WarningCollectorFactory warningCollectorFactory)
7888
{
7989
SessionBuilder sessionBuilder = Session.builder(sessionPropertyManager)
8090
.setQueryId(queryId)
@@ -128,11 +138,7 @@ else if (context.getTimeZoneId() != null) {
128138
WarningCollector warningCollector = warningCollectorFactory.create(sessionBuilder.getSystemProperty(WARNING_HANDLING, WarningHandlingLevel.class));
129139
sessionBuilder.setWarningCollector(warningCollector);
130140

131-
Session session = sessionBuilder.build();
132-
if (context.getTransactionId().isPresent()) {
133-
session = session.beginTransactionId(context.getTransactionId().get(), transactionManager, accessControl);
134-
}
135-
return session;
141+
return sessionBuilder;
136142
}
137143

138144
private Identity authenticateIdentity(QueryId queryId, SessionContext context)

0 commit comments

Comments
 (0)