[#1277] Improvement(column): Add column attribute autoIncrement of th…

…e increment interface.
apache · Jan 5, 2024 · 9b80ba8 · 9b80ba8
1 parent 26731a9
commit 9b80ba8
Show file tree

Hide file tree

Showing 9 changed files with 125 additions and 18 deletions.
diff --git a/api/src/main/java/com/datastrato/gravitino/rel/Column.java b/api/src/main/java/com/datastrato/gravitino/rel/Column.java
@@ -30,5 +30,8 @@ public interface Column {
   /** @return True if this column may produce null values. Default is true. */
   boolean nullable();
 
+  /** @return True if this column is an auto-increment column. Default is false. */
+  boolean autoIncrement();
+
   // TODO. Support column default value. @Jerry
 }
diff --git a/api/src/main/java/com/datastrato/gravitino/rel/types/Types.java b/api/src/main/java/com/datastrato/gravitino/rel/types/Types.java
@@ -923,4 +923,12 @@ public int hashCode() {
       return Arrays.hashCode(types);
     }
   }
+
+  /**
+   * @param dataType The data type to check.
+   * @return True if the given data type is allowed to be an auto-increment column.
+   */
+  public static boolean allowAutoIncrement(Type dataType) {
+    return dataType instanceof IntegerType || dataType instanceof LongType;
+  }
 }
diff --git a/api/src/test/java/com/datastrato/gravitino/rel/TestTransforms.java b/api/src/test/java/com/datastrato/gravitino/rel/TestTransforms.java
@@ -50,6 +50,11 @@ public String comment() {
           public boolean nullable() {
             return true;
           }
+
+          @Override
+          public boolean autoIncrement() {
+            return false;
+          }
         };
     String[] fieldName = new String[] {column.name()};
 
@@ -97,6 +102,11 @@ public String comment() {
           public boolean nullable() {
             return true;
           }
+
+          @Override
+          public boolean autoIncrement() {
+            return false;
+          }
         };
     // partition by foo(col_1, 'bar')
     NamedReference.FieldReference arg1 = field(column.name());

