[#1759] feat(TableChange) : TableChange support index. (#2359)

### What changes were proposed in this pull request? TableChange support index. The current PR does not support the extension of `AlterIndex`. Currently, only `AddIndex` and `DeleteIndex` are supported. ### Why are the changes needed? Fix: #1758 ### Does this PR introduce _any_ user-facing change? Add Index TableChange ### How was this patch tested? IT,UT --------- Co-authored-by: Clearvive <[email protected]>
apache · Feb 28, 2024 · 52ec733 · 52ec733
1 parent ee9788a
commit 52ec733
Show file tree

Hide file tree

Showing 11 changed files with 651 additions and 23 deletions.
diff --git a/api/src/main/java/com/datastrato/gravitino/rel/TableChange.java b/api/src/main/java/com/datastrato/gravitino/rel/TableChange.java
@@ -20,6 +20,7 @@
 
 package com.datastrato.gravitino.rel;
 
+import com.datastrato.gravitino.rel.indexes.Index;
 import com.datastrato.gravitino.rel.types.Type;
 import java.util.Arrays;
 import java.util.Objects;
@@ -284,6 +285,30 @@ static TableChange updateColumnNullability(String[] fieldName, boolean nullable)
     return new UpdateColumnNullability(fieldName, nullable);
   }
 
+  /**
+   * Create a TableChange for adding an index.
+   *
+   * @param type The type of the index.
+   * @param name The name of the index.
+   * @param fieldNames The field names of the index.
+   * @return A TableChange for the add index.
+   */
+  static TableChange addIndex(Index.IndexType type, String name, String[][] fieldNames) {
+    return new AddIndex(type, name, fieldNames);
+  }
+
+  /**
+   * Create a TableChange for deleting an index.
+   *
+   * @param name The name of the index to be dropped.
+   * @param ifExists If true, silence the error if column does not exist during drop. Otherwise, an
+   *     {@link IllegalArgumentException} will be thrown.
+   * @return A TableChange for the delete index.
+   */
+  static TableChange deleteIndex(String name, Boolean ifExists) {
+    return new DeleteIndex(name, ifExists);
+  }
+
   /** A TableChange to rename a table. */
   final class RenameTable implements TableChange {
     private final String newName;
@@ -529,6 +554,103 @@ public String toString() {
     }
   }
 
