Make list*TablePrivileges to return Stream

This way table privileges are enumerated lazily. Extracted-From: prestodb/presto#10904
trinodb · Jan 29, 2019 · 1d11c06 · 1d11c06
1 parent 0c2fa44
commit 1d11c06
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 25 deletions.
diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/HiveMetadata.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/HiveMetadata.java
@@ -1845,7 +1845,8 @@ public void revokeTablePrivileges(ConnectorSession session, SchemaTableName sche
     @Override
     public List<GrantInfo> listTablePrivileges(ConnectorSession session, SchemaTablePrefix schemaTablePrefix)
     {
-        Set<PrestoPrincipal> principals = listEnabledPrincipals(metastore, session.getIdentity());
+        Set<PrestoPrincipal> principals = listEnabledPrincipals(metastore, session.getIdentity())
+                .collect(toImmutableSet());
         ImmutableList.Builder<GrantInfo> result = ImmutableList.builder();
         for (SchemaTableName tableName : listTables(session, schemaTablePrefix)) {
             for (PrestoPrincipal grantee : principals) {

diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/thrift/ThriftMetastoreUtil.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/thrift/ThriftMetastoreUtil.java
@@ -274,40 +274,32 @@ public static Stream<String> listApplicableRoles(SemiTransactionalHiveMetastore
                 .map(RoleGrant::getRoleName);
     }
 
-    public static Set<PrestoPrincipal> listEnabledPrincipals(SemiTransactionalHiveMetastore metastore, ConnectorIdentity identity)
+    public static Stream<PrestoPrincipal> listEnabledPrincipals(SemiTransactionalHiveMetastore metastore, ConnectorIdentity identity)
     {
-        ImmutableSet.Builder<PrestoPrincipal> principals = ImmutableSet.builder();
-        PrestoPrincipal userPrincipal = new PrestoPrincipal(USER, identity.getUser());
-        principals.add(userPrincipal);
-        listEnabledRoles(identity, metastore::listRoleGrants)
-                .map(role -> new PrestoPrincipal(ROLE, role))
-                .forEach(principals::add);
-        return principals.build();
+        return Stream.concat(
+                Stream.of(new PrestoPrincipal(USER, identity.getUser())),
+                listEnabledRoles(identity, metastore::listRoleGrants)
+                        .map(role -> new PrestoPrincipal(ROLE, role)));
     }
 
-    public static Set<HivePrivilegeInfo> listEnabledTablePrivileges(SemiTransactionalHiveMetastore metastore, String databaseName, String tableName, ConnectorIdentity identity)
+    public static Stream<HivePrivilegeInfo> listEnabledTablePrivileges(SemiTransactionalHiveMetastore metastore, String databaseName, String tableName, ConnectorIdentity identity)
     {
         return listTablePrivileges(metastore, databaseName, tableName, listEnabledPrincipals(metastore, identity));
     }
 
-    public static Set<HivePrivilegeInfo> listApplicableTablePrivileges(SemiTransactionalHiveMetastore metastore, String databaseName, String tableName, String user)
+    public static Stream<HivePrivilegeInfo> listApplicableTablePrivileges(SemiTransactionalHiveMetastore metastore, String databaseName, String tableName, String user)
     {
-        ImmutableSet.Builder<PrestoPrincipal> principals = ImmutableSet.builder();
         PrestoPrincipal userPrincipal = new PrestoPrincipal(USER, user);
-        principals.add(userPrincipal);
-        listApplicableRoles(metastore, userPrincipal)
-                .map(role -> new PrestoPrincipal(ROLE, role))
-                .forEach(principals::add);
-        return listTablePrivileges(metastore, databaseName, tableName, principals.build());
+        Stream<PrestoPrincipal> principals = Stream.concat(
+                Stream.of(userPrincipal),
+                listApplicableRoles(metastore, userPrincipal)
+                .map(role -> new PrestoPrincipal(ROLE, role)));
+        return listTablePrivileges(metastore, databaseName, tableName, principals);
     }
 
-    private static Set<HivePrivilegeInfo> listTablePrivileges(SemiTransactionalHiveMetastore metastore, String databaseName, String tableName, Set<PrestoPrincipal> principals)
+    private static Stream<HivePrivilegeInfo> listTablePrivileges(SemiTransactionalHiveMetastore metastore, String databaseName, String tableName, Stream<PrestoPrincipal> principals)
     {
-        ImmutableSet.Builder<HivePrivilegeInfo> result = ImmutableSet.builder();
-        for (PrestoPrincipal current : principals) {
-            result.addAll(metastore.listTablePrivileges(databaseName, tableName, current));
-        }
-        return result.build();
+        return principals.flatMap(principal -> metastore.listTablePrivileges(databaseName, tableName, principal).stream());
     }
 
     public static boolean isRoleEnabled(ConnectorIdentity identity, Function<PrestoPrincipal, Set<RoleGrant>> listRoleGrants, String role)

diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/security/SqlStandardAccessControl.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/security/SqlStandardAccessControl.java
@@ -402,7 +402,6 @@ private boolean checkTablePermission(
 
         SemiTransactionalHiveMetastore metastore = metastoreProvider.apply(((HiveTransactionHandle) transaction));
         return listEnabledTablePrivileges(metastore, tableName.getSchemaName(), tableName.getTableName(), identity)
-                .stream()
                 .filter(privilegeInfo -> !grantOptionRequired || privilegeInfo.isGrantOption())
                 .anyMatch(privilegeInfo -> privilegeInfo.getHivePrivilege().equals(requiredPrivilege));
     }
@@ -419,7 +418,6 @@ private boolean hasGrantOptionForPrivilege(ConnectorTransactionHandle transactio
                 tableName.getSchemaName(),
                 tableName.getTableName(),
                 identity.getUser())
-                .stream()
                 .anyMatch(privilegeInfo -> privilegeInfo.getHivePrivilege().equals(toHivePrivilege(privilege)) && privilegeInfo.isGrantOption());
     }