diff --git a/common/src/main/java/com/datastrato/gravitino/dto/rel/ColumnDTO.java b/common/src/main/java/com/datastrato/gravitino/dto/rel/ColumnDTO.java
@@ -33,18 +33,10 @@ public class ColumnDTO implements Column {
   @JsonProperty("nullable")
   private boolean nullable = true;
 
-  private ColumnDTO() {}
+  @JsonProperty("autoIncrement")
+  private boolean autoIncrement = false;
 
-  /**
-   * Constructs a Column DTO.
-   *
-   * @param name The name of the column.
-   * @param dataType The data type of the column.
-   * @param comment The comment associated with the column.
-   */
-  private ColumnDTO(String name, Type dataType, String comment) {
-    this(name, dataType, comment, true);
-  }
+  private ColumnDTO() {}
 
   /**
    * Constructs a Column DTO.
@@ -53,12 +45,15 @@ private ColumnDTO(String name, Type dataType, String comment) {
    * @param dataType The data type of the column.
    * @param comment The comment associated with the column.
    * @param nullable Whether the column value can be null.
+   * @param autoIncrement Whether the column is an auto-increment column.
    */
-  private ColumnDTO(String name, Type dataType, String comment, boolean nullable) {
+  private ColumnDTO(
+      String name, Type dataType, String comment, boolean nullable, boolean autoIncrement) {
     this.name = name;
     this.dataType = dataType;
     this.comment = comment;
     this.nullable = nullable;
+    this.autoIncrement = autoIncrement;
   }
 
   @Override
@@ -81,6 +76,11 @@ public boolean nullable() {
     return nullable;
   }
 
+  @Override
+  public boolean autoIncrement() {
+    return autoIncrement;
+  }
+
   /**
    * Creates a new Builder to build a Column DTO.
    *
@@ -101,6 +101,7 @@ public static class Builder<S extends Builder> {
     protected Type dataType;
     protected String comment;
     protected boolean nullable = true;
+    protected boolean autoIncrement = false;
 
     public Builder() {}
 
@@ -148,6 +149,17 @@ public S withNullable(boolean nullable) {
       return (S) this;
     }
 
+    /**
+     * Sets whether the column is an auto-increment column.
+     *
+     * @param autoIncrement Whether the column is an auto-increment column.
+     * @return The Builder instance.
+     */
+    public S withAutoIncrement(boolean autoIncrement) {
+      this.autoIncrement = autoIncrement;
+      return (S) this;
+    }
+
     /**
      * Builds a Column DTO based on the provided builder parameters.
      *
@@ -157,7 +169,20 @@ public S withNullable(boolean nullable) {
     public ColumnDTO build() {
       Preconditions.checkNotNull(name, "Column name cannot be null");
       Preconditions.checkNotNull(dataType, "Column data type cannot be null");
-      return new ColumnDTO(name, dataType, comment, nullable);
+      return new ColumnDTO(name, dataType, comment, nullable, autoIncrement);
+    }
+  }
+
+  public void validate() throws IllegalArgumentException {
+    if (autoIncrement()) {
+      // TODO This part of the code will be deleted after underlying support.
+      throw new UnsupportedOperationException("Auto-increment column is not supported yet.");
+    }
+    if (name() == null || name().isEmpty()) {
+      throw new IllegalArgumentException("Column name cannot be null or empty.");
+    }
+    if (dataType() == null) {
+      throw new IllegalArgumentException("Column data type cannot be null.");
     }
   }
 }
diff --git a/common/src/main/java/com/datastrato/gravitino/dto/requests/TableCreateRequest.java b/common/src/main/java/com/datastrato/gravitino/dto/requests/TableCreateRequest.java
@@ -13,7 +13,9 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Preconditions;
 import java.util.Arrays;
+import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import lombok.Builder;
 import lombok.EqualsAndHashCode;
@@ -108,5 +110,17 @@ public void validate() throws IllegalArgumentException {
     if (partitioning != null) {
       Arrays.stream(partitioning).forEach(p -> p.validate(columns));
     }
+
+    List<ColumnDTO> autoIncrementCols =
+        Arrays.stream(columns)
+            .peek(ColumnDTO::validate)
+            .filter(ColumnDTO::autoIncrement)
+            .collect(Collectors.toList());
+    String autoIncrementColsStr =
+        autoIncrementCols.stream().map(ColumnDTO::name).collect(Collectors.joining(",", "[", "]"));
+    Preconditions.checkArgument(
+        autoIncrementCols.size() <= 1,
+        "Only one column can be auto-incremented. There are multiple auto-increment columns in your table: "
+            + autoIncrementColsStr);
   }
 }
diff --git a/common/src/main/java/com/datastrato/gravitino/dto/util/DTOConverters.java b/common/src/main/java/com/datastrato/gravitino/dto/util/DTOConverters.java
@@ -97,6 +97,7 @@ public static ColumnDTO toDTO(Column column) {
         .withDataType(column.dataType())
         .withComment(column.comment())
         .withNullable(column.nullable())
+        .withAutoIncrement(column.autoIncrement())
         .build();
   }
 

diff --git a/common/src/test/java/com/datastrato/gravitino/json/TestDTOJsonSerDe.java b/common/src/test/java/com/datastrato/gravitino/json/TestDTOJsonSerDe.java
@@ -42,7 +42,8 @@ public class TestDTOJsonSerDe {
 
   private final String metalakeJson = "{\"name\":%s,\"comment\":%s,\"properties\":%s,\"audit\":%s}";
 
-  private final String columnJson = "{\"name\":%s,\"type\":%s,\"comment\":%s,\"nullable\":%s}";
+  private final String columnJson =
+      "{\"name\":%s,\"type\":%s,\"comment\":%s,\"nullable\":%s,\"autoIncrement\":%s}";
 
   private final String tableJson =
       "{\"name\":%s,\"comment\":%s,\"columns\":[%s],\"properties\":%s,\"audit\":%s,\"distribution\":%s,\"sortOrders\":%s,\"partitioning\":%s}";
@@ -183,7 +184,8 @@ public void testColumnDTOSerDe() throws Exception {
             withQuotes(name),
             withQuotes(type.simpleString()),
             withQuotes(comment),
-            column.nullable());
+            column.nullable(),
+            column.autoIncrement());
     Assertions.assertEquals(expectedJson, serJson);
     ColumnDTO deserColumn = JsonUtils.objectMapper().readValue(serJson, ColumnDTO.class);
     Assertions.assertEquals(column, deserColumn);
@@ -234,7 +236,8 @@ public void testTableDTOSerDe() throws Exception {
                 withQuotes(name),
                 withQuotes(type.simpleString()),
                 withQuotes(comment),
-                column.nullable()),
+                column.nullable(),
+                column.autoIncrement()),
             JsonUtils.objectMapper().writeValueAsString(properties),
             String.format(auditJson, withQuotes(creator), withQuotes(now.toString()), null, null),
             null,

diff --git a/core/src/main/java/com/datastrato/gravitino/catalog/rel/BaseColumn.java b/core/src/main/java/com/datastrato/gravitino/catalog/rel/BaseColumn.java
@@ -24,6 +24,8 @@ public abstract class BaseColumn implements Column {
 
   protected boolean nullable;
 
+  protected boolean autoIncrement;
+
   /**
    * Returns the name of the column.
    *
@@ -61,6 +63,12 @@ public boolean nullable() {
     return nullable;
   }
 
+  /** @return True if this column is an auto-increment column. Default is false. */
+  @Override
+  public boolean autoIncrement() {
+    return autoIncrement;
+  }
+
   /**
    * Builder interface for creating instances of {@link BaseColumn}.
    *
@@ -76,6 +84,8 @@ interface Builder<SELF extends BaseColumn.Builder<SELF, T>, T extends BaseColumn
 
     SELF withNullable(boolean nullable);
 
+    SELF withAutoIncrement(boolean autoIncrement);
+
     T build();
   }
 
@@ -92,6 +102,7 @@ public abstract static class BaseColumnBuilder<
     protected String comment;
     protected Type dataType;
     protected boolean nullable = true;
+    protected boolean autoIncrement = false;
 
     /**
      * Sets the name of the column.
@@ -141,6 +152,18 @@ public SELF withNullable(boolean nullable) {
       return self();
     }
 
+    /**
+     * Sets whether the column is an auto-increment column.
+     *
+     * @param autoIncrement Whether the column is an auto-increment column.
+     * @return The builder instance.
+     */
+    @Override
+    public SELF withAutoIncrement(boolean autoIncrement) {
+      this.autoIncrement = autoIncrement;
+      return self();
+    }
+
     /**
      * Builds the instance of the column with the provided attributes.
      *

diff --git a/server/src/test/java/com/datastrato/gravitino/server/web/rest/TestTableOperations.java b/server/src/test/java/com/datastrato/gravitino/server/web/rest/TestTableOperations.java
@@ -310,7 +310,9 @@ public void testCreateTable() {
   public void testCreatePartitionedTable() {
     Column[] columns =
         new Column[] {
-          mockColumn("col1", Types.StringType.get()), mockColumn("col2", Types.ByteType.get())
+          mockColumn("col1", Types.StringType.get()),
+          mockColumn("col2", Types.ByteType.get()),
+          mockColumn("col3", Types.IntegerType.get(), "test", false, false)
         };
     Partitioning[] partitioning =
         new Partitioning[] {IdentityPartitioningDTO.of(columns[0].name())};
@@ -345,14 +347,21 @@ public void testCreatePartitionedTable() {
     Assertions.assertEquals(ImmutableMap.of("k1", "v1"), tableDTO.properties());
 
     Column[] columnDTOs = tableDTO.columns();
-    Assertions.assertEquals(2, columnDTOs.length);
+    Assertions.assertEquals(3, columnDTOs.length);
     Assertions.assertEquals(columns[0].name(), columnDTOs[0].name());
     Assertions.assertEquals(columns[0].dataType(), columnDTOs[0].dataType());
     Assertions.assertEquals(columns[0].comment(), columnDTOs[0].comment());
+    Assertions.assertEquals(columns[0].autoIncrement(), columnDTOs[0].autoIncrement());
 
     Assertions.assertEquals(columns[1].name(), columnDTOs[1].name());
     Assertions.assertEquals(columns[1].dataType(), columnDTOs[1].dataType());
     Assertions.assertEquals(columns[1].comment(), columnDTOs[1].comment());
+    Assertions.assertEquals(columns[1].autoIncrement(), columnDTOs[1].autoIncrement());
+
+    Assertions.assertEquals(columns[2].name(), columnDTOs[2].name());
+    Assertions.assertEquals(columns[2].dataType(), columnDTOs[2].dataType());
+    Assertions.assertEquals(columns[2].comment(), columnDTOs[2].comment());
+    Assertions.assertEquals(columns[2].autoIncrement(), columnDTOs[2].autoIncrement());
 
     Assertions.assertArrayEquals(partitioning, tableDTO.partitioning());
 
@@ -777,6 +786,17 @@ private static Column mockColumn(String name, Type type, String comment, boolean
     return column;
   }
 
+  private static Column mockColumn(
+      String name, Type type, String comment, boolean nullable, boolean autoIncrement) {
+    Column column = mock(Column.class);
+    when(column.name()).thenReturn(name);
+    when(column.dataType()).thenReturn(type);
+    when(column.comment()).thenReturn(comment);
+    when(column.nullable()).thenReturn(nullable);
+    when(column.autoIncrement()).thenReturn(autoIncrement);
+    return column;
+  }
+
   private static Table mockTable(
       String tableName, Column[] columns, String comment, Map<String, String> properties) {
     return mockTable(tableName, columns, comment, properties, new Transform[0]);