+  /**
+   * A TableChange to add an index. Add an index key based on the type and field name passed in as
+   * well as the name.
+   */
+  final class AddIndex implements TableChange {
+
+    private final Index.IndexType type;
+    private final String name;
+
+    private final String[][] fieldNames;
+
+    /**
+     * @param type The type of the index.
+     * @param name The name of the index.
+     * @param fieldNames The field names of the index.
+     */
+    public AddIndex(Index.IndexType type, String name, String[][] fieldNames) {
+      this.type = type;
+      this.name = name;
+      this.fieldNames = fieldNames;
+    }
+
+    /** @return The type of the index. */
+    public Index.IndexType getType() {
+      return type;
+    }
+
+    /** @return The name of the index. */
+    public String getName() {
+      return name;
+    }
+
+    /** @return The field names of the index. */
+    public String[][] getFieldNames() {
+      return fieldNames;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+      AddIndex addIndex = (AddIndex) o;
+      return type == addIndex.type
+          && Objects.equals(name, addIndex.name)
+          && Arrays.deepEquals(fieldNames, addIndex.fieldNames);
+    }
+
+    @Override
+    public int hashCode() {
+      int result = Objects.hash(type, name);
+      result = 31 * result + Arrays.hashCode(fieldNames);
+      return result;
+    }
+  }
+
+  /**
+   * A TableChange to delete an index.
+   *
+   * <p>If the index does not exist, the change must result in an {@link IllegalArgumentException}.
+   */
+  final class DeleteIndex implements TableChange {
+    private final String name;
+    private final boolean ifExists;
+
+    /**
+     * @param name name of the index.
+     * @param ifExists If true, silence the error if index does not exist during drop.
+     */
+    public DeleteIndex(String name, boolean ifExists) {
+      this.name = name;
+      this.ifExists = ifExists;
+    }
+
+    /** @return The name of the index to be deleted. */
+    public String getName() {
+      return name;
+    }
+
+    /** @return If true, silence the error if index does not exist during drop. */
+    public boolean isIfExists() {
+      return ifExists;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+      DeleteIndex that = (DeleteIndex) o;
+      return Objects.equals(name, that.name) && Objects.equals(ifExists, that.ifExists);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(name, ifExists);
+    }
+  }
+
   /**
    * The interface for all column positions. Column positions are used to specify the position of a
    * column when adding a new column to a table.

diff --git a/.../src/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlTableOperations.java b/.../src/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlTableOperations.java
@@ -20,6 +20,7 @@
 import com.datastrato.gravitino.rel.expressions.transforms.Transform;
 import com.datastrato.gravitino.rel.indexes.Index;
 import com.datastrato.gravitino.rel.indexes.Indexes;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
@@ -170,17 +171,7 @@ private static void validateIncrementCol(JdbcColumn[] columns, Index[] indexes)
 
   public static void appendIndexesSql(Index[] indexes, StringBuilder sqlBuilder) {
     for (Index index : indexes) {
-      String fieldStr =
-          Arrays.stream(index.fieldNames())
-              .map(
-                  colNames -> {
-                    if (colNames.length > 1) {
-                      throw new IllegalArgumentException(
-                          "Index does not support complex fields in MySQL");
-                    }
-                    return BACK_QUOTE + colNames[0] + BACK_QUOTE;
-                  })
-              .collect(Collectors.joining(", "));
+      String fieldStr = getIndexFieldStr(index.fieldNames());
       sqlBuilder.append(",\n");
       switch (index.type()) {
         case PRIMARY_KEY:
@@ -204,6 +195,19 @@ public static void appendIndexesSql(Index[] indexes, StringBuilder sqlBuilder) {
     }
   }
 
+  private static String getIndexFieldStr(String[][] fieldNames) {
+    return Arrays.stream(fieldNames)
+        .map(
+            colNames -> {
+              if (colNames.length > 1) {
+                throw new IllegalArgumentException(
+                    "Index does not support complex fields in MySQL");
+              }
+              return BACK_QUOTE + colNames[0] + BACK_QUOTE;
+            })
+        .collect(Collectors.joining(", "));
+  }
+
   @Override
   protected boolean getAutoIncrementInfo(ResultSet resultSet) throws SQLException {
     return "YES".equalsIgnoreCase(resultSet.getString("IS_AUTOINCREMENT"));
@@ -319,6 +323,11 @@ protected String generateAlterTableSql(
         alterSql.add(
             updateColumnNullabilityDefinition(
                 (TableChange.UpdateColumnNullability) change, lazyLoadTable));
+      } else if (change instanceof TableChange.AddIndex) {
+        alterSql.add(addIndexDefinition((TableChange.AddIndex) change));
+      } else if (change instanceof TableChange.DeleteIndex) {
+        lazyLoadTable = getOrCreateTable(databaseName, tableName, lazyLoadTable);
+        alterSql.add(deleteIndexDefinition(lazyLoadTable, (TableChange.DeleteIndex) change));
       } else {
         throw new IllegalArgumentException(
             "Unsupported table change type: " + change.getClass().getName());
@@ -355,6 +364,18 @@ protected String generateAlterTableSql(
     return result;
   }
 
+  @VisibleForTesting
+  static String deleteIndexDefinition(
+      JdbcTable lazyLoadTable, TableChange.DeleteIndex deleteIndex) {
+    if (deleteIndex.isIfExists()) {
+      if (Arrays.stream(lazyLoadTable.index())
+          .anyMatch(index -> index.name().equals(deleteIndex.getName()))) {
+        throw new IllegalArgumentException("Index does not exist");
+      }
+    }
+    return "DROP INDEX " + BACK_QUOTE + deleteIndex.getName() + BACK_QUOTE;
+  }
+
   private String updateColumnNullabilityDefinition(
       TableChange.UpdateColumnNullability change, JdbcTable table) {
     validateUpdateColumnNullable(change, table);
@@ -376,6 +397,33 @@ private String updateColumnNullabilityDefinition(
         + appendColumnDefinition(updateColumn, new StringBuilder());
   }
 
+  @VisibleForTesting
+  static String addIndexDefinition(TableChange.AddIndex addIndex) {
+    StringBuilder sqlBuilder = new StringBuilder();
+    sqlBuilder.append("ADD ");
+    switch (addIndex.getType()) {
+      case PRIMARY_KEY:
+        if (null != addIndex.getName()
+            && !StringUtils.equalsIgnoreCase(
+                addIndex.getName(), Indexes.DEFAULT_MYSQL_PRIMARY_KEY_NAME)) {
+          throw new IllegalArgumentException("Primary key name must be PRIMARY in MySQL");
+        }
+        sqlBuilder.append("PRIMARY KEY ");
+        break;
+      case UNIQUE_KEY:
+        sqlBuilder
+            .append("UNIQUE INDEX ")
+            .append(BACK_QUOTE)
+            .append(addIndex.getName())
+            .append(BACK_QUOTE);
+        break;
+      default:
+        break;
+    }
+    sqlBuilder.append(" (").append(getIndexFieldStr(addIndex.getFieldNames())).append(")");
+    return sqlBuilder.toString();
+  }
+
   private String generateTableProperties(List<TableChange.SetProperty> setProperties) {
     return setProperties.stream()
         .map(

diff --git a/.../test/java/com/datastrato/gravitino/catalog/mysql/operation/TestMysqlTableOperations.java b/.../test/java/com/datastrato/gravitino/catalog/mysql/operation/TestMysqlTableOperations.java
@@ -4,6 +4,7 @@
  */
 package com.datastrato.gravitino.catalog.mysql.operation;
 
