[apache#1707] feat(mysql): Support mysql index.

Clearvive · Jan 26, 2024 · e15078b · e15078b
1 parent cf892cd
commit e15078b
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 48 deletions.
diff --git a/api/src/main/java/com/datastrato/gravitino/rel/indexes/Indexes.java b/api/src/main/java/com/datastrato/gravitino/rel/indexes/Indexes.java
@@ -9,6 +9,7 @@ public class Indexes {
 
   public static final Index[] EMPTY_INDEXES = new Index[0];
 
+  /** My SQL does not support setting the name of the primary key, so the default name is used. */
   public static final String DEFAULT_MYSQL_PRIMARY_KEY_NAME = "PRIMARY";
 
   /**
@@ -23,14 +24,13 @@ public static Index unique(String name, String[][] fieldNames) {
   }
 
   /**
-   * Mysql only supports single-column primary key, and the name of the primary key is fixed, which
-   * is PRIMARY
+   * To create a My SQL primary key, you need to use the default primary key name.
    *
-   * @param fieldName The field name under the table contained in the index.
+   * @param fieldNames T.
    * @return The primary key index
    */
-  public static Index createMysqlPrimaryKey(String fieldName) {
-    return primary(DEFAULT_MYSQL_PRIMARY_KEY_NAME, new String[][] {{fieldName}});
+  public static Index createMysqlPrimaryKey(String[][] fieldNames) {
+    return primary(DEFAULT_MYSQL_PRIMARY_KEY_NAME, fieldNames);
   }
 
   /**

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
@@ -95,14 +95,27 @@ public JdbcTable load(String databaseName, String tableName) throws NoSuchTableE
 
       List<IndexDTO> indexes = new ArrayList<>();
 
+      Map<String, Set<String>> primaryKeyGroupByName = new HashMap<>();
       // 4.Get primary key information
       ResultSet primaryKeys = metaData.getPrimaryKeys(databaseName, null, tableName);
       while (primaryKeys.next()) {
+        String columnName = primaryKeys.getString("COLUMN_NAME");
+        primaryKeyGroupByName.compute(
+            primaryKeys.getString("PK_NAME"),
+            (k, v) -> {
+              if (v == null) {
+                v = new HashSet<>();
+              }
+              v.add(columnName);
+              return v;
+            });
+      }
+      for (Map.Entry<String, Set<String>> entry : primaryKeyGroupByName.entrySet()) {
         indexes.add(
             IndexDTO.builder()
                 .withIndexType(Index.IndexType.PRIMARY_KEY)
-                .withName(primaryKeys.getString("PK_NAME"))
-                .withFieldNames(new String[][] {{primaryKeys.getString("COLUMN_NAME")}})
+                .withName(entry.getKey())
+                .withFieldNames(convertIndexFieldNames(entry.getValue().toArray(new String[0])))
                 .build());
       }
 
@@ -130,7 +143,7 @@ public JdbcTable load(String databaseName, String tableName) throws NoSuchTableE
             IndexDTO.builder()
                 .withIndexType(Index.IndexType.UNIQUE_KEY)
                 .withName(entry.getKey())
-                .withFieldNames(new String[][] {entry.getValue().toArray(new String[0])})
+                .withFieldNames(convertIndexFieldNames(entry.getValue().toArray(new String[0])))
                 .build());
       }
 
@@ -157,6 +170,12 @@ public JdbcTable load(String databaseName, String tableName) throws NoSuchTableE
     }
   }
 
+  private String[][] convertIndexFieldNames(String[] fieldNames) {
+    return Arrays.stream(fieldNames)
+        .map(colName -> new String[] {colName})
+        .toArray(String[][]::new);
+  }
+
   private Map<String, String> loadTablePropertiesFromSql(Connection connection, String tableName)
       throws SQLException {
     try (PreparedStatement statement = connection.prepareStatement("SHOW TABLE STATUS LIKE ?")) {
@@ -251,42 +270,36 @@ protected String generateCreateTableSql(
       }
     }
 
-    if (indexes.length > 0) {
-      for (Index index : indexes) {
-        sqlBuilder.append(",\n");
-        switch (index.type()) {
-          case PRIMARY_KEY:
-            if (index.fieldNames().length != 1 || index.fieldNames()[0].length != 1) {
-              throw new IllegalArgumentException(
-                  "Primary key does not support complex fields in MySQL");
-            }
-            if (!"PRIMARY".equals(index.name().toUpperCase(Locale.ROOT))) {
-              throw new IllegalArgumentException("Primary key name must be PRIMARY in MySQL");
-            }
-            sqlBuilder
-                .append("CONSTRAINT ")
-                .append("PRIMARY KEY (")
-                .append(index.fieldNames()[0][0])
-                .append(")");
-            break;
-          case UNIQUE_KEY:
-            if (index.fieldNames().length != 1) {
-              throw new IllegalArgumentException(
-                  "Unique key does not support complex fields in MySQL");
-            }
-
-            String fields = String.join(", ", index.fieldNames()[0]);
-            sqlBuilder
-                .append("CONSTRAINT ")
-                .append(index.name())
-                .append(" UNIQUE (")
-                .append(fields)
-                .append(")");
-            break;
-          default:
-            throw new UnsupportedOperationException(
-                "Currently we do not support index in MySQL, index type: " + index.type());
-        }
+    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 colNames[0];
+                  })
+              .collect(Collectors.joining(", "));
+      sqlBuilder.append(",\n");
+      switch (index.type()) {
+        case PRIMARY_KEY:
+          if (!"PRIMARY".equals(index.name().toUpperCase(Locale.ROOT))) {
+            throw new IllegalArgumentException("Primary key name must be PRIMARY in MySQL");
+          }
+          sqlBuilder.append("CONSTRAINT ").append("PRIMARY KEY (").append(fieldStr).append(")");
+          break;
+        case UNIQUE_KEY:
+          sqlBuilder
+              .append("CONSTRAINT ")
+              .append(index.name())
+              .append(" UNIQUE (")
+              .append(fieldStr)
+              .append(")");
+          break;
+        default:
+          throw new IllegalArgumentException("MySQL doesn't support index : " + index.type());
       }
     }
 

diff --git a/...om/datastrato/gravitino/integration/test/catalog/jdbc/mysql/TestMysqlTableOperations.java b/...om/datastrato/gravitino/integration/test/catalog/jdbc/mysql/TestMysqlTableOperations.java
@@ -67,7 +67,7 @@ public void testOperationTable() {
     Map<String, String> properties = new HashMap<>();
     properties.put(MYSQL_AUTO_INCREMENT_OFFSET_KEY, "10");
 
-    Index[] indexes = new Index[] {Indexes.unique("test", new String[][] {{"col_1", "col_2"}})};
+    Index[] indexes = new Index[] {Indexes.unique("test", new String[][] {{"col_1"}, {"col_2"}})};
     // create table
     TABLE_OPERATIONS.create(
         TEST_DB_NAME,
@@ -190,8 +190,8 @@ public void testAlterTable() {
 
     Index[] indexes =
         new Index[] {
-          Indexes.createMysqlPrimaryKey("col_1"),
-          Indexes.unique("uk_2", new String[][] {{"col_1", "col_2"}})
+          Indexes.createMysqlPrimaryKey(new String[][] {{"col_1"}, {"col_2"}}),
+          Indexes.unique("uk_2", new String[][] {{"col_1"}, {"col_2"}})
         };
     // create table
     TABLE_OPERATIONS.create(
@@ -401,7 +401,7 @@ public void testCreateAndLoadTable() {
 
     Index[] indexes =
         new Index[] {
-          Indexes.createMysqlPrimaryKey("col_2"),
+          Indexes.createMysqlPrimaryKey(new String[][] {{"col_2"}}),
           Indexes.unique("uk_col_4", new String[][] {{"col_4"}})
         };
     // create table