change test way

Author: swuferhong

fluss-client/src/test/java/org/apache/fluss/client/admin/RebalanceITCase.java

@@ -40,7 +40,6 @@
 import java.util.Collections;
 
 import static org.apache.fluss.record.TestData.DATA1_SCHEMA;
-import static org.apache.fluss.record.TestData.DATA1_TABLE_DESCRIPTOR;
 import static org.apache.fluss.testutils.common.CommonTestUtils.retry;
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -79,26 +78,38 @@ protected void teardown() throws Exception {
     void testRebalanceForLogTable() throws Exception {
         String dbName = "db-balance";
         admin.createDatabase(dbName, DatabaseDescriptor.EMPTY, false).get();
-        // create somne none partitioned log table.
+
+        TableDescriptor logDescriptor =
+                TableDescriptor.builder()
+                        .schema(DATA1_SCHEMA)
+                        .distributedBy(3)
+                        .property(
+                                ConfigOptions.TABLE_GENERATE_UNBALANCE_TABLE_ASSIGNMENT.key(),
+                                "true")
+                        .build();
+        // create some none partitioned log table.
         for (int i = 0; i < 6; i++) {
             long tableId =
                     createTable(
                             new TablePath(dbName, "test-rebalance_table-" + i),
-                            DATA1_TABLE_DESCRIPTOR,
+                            logDescriptor,
                             false);
             FLUSS_CLUSTER_EXTENSION.waitUntilTableReady(tableId);
         }
 
         // create some partitioned table with partition.
-        TableDescriptor descriptor =
+        TableDescriptor partitionedDescriptor =
                 TableDescriptor.builder()
                         .schema(DATA1_SCHEMA)
                         .distributedBy(3)
+                        .property(
+                                ConfigOptions.TABLE_GENERATE_UNBALANCE_TABLE_ASSIGNMENT.key(),
+                                "true")
                         .partitionedBy("b")
                         .build();
         for (int i = 0; i < 3; i++) {
             TablePath tablePath = new TablePath(dbName, "test-rebalance_partitioned_table-" + i);
-            long tableId = createTable(tablePath, descriptor, false);
+            long tableId = createTable(tablePath, partitionedDescriptor, false);
             for (int j = 0; j < 2; j++) {
                 PartitionSpec partitionSpec =
                         new PartitionSpec(Collections.singletonMap("b", String.valueOf(j)));
@@ -185,9 +196,6 @@ void testRebalanceForLogTable() throws Exception {
 
     private static Configuration initConfig() {
         Configuration configuration = new Configuration();
-        // As we want to test rebalance, we need to set this option to true for generate unbalance
-        // buckets location in server.
-        configuration.set(ConfigOptions.SERVER_GENERATE_UNBALANCE_ASSIGNMENT_FOR_TEST, true);
         configuration.set(ConfigOptions.DEFAULT_REPLICATION_FACTOR, 3);
         configuration.set(ConfigOptions.DEFAULT_BUCKET_NUMBER, 3);
         return configuration;

fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java

@@ -481,15 +481,6 @@ public class ConfigOptions {
                     .withDescription(
                             "Defines how long the buffer pool will block when waiting for segments to become available.");
 
-    public static final ConfigOption<Boolean> SERVER_GENERATE_UNBALANCE_ASSIGNMENT_FOR_TEST =
-            key("server.generate-unbalance-assignment-for-test")
-                    .booleanType()
-                    .defaultValue(false)
-                    .withDescription(
-                            "Whether to generate unbalance table or partition assignment. This parameter is only used "
-                                    + "for itCase. If set to true, the assignment will always be [0,1,2] as replica factor "
-                                    + "set as 3 even if there are tabletServers more than 3.");
-
     // ------------------------------------------------------------------
     // ZooKeeper Settings
     // ------------------------------------------------------------------
@@ -1403,6 +1394,14 @@ public class ConfigOptions {
                                     + "The `disable` behavior rejects delete requests with a clear error message. "
                                     + "For tables with FIRST_ROW or VERSIONED merge engines, this option defaults to `ignore`.");
 
+    public static final ConfigOption<Boolean> TABLE_GENERATE_UNBALANCE_TABLE_ASSIGNMENT =
+            key("table.generate-unbalance-table-assignment")
+                    .booleanType()
+                    .defaultValue(false)
+                    .withDescription(
+                            "This parameter is only used for test. If set to true, the assignment will always be [0,1,2] "
+                                    + "as replica factor set as 3 even if there are tabletServers more than 3.");
+
     // ------------------------------------------------------------------------
     //  ConfigOptions for Kv
     // ------------------------------------------------------------------------

fluss-common/src/main/java/org/apache/fluss/config/TableConfig.java

@@ -126,4 +126,9 @@ public ArrowCompressionInfo getArrowCompressionInfo() {
     public AutoPartitionStrategy getAutoPartitionStrategy() {
         return AutoPartitionStrategy.from(config);
     }
+
+    /** Whether to generate unbalance assignment fot this table. */
+    public boolean generateUnbalanceAssignment() {
+        return config.get(ConfigOptions.TABLE_GENERATE_UNBALANCE_TABLE_ASSIGNMENT);
+    }
 }

fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorService.java

@@ -176,9 +176,6 @@ public final class CoordinatorService extends RpcServiceBase implements Coordina
     private final LakeCatalogDynamicLoader lakeCatalogDynamicLoader;
     private final Supplier<RebalanceManager> rebalanceManagerSupplier;
 
-    // This parameter is only used for testing.
-    private final boolean generateUnBalanceAssignment;
-
     public CoordinatorService(
             Configuration conf,
             FileSystem remoteFileSystem,
@@ -201,8 +198,6 @@ public CoordinatorService(
         this.defaultReplicationFactor = conf.getInt(ConfigOptions.DEFAULT_REPLICATION_FACTOR);
         this.logTableAllowCreation = conf.getBoolean(ConfigOptions.LOG_TABLE_ALLOW_CREATION);
         this.kvTableAllowCreation = conf.getBoolean(ConfigOptions.KV_TABLE_ALLOW_CREATION);
-        this.generateUnBalanceAssignment =
-                conf.getBoolean(ConfigOptions.SERVER_GENERATE_UNBALANCE_ASSIGNMENT_FOR_TEST);
         this.eventManagerSupplier =
                 () -> coordinatorEventProcessorSupplier.get().getCoordinatorEventManager();
         this.coordinatorEpochSupplier =
@@ -303,12 +298,22 @@ public CompletableFuture<CreateTableResponse> createTable(CreateTableRequest req
 
         // first, generate the assignment
         TableAssignment tableAssignment = null;
+        Map<String, String> properties = tableDescriptor.getProperties();
+        boolean generateUnbalanceAssignment;
+        if (properties.containsKey(ConfigOptions.TABLE_GENERATE_UNBALANCE_TABLE_ASSIGNMENT.key())) {
+            generateUnbalanceAssignment =
+                    Boolean.parseBoolean(
+                            properties.get(
+                                    ConfigOptions.TABLE_GENERATE_UNBALANCE_TABLE_ASSIGNMENT.key()));
+        } else {
+            generateUnbalanceAssignment = false;
+        }
         // only when it's no partitioned table do we generate the assignment for it
         if (!tableDescriptor.isPartitioned()) {
             // the replication factor must be set now
             int replicaFactor = tableDescriptor.getReplicationFactor();
             TabletServerInfo[] servers = metadataCache.getLiveServers();
-            if (generateUnBalanceAssignment) {
+            if (generateUnbalanceAssignment) {
                 // this branch is only used for testing.
                 tableAssignment =
                         new TableAssignment(
@@ -509,7 +514,9 @@ public CompletableFuture<CreatePartitionResponse> createPartition(
         int replicaFactor = table.getTableConfig().getReplicationFactor();
         TabletServerInfo[] servers = metadataCache.getLiveServers();
         Map<Integer, BucketAssignment> bucketAssignments;
-        if (generateUnBalanceAssignment) {
+
+        boolean generateUnbalanceAssignment = table.getTableConfig().generateUnbalanceAssignment();
+        if (generateUnbalanceAssignment) {
             // This branch is only used for testing.
             bucketAssignments = generateUnBalanceAssignment(table.bucketCount, replicaFactor);
         } else {