This is an automated email from the ASF dual-hosted git repository.
panjuan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shardingsphere.git
The following commit(s) were added to refs/heads/master by this push:
new 4204651 Add PostgreSQL privilege loader (#9924)
4204651 is described below
commit 420465121ba052b648beaa66905161368f0b6adc
Author: Zhu jun <[email protected]>
AuthorDate: Tue Apr 6 12:45:15 2021 +0800
Add PostgreSQL privilege loader (#9924)
* add postgresql privilege loader
* add postgresql privilege loader
* add postgresql privilege loader
* fix testcase
* fix review
* fix review
* fix review
* fix review
* fix review
* fix review
* fix
---
.../loader/dialect/PostgreSQLPrivilegeLoader.java | 191 +++++++++++++++++++++
.../auth/model/privilege/PrivilegeType.java | 9 +-
...ra.metadata.auth.builder.loader.PrivilegeLoader | 1 +
.../dialect/PostgreSQLPrivilegeLoaderTest.java | 128 ++++++++++++++
4 files changed, 328 insertions(+), 1 deletion(-)
diff --git
a/shardingsphere-infra/shardingsphere-infra-common/src/main/java/org/apache/shardingsphere/infra/metadata/auth/builder/loader/dialect/PostgreSQLPrivilegeLoader.java
b/shardingsphere-infra/shardingsphere-infra-common/src/main/java/org/apache/shardingsphere/infra/metadata/auth/builder/loader/dialect/PostgreSQLPrivilegeLoader.java
new file mode 100644
index 0000000..f56f158
--- /dev/null
+++
b/shardingsphere-infra/shardingsphere-infra-common/src/main/java/org/apache/shardingsphere/infra/metadata/auth/builder/loader/dialect/PostgreSQLPrivilegeLoader.java
@@ -0,0 +1,191 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.shardingsphere.infra.metadata.auth.builder.loader.dialect;
+
+import
org.apache.shardingsphere.infra.metadata.auth.builder.loader.PrivilegeLoader;
+import
org.apache.shardingsphere.infra.metadata.auth.model.privilege.PrivilegeType;
+import
org.apache.shardingsphere.infra.metadata.auth.model.privilege.ShardingSpherePrivilege;
+import
org.apache.shardingsphere.infra.metadata.auth.model.privilege.database.SchemaPrivilege;
+import
org.apache.shardingsphere.infra.metadata.auth.model.privilege.database.TablePrivilege;
+import org.apache.shardingsphere.infra.metadata.auth.model.user.Grantee;
+import
org.apache.shardingsphere.infra.metadata.auth.model.user.ShardingSphereUser;
+
+import javax.sql.DataSource;
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Collectors;
+
+/**
+ * PostgreSQL privilege loader.
+ */
+public final class PostgreSQLPrivilegeLoader implements PrivilegeLoader {
+
+ private static final String ROLES_SQL = "select * from pg_roles WHERE
rolname IN (%s)";
+
+ private static final String TABLE_PRIVILEGE_SQL = "SELECT grantor,
grantee, table_catalog, table_name, privilege_type, is_grantable from
information_schema.table_privileges WHERE grantee IN (%s)";
+
+ @Override
+ public Map<ShardingSphereUser, ShardingSpherePrivilege> load(final
Collection<ShardingSphereUser> users, final DataSource dataSource) throws
SQLException {
+ Map<ShardingSphereUser, ShardingSpherePrivilege> result = new
LinkedHashMap<>();
+ users.forEach(user -> result.put(user, new ShardingSpherePrivilege()));
+ fillTablePrivilege(result, dataSource, users);
+ fillRolePrivilege(result, dataSource, users);
+ return result;
+ }
+
+ private void fillTablePrivilege(final Map<ShardingSphereUser,
ShardingSpherePrivilege> privileges, final DataSource dataSource, final
Collection<ShardingSphereUser> users) throws SQLException {
+ Map<ShardingSphereUser, Map<String, Map<String, List<PrivilegeType>>>>
privilegeCache = new HashMap<>();
+ try (Connection connection = dataSource.getConnection()) {
+ Statement statement = connection.createStatement();
+ try (ResultSet resultSet =
statement.executeQuery(getTablePrivilegeSQL(users))) {
+ while (resultSet.next()) {
+ collectPrivilege(privilegeCache, resultSet);
+ }
+ }
+ }
+ fillTablePrivilege(privilegeCache, privileges);
+ }
+
+ private void fillTablePrivilege(final Map<ShardingSphereUser, Map<String,
Map<String, List<PrivilegeType>>>> privilegeCache, final
Map<ShardingSphereUser, ShardingSpherePrivilege> privileges) {
+ for (ShardingSphereUser user : privilegeCache.keySet()) {
+ for (String db : privilegeCache.get(user).keySet()) {
+ for (String tableName :
privilegeCache.get(user).get(db).keySet()) {
+ TablePrivilege tablePrivilege = new
TablePrivilege(tableName, privilegeCache.get(user).get(db).get(tableName));
+ ShardingSpherePrivilege privilege = privileges.get(user);
+ if
(!privilege.getDatabasePrivilege().getSpecificPrivileges().containsKey(db)) {
+
privilege.getDatabasePrivilege().getSpecificPrivileges().put(db, new
SchemaPrivilege(db));
+ }
+
privilege.getDatabasePrivilege().getSpecificPrivileges().get(db).getSpecificPrivileges().put(tableName,
tablePrivilege);
+ }
+ }
+ }
+ }
+
+ private void collectPrivilege(final Map<ShardingSphereUser, Map<String,
Map<String, List<PrivilegeType>>>> privilegeCache, final ResultSet resultSet)
throws SQLException {
+ String db = resultSet.getString("table_catalog");
+ String tableName = resultSet.getString("table_name");
+ String privilegeType = resultSet.getString("privilege_type");
+ Boolean hasPrivilege =
resultSet.getString("is_grantable").equalsIgnoreCase("TRUE");
+ String grantee = resultSet.getString("grantee");
+ if (hasPrivilege) {
+ privilegeCache
+ .computeIfAbsent(new ShardingSphereUser(grantee, "", ""),
k -> new HashMap<>())
+ .computeIfAbsent(db, k -> new HashMap<>())
+ .computeIfAbsent(tableName, k -> new ArrayList<>())
+ .add(getPrivilegeType(privilegeType));
+ }
+ }
+
+ private void fillRolePrivilege(final Map<ShardingSphereUser,
ShardingSpherePrivilege> privileges, final DataSource dataSource, final
Collection<ShardingSphereUser> users) throws SQLException {
+ try (Connection connection = dataSource.getConnection()) {
+ Statement statement = connection.createStatement();
+ try (ResultSet resultSet =
statement.executeQuery(getRolePrivilegeSQL(users))) {
+ while (resultSet.next()) {
+ fillRolePrivilege(privileges, resultSet);
+ }
+ }
+ }
+ }
+
+ private void fillRolePrivilege(final Map<ShardingSphereUser,
ShardingSpherePrivilege> privileges, final ResultSet resultSet) throws
SQLException {
+ Optional<ShardingSphereUser> user = getShardingSphereUser(privileges,
resultSet);
+ if (user.isPresent()) {
+
privileges.get(user.get()).getAdministrativePrivilege().getPrivileges().addAll(loadRolePrivileges(resultSet));
+ }
+ }
+
+ private Optional<ShardingSphereUser> getShardingSphereUser(final
Map<ShardingSphereUser, ShardingSpherePrivilege> privileges, final ResultSet
resultSet) throws SQLException {
+ Grantee grantee = new Grantee(resultSet.getString("rolname"), "");
+ return privileges.keySet().stream().filter(each ->
each.getGrantee().equals(grantee)).findFirst();
+ }
+
+ private Collection<PrivilegeType> loadRolePrivileges(final ResultSet
resultSet) throws SQLException {
+ Collection<PrivilegeType> result = new LinkedList<>();
+ addToPrivilegeTypesIfPresent(resultSet.getBoolean("rolsuper"),
PrivilegeType.SUPER, result);
+ addToPrivilegeTypesIfPresent(resultSet.getBoolean("rolcreaterole"),
PrivilegeType.CREATE_ROLE, result);
+ addToPrivilegeTypesIfPresent(resultSet.getBoolean("rolcreatedb"),
PrivilegeType.CREATE_DATABASE, result);
+ addToPrivilegeTypesIfPresent(resultSet.getBoolean("rolreplication"),
PrivilegeType.REPL_CLIENT, result);
+ addToPrivilegeTypesIfPresent(resultSet.getBoolean("rolinherit"),
PrivilegeType.INHERIT, result);
+ addToPrivilegeTypesIfPresent(resultSet.getBoolean("rolcanlogin"),
PrivilegeType.CAN_LOGIN, result);
+ return result;
+ }
+
+ private String getTablePrivilegeSQL(final Collection<ShardingSphereUser>
users) {
+ String userList = users.stream().map(each -> String.format("'%s'",
each.getGrantee().getUsername()))
+ .collect(Collectors.joining(", "));
+ return String.format(TABLE_PRIVILEGE_SQL, userList);
+ }
+
+ private String getRolePrivilegeSQL(final Collection<ShardingSphereUser>
users) {
+ String userList = users.stream().map(each -> String.format("'%s'",
each.getGrantee().getUsername()))
+ .collect(Collectors.joining(", "));
+ return String.format(ROLES_SQL, userList);
+ }
+
+ private PrivilegeType getPrivilegeType(final String privilege) {
+ switch (privilege) {
+ case "SELECT":
+ return PrivilegeType.SELECT;
+ case "INSERT":
+ return PrivilegeType.INSERT;
+ case "UPDATE":
+ return PrivilegeType.UPDATE;
+ case "DELETE":
+ return PrivilegeType.DELETE;
+ case "TRUNCATE":
+ return PrivilegeType.TRUNCATE;
+ case "REFERENCES":
+ return PrivilegeType.REFERENCES;
+ case "TRIGGER":
+ return PrivilegeType.TRIGGER;
+ case "CREATE":
+ return PrivilegeType.CREATE;
+ case "EXECUTE":
+ return PrivilegeType.EXECUTE;
+ case "USAGE":
+ return PrivilegeType.USAGE;
+ case "CONNECT":
+ return PrivilegeType.CONNECT;
+ case "TEMPORARY":
+ return PrivilegeType.TEMPORARY;
+ default:
+ throw new UnsupportedOperationException(privilege);
+ }
+ }
+
+ private void addToPrivilegeTypesIfPresent(final boolean hasPrivilege,
final PrivilegeType privilegeType, final Collection<PrivilegeType> target) {
+ if (hasPrivilege) {
+ target.add(privilegeType);
+ }
+ }
+
+ @Override
+ public String getDatabaseType() {
+ return "PostgreSQL";
+ }
+}
diff --git
a/shardingsphere-infra/shardingsphere-infra-common/src/main/java/org/apache/shardingsphere/infra/metadata/auth/model/privilege/PrivilegeType.java
b/shardingsphere-infra/shardingsphere-infra-common/src/main/java/org/apache/shardingsphere/infra/metadata/auth/model/privilege/PrivilegeType.java
index a18b959..a16040f 100644
---
a/shardingsphere-infra/shardingsphere-infra-common/src/main/java/org/apache/shardingsphere/infra/metadata/auth/model/privilege/PrivilegeType.java
+++
b/shardingsphere-infra/shardingsphere-infra-common/src/main/java/org/apache/shardingsphere/infra/metadata/auth/model/privilege/PrivilegeType.java
@@ -58,5 +58,12 @@ public enum PrivilegeType {
CREATE_USER,
CREATE_TABLESPACE,
CREATE_ROLE,
- DROP_ROLE
+ DROP_ROLE,
+ TRUNCATE,
+ USAGE,
+ CONNECT,
+ TEMPORARY,
+ CREATE_DATABASE,
+ INHERIT,
+ CAN_LOGIN
}
diff --git
a/shardingsphere-infra/shardingsphere-infra-common/src/main/resources/META-INF/services/org.apache.shardingsphere.infra.metadata.auth.builder.loader.PrivilegeLoader
b/shardingsphere-infra/shardingsphere-infra-common/src/main/resources/META-INF/services/org.apache.shardingsphere.infra.metadata.auth.builder.loader.PrivilegeLoader
index be5992f..ec7c809 100644
---
a/shardingsphere-infra/shardingsphere-infra-common/src/main/resources/META-INF/services/org.apache.shardingsphere.infra.metadata.auth.builder.loader.PrivilegeLoader
+++
b/shardingsphere-infra/shardingsphere-infra-common/src/main/resources/META-INF/services/org.apache.shardingsphere.infra.metadata.auth.builder.loader.PrivilegeLoader
@@ -16,3 +16,4 @@
#
org.apache.shardingsphere.infra.metadata.auth.builder.loader.dialect.MySQLPrivilegeLoader
+org.apache.shardingsphere.infra.metadata.auth.builder.loader.dialect.PostgreSQLPrivilegeLoader
diff --git
a/shardingsphere-infra/shardingsphere-infra-common/src/test/java/org/apache/shardingsphere/infra/metadata/auth/builder/loader/dialect/PostgreSQLPrivilegeLoaderTest.java
b/shardingsphere-infra/shardingsphere-infra-common/src/test/java/org/apache/shardingsphere/infra/metadata/auth/builder/loader/dialect/PostgreSQLPrivilegeLoaderTest.java
new file mode 100644
index 0000000..156c1c6
--- /dev/null
+++
b/shardingsphere-infra/shardingsphere-infra-common/src/test/java/org/apache/shardingsphere/infra/metadata/auth/builder/loader/dialect/PostgreSQLPrivilegeLoaderTest.java
@@ -0,0 +1,128 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.shardingsphere.infra.metadata.auth.builder.loader.dialect;
+
+import
org.apache.shardingsphere.infra.metadata.auth.builder.loader.PrivilegeLoader;
+import
org.apache.shardingsphere.infra.metadata.auth.model.privilege.PrivilegeType;
+import
org.apache.shardingsphere.infra.metadata.auth.model.privilege.ShardingSpherePrivilege;
+import
org.apache.shardingsphere.infra.metadata.auth.model.privilege.database.SchemaPrivilege;
+import
org.apache.shardingsphere.infra.metadata.auth.model.user.ShardingSphereUser;
+import org.apache.shardingsphere.infra.spi.ShardingSphereServiceLoader;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.junit.MockitoJUnitRunner;
+
+import javax.sql.DataSource;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.LinkedList;
+import java.util.Map;
+import java.util.concurrent.CopyOnWriteArraySet;
+import java.util.stream.Collectors;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThat;
+import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+@RunWith(MockitoJUnitRunner.class)
+public final class PostgreSQLPrivilegeLoaderTest {
+
+ @BeforeClass
+ public static void setUp() {
+ ShardingSphereServiceLoader.register(PrivilegeLoader.class);
+ }
+
+ @Test
+ public void assertLoad() throws SQLException {
+ Collection<ShardingSphereUser> users = createUsers();
+ DataSource dataSource = mockDataSource(users);
+ assertPrivilege(getPrivilegeLoader().load(users, dataSource));
+ }
+
+ private void assertPrivilege(final Map<ShardingSphereUser,
ShardingSpherePrivilege> actual) {
+ assertThat(actual.size(), is(1));
+ ShardingSphereUser user = new ShardingSphereUser("postgres", "", "");
+
assertThat(actual.get(user).getDatabasePrivilege().getGlobalPrivileges().size(),
is(0));
+
assertThat(actual.get(user).getDatabasePrivilege().getSpecificPrivileges().size(),
is(1));
+ Collection<PrivilegeType> expectedSpecificPrivilege = new
CopyOnWriteArraySet(Arrays.asList(PrivilegeType.INSERT, PrivilegeType.SELECT,
PrivilegeType.UPDATE,
+ PrivilegeType.DELETE));
+ SchemaPrivilege schemaPrivilege =
actual.get(user).getDatabasePrivilege().getSpecificPrivileges().get("db0");
+
assertThat(schemaPrivilege.getSpecificPrivileges().get("t_order").hasPrivileges(expectedSpecificPrivilege),
is(true));
+
assertThat(actual.get(user).getAdministrativePrivilege().getPrivileges().size(),
is(4));
+ Collection<PrivilegeType> expectedAdministrativePrivilege = new
CopyOnWriteArraySet(Arrays.asList(PrivilegeType.SUPER,
PrivilegeType.CREATE_ROLE,
+ PrivilegeType.CREATE_DATABASE, PrivilegeType.CAN_LOGIN));
+
assertEquals(actual.get(user).getAdministrativePrivilege().getPrivileges(),
expectedAdministrativePrivilege);
+ }
+
+ private Collection<ShardingSphereUser> createUsers() {
+ LinkedList<ShardingSphereUser> users = new LinkedList<>();
+ users.add(new ShardingSphereUser("postgres", "", ""));
+ return users;
+ }
+
+ private DataSource mockDataSource(final Collection<ShardingSphereUser>
users) throws SQLException {
+ ResultSet tablePrivilegeResultSet = mockTablePrivilegeResultSet();
+ DataSource dataSource = mock(DataSource.class, RETURNS_DEEP_STUBS);
+ String tablePrivilegeSql = "SELECT grantor, grantee, table_catalog,
table_name, privilege_type, is_grantable from
information_schema.table_privileges WHERE grantee IN (%s)";
+ String userList = users.stream().map(item -> String.format("'%s'",
item.getGrantee().getUsername(),
item.getGrantee().getHostname())).collect(Collectors.joining(", "));
+
when(dataSource.getConnection().createStatement().executeQuery(String.format(tablePrivilegeSql,
userList))).thenReturn(tablePrivilegeResultSet);
+ ResultSet rolePrivilegeResultSet = mockRolePrivilegeResultSet();
+ String rolePrivilegeSql = "select * from pg_roles WHERE rolname IN
(%s)";
+
when(dataSource.getConnection().createStatement().executeQuery(String.format(rolePrivilegeSql,
userList))).thenReturn(rolePrivilegeResultSet);
+ return dataSource;
+ }
+
+ private ResultSet mockTablePrivilegeResultSet() throws SQLException {
+ ResultSet result = mock(ResultSet.class, RETURNS_DEEP_STUBS);
+ when(result.next()).thenReturn(true, true, true, true, true, true,
true, false);
+ when(result.getString("table_catalog")).thenReturn("db0");
+ when(result.getString("table_name")).thenReturn("t_order");
+ when(result.getString("privilege_type")).thenReturn("INSERT",
"SELECT", "UPDATE", "DELETE", "TRUNCATE", "REFERENCES", "TRIGGER");
+ when(result.getString("is_grantable")).thenReturn("TRUE", "TRUE",
"TRUE", "TRUE", "TRUE", "TRUE", "TRUE");
+ when(result.getString("grantee")).thenReturn("postgres");
+ return result;
+ }
+
+ private ResultSet mockRolePrivilegeResultSet() throws SQLException {
+ ResultSet result = mock(ResultSet.class, RETURNS_DEEP_STUBS);
+ when(result.next()).thenReturn(true, false);
+ when(result.getString("rolname")).thenReturn("postgres");
+ when(result.getBoolean("rolsuper")).thenReturn(true);
+ when(result.getBoolean("rolcreaterole")).thenReturn(true);
+ when(result.getBoolean("rolcreatedb")).thenReturn(true);
+ when(result.getBoolean("rolreplication")).thenReturn(false);
+ when(result.getBoolean("rolinherit")).thenReturn(false);
+ when(result.getBoolean("rolcanlogin")).thenReturn(true);
+ return result;
+ }
+
+ private PrivilegeLoader getPrivilegeLoader() {
+ for (PrivilegeLoader each :
ShardingSphereServiceLoader.getSingletonServiceInstances(PrivilegeLoader.class))
{
+ if ("PostgreSQL".equals(each.getDatabaseType())) {
+ return each;
+ }
+ }
+ throw new IllegalStateException("Can not find
PostgreSQLPrivilegeLoader");
+ }
+}
