Skip to content

Commit e2ee434

Browse files
kevintang2022rschlussel
kevintang2022
authored and
rschlussel
committed
Add checkQueryIntegrity to DispatchManager and Add schema/catalog to AccessControlContext
Add call to check query integrity Checkstyle Fix constrcutor usage Fix usages of constructor Whitespace Add test coverage for access control context Rename to schema from sqlnamespace Fix catalog name
1 parent 15afc5c commit e2ee434

File tree

9 files changed

+65
-15
lines changed

9 files changed

+65
-15
lines changed

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

+1-1
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(), Optional.empty());
51+
public static final AccessControlContext CONTEXT = new AccessControlContext(new QueryId("query_id"), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty(), Optional.empty(), Optional.empty());
5252

5353
@Test
5454
public void testTablePriviledgesRolesNotAllowed()

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,7 @@ public Session(
176176
this.warningCollector = requireNonNull(warningCollector, "warningCollector is null");
177177
this.runtimeStats = requireNonNull(runtimeStats, "runtimeStats is null");
178178
this.queryType = requireNonNull(queryType, "queryType is null");
179-
this.context = new AccessControlContext(queryId, clientInfo, clientTags, source, warningCollector, runtimeStats, queryType);
179+
this.context = new AccessControlContext(queryId, clientInfo, clientTags, source, warningCollector, runtimeStats, queryType, catalog, schema);
180180
}
181181

182182
public QueryId getQueryId()

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

+15
Original file line numberDiff line numberDiff line change
@@ -32,11 +32,13 @@
3232
import com.facebook.presto.server.SessionSupplier;
3333
import com.facebook.presto.spi.PrestoException;
3434
import com.facebook.presto.spi.QueryId;
35+
import com.facebook.presto.spi.WarningCollector;
3536
import com.facebook.presto.spi.analyzer.AnalyzerOptions;
3637
import com.facebook.presto.spi.analyzer.QueryPreparerProvider;
3738
import com.facebook.presto.spi.resourceGroups.SelectionContext;
3839
import com.facebook.presto.spi.resourceGroups.SelectionCriteria;
3940
import com.facebook.presto.spi.security.AccessControl;
41+
import com.facebook.presto.spi.security.AccessControlContext;
4042
import com.facebook.presto.sql.analyzer.QueryPreparerProviderManager;
4143
import com.facebook.presto.transaction.TransactionManager;
4244
import com.google.common.util.concurrent.AbstractFuture;
@@ -271,6 +273,19 @@ private <C> void createQueryInternal(QueryId queryId, String slug, int retryCoun
271273

272274
// decode session
273275
sessionBuilder = sessionSupplier.createSessionBuilder(queryId, sessionContext, warningCollectorFactory);
276+
277+
AccessControlContext accessControlContext = new AccessControlContext(
278+
queryId,
279+
Optional.ofNullable(sessionContext.getClientInfo()),
280+
sessionContext.getClientTags(),
281+
Optional.ofNullable(sessionContext.getSource()),
282+
WarningCollector.NOOP,
283+
sessionContext.getRuntimeStats(),
284+
Optional.empty(),
285+
Optional.ofNullable(sessionContext.getCatalog()),
286+
Optional.ofNullable(sessionContext.getSchema()));
287+
288+
accessControl.checkQueryIntegrity(sessionContext.getIdentity(), accessControlContext, query);
274289
session = sessionBuilder.build();
275290

276291
// prepare query

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

+6-2
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,9 @@ public static void checkPermissions(AccessControl accessControl, SecurityConfig
4747
Optional.ofNullable(sessionContext.getSource()),
4848
WarningCollector.NOOP,
4949
sessionContext.getRuntimeStats(),
50-
Optional.empty()),
50+
Optional.empty(),
51+
Optional.ofNullable(sessionContext.getCatalog()),
52+
Optional.ofNullable(sessionContext.getSchema())),
5153
identity.getPrincipal(),
5254
identity.getUser());
5355
}
@@ -73,7 +75,9 @@ public static Optional<AuthorizedIdentity> getAuthorizedIdentity(AccessControl a
7375
Optional.ofNullable(sessionContext.getSource()),
7476
WarningCollector.NOOP,
7577
sessionContext.getRuntimeStats(),
76-
Optional.empty()),
78+
Optional.empty(),
79+
Optional.ofNullable(sessionContext.getCatalog()),
80+
Optional.ofNullable(sessionContext.getSchema())),
7781
identity.getUser(),
7882
sessionContext.getCertificates());
7983
return Optional.of(authorizedIdentity);

‎presto-main/src/test/java/com/facebook/presto/security/TestAccessControlManager.java

