[#6306] fix(authz): Fix the OOM of JdbcAuthorizationPlugin

apache · Jan 17, 2025 · 92f9037 · 92f9037
1 parent e6225a0
commit 92f9037
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 15 deletions.
diff --git a/...common/src/main/java/org/apache/gravitino/authorization/jdbc/JdbcAuthorizationPlugin.java b/...common/src/main/java/org/apache/gravitino/authorization/jdbc/JdbcAuthorizationPlugin.java
@@ -110,8 +110,13 @@ public Boolean onMetadataUpdated(MetadataObjectChange... changes) throws Runtime
   @Override
   public Boolean onRoleCreated(Role role) throws AuthorizationPluginException {
     List<String> sqls = getCreateRoleSQL(role.name());
+    boolean createdNewly = false;
     for (String sql : sqls) {
-      executeUpdateSQL(sql, "already exists");
+      createdNewly = executeUpdateSQL(sql, "already exists");
+    }
+
+    if (!createdNewly) {
+      return true;
     }
 
     if (role.securableObjects() != null) {
@@ -140,7 +145,6 @@ public Boolean onRoleDeleted(Role role) throws AuthorizationPluginException {
   @Override
   public Boolean onRoleUpdated(Role role, RoleChange... changes)
       throws AuthorizationPluginException {
-    onRoleCreated(role);
     for (RoleChange change : changes) {
       if (change instanceof RoleChange.AddSecurableObject) {
         SecurableObject object = ((RoleChange.AddSecurableObject) change).getSecurableObject();
@@ -381,14 +385,15 @@ protected AuthorizationPluginException toAuthorizationPluginException(SQLExcepti
         "JDBC authorization plugin fail to execute SQL, error code: %d", se.getErrorCode());
   }
 
-  public void executeUpdateSQL(String sql, String ignoreErrorMsg) {
+  public boolean executeUpdateSQL(String sql, String ignoreErrorMsg) {
     try (final Connection connection = getConnection()) {
       try (final Statement statement = connection.createStatement()) {
         statement.executeUpdate(sql);
+        return true;
       }
     } catch (SQLException se) {
       if (ignoreErrorMsg != null && se.getMessage().contains(ignoreErrorMsg)) {
-        return;
+        return false;
       }
       LOG.error("JDBC authorization plugin exception: ", se);
       throw toAuthorizationPluginException(se);

diff --git a/...on/src/test/java/org/apache/gravitino/authorization/jdbc/TestJdbcAuthorizationPlugin.java b/...on/src/test/java/org/apache/gravitino/authorization/jdbc/TestJdbcAuthorizationPlugin.java
@@ -74,9 +74,10 @@ public List<String> getSetOwnerSQL(
           return Collections.emptyList();
         }
 
-        public void executeUpdateSQL(String sql, String ignoreErrorMsg) {
+        public boolean executeUpdateSQL(String sql, String ignoreErrorMsg) {
           Assertions.assertEquals(expectSQLs.get(currentSQLIndex), sql);
           currentSQLIndex++;
+          return true;
         }
       };
 
@@ -148,23 +149,21 @@ public void testPermissionManagement() {
 
     // Test metalake object and different role change
     resetSQLIndex();
-    expectSQLs = Lists.newArrayList("CREATE ROLE tmp", "GRANT SELECT ON TABLE *.* TO ROLE tmp");
+    expectSQLs = Lists.newArrayList("GRANT SELECT ON TABLE *.* TO ROLE tmp");
     SecurableObject metalakeObject =
         SecurableObjects.ofMetalake("metalake", Lists.newArrayList(Privileges.SelectTable.allow()));
     RoleChange roleChange = RoleChange.addSecurableObject("tmp", metalakeObject);
     plugin.onRoleUpdated(role, roleChange);
 
     resetSQLIndex();
-    expectSQLs = Lists.newArrayList("CREATE ROLE tmp", "REVOKE SELECT ON TABLE *.* FROM ROLE tmp");
+    expectSQLs = Lists.newArrayList("REVOKE SELECT ON TABLE *.* FROM ROLE tmp");
     roleChange = RoleChange.removeSecurableObject("tmp", metalakeObject);
     plugin.onRoleUpdated(role, roleChange);
 
     resetSQLIndex();
     expectSQLs =
         Lists.newArrayList(
-            "CREATE ROLE tmp",
-            "REVOKE SELECT ON TABLE *.* FROM ROLE tmp",
-            "GRANT CREATE ON TABLE *.* TO ROLE tmp");
+            "REVOKE SELECT ON TABLE *.* FROM ROLE tmp", "GRANT CREATE ON TABLE *.* TO ROLE tmp");
     SecurableObject newMetalakeObject =
         SecurableObjects.ofMetalake("metalake", Lists.newArrayList(Privileges.CreateTable.allow()));
     roleChange = RoleChange.updateSecurableObject("tmp", metalakeObject, newMetalakeObject);
@@ -175,7 +174,7 @@ public void testPermissionManagement() {
     SecurableObject catalogObject =
         SecurableObjects.ofCatalog("catalog", Lists.newArrayList(Privileges.SelectTable.allow()));
     roleChange = RoleChange.addSecurableObject("tmp", catalogObject);
-    expectSQLs = Lists.newArrayList("CREATE ROLE tmp", "GRANT SELECT ON TABLE *.* TO ROLE tmp");
+    expectSQLs = Lists.newArrayList("GRANT SELECT ON TABLE *.* TO ROLE tmp");
     plugin.onRoleUpdated(role, roleChange);
 
     // Test schema object
@@ -184,8 +183,7 @@ public void testPermissionManagement() {
         SecurableObjects.ofSchema(
             catalogObject, "schema", Lists.newArrayList(Privileges.SelectTable.allow()));
     roleChange = RoleChange.addSecurableObject("tmp", schemaObject);
-    expectSQLs =
-        Lists.newArrayList("CREATE ROLE tmp", "GRANT SELECT ON TABLE schema.* TO ROLE tmp");
+    expectSQLs = Lists.newArrayList("GRANT SELECT ON TABLE schema.* TO ROLE tmp");
     plugin.onRoleUpdated(role, roleChange);
 
     // Test table object
@@ -194,8 +192,18 @@ public void testPermissionManagement() {
         SecurableObjects.ofTable(
             schemaObject, "table", Lists.newArrayList(Privileges.SelectTable.allow()));
     roleChange = RoleChange.addSecurableObject("tmp", tableObject);
-    expectSQLs =
-        Lists.newArrayList("CREATE ROLE tmp", "GRANT SELECT ON TABLE schema.table TO ROLE tmp");
+    expectSQLs = Lists.newArrayList("GRANT SELECT ON TABLE schema.table TO ROLE tmp");
+    plugin.onRoleUpdated(role, roleChange);
+
+    // Test the role with objects
+    resetSQLIndex();
+    role =
+        RoleEntity.builder()
+            .withId(-1L)
+            .withName("tmp")
+            .withSecurableObjects(Lists.newArrayList(tableObject))
+            .withAuditInfo(AuditInfo.EMPTY)
+            .build();
     plugin.onRoleUpdated(role, roleChange);
   }