From 08b2394b5044b0c5c3d1467cf023b1f180015d28 Mon Sep 17 00:00:00 2001 From: Clearvive Date: Tue, 27 Feb 2024 14:20:20 +0800 Subject: [PATCH] [#1759] feat(TableChange) : TableChange support index. --- .../datastrato/gravitino/rel/TableChange.java | 113 ++++++++++++++++++ .../mysql/operation/MysqlTableOperations.java | 65 ++++++++-- .../operation/TestMysqlTableOperations.java | 16 +++ .../operation/PostgreSqlTableOperations.java | 61 ++++++++-- .../TestPostgreSqlTableOperations.java | 33 +++++ .../gravitino/client/DTOConverters.java | 9 ++ .../dto/requests/TableUpdateRequest.java | 95 ++++++++++++++- .../dto/requests/TestTableUpdatesRequest.java | 29 +++++ .../catalog/jdbc/mysql/CatalogMysqlIT.java | 55 +++++++++ .../jdbc/postgresql/CatalogPostgreSqlIT.java | 56 +++++++++ .../server/web/rest/TestTableOperations.java | 17 +++ 11 files changed, 525 insertions(+), 24 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 index cdd68468fa8..4ec15d0934e 100644 --- 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 + */ + 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,94 @@ public String toString() { } } + /** + * A TableChange to add a 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; + + 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.equals(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. + * + *

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; + + 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. Otherwise, an */ + public Boolean getIfExists() { + 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/catalogs/catalog-jdbc-mysql/src/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlTableOperations.java b/catalogs/catalog-jdbc-mysql/src/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlTableOperations.java index 415061d61f2..c313c86e2d2 100644 --- a/catalogs/catalog-jdbc-mysql/src/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlTableOperations.java +++ b/catalogs/catalog-jdbc-mysql/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,14 @@ protected String generateAlterTableSql( return result; } + private String deleteIndexDefinition(JdbcTable lazyLoadTable, TableChange.DeleteIndex deleteIndex) { + if (deleteIndex.getIfExists()) { + Arrays.stream(lazyLoadTable.index()).filter(index -> index.name().equals(deleteIndex.getName())) + .findAny().orElseThrow(() -> new IllegalArgumentException("Index does not exist")); + } + return "DROP INDEX " + deleteIndex.getName(); + } + private String updateColumnNullabilityDefinition( TableChange.UpdateColumnNullability change, JdbcTable table) { validateUpdateColumnNullable(change, table); @@ -376,6 +393,32 @@ private String updateColumnNullabilityDefinition( + appendColumnDefinition(updateColumn, new StringBuilder()); } + @VisibleForTesting + static String addIndexDefinition(TableChange.AddIndex addIndex) { + StringBuilder sqlBuilder = new StringBuilder(); + sqlBuilder.append("ADD").append(SPACE); + 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(addIndex.getName()); + break; + default: + break; + } + sqlBuilder + .append(" (") + .append(getIndexFieldStr(addIndex.getFieldNames())) + .append(")"); + return sqlBuilder.toString(); + } + private String generateTableProperties(List setProperties) { return setProperties.stream() .map( diff --git a/catalogs/catalog-jdbc-mysql/src/test/java/com/datastrato/gravitino/catalog/mysql/operation/TestMysqlTableOperations.java b/catalogs/catalog-jdbc-mysql/src/test/java/com/datastrato/gravitino/catalog/mysql/operation/TestMysqlTableOperations.java index 4462b66c2d6..ce6e02a782b 100644 --- a/catalogs/catalog-jdbc-mysql/src/test/java/com/datastrato/gravitino/catalog/mysql/operation/TestMysqlTableOperations.java +++ b/catalogs/catalog-jdbc-mysql/src/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,19 @@ public void testAppendIndexesBuilder() { + "CONSTRAINT `uk_3` UNIQUE (`col_4`, `col_5`, `col_6`, `col_7`)"; Assertions.assertEquals(expectedStr, sql.toString()); } + + @Test + public void testCreateIndexDefinition() { + 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); + } } diff --git a/catalogs/catalog-jdbc-postgresql/src/main/java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlTableOperations.java b/catalogs/catalog-jdbc-postgresql/src/main/java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlTableOperations.java index b0462a68d70..d75030a8845 100644 --- a/catalogs/catalog-jdbc-postgresql/src/main/java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlTableOperations.java +++ b/catalogs/catalog-jdbc-postgresql/src/main/java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlTableOperations.java @@ -143,17 +143,7 @@ protected String generateCreateTableSql( @VisibleForTesting 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 PostgreSQL"); - } - return PG_QUOTE + colNames[0] + PG_QUOTE; - }) - .collect(Collectors.joining(", ")); + String fieldStr = getIndexFieldStr(index.fieldNames()); sqlBuilder.append(",").append(NEW_LINE); switch (index.type()) { case PRIMARY_KEY: @@ -174,6 +164,19 @@ 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 PostgreSQL"); + } + return PG_QUOTE + colNames[0] + PG_QUOTE; + }) + .collect(Collectors.joining(", ")); + } + private void appendColumnDefinition(JdbcColumn column, StringBuilder sqlBuilder) { // Add data type sqlBuilder @@ -270,7 +273,11 @@ protected String generateAlterTableSql( validateUpdateColumnNullable(updateColumnNullability, lazyLoadTable); alterSql.add(updateColumnNullabilityDefinition(updateColumnNullability, tableName)); - } else { + } else if (change instanceof TableChange.AddIndex) { + alterSql.add(addIndexDefinition(tableName, (TableChange.AddIndex) change)); + } else if (change instanceof TableChange.DeleteIndex) { + alterSql.add(deleteIndexDefinition(tableName, (TableChange.DeleteIndex) change)); + } else { throw new IllegalArgumentException( "Unsupported table change type: " + change.getClass().getName()); } @@ -287,6 +294,36 @@ protected String generateAlterTableSql( return result; } + @VisibleForTesting + static String deleteIndexDefinition(String tableName, TableChange.DeleteIndex deleteIndex) { + StringBuilder sqlBuilder = new StringBuilder(); + sqlBuilder.append("ALTER TABLE ").append(tableName).append(" DROP CONSTRAINT ").append(deleteIndex.getName()).append(";\n"); + if (deleteIndex.getIfExists()) { + sqlBuilder.append("DROP INDEX IF EXISTS ").append(PG_QUOTE).append(deleteIndex.getName()).append(PG_QUOTE).append(";"); + } else { + sqlBuilder.append("DROP INDEX ").append(PG_QUOTE).append(deleteIndex.getName()).append(PG_QUOTE).append(";"); + } + return sqlBuilder.toString(); + } + + @VisibleForTesting + static String addIndexDefinition(String tableName, TableChange.AddIndex addIndex) { + StringBuilder sqlBuilder = new StringBuilder(); + sqlBuilder.append("ALTER TABLE ").append(PG_QUOTE).append(tableName).append(PG_QUOTE).append(" ADD CONSTRAINT ").append(PG_QUOTE).append(addIndex.getName()).append(PG_QUOTE); + switch (addIndex.getType()) { + case PRIMARY_KEY: + sqlBuilder.append(" PRIMARY KEY "); + break; + case UNIQUE_KEY: + sqlBuilder.append(" UNIQUE "); + break; + default: + throw new IllegalArgumentException("Unsupported index type: " + addIndex.getType()); + } + sqlBuilder.append("(").append(getIndexFieldStr(addIndex.getFieldNames())).append(");"); + return sqlBuilder.toString(); + } + private String updateColumnNullabilityDefinition( TableChange.UpdateColumnNullability updateColumnNullability, String tableName) { if (updateColumnNullability.fieldName().length > 1) { diff --git a/catalogs/catalog-jdbc-postgresql/src/test/java/com/datastrato/gravitino/catalog/postgresql/operation/TestPostgreSqlTableOperations.java b/catalogs/catalog-jdbc-postgresql/src/test/java/com/datastrato/gravitino/catalog/postgresql/operation/TestPostgreSqlTableOperations.java index 5b442c4e6c1..97526a0df7b 100644 --- a/catalogs/catalog-jdbc-postgresql/src/test/java/com/datastrato/gravitino/catalog/postgresql/operation/TestPostgreSqlTableOperations.java +++ b/catalogs/catalog-jdbc-postgresql/src/test/java/com/datastrato/gravitino/catalog/postgresql/operation/TestPostgreSqlTableOperations.java @@ -4,6 +4,7 @@ */ package com.datastrato.gravitino.catalog.postgresql.operation; +import com.datastrato.gravitino.rel.TableChange; import com.datastrato.gravitino.rel.indexes.Index; import com.datastrato.gravitino.rel.indexes.Indexes; import org.apache.commons.lang3.StringUtils; @@ -64,4 +65,36 @@ public void testAppendIndexesSql() { illegalArgumentException.getMessage(), "Index does not support complex fields in PostgreSQL")); } + + @Test + public void testOperationIndexDefinition() { + String tableName = "t1"; + // Test add index definition success. + TableChange.AddIndex addIndex = new TableChange.AddIndex(Index.IndexType.PRIMARY_KEY, "test_pk", new String[][]{{"col_1"}}); + String result = PostgreSqlTableOperations.addIndexDefinition(tableName, addIndex); + Assertions.assertEquals( + "ALTER TABLE \"t1\" ADD CONSTRAINT \"test_pk\" PRIMARY KEY (\"col_1\");", + result); + + addIndex = new TableChange.AddIndex(Index.IndexType.UNIQUE_KEY, "test_uk", new String[][]{{"col_1"},{"col_2"}}); + result = PostgreSqlTableOperations.addIndexDefinition(tableName, addIndex); + Assertions.assertEquals( + "ALTER TABLE \"t1\" ADD CONSTRAINT \"test_uk\" UNIQUE (\"col_1\", \"col_2\");", + result); + + // Test delete index definition. + TableChange.DeleteIndex deleteIndex = new TableChange.DeleteIndex("test_pk", false); + result = PostgreSqlTableOperations.deleteIndexDefinition(tableName, deleteIndex); + Assertions.assertEquals( + "ALTER TABLE t1 DROP CONSTRAINT test_pk;\n" + + "DROP INDEX \"test_pk\";", + result); + + deleteIndex = new TableChange.DeleteIndex("test_2_pk", true); + result = PostgreSqlTableOperations.deleteIndexDefinition(tableName, deleteIndex); + Assertions.assertEquals( + "ALTER TABLE t1 DROP CONSTRAINT test_2_pk;\n" + + "DROP INDEX IF EXISTS \"test_2_pk\";", + result); + } } diff --git a/clients/client-java/src/main/java/com/datastrato/gravitino/client/DTOConverters.java b/clients/client-java/src/main/java/com/datastrato/gravitino/client/DTOConverters.java index 0be41467554..9e4a559963f 100644 --- a/clients/client-java/src/main/java/com/datastrato/gravitino/client/DTOConverters.java +++ b/clients/client-java/src/main/java/com/datastrato/gravitino/client/DTOConverters.java @@ -147,6 +147,15 @@ static TableUpdateRequest toTableUpdateRequest(TableChange change) { } else if (change instanceof TableChange.ColumnChange) { return toColumnUpdateRequest((TableChange.ColumnChange) change); + } else if (change instanceof TableChange.AddIndex) { + return new TableUpdateRequest.AddTableIndexRequest( + ((TableChange.AddIndex) change).getType(), + ((TableChange.AddIndex) change).getName(), + ((TableChange.AddIndex) change).getFieldNames()); + } else if (change instanceof TableChange.DeleteIndex) { + return new TableUpdateRequest.DeleteTableIndexRequest( + ((TableChange.DeleteIndex) change).getName(), + ((TableChange.DeleteIndex) change).getIfExists()); } else { throw new IllegalArgumentException( "Unknown change type: " + change.getClass().getSimpleName()); diff --git a/common/src/main/java/com/datastrato/gravitino/dto/requests/TableUpdateRequest.java b/common/src/main/java/com/datastrato/gravitino/dto/requests/TableUpdateRequest.java index a346065ec7c..ed07c06ba2f 100644 --- a/common/src/main/java/com/datastrato/gravitino/dto/requests/TableUpdateRequest.java +++ b/common/src/main/java/com/datastrato/gravitino/dto/requests/TableUpdateRequest.java @@ -6,6 +6,8 @@ import com.datastrato.gravitino.json.JsonUtils; import com.datastrato.gravitino.rel.TableChange; +import com.datastrato.gravitino.rel.indexes.Index; +import com.datastrato.gravitino.rel.indexes.Indexes; import com.datastrato.gravitino.rel.types.Type; import com.datastrato.gravitino.rest.RESTRequest; import com.fasterxml.jackson.annotation.JsonIgnoreProperties; @@ -54,7 +56,13 @@ name = "updateColumnNullability"), @JsonSubTypes.Type( value = TableUpdateRequest.DeleteTableColumnRequest.class, - name = "deleteColumn") + name = "deleteColumn"), + @JsonSubTypes.Type( + value = TableUpdateRequest.AddTableIndexRequest.class, + name = "addTableIndex"), + @JsonSubTypes.Type( + value = TableUpdateRequest.DeleteTableIndexRequest.class, + name = "deleteTableIndex") }) public interface TableUpdateRequest extends RESTRequest { @@ -672,4 +680,89 @@ public TableChange tableChange() { return TableChange.deleteColumn(fieldName, ifExists); } } + + @EqualsAndHashCode + @ToString + class AddTableIndexRequest implements TableUpdateRequest { + + @JsonProperty("index") + @JsonSerialize(using = JsonUtils.IndexSerializer.class) + @JsonDeserialize(using = JsonUtils.IndexDeserializer.class) + private Index index; + + /** Default constructor for Jackson deserialization. */ + public AddTableIndexRequest() { + } + + /** + * The constructor of the add table index request. + * @param type The type of the index + * @param name The name of the index + * @param fieldNames The field names under the table contained in the index. + */ + public AddTableIndexRequest(Index.IndexType type, String name, String[][] fieldNames) { + this.index = Indexes.of(type, name, fieldNames); + } + + /** + * Validates the request. + * + * @throws IllegalArgumentException If the request is invalid, this exception is thrown. + */ + @Override + public void validate() throws IllegalArgumentException { + Preconditions.checkNotNull(index, "Index cannot be null"); + Preconditions.checkArgument(index.type() != null, "Index type cannot be null"); + Preconditions.checkArgument( + index.fieldNames() != null && index.fieldNames().length > 0, + "The index must be set with corresponding column names"); + } + + /** @return An instance of TableChange. */ + @Override + public TableChange tableChange() { + return TableChange.addIndex(index.type(), index.name(), index.fieldNames()); + } + } + + @EqualsAndHashCode + @ToString + class DeleteTableIndexRequest implements TableUpdateRequest { + + @JsonProperty("name") + private String name; + + @JsonProperty("ifExists") + private Boolean ifExists; + + /** Default constructor for Jackson deserialization. */ + public DeleteTableIndexRequest() { + } + + /** + * The constructor of the delete table index request. + * @param name The name of the index + * @param ifExists Whether to delete the index if it exists + */ + public DeleteTableIndexRequest(String name, Boolean ifExists) { + this.name = name; + this.ifExists = ifExists; + } + + /** + * Validates the request. + * + * @throws IllegalArgumentException If the request is invalid, this exception is thrown. + */ + @Override + public void validate() throws IllegalArgumentException { + Preconditions.checkNotNull(name, "Index name cannot be null"); + } + + /** @return An instance of TableChange. */ + @Override + public TableChange tableChange() { + return TableChange.deleteIndex(name, ifExists); + } + } } diff --git a/common/src/test/java/com/datastrato/gravitino/dto/requests/TestTableUpdatesRequest.java b/common/src/test/java/com/datastrato/gravitino/dto/requests/TestTableUpdatesRequest.java index 42c6fc8a95a..d6b3c9ecd89 100644 --- a/common/src/test/java/com/datastrato/gravitino/dto/requests/TestTableUpdatesRequest.java +++ b/common/src/test/java/com/datastrato/gravitino/dto/requests/TestTableUpdatesRequest.java @@ -6,6 +6,8 @@ import com.datastrato.gravitino.json.JsonUtils; import com.datastrato.gravitino.rel.TableChange; +import com.datastrato.gravitino.rel.indexes.Index; +import com.datastrato.gravitino.rel.indexes.Indexes; import com.datastrato.gravitino.rel.types.Types; import com.fasterxml.jackson.core.JsonProcessingException; import com.google.common.collect.ImmutableList; @@ -193,4 +195,31 @@ public void testAddTableColumnRequest() throws JsonProcessingException { .getPosition() instanceof TableChange.Default); } + + @Test + public void testOperationTableIndexRequest() throws JsonProcessingException { + //check add index request + TableUpdateRequest tableUpdateRequest = new TableUpdateRequest.AddTableIndexRequest( + Index.IndexType.PRIMARY_KEY, Indexes.DEFAULT_MYSQL_PRIMARY_KEY_NAME, new String[][] {{"column1"}}); + String jsonString = JsonUtils.objectMapper().writeValueAsString(tableUpdateRequest); + String expected = + "{\"@type\":\"addTableIndex\",\"index\":{\"indexType\":\"PRIMARY_KEY\",\"name\":\"PRIMARY\",\"fieldNames\":[[\"column1\"]]}}"; + Assertions.assertEquals( + JsonUtils.objectMapper().readTree(expected), JsonUtils.objectMapper().readTree(jsonString)); + + tableUpdateRequest = new TableUpdateRequest.AddTableIndexRequest( + Index.IndexType.UNIQUE_KEY, "uk_2", new String[][] {{"column2"}}); + jsonString = JsonUtils.objectMapper().writeValueAsString(tableUpdateRequest); + expected = + "{\"@type\":\"addTableIndex\",\"index\":{\"indexType\":\"UNIQUE_KEY\",\"name\":\"uk_2\",\"fieldNames\":[[\"column2\"]]}}"; + Assertions.assertEquals( + JsonUtils.objectMapper().readTree(expected), JsonUtils.objectMapper().readTree(jsonString)); + + //check delete index request + TableUpdateRequest.DeleteTableIndexRequest deleteTableIndexRequest = new TableUpdateRequest.DeleteTableIndexRequest("uk_2", true); + jsonString = JsonUtils.objectMapper().writeValueAsString(deleteTableIndexRequest); + expected = "{\"@type\":\"deleteTableIndex\",\"name\":\"uk_2\",\"ifExists\":true}"; + Assertions.assertEquals( + JsonUtils.objectMapper().readTree(expected), JsonUtils.objectMapper().readTree(jsonString)); + } } diff --git a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/mysql/CatalogMysqlIT.java b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/mysql/CatalogMysqlIT.java index 276d5405c2d..a8b4e39e357 100644 --- a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/mysql/CatalogMysqlIT.java +++ b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/mysql/CatalogMysqlIT.java @@ -1301,4 +1301,59 @@ void testUnparsedTypeConverter() { .loadTable(NameIdentifier.of(metalakeName, catalogName, schemaName, tableName)); Assertions.assertEquals(Types.UnparsedType.of("BIT"), loadedTable.columns()[0].dataType()); } + + @Test + void testOperationTableIndex() { + String tableName = GravitinoITUtils.genRandomName("test_add_index"); + Column col1 = Column.of("col_1", Types.LongType.get(), "id", false, false, null); + Column col2 = Column.of("col_2", Types.VarCharType.of(255), "code", false, false, null); + Column col3 = Column.of("col_3", Types.VarCharType.of(255), "config", false, false, null); + Column[] newColumns = new Column[] {col1, col2, col3}; + TableCatalog tableCatalog = catalog.asTableCatalog(); + tableCatalog.createTable( + NameIdentifier.of(metalakeName, catalogName, schemaName, tableName), + newColumns, + table_comment, + createProperties(), + Transforms.EMPTY_TRANSFORM, + Distributions.NONE, + new SortOrder[0], + Indexes.EMPTY_INDEXES); + + //add index test. + tableCatalog.alterTable( + NameIdentifier.of(metalakeName, catalogName, schemaName, tableName), + TableChange.addIndex(Index.IndexType.UNIQUE_KEY, "u1_key", new String[][]{{"col_2"}, {"col_3"}}), + TableChange.addIndex(Index.IndexType.PRIMARY_KEY, Indexes.DEFAULT_MYSQL_PRIMARY_KEY_NAME, new String[][]{{"col_1"}})); + + Table table = tableCatalog.loadTable(NameIdentifier.of(metalakeName, catalogName, schemaName, tableName)); + Index[] indexes = new Index[] {Indexes.unique("u1_key", new String[][] {{"col_2"}, {"col_3"}}),Indexes.createMysqlPrimaryKey(new String[][] {{"col_1"}})}; + assertionsTableInfo(tableName, table_comment, Arrays.asList(newColumns), createProperties(), indexes, table); + + //delete index and add new column and index. + tableCatalog.alterTable( + NameIdentifier.of(metalakeName, catalogName, schemaName, tableName), + TableChange.deleteIndex("u1_key", false), + TableChange.addColumn(new String[]{"col_4"}, Types.VarCharType.of(255), TableChange.ColumnPosition.defaultPos()), + TableChange.addIndex(Index.IndexType.UNIQUE_KEY, "u2_key", new String[][]{{"col_4"}})); + + indexes = new Index[] {Indexes.createMysqlPrimaryKey(new String[][] {{"col_1"}}),Indexes.unique("u2_key", new String[][] {{"col_4"}})}; + table = tableCatalog.loadTable(NameIdentifier.of(metalakeName, catalogName, schemaName, tableName)); + Column col4 = Column.of("col_4", Types.VarCharType.of(255), null, true, false, null); + newColumns = new Column[] {col1, col2, col3 , col4}; + assertionsTableInfo(tableName, table_comment, Arrays.asList(newColumns), createProperties(), indexes, table); + + //Add a previously existing index + tableCatalog.alterTable( + NameIdentifier.of(metalakeName, catalogName, schemaName, tableName), + TableChange.addIndex(Index.IndexType.UNIQUE_KEY, "u1_key", new String[][]{{"col_2"}, {"col_3"}}), + TableChange.addIndex(Index.IndexType.UNIQUE_KEY, "u3_key", new String[][]{{"col_1"}, {"col_4"}})); + + indexes = new Index[]{Indexes.createMysqlPrimaryKey(new String[][]{{"col_1"}}), + Indexes.unique("u2_key", new String[][]{{"col_4"}}), + Indexes.unique("u1_key", new String[][]{{"col_2"}, {"col_3"}}), + Indexes.unique("u3_key", new String[][]{{"col_1"}, {"col_4"}})}; + table = tableCatalog.loadTable(NameIdentifier.of(metalakeName, catalogName, schemaName, tableName)); + assertionsTableInfo(tableName, table_comment, Arrays.asList(newColumns), createProperties(), indexes, table); + } } diff --git a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/CatalogPostgreSqlIT.java b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/CatalogPostgreSqlIT.java index 844919d3e98..a72cac52645 100644 --- a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/CatalogPostgreSqlIT.java +++ b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/CatalogPostgreSqlIT.java @@ -1193,4 +1193,60 @@ void testUnparsedTypeConverter() { .loadTable(NameIdentifier.of(metalakeName, catalogName, schemaName, tableName)); Assertions.assertEquals(Types.UnparsedType.of("bit"), loadedTable.columns()[0].dataType()); } + + @Test + void testOperationTableIndex() { + String tableName = GravitinoITUtils.genRandomName("test_add_index"); + Column col1 = Column.of("col_1", Types.LongType.get(), "id", false, false, null); + Column col2 = Column.of("col_2", Types.VarCharType.of(255), "code", false, false, null); + Column col3 = Column.of("col_3", Types.VarCharType.of(255), "config", false, false, null); + Column[] newColumns = new Column[] {col1, col2, col3}; + TableCatalog tableCatalog = catalog.asTableCatalog(); + tableCatalog.createTable( + NameIdentifier.of(metalakeName, catalogName, schemaName, tableName), + newColumns, + table_comment, + createProperties(), + Transforms.EMPTY_TRANSFORM, + Distributions.NONE, + new SortOrder[0], + Indexes.EMPTY_INDEXES); + + //add index test. + tableCatalog.alterTable( + NameIdentifier.of(metalakeName, catalogName, schemaName, tableName), + TableChange.addIndex(Index.IndexType.UNIQUE_KEY, "u1_key", new String[][]{{"col_2"}, {"col_3"}}), + TableChange.addIndex(Index.IndexType.PRIMARY_KEY, "pk1_key", new String[][]{{"col_1"}})); + + Table table = tableCatalog.loadTable(NameIdentifier.of(metalakeName, catalogName, schemaName, tableName)); + Index[] indexes = new Index[]{Indexes.unique("u1_key", new String[][]{{"col_2"}, {"col_3"}}), Indexes.primary("pk1_key", new String[][]{{"col_1"}})}; + assertionsTableInfo(tableName, table_comment, Arrays.asList(newColumns), createProperties(), indexes, table); + + //delete index and add new column and index. + tableCatalog.alterTable( + NameIdentifier.of(metalakeName, catalogName, schemaName, tableName), + TableChange.deleteIndex("u1_key", true), + TableChange.addColumn(new String[]{"col_4"}, Types.VarCharType.of(255), TableChange.ColumnPosition.defaultPos()), + TableChange.addIndex(Index.IndexType.UNIQUE_KEY, "u2_key", new String[][]{{"col_4"}})); + + indexes = new Index[] {Indexes.primary("pk1_key",new String[][] {{"col_1"}}),Indexes.unique("u2_key", new String[][] {{"col_4"}})}; + table = tableCatalog.loadTable(NameIdentifier.of(metalakeName, catalogName, schemaName, tableName)); + Column col4 = Column.of("col_4", Types.VarCharType.of(255), null, true, false, null); + newColumns = new Column[] {col1, col2, col3 , col4}; + assertionsTableInfo(tableName, table_comment, Arrays.asList(newColumns), createProperties(), indexes, table); + + + //Add a previously existing index + tableCatalog.alterTable( + NameIdentifier.of(metalakeName, catalogName, schemaName, tableName), + TableChange.addIndex(Index.IndexType.UNIQUE_KEY, "u1_key", new String[][]{{"col_2"}, {"col_3"}}), + TableChange.addIndex(Index.IndexType.UNIQUE_KEY, "u3_key", new String[][]{{"col_1"}, {"col_4"}})); + + indexes = new Index[]{Indexes.primary("pk1_key", new String[][]{{"col_1"}}), + Indexes.unique("u2_key", new String[][]{{"col_4"}}), + Indexes.unique("u1_key", new String[][]{{"col_2"}, {"col_3"}}), + Indexes.unique("u3_key", new String[][]{{"col_1"}, {"col_4"}})}; + table = tableCatalog.loadTable(NameIdentifier.of(metalakeName, catalogName, schemaName, tableName)); + assertionsTableInfo(tableName, table_comment, Arrays.asList(newColumns), createProperties(), indexes, table); + } } 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 index 59fc3dcdb78..49e6fa1cb6a 100644 --- 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 @@ -42,6 +42,8 @@ import com.datastrato.gravitino.rel.expressions.sorts.SortDirection; import com.datastrato.gravitino.rel.expressions.sorts.SortOrder; import com.datastrato.gravitino.rel.expressions.transforms.Transform; +import com.datastrato.gravitino.rel.indexes.Index; +import com.datastrato.gravitino.rel.indexes.Indexes; import com.datastrato.gravitino.rel.types.Type; import com.datastrato.gravitino.rel.types.Types; import com.datastrato.gravitino.rest.RESTUtils; @@ -617,6 +619,20 @@ public void testUpdateTableColumnPosition() { testAlterTableRequest(req, table); } + @Test + public void testAddTableIndex() { + TableUpdateRequest.AddTableIndexRequest req = + new TableUpdateRequest.AddTableIndexRequest( + Index.IndexType.PRIMARY_KEY, Indexes.DEFAULT_MYSQL_PRIMARY_KEY_NAME, new String[][] {{"col1"}}); + Column[] columns = + new Column[] { + mockColumn("col2", Types.ByteType.get()), mockColumn("col1", Types.StringType.get()) + }; + Table table = + mockTable("table1", columns, "mock comment", ImmutableMap.of("k1", "v1"), new Transform[0]); + testAlterTableRequest(req, table); + } + @Test public void testDropTable() { when(dispatcher.dropTable(any())).thenReturn(true); @@ -762,6 +778,7 @@ private void testAlterTableRequest(TableUpdateRequest req, Table updatedTable) { Assertions.assertEquals(tableDTO.distribution(), updatedTable.distribution()); Assertions.assertArrayEquals(tableDTO.sortOrder(), updatedTable.sortOrder()); + Assertions.assertArrayEquals(tableDTO.index(), updatedTable.index()); } private static String tablePath(String metalake, String catalog, String schema) {