Fix RegionGroup limits for same-name tree and table databases

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ConfigManager.java

@@ -313,6 +313,7 @@
 import java.util.stream.Collectors;
 
 import static org.apache.iotdb.commons.conf.IoTDBConstant.ONE_LEVEL_PATH_WILDCARD;
+import static org.apache.iotdb.commons.schema.table.Audit.TABLE_MODEL_AUDIT_DATABASE;
 import static org.apache.iotdb.commons.schema.table.Audit.TREE_MODEL_AUDIT_DATABASE;
 
 /** Entry of all management, AssignPartitionManager, AssignRegionManager. */
@@ -826,8 +827,9 @@ public synchronized TSStatus deleteDatabases(final TDeleteDatabasesReq tDeleteRe
               .getMatchedDatabaseSchemasByName(
                   deletedPaths, tDeleteReq.isSetIsTableModel() && tDeleteReq.isIsTableModel());
 
-      // remove root.__audit
+      // remove audit databases
       deleteDatabaseSchemaMap.remove(TREE_MODEL_AUDIT_DATABASE);
+      deleteDatabaseSchemaMap.remove(TABLE_MODEL_AUDIT_DATABASE);
 
       if (deleteDatabaseSchemaMap.isEmpty()) {
         return RpcUtils.getStatus(

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/schema/ClusterSchemaInfo.java

@@ -607,7 +607,7 @@ public TDatabaseSchema getMatchedDatabaseSchemaByName(final String database)
    * Only leader use this interface. Get the matched DatabaseSchemas.
    *
    * @param rawPathList Databases' path patterns or full paths
-   * @return All DatabaseSchemas that matches to the specific Database patterns
+   * @return All matching DatabaseSchemas keyed by their database names
    */
   public Map<String, TDatabaseSchema> getMatchedDatabaseSchemasByName(
       final List<String> rawPathList, final Boolean isTableModel) {
@@ -638,8 +638,10 @@ private void enrichSchemaMap(
       final List<PartialPath> matchedPaths =
           mTree.getMatchedDatabases(patternPath, ALL_MATCH_SCOPE, false);
       for (final PartialPath path : matchedPaths) {
-        schemaMap.put(
-            path.getFullPath(), mTree.getDatabaseNodeByPath(path).getAsMNode().getDatabaseSchema());
+        final TDatabaseSchema schema =
+            mTree.getDatabaseNodeByPath(path).getAsMNode().getDatabaseSchema();
+        // Table database "pg" and tree database "root.pg" share the same internal path.
+        schemaMap.put(schema.getName(), schema);
       }
     }
   }

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/ClusterSchemaManagerTest.java

@@ -18,20 +18,29 @@
  */
 package org.apache.iotdb.confignode.manager;
 
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.commons.exception.table.ColumnInAlterException;
 import org.apache.iotdb.commons.exception.table.ColumnInDeletionException;
 import org.apache.iotdb.commons.exception.table.TableInDeletionException;
 import org.apache.iotdb.commons.schema.table.TsTable;
 import org.apache.iotdb.commons.schema.table.TsTableInternalRPCUtil;
 import org.apache.iotdb.commons.schema.table.column.FieldColumnSchema;
 import org.apache.iotdb.commons.schema.table.column.TsTableColumnSchema;
+import org.apache.iotdb.confignode.conf.ConfigNodeConfig;
+import org.apache.iotdb.confignode.conf.ConfigNodeDescriptor;
 import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlanType;
+import org.apache.iotdb.confignode.consensus.request.write.database.AdjustMaxRegionGroupNumPlan;
 import org.apache.iotdb.confignode.consensus.request.write.database.DatabaseSchemaPlan;
 import org.apache.iotdb.confignode.consensus.request.write.table.CommitCreateTablePlan;
 import org.apache.iotdb.confignode.consensus.request.write.table.PreAlterColumnDataTypePlan;
 import org.apache.iotdb.confignode.consensus.request.write.table.PreCreateTablePlan;
 import org.apache.iotdb.confignode.consensus.request.write.table.PreDeleteColumnPlan;
 import org.apache.iotdb.confignode.consensus.request.write.table.PreDeleteTablePlan;
+import org.apache.iotdb.confignode.manager.consensus.ConsensusManager;
+import org.apache.iotdb.confignode.manager.node.NodeManager;
+import org.apache.iotdb.confignode.manager.partition.PartitionManager;
+import org.apache.iotdb.confignode.manager.partition.RegionGroupExtensionPolicy;
 import org.apache.iotdb.confignode.manager.schema.ClusterSchemaManager;
 import org.apache.iotdb.confignode.manager.schema.ClusterSchemaQuotaStatistics;
 import org.apache.iotdb.confignode.persistence.schema.ClusterSchemaInfo;
@@ -47,6 +56,7 @@
 import org.mockito.Mockito;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -70,6 +80,86 @@ public void testCalcMaxRegionGroupNum() {
     Assert.assertEquals(20, ClusterSchemaManager.calcMaxRegionGroupNum(3, 1.0, 120, 2, 3, 5));
   }
 
+  @Test
+  public void testAdjustMaxRegionGroupNumWithMixedModels() throws Exception {
+    final ConfigNodeConfig conf = ConfigNodeDescriptor.getInstance().getConf();
+    final int originalSchemaQuota = conf.getSchemaRegionPerDataNode();
+    final int originalDataQuota = conf.getDataRegionPerDataNode();
+    final RegionGroupExtensionPolicy originalSchemaPolicy =
+        conf.getSchemaRegionGroupExtensionPolicy();
+    final RegionGroupExtensionPolicy originalDataPolicy = conf.getDataRegionGroupExtensionPolicy();
+    try {
+      conf.setSchemaRegionPerDataNode(12);
+      conf.setDataRegionPerDataNode(12);
+      conf.setSchemaRegionGroupExtensionPolicy(RegionGroupExtensionPolicy.PROACTIVE);
+      conf.setDataRegionGroupExtensionPolicy(RegionGroupExtensionPolicy.PROACTIVE);
+
+      final ClusterSchemaInfo clusterSchemaInfo = new ClusterSchemaInfo();
+      final IManager configManager = Mockito.mock(IManager.class);
+      final NodeManager nodeManager = Mockito.mock(NodeManager.class);
+      final PartitionManager partitionManager = Mockito.mock(PartitionManager.class);
+      final ConsensusManager consensusManager = Mockito.mock(ConsensusManager.class);
+      Mockito.when(configManager.getNodeManager()).thenReturn(nodeManager);
+      Mockito.when(configManager.getPartitionManager()).thenReturn(partitionManager);
+      Mockito.when(configManager.getConsensusManager()).thenReturn(consensusManager);
+      Mockito.when(nodeManager.getRegisteredDataNodeCount()).thenReturn(1);
+      Mockito.when(nodeManager.getDataNodeCpuCoreCount()).thenReturn(1);
+      Mockito.when(
+              partitionManager.getRegionGroupCount(
+                  Mockito.anyString(), Mockito.any(TConsensusGroupType.class)))
+          .thenReturn(0);
+      Mockito.when(consensusManager.write(Mockito.any(AdjustMaxRegionGroupNumPlan.class)))
+          .thenAnswer(
+              invocation -> {
+                final TSStatus status =
+                    clusterSchemaInfo.adjustMaxRegionGroupCount(invocation.getArgument(0));
+                Assert.assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), status.getCode());
+                return status;
+              });
+      final ClusterSchemaManager clusterSchemaManager =
+          new ClusterSchemaManager(
+              configManager, clusterSchemaInfo, Mockito.mock(ClusterSchemaQuotaStatistics.class));
+
+      final List<String> databases = Arrays.asList("pg", "root.pg", "root.control");
+      final int[] expectedLimits = {12, 6, 4};
+      for (int i = 0; i < databases.size(); i++) {
+        final String database = databases.get(i);
+        final TDatabaseSchema databaseSchema =
+            new TDatabaseSchema(database)
+                .setIsTableModel(i == 0)
+                .setSchemaReplicationFactor(1)
+                .setDataReplicationFactor(1)
+                .setMinSchemaRegionGroupNum(1)
+                .setMaxSchemaRegionGroupNum(1)
+                .setMinDataRegionGroupNum(2)
+                .setMaxDataRegionGroupNum(2);
+        Assert.assertEquals(
+            TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+            clusterSchemaInfo
+                .createDatabase(
+                    new DatabaseSchemaPlan(ConfigPhysicalPlanType.CreateDatabase, databaseSchema))
+                .getCode());
+        Mockito.when(partitionManager.isDatabaseExist(database)).thenReturn(true);
+
+        clusterSchemaManager.adjustMaxRegionGroupNum();
+
+        for (final String createdDatabase : databases.subList(0, i + 1)) {
+          final TDatabaseSchema adjustedSchema =
+              clusterSchemaInfo.getMatchedDatabaseSchemaByName(createdDatabase);
+          Assert.assertEquals(
+              createdDatabase, expectedLimits[i], adjustedSchema.getMaxSchemaRegionGroupNum());
+          Assert.assertEquals(
+              createdDatabase, expectedLimits[i], adjustedSchema.getMaxDataRegionGroupNum());
+        }
+      }
+    } finally {
+      conf.setSchemaRegionPerDataNode(originalSchemaQuota);
+      conf.setDataRegionPerDataNode(originalDataQuota);
+      conf.setSchemaRegionGroupExtensionPolicy(originalSchemaPolicy);
+      conf.setDataRegionGroupExtensionPolicy(originalDataPolicy);
+    }
+  }
+
   @Test
   public void testNeedLastCacheDefaultsToTrueWhenUnset() {
     final TDatabaseSchema unsetSchema = new TDatabaseSchema();

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/ConfigManagerDatabaseDeletionTest.java

@@ -0,0 +1,65 @@
+/*
+ * 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.iotdb.confignode.manager;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.confignode.manager.schema.ClusterSchemaManager;
+import org.apache.iotdb.confignode.rpc.thrift.TDatabaseSchema;
+import org.apache.iotdb.confignode.rpc.thrift.TDeleteDatabasesReq;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import org.junit.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.apache.iotdb.commons.schema.table.Audit.TABLE_MODEL_AUDIT_DATABASE;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class ConfigManagerDatabaseDeletionTest {
+
+  @Test
+  public void testTableAuditDatabaseIsExcludedFromDeletion() {
+    final ConfigManager manager = mock(ConfigManager.class);
+    final ClusterSchemaManager schemaManager = mock(ClusterSchemaManager.class);
+    when(manager.confirmLeader())
+        .thenReturn(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()));
+    when(manager.getClusterSchemaManager()).thenReturn(schemaManager);
+
+    final Map<String, TDatabaseSchema> schemas = new HashMap<>();
+    schemas.put(
+        TABLE_MODEL_AUDIT_DATABASE,
+        new TDatabaseSchema(TABLE_MODEL_AUDIT_DATABASE).setIsTableModel(true));
+    when(schemaManager.getMatchedDatabaseSchemasByName(
+            Collections.singletonList(TABLE_MODEL_AUDIT_DATABASE), true))
+        .thenReturn(schemas);
+    final TDeleteDatabasesReq request =
+        new TDeleteDatabasesReq(Collections.singletonList(TABLE_MODEL_AUDIT_DATABASE))
+            .setIsTableModel(true);
+    when(manager.deleteDatabases(request)).thenCallRealMethod();
+
+    assertEquals(
+        TSStatusCode.PATH_NOT_EXIST.getStatusCode(), manager.deleteDatabases(request).getCode());
+    assertTrue(schemas.isEmpty());
+  }
+}