+import com.datastrato.gravitino.rel.TableChange;
 import com.datastrato.gravitino.rel.indexes.Index;
 import com.datastrato.gravitino.rel.indexes.Indexes;
 import org.junit.jupiter.api.Assertions;
@@ -48,4 +49,36 @@ public void testAppendIndexesBuilder() {
             + "CONSTRAINT `uk_3` UNIQUE (`col_4`, `col_5`, `col_6`, `col_7`)";
     Assertions.assertEquals(expectedStr, sql.toString());
   }
+
+  @Test
+  public void testOperationIndexDefinition() {
+    TableChange.AddIndex failIndex =
+        new TableChange.AddIndex(Index.IndexType.PRIMARY_KEY, "pk_1", new String[][] {{"col_1"}});
+    IllegalArgumentException illegalArgumentException =
+        Assertions.assertThrows(
+            IllegalArgumentException.class,
+            () -> MysqlTableOperations.addIndexDefinition(failIndex));
+    Assertions.assertTrue(
+        illegalArgumentException
+            .getMessage()
+            .contains("Primary key name must be PRIMARY in MySQL"));
+
+    TableChange.AddIndex successIndex =
+        new TableChange.AddIndex(
+            Index.IndexType.UNIQUE_KEY, "uk_1", new String[][] {{"col_1"}, {"col_2"}});
+    String sql = MysqlTableOperations.addIndexDefinition(successIndex);
+    Assertions.assertEquals("ADD UNIQUE INDEX `uk_1` (`col_1`, `col_2`)", sql);
+
+    successIndex =
+        new TableChange.AddIndex(
+            Index.IndexType.PRIMARY_KEY,
+            Indexes.DEFAULT_MYSQL_PRIMARY_KEY_NAME,
+            new String[][] {{"col_1"}, {"col_2"}});
+    sql = MysqlTableOperations.addIndexDefinition(successIndex);
+    Assertions.assertEquals("ADD PRIMARY KEY  (`col_1`, `col_2`)", sql);
+
+    TableChange.DeleteIndex deleteIndex = new TableChange.DeleteIndex("uk_1", false);
+    sql = MysqlTableOperations.deleteIndexDefinition(null, deleteIndex);
+    Assertions.assertEquals("DROP INDEX `uk_1`", sql);
+  }
 }