Skip to content

Commit

Permalink
[apache#1759] feat(TableChange) : TableChange support index.
Browse files Browse the repository at this point in the history
  • Loading branch information
@Clearvive
Clearvive authored and Clearvive committed Feb 27, 2024
1 parent 989ad92 commit b2592af
Show file tree
Hide file tree
Showing 11 changed files with 612 additions and 23 deletions.
113 changes: 113 additions & 0 deletions api/src/main/java/com/datastrato/gravitino/rel/TableChange.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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.
*
* <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;

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.
Expand Down
65 changes: 54 additions & 11 deletions .../src/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlTableOperations.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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:
Expand All @@ -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"));
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -355,6 +364,17 @@ 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);
Expand All @@ -376,6 +396,29 @@ 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<TableChange.SetProperty> setProperties) {
return setProperties.stream()
.map(
Expand Down
29 changes: 29 additions & 0 deletions .../test/java/com/datastrato/gravitino/catalog/mysql/operation/TestMysqlTableOperations.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -48,4 +49,32 @@ 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);
}
}
82 changes: 71 additions & 11 deletions ...java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlTableOperations.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -270,6 +273,10 @@ protected String generateAlterTableSql(
validateUpdateColumnNullable(updateColumnNullability, lazyLoadTable);

alterSql.add(updateColumnNullabilityDefinition(updateColumnNullability, tableName));
} 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());
Expand All @@ -287,6 +294,59 @@ 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) {
Expand Down
Loading

0 comments on commit b2592af

Please sign in to comment.