+8-8
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ public void testInitializing()
8282
AccessControlManager accessControlManager = new AccessControlManager(createTestTransactionManager());
8383
accessControlManager.checkCanSetUser(
8484
new Identity(USER_NAME, Optional.of(PRINCIPAL)),
85-
new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty()),
85+
new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty(), Optional.empty(), Optional.empty()),
8686
Optional.empty(),
8787
"foo");
8888
}
@@ -94,7 +94,7 @@ public void testNoneSystemAccessControl()
9494
accessControlManager.setSystemAccessControl(AllowAllSystemAccessControl.NAME, ImmutableMap.of());
9595
accessControlManager.checkCanSetUser(
9696
new Identity(USER_NAME, Optional.of(PRINCIPAL)),
97-
new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty()),
97+
new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty(), Optional.empty(), Optional.empty()),
9898
Optional.empty(),
9999
USER_NAME);
100100
}
@@ -106,7 +106,7 @@ public void testReadOnlySystemAccessControl()
106106
QualifiedObjectName tableName = new QualifiedObjectName("catalog", "schema", "table");
107107
TransactionManager transactionManager = createTestTransactionManager();
108108
AccessControlManager accessControlManager = new AccessControlManager(transactionManager);
109-
AccessControlContext context = new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty());
109+
AccessControlContext context = new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty(), Optional.empty(), Optional.empty());
110110

111111
accessControlManager.setSystemAccessControl(ReadOnlySystemAccessControl.NAME, ImmutableMap.of());
112112
accessControlManager.checkCanSetUser(identity, context, Optional.of(PRINCIPAL), USER_NAME);
@@ -149,7 +149,7 @@ public void testSetAccessControl()
149149

150150
accessControlManager.checkCanSetUser(
151151
new Identity(USER_NAME, Optional.of(PRINCIPAL)),
152-
new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty()),
152+
new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty(), Optional.empty(), Optional.empty()),
153153
Optional.of(PRINCIPAL),
154154
USER_NAME);
155155
assertEquals(accessControlFactory.getCheckedUserName(), USER_NAME);
@@ -160,7 +160,7 @@ public void testSetAccessControl()
160160
public void testCheckQueryIntegrity()
161161
{
162162
AccessControlManager accessControlManager = new AccessControlManager(createTestTransactionManager());
163-
AccessControlContext context = new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty());
163+
AccessControlContext context = new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty(), Optional.empty(), Optional.empty());
164164

165165
TestSystemAccessControlFactory accessControlFactory = new TestSystemAccessControlFactory("test");
166166
accessControlManager.addSystemAccessControlFactory(accessControlFactory);
@@ -210,7 +210,7 @@ public void testNoCatalogAccessControl()
210210
transaction(transactionManager, accessControlManager)
211211
.execute(transactionId -> {
212212
accessControlManager.checkCanSelectFromColumns(transactionId, new Identity(USER_NAME, Optional.of(PRINCIPAL)),
213-
new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty()),
213+
new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty(), Optional.empty(), Optional.empty()),
214214
new QualifiedObjectName("catalog", "schema", "table"), ImmutableSet.of(new Subfield("column")));
215215
});
216216
}
@@ -232,7 +232,7 @@ public void testDenyCatalogAccessControl()
232232
transaction(transactionManager, accessControlManager)
233233
.execute(transactionId -> {
234234
accessControlManager.checkCanSelectFromColumns(transactionId, new Identity(USER_NAME, Optional.of(PRINCIPAL)),
235-
new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty()),
235+
new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty(), Optional.empty(), Optional.empty()),
236236
new QualifiedObjectName("catalog", "schema", "table"), ImmutableSet.of(new Subfield("column")));
237237
});
238238
}
@@ -254,7 +254,7 @@ public void testDenySystemAccessControl()
254254
transaction(transactionManager, accessControlManager)
255255
.execute(transactionId -> {
256256
accessControlManager.checkCanSelectFromColumns(transactionId, new Identity(USER_NAME, Optional.of(PRINCIPAL)),
257-
new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty()),
257+
new AccessControlContext(new QueryId(QUERY_ID), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty(), Optional.empty(), Optional.empty()),
258258
new QualifiedObjectName("secured_catalog", "schema", "table"), ImmutableSet.of(new Subfield("column")));
259259
});
260260
}

‎presto-main/src/test/java/com/facebook/presto/security/TestFileBasedSystemAccessControl.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ public class TestFileBasedSystemAccessControl
7070
private static final QualifiedObjectName aliceTable = new QualifiedObjectName("alice-catalog", "schema", "table");
7171
private static final QualifiedObjectName aliceView = new QualifiedObjectName("alice-catalog", "schema", "view");
7272
private static final CatalogSchemaName aliceSchema = new CatalogSchemaName("alice-catalog", "schema");
73-
private static final AccessControlContext context = new AccessControlContext(new QueryId("query_id"), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty());
73+
private static final AccessControlContext context = new AccessControlContext(new QueryId("query_id"), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty(), Optional.empty(), Optional.empty());
7474
@Test
7575
public void testCanSetUserOperations() throws IOException
7676
{

‎presto-plugin-toolkit/src/test/java/com/facebook/presto/plugin/base/security/TestFileBasedAccessControl.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@
4343
public class TestFileBasedAccessControl
4444
{
4545
public static final ConnectorTransactionHandle TRANSACTION_HANDLE = new ConnectorTransactionHandle() {};
46-
public static final AccessControlContext CONTEXT = new AccessControlContext(new QueryId("query_id"), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty());
46+
public static final AccessControlContext CONTEXT = new AccessControlContext(new QueryId("query_id"), Optional.empty(), Collections.emptySet(), Optional.empty(), WarningCollector.NOOP, new RuntimeStats(), Optional.empty(), Optional.empty(), Optional.empty());
4747

4848
@Test
4949
public void testSchemaRules()

‎presto-spi/src/main/java/com/facebook/presto/spi/security/AccessControlContext.java

+24-1
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,19 @@ public class AccessControlContext
3232
private final WarningCollector warningCollector;
3333
private final RuntimeStats runtimeStats;
3434
private final Optional<QueryType> queryType;
35+
private final Optional<String> catalog;
36+
private final Optional<String> schema;
3537

36-
public AccessControlContext(QueryId queryId, Optional<String> clientInfo, Set<String> clientTags, Optional<String> source, WarningCollector warningCollector, RuntimeStats runtimeStats, Optional<QueryType> queryType)
38+
public AccessControlContext(
39+
QueryId queryId,
40+
Optional<String> clientInfo,
41+
Set<String> clientTags,
42+
Optional<String> source,
43+
WarningCollector warningCollector,
44+
RuntimeStats runtimeStats,
45+
Optional<QueryType> queryType,
46+
Optional<String> catalog,
47+
Optional<String> schema)
3748
{
3849
this.queryId = requireNonNull(queryId, "queryId is null");
3950
this.clientInfo = requireNonNull(clientInfo, "clientInfo is null");
@@ -42,6 +53,8 @@ public AccessControlContext(QueryId queryId, Optional<String> clientInfo, Set<St
4253
this.warningCollector = requireNonNull(warningCollector, "warningCollector is null");
4354
this.runtimeStats = requireNonNull(runtimeStats, "runtimeStats is null");
4455
this.queryType = requireNonNull(queryType, "queryType is null");
56+
this.catalog = requireNonNull(catalog, "catalog is null");
57+
this.schema = requireNonNull(schema, "schema is null");
4558
}
4659

4760
public QueryId getQueryId()
@@ -78,4 +91,14 @@ public Optional<QueryType> getQueryType()
7891
{
7992
return queryType;
8093
}
94+
95+
public Optional<String> getCatalog()
96+
{
97+
return catalog;
98+
}
99+
100+
public Optional<String> getSchema()
101+
{
102+
return schema;
103+
}
81104
}

‎presto-tests/src/test/java/com/facebook/presto/tests/TestQueryManager.java

+8
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,8 @@ public void testQueryCpuLimit()
220220
BasicQueryInfo queryInfo = queryManager.getQueryInfo(queryId);
221221
assertEquals(queryInfo.getState(), FAILED);
222222
assertEquals(queryInfo.getErrorCode(), EXCEEDED_CPU_LIMIT.toErrorCode());
223+
assertEquals(queryManager.getQuerySession(queryId).getAccessControlContext().getSchema(), TEST_SESSION.getSchema());
224+
assertEquals(queryManager.getQuerySession(queryId).getAccessControlContext().getCatalog(), TEST_SESSION.getCatalog());
223225
}
224226
}
225227

@@ -234,6 +236,8 @@ public void testQueryScanExceeded()
234236
BasicQueryInfo queryInfo = queryManager.getQueryInfo(queryId);
235237
assertEquals(queryInfo.getState(), FAILED);
236238
assertEquals(queryInfo.getErrorCode(), EXCEEDED_SCAN_RAW_BYTES_READ_LIMIT.toErrorCode());
239+
assertEquals(queryManager.getQuerySession(queryId).getAccessControlContext().getSchema(), TEST_SESSION.getSchema());
240+
assertEquals(queryManager.getQuerySession(queryId).getAccessControlContext().getCatalog(), TEST_SESSION.getCatalog());
237241
}
238242
}
239243

@@ -248,6 +252,8 @@ public void testQueryOutputPositionsExceeded()
248252
BasicQueryInfo queryInfo = queryManager.getQueryInfo(queryId);
249253
assertEquals(queryInfo.getState(), FAILED);
250254
assertEquals(queryInfo.getErrorCode(), EXCEEDED_OUTPUT_POSITIONS_LIMIT.toErrorCode());
255+
assertEquals(queryManager.getQuerySession(queryId).getAccessControlContext().getSchema(), TEST_SESSION.getSchema());
256+
assertEquals(queryManager.getQuerySession(queryId).getAccessControlContext().getCatalog(), TEST_SESSION.getCatalog());
251257
}
252258
}
253259

@@ -262,6 +268,8 @@ public void testQueryOutputSizeExceeded()
262268
BasicQueryInfo queryInfo = queryManager.getQueryInfo(queryId);
263269
assertEquals(queryInfo.getState(), FAILED);
264270
assertEquals(queryInfo.getErrorCode(), EXCEEDED_OUTPUT_SIZE_LIMIT.toErrorCode());
271+
assertEquals(queryManager.getQuerySession(queryId).getAccessControlContext().getSchema(), TEST_SESSION.getSchema());
272+
assertEquals(queryManager.getQuerySession(queryId).getAccessControlContext().getCatalog(), TEST_SESSION.getCatalog());
265273
}
266274
}
267275

0 commit comments

Comments
 (0)