[#11829] fix(doris): support replication allocation property

catalogs/catalog-jdbc-doris/src/main/java/org/apache/gravitino/catalog/doris/DorisTablePropertiesMetadata.java

@@ -18,6 +18,8 @@
  */
 package org.apache.gravitino.catalog.doris;
 
+import static org.apache.gravitino.connector.PropertyEntry.stringOptionalPropertyEntry;
+
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Maps;
 import java.util.List;
@@ -28,6 +30,7 @@
 public class DorisTablePropertiesMetadata extends JdbcTablePropertiesMetadata {
 
   public static final String REPLICATION_FACTOR = "replication_num";
+  public static final String REPLICATION_ALLOCATION = "replication_allocation";
   public static final int DEFAULT_REPLICATION_FACTOR = 1;
   public static final int DEFAULT_REPLICATION_FACTOR_IN_SERVER_SIDE = 3;
 
@@ -42,6 +45,12 @@ public class DorisTablePropertiesMetadata extends JdbcTablePropertiesMetadata {
                     + " the default value will be used",
                 false /* immutable */,
                 DEFAULT_REPLICATION_FACTOR, /* default value */
+                false /* hidden */),
+            stringOptionalPropertyEntry(
+                REPLICATION_ALLOCATION,
+                "The replication allocation policy for the table.",
+                false /* immutable */,
+                null /* default value */,
                 false /* hidden */));
 
     PROPERTIES_METADATA = Maps.uniqueIndex(propertyEntries, PropertyEntry::getName);

catalogs/catalog-jdbc-doris/src/main/java/org/apache/gravitino/catalog/doris/operation/DorisTableOperations.java

@@ -20,6 +20,7 @@
 
 import static org.apache.gravitino.catalog.doris.DorisCatalog.DORIS_TABLE_PROPERTIES_META;
 import static org.apache.gravitino.catalog.doris.DorisTablePropertiesMetadata.DEFAULT_REPLICATION_FACTOR_IN_SERVER_SIDE;
+import static org.apache.gravitino.catalog.doris.DorisTablePropertiesMetadata.REPLICATION_ALLOCATION;
 import static org.apache.gravitino.catalog.doris.DorisTablePropertiesMetadata.REPLICATION_FACTOR;
 import static org.apache.gravitino.catalog.doris.utils.DorisUtils.generatePartitionSqlFragment;
 import static org.apache.gravitino.rel.Column.DEFAULT_VALUE_NOT_SET;
@@ -160,7 +161,8 @@ Map<String, String> appendNecessaryProperties(Map<String, String> properties) {
 
     // If the backend server is less than DEFAULT_REPLICATION_FACTOR_IN_SERVER_SIDE (3), we need to
     // set the property 'replication_num' to 1 explicitly.
-    if (!resultMap.containsKey(REPLICATION_FACTOR)) {
+    if (!resultMap.containsKey(REPLICATION_FACTOR)
+        && !resultMap.containsKey(REPLICATION_ALLOCATION)) {
       // Try to check the number of backend servers using `show backends`, this SQL is supported by
       // all versions of Doris
       String query = "show backends";

catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/TestDorisCatalog.java

@@ -19,6 +19,7 @@
 
 package org.apache.gravitino.catalog.doris;
 
+import static org.apache.gravitino.catalog.doris.DorisTablePropertiesMetadata.REPLICATION_ALLOCATION;
 import static org.apache.gravitino.catalog.doris.DorisTablePropertiesMetadata.REPLICATION_FACTOR;
 
 import java.util.Map;
@@ -41,5 +42,11 @@ void testDorisTablePropertiesMetadata() {
     Assertions.assertEquals(
         DorisTablePropertiesMetadata.DEFAULT_REPLICATION_FACTOR, propertyEntry.getDefaultValue());
     Assertions.assertFalse(propertyEntry.isHidden());
+
+    PropertyEntry<?> replicationAllocationEntry = propertyEntryMap.get(REPLICATION_ALLOCATION);
+    Assertions.assertEquals(REPLICATION_ALLOCATION, replicationAllocationEntry.getName());
+    Assertions.assertFalse(replicationAllocationEntry.isImmutable());
+    Assertions.assertNull(replicationAllocationEntry.getDefaultValue());
+    Assertions.assertFalse(replicationAllocationEntry.isHidden());
   }
 }

catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/operation/TestDorisTableOperationsSqlGeneration.java

@@ -18,7 +18,16 @@
  */
 package org.apache.gravitino.catalog.doris.operation;
 
+import static org.apache.gravitino.catalog.doris.DorisTablePropertiesMetadata.REPLICATION_ALLOCATION;
+import static org.apache.gravitino.catalog.doris.DorisTablePropertiesMetadata.REPLICATION_FACTOR;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.Statement;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import javax.sql.DataSource;
 import org.apache.gravitino.catalog.doris.converter.DorisTypeConverter;
 import org.apache.gravitino.catalog.jdbc.JdbcColumn;
 import org.apache.gravitino.catalog.jdbc.converter.JdbcColumnDefaultValueConverter;
@@ -43,6 +52,10 @@ public TestableDorisTableOperations() {
       super.columnDefaultValueConverter = new JdbcColumnDefaultValueConverter();
     }
 
+    public void setDataSource(DataSource dataSource) {
+      super.dataSource = dataSource;
+    }
+
     public String createTableSql(
         String tableName, JdbcColumn[] columns, Distribution distribution) {
       return generateCreateTableSql(
@@ -107,4 +120,45 @@ public void testCreateTableWithNonEmptyStringDefaultValue() {
         sql.contains("DEFAULT " + converter.fromGravitino(col1.defaultValue())),
         "Should contain DEFAULT value but was: " + sql);
   }
+
+  @Test
+  public void testAppendNecessaryPropertiesAddsReplicationNumWhenBackendsAreNotEnough()
+      throws Exception {
+    TestableDorisTableOperations ops = new TestableDorisTableOperations();
+    ops.setDataSource(mockBackendDataSource(1));
+
+    Map<String, String> properties = ops.appendNecessaryProperties(Collections.emptyMap());
+
+    Assertions.assertEquals("1", properties.get(REPLICATION_FACTOR));
+  }
+
+  @Test
+  public void testAppendNecessaryPropertiesKeepsReplicationAllocation() throws Exception {
+    TestableDorisTableOperations ops = new TestableDorisTableOperations();
+    ops.setDataSource(mockBackendDataSource(1));
+
+    Map<String, String> properties = new HashMap<>();
+    properties.put(REPLICATION_ALLOCATION, "tag.location.default: 1");
+
+    Map<String, String> result = ops.appendNecessaryProperties(properties);
+
+    Assertions.assertEquals("tag.location.default: 1", result.get(REPLICATION_ALLOCATION));
+    Assertions.assertFalse(result.containsKey(REPLICATION_FACTOR));
+  }
+
+  private static DataSource mockBackendDataSource(int aliveBackendCount) throws Exception {
+    DataSource dataSource = Mockito.mock(DataSource.class);
+    Connection connection = Mockito.mock(Connection.class);
+    Statement statement = Mockito.mock(Statement.class);
+    ResultSet resultSet = Mockito.mock(ResultSet.class);
+
+    Mockito.when(dataSource.getConnection()).thenReturn(connection);
+    Mockito.when(connection.createStatement()).thenReturn(statement);
+    Mockito.when(statement.executeQuery("show backends")).thenReturn(resultSet);
+    int[] remainingAliveBackends = new int[] {aliveBackendCount};
+    Mockito.when(resultSet.next()).thenAnswer(invocation -> remainingAliveBackends[0]-- > 0);
+    Mockito.when(resultSet.getString("Alive")).thenReturn("true");
+
+    return dataSource;
+  }
 }