From 204ae9a536ce3cfa3783b44d476b47ab9389dac6 Mon Sep 17 00:00:00 2001 From: Clearvive Date: Tue, 30 Jan 2024 20:06:09 +0800 Subject: [PATCH] [#1736] feat(postgresql): Support PostgreSQL index. --- .../catalog/jdbc/bean/JdbcIndexBean.java | 76 +++++++++++ .../jdbc/operation/JdbcTableOperations.java | 75 ++++++++++- .../mysql/operation/MysqlTableOperations.java | 41 ------ .../catalog-jdbc-postgresql/build.gradle.kts | 6 + .../operation/PostgreSqlTableOperations.java | 49 +++++++ .../TestPostgreSqlTableOperations.java | 67 ++++++++++ .../catalog/jdbc/mysql/CatalogMysqlIT.java | 5 + .../jdbc/postgresql/CatalogPostgreSqlIT.java | 123 ++++++++++++++++++ ....java => PostgreSqlTableOperationsIT.java} | 85 +++++++++++- .../integration/test/util/AbstractIT.java | 10 +- 10 files changed, 489 insertions(+), 48 deletions(-) create mode 100644 catalogs/catalog-jdbc-common/src/main/java/com/datastrato/gravitino/catalog/jdbc/bean/JdbcIndexBean.java create mode 100644 catalogs/catalog-jdbc-postgresql/src/test/java/com/datastrato/gravitino/catalog/postgresql/operation/TestPostgreSqlTableOperations.java rename integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/{TestPostgreSqlTableOperations.java => PostgreSqlTableOperationsIT.java} (85%) diff --git a/catalogs/catalog-jdbc-common/src/main/java/com/datastrato/gravitino/catalog/jdbc/bean/JdbcIndexBean.java b/catalogs/catalog-jdbc-common/src/main/java/com/datastrato/gravitino/catalog/jdbc/bean/JdbcIndexBean.java new file mode 100644 index 00000000000..1fedb87ca8a --- /dev/null +++ b/catalogs/catalog-jdbc-common/src/main/java/com/datastrato/gravitino/catalog/jdbc/bean/JdbcIndexBean.java @@ -0,0 +1,76 @@ +/* + * Copyright 2024 Datastrato Pvt Ltd. + * This software is licensed under the Apache License version 2. + */ +package com.datastrato.gravitino.catalog.jdbc.bean; + +import com.datastrato.gravitino.rel.indexes.Index; +import java.util.Objects; + +/** Store JDBC index information. */ +public class JdbcIndexBean { + + private final Index.IndexType indexType; + + private final String colName; + + private final String name; + + /** Used for sorting */ + private final int order; + + public JdbcIndexBean(Index.IndexType indexType, String colName, String name, int order) { + this.indexType = indexType; + this.colName = colName; + this.name = name; + this.order = order; + } + + public Index.IndexType getIndexType() { + return indexType; + } + + public String getColName() { + return colName; + } + + public String getName() { + return name; + } + + public int getOrder() { + return order; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + JdbcIndexBean that = (JdbcIndexBean) o; + return order == that.order + && indexType == that.indexType + && Objects.equals(colName, that.colName) + && Objects.equals(name, that.name); + } + + @Override + public int hashCode() { + return Objects.hash(indexType, colName, name, order); + } + + @Override + public String toString() { + return "JdbcIndexBean{" + + "indexType=" + + indexType + + ", colName='" + + colName + + '\'' + + ", name='" + + name + + '\'' + + ", order=" + + order + + '}'; + } +} diff --git a/catalogs/catalog-jdbc-common/src/main/java/com/datastrato/gravitino/catalog/jdbc/operation/JdbcTableOperations.java b/catalogs/catalog-jdbc-common/src/main/java/com/datastrato/gravitino/catalog/jdbc/operation/JdbcTableOperations.java index 697c28004f0..233de2a265e 100644 --- a/catalogs/catalog-jdbc-common/src/main/java/com/datastrato/gravitino/catalog/jdbc/operation/JdbcTableOperations.java +++ b/catalogs/catalog-jdbc-common/src/main/java/com/datastrato/gravitino/catalog/jdbc/operation/JdbcTableOperations.java @@ -6,6 +6,7 @@ import com.datastrato.gravitino.catalog.jdbc.JdbcColumn; import com.datastrato.gravitino.catalog.jdbc.JdbcTable; +import com.datastrato.gravitino.catalog.jdbc.bean.JdbcIndexBean; import com.datastrato.gravitino.catalog.jdbc.converter.JdbcExceptionConverter; import com.datastrato.gravitino.catalog.jdbc.converter.JdbcTypeConverter; import com.datastrato.gravitino.catalog.jdbc.utils.JdbcConnectorUtils; @@ -16,6 +17,7 @@ import com.datastrato.gravitino.rel.TableChange; 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.collect.Lists; import java.sql.Connection; import java.sql.DatabaseMetaData; @@ -23,8 +25,11 @@ import java.sql.SQLException; import java.util.ArrayList; import java.util.Collections; +import java.util.Comparator; import java.util.List; import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; import javax.sql.DataSource; import org.apache.commons.lang3.StringUtils; import org.slf4j.Logger; @@ -239,7 +244,75 @@ protected void correctJdbcTableFields( protected List getIndexes(String databaseName, String tableName, DatabaseMetaData metaData) throws SQLException { - return Collections.emptyList(); + List indexes = new ArrayList<>(); + + // Get primary key information + ResultSet primaryKeys = getPrimaryKeys(databaseName, tableName, metaData); + List jdbcIndexBeans = new ArrayList<>(); + while (primaryKeys.next()) { + jdbcIndexBeans.add( + new JdbcIndexBean( + Index.IndexType.PRIMARY_KEY, + primaryKeys.getString("COLUMN_NAME"), + primaryKeys.getString("PK_NAME"), + primaryKeys.getInt("KEY_SEQ"))); + } + + Set primaryIndexNames = + jdbcIndexBeans.stream().map(JdbcIndexBean::getName).collect(Collectors.toSet()); + + // Get unique key information + ResultSet indexInfo = getIndexInfo(databaseName, tableName, metaData); + while (indexInfo.next()) { + String indexName = indexInfo.getString("INDEX_NAME"); + if (!indexInfo.getBoolean("NON_UNIQUE") && !primaryIndexNames.contains(indexName)) { + jdbcIndexBeans.add( + new JdbcIndexBean( + Index.IndexType.UNIQUE_KEY, + indexInfo.getString("COLUMN_NAME"), + indexName, + indexInfo.getInt("ORDINAL_POSITION"))); + } + } + + // Assemble into Index + Map> indexBeanGroupByIndexType = + jdbcIndexBeans.stream().collect(Collectors.groupingBy(JdbcIndexBean::getIndexType)); + + for (Map.Entry> entry : + indexBeanGroupByIndexType.entrySet()) { + // Group by index Name + Map> indexBeanGroupByName = + entry.getValue().stream().collect(Collectors.groupingBy(JdbcIndexBean::getName)); + for (Map.Entry> indexEntry : indexBeanGroupByName.entrySet()) { + List colNames = + indexEntry.getValue().stream() + .sorted(Comparator.comparingInt(JdbcIndexBean::getOrder)) + .map(JdbcIndexBean::getColName) + .collect(Collectors.toList()); + String[][] colStrArrays = convertIndexFieldNames(colNames); + if (entry.getKey() == Index.IndexType.PRIMARY_KEY) { + indexes.add(Indexes.primary(indexEntry.getKey(), colStrArrays)); + } else { + indexes.add(Indexes.unique(indexEntry.getKey(), colStrArrays)); + } + } + } + return indexes; + } + + protected ResultSet getIndexInfo(String databaseName, String tableName, DatabaseMetaData metaData) + throws SQLException { + return metaData.getIndexInfo(databaseName, null, tableName, false, false); + } + + protected ResultSet getPrimaryKeys( + String databaseName, String tableName, DatabaseMetaData metaData) throws SQLException { + return metaData.getPrimaryKeys(databaseName, null, tableName); + } + + protected String[][] convertIndexFieldNames(List fieldNames) { + return fieldNames.stream().map(colName -> new String[] {colName}).toArray(String[][]::new); } protected abstract String generateCreateTableSql( 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 f8b63238d3a..ae7dd8df682 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,10 +20,7 @@ import com.datastrato.gravitino.rel.indexes.Index; import com.datastrato.gravitino.rel.indexes.Indexes; import com.google.common.base.Preconditions; -import com.google.common.collect.HashMultimap; -import com.google.common.collect.SetMultimap; import java.sql.Connection; -import java.sql.DatabaseMetaData; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; @@ -35,7 +32,6 @@ import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.Set; import java.util.stream.Collectors; import org.apache.commons.collections4.CollectionUtils; import org.apache.commons.collections4.MapUtils; @@ -215,43 +211,6 @@ public static void appendIndexesSql(Index[] indexes, StringBuilder sqlBuilder) { } } - protected List getIndexes(String databaseName, String tableName, DatabaseMetaData metaData) - throws SQLException { - List indexes = new ArrayList<>(); - - // Get primary key information - SetMultimap primaryKeyGroupByName = HashMultimap.create(); - ResultSet primaryKeys = metaData.getPrimaryKeys(databaseName, null, tableName); - while (primaryKeys.next()) { - String columnName = primaryKeys.getString("COLUMN_NAME"); - primaryKeyGroupByName.put(primaryKeys.getString("PK_NAME"), columnName); - } - for (String key : primaryKeyGroupByName.keySet()) { - indexes.add(Indexes.primary(key, convertIndexFieldNames(primaryKeyGroupByName.get(key)))); - } - - // Get unique key information - SetMultimap indexGroupByName = HashMultimap.create(); - ResultSet indexInfo = metaData.getIndexInfo(databaseName, null, tableName, false, false); - while (indexInfo.next()) { - String indexName = indexInfo.getString("INDEX_NAME"); - if (!indexInfo.getBoolean("NON_UNIQUE") - && !StringUtils.equalsIgnoreCase(Indexes.DEFAULT_MYSQL_PRIMARY_KEY_NAME, indexName)) { - String columnName = indexInfo.getString("COLUMN_NAME"); - indexGroupByName.put(indexName, columnName); - } - } - for (String key : indexGroupByName.keySet()) { - indexes.add(Indexes.unique(key, convertIndexFieldNames(indexGroupByName.get(key)))); - } - - return indexes; - } - - private String[][] convertIndexFieldNames(Set fieldNames) { - return fieldNames.stream().map(colName -> new String[] {colName}).toArray(String[][]::new); - } - @Override protected boolean getAutoIncrementInfo(ResultSet resultSet) throws SQLException { return "YES".equalsIgnoreCase(resultSet.getString("IS_AUTOINCREMENT")); diff --git a/catalogs/catalog-jdbc-postgresql/build.gradle.kts b/catalogs/catalog-jdbc-postgresql/build.gradle.kts index c1924569197..d9c2f4f5986 100644 --- a/catalogs/catalog-jdbc-postgresql/build.gradle.kts +++ b/catalogs/catalog-jdbc-postgresql/build.gradle.kts @@ -20,6 +20,12 @@ dependencies { implementation(libs.commons.lang3) implementation(libs.commons.collections4) implementation(libs.jsqlparser) + + testImplementation(libs.guava) + testImplementation(libs.commons.lang3) + testImplementation(libs.junit.jupiter.api) + testImplementation(libs.junit.jupiter.params) + testRuntimeOnly(libs.junit.jupiter.engine) } tasks { 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 665fd6f5f3e..e78ebc68cb9 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 @@ -16,6 +16,7 @@ import com.datastrato.gravitino.rel.expressions.transforms.Transform; import com.datastrato.gravitino.rel.indexes.Index; import com.datastrato.gravitino.rel.types.Types; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import java.sql.Connection; import java.sql.DatabaseMetaData; @@ -25,6 +26,7 @@ import java.util.Arrays; import java.util.List; import java.util.Map; +import java.util.stream.Collectors; import javax.sql.DataSource; import org.apache.commons.collections4.MapUtils; import org.apache.commons.lang3.ArrayUtils; @@ -82,6 +84,7 @@ protected String generateCreateTableSql( sqlBuilder.append(",\n"); } } + appendIndexesSql(indexes, sqlBuilder); sqlBuilder.append("\n)"); // Add table properties if any if (MapUtils.isNotEmpty(properties)) { @@ -120,6 +123,40 @@ protected String generateCreateTableSql( return result; } + @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(", ")); + sqlBuilder.append(",\n"); + switch (index.type()) { + case PRIMARY_KEY: + if (StringUtils.isNotEmpty(index.name())) { + sqlBuilder.append("CONSTRAINT ").append(PG_QUOTE).append(index.name()).append(PG_QUOTE); + } + sqlBuilder.append(" PRIMARY KEY (").append(fieldStr).append(")"); + break; + case UNIQUE_KEY: + if (StringUtils.isNotEmpty(index.name())) { + sqlBuilder.append("CONSTRAINT ").append(PG_QUOTE).append(index.name()).append(PG_QUOTE); + } + sqlBuilder.append(" UNIQUE (").append(fieldStr).append(")"); + break; + default: + throw new IllegalArgumentException("PostgreSQL doesn't support index : " + index.type()); + } + } + } + private StringBuilder appendColumnDefinition(JdbcColumn column, StringBuilder sqlBuilder) { // Add data type sqlBuilder @@ -440,6 +477,18 @@ private String updateColumnCommentFieldDefinition( + "';"; } + @Override + protected ResultSet getIndexInfo(String schemaName, String tableName, DatabaseMetaData metaData) + throws SQLException { + return metaData.getIndexInfo(database, schemaName, tableName, false, false); + } + + @Override + protected ResultSet getPrimaryKeys(String schemaName, String tableName, DatabaseMetaData metaData) + throws SQLException { + return metaData.getPrimaryKeys(database, schemaName, tableName); + } + @Override protected Connection getConnection(String schema) throws SQLException { Connection connection = dataSource.getConnection(); 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 new file mode 100644 index 00000000000..e779173ce36 --- /dev/null +++ b/catalogs/catalog-jdbc-postgresql/src/test/java/com/datastrato/gravitino/catalog/postgresql/operation/TestPostgreSqlTableOperations.java @@ -0,0 +1,67 @@ +/* + * Copyright 2024 Datastrato Pvt Ltd. + * This software is licensed under the Apache License version 2. + */ +package com.datastrato.gravitino.catalog.postgresql.operation; + +import com.datastrato.gravitino.rel.indexes.Index; +import com.datastrato.gravitino.rel.indexes.Indexes; +import org.apache.commons.lang3.StringUtils; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** Unit test for {@link PostgreSqlTableOperations}. */ +public class TestPostgreSqlTableOperations { + + @Test + public void testAppendIndexesSql() { + // Test append index sql success. + Index[] indexes = + new Index[] { + Indexes.primary("test_pk", new String[][] {{"col_1"}, {"col_2"}}), + Indexes.unique("u1_key", new String[][] {{"col_2"}, {"col_3"}}), + Indexes.unique("u2_key", new String[][] {{"col_3"}, {"col_4"}}) + }; + StringBuilder successBuilder = new StringBuilder(); + PostgreSqlTableOperations.appendIndexesSql(indexes, successBuilder); + Assertions.assertEquals( + ",\n" + + "CONSTRAINT \"test_pk\" PRIMARY KEY (\"col_1\", \"col_2\"),\n" + + "CONSTRAINT \"u1_key\" UNIQUE (\"col_2\", \"col_3\"),\n" + + "CONSTRAINT \"u2_key\" UNIQUE (\"col_3\", \"col_4\")", + successBuilder.toString()); + + // Test append index sql not have name. + indexes = + new Index[] { + Indexes.primary(null, new String[][] {{"col_1"}, {"col_2"}}), + Indexes.unique(null, new String[][] {{"col_2"}, {"col_3"}}), + Indexes.unique(null, new String[][] {{"col_3"}, {"col_4"}}) + }; + successBuilder = new StringBuilder(); + PostgreSqlTableOperations.appendIndexesSql(indexes, successBuilder); + Assertions.assertEquals( + ",\n" + + " PRIMARY KEY (\"col_1\", \"col_2\"),\n" + + " UNIQUE (\"col_2\", \"col_3\"),\n" + + " UNIQUE (\"col_3\", \"col_4\")", + successBuilder.toString()); + + // Test append index sql failed. + IllegalArgumentException illegalArgumentException = + Assertions.assertThrows( + IllegalArgumentException.class, + () -> + PostgreSqlTableOperations.appendIndexesSql( + new Index[] { + Indexes.primary("test_pk", new String[][] {{"col_1", "col_2"}}), + Indexes.unique("u1_key", new String[][] {{"col_2", "col_3"}}), + Indexes.unique("u2_key", new String[][] {{"col_3", "col_4"}}) + }, + new StringBuilder())); + Assertions.assertTrue( + StringUtils.contains( + illegalArgumentException.getMessage(), + "Index does not support complex fields in PostgreSQL")); + } +} 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 6f56e4d1e4a..f7e97f48259 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 @@ -754,6 +754,11 @@ public void testAutoIncrement() { col4, col5 }; + indexes = + new Index[] { + Indexes.createMysqlPrimaryKey(new String[][] {{"col_1_1"}, {"col_2"}}), + Indexes.unique("u1_key", new String[][] {{"col_2"}, {"col_3"}}) + }; assertionsTableInfo( tableName, table_comment, Arrays.asList(alterColumns), properties, 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 c9bfd33f74d..23477a5bd61 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 @@ -4,6 +4,8 @@ */ package com.datastrato.gravitino.integration.test.catalog.jdbc.postgresql; +import static org.junit.jupiter.api.Assertions.assertThrows; + import com.datastrato.gravitino.Catalog; import com.datastrato.gravitino.NameIdentifier; import com.datastrato.gravitino.Namespace; @@ -29,6 +31,7 @@ import com.datastrato.gravitino.rel.expressions.sorts.SortOrder; import com.datastrato.gravitino.rel.expressions.transforms.Transform; import com.datastrato.gravitino.rel.expressions.transforms.Transforms; +import com.datastrato.gravitino.rel.indexes.Index; import com.datastrato.gravitino.rel.indexes.Indexes; import com.datastrato.gravitino.rel.types.Types; import com.google.common.collect.ImmutableMap; @@ -594,4 +597,124 @@ public void testBackQuoteTable() { Assertions.assertDoesNotThrow(() -> tableCatalog.dropTable(tableIdentifier)); } + + @Test + void testCreateIndexTable() { + Column col1 = Column.of("col_1", Types.LongType.get(), "id", false, false, null); + Column col2 = Column.of("col_2", Types.VarCharType.of(100), "yes", false, false, null); + Column col3 = Column.of("col_3", Types.DateType.get(), "comment", false, false, null); + Column col4 = Column.of("col_4", Types.VarCharType.of(255), "code", false, false, null); + Column col5 = Column.of("col_5", Types.VarCharType.of(255), "config", false, false, null); + Column[] newColumns = new Column[] {col1, col2, col3, col4, col5}; + + Index[] indexes = + new Index[] { + Indexes.primary("k1_pk", new String[][] {{"col_2"}, {"col_1"}}), + Indexes.unique("u1_key", new String[][] {{"col_2"}, {"col_3"}}), + Indexes.unique("u2_key", new String[][] {{"col_3"}, {"col_4"}}), + Indexes.unique("u3_key", new String[][] {{"col_5"}, {"col_4"}}), + Indexes.unique("u4_key", new String[][] {{"col_2"}, {"col_4"}, {"col_3"}}), + Indexes.unique("u5_key", new String[][] {{"col_5"}, {"col_3"}, {"col_2"}}), + Indexes.unique("u6_key", new String[][] {{"col_1"}, {"col_3"}, {"col_2"}, {"col_4"}}), + }; + + NameIdentifier tableIdentifier = + NameIdentifier.of(metalakeName, catalogName, schemaName, tableName); + + // Test create many indexes with name success. + Map properties = createProperties(); + TableCatalog tableCatalog = catalog.asTableCatalog(); + Table createdTable = + tableCatalog.createTable( + tableIdentifier, + newColumns, + table_comment, + properties, + Transforms.EMPTY_TRANSFORM, + Distributions.NONE, + new SortOrder[0], + indexes); + assertionsTableInfo( + tableName, table_comment, Arrays.asList(newColumns), properties, indexes, createdTable); + Table table = tableCatalog.loadTable(tableIdentifier); + assertionsTableInfo( + tableName, table_comment, Arrays.asList(newColumns), properties, indexes, table); + + // Test create index complex fields fail. + IllegalArgumentException illegalArgumentException = + assertThrows( + IllegalArgumentException.class, + () -> { + tableCatalog.createTable( + NameIdentifier.of(metalakeName, catalogName, schemaName, "test_failed"), + newColumns, + table_comment, + properties, + Transforms.EMPTY_TRANSFORM, + Distributions.NONE, + new SortOrder[0], + new Index[] {Indexes.createMysqlPrimaryKey(new String[][] {{"col_1", "col_2"}})}); + }); + Assertions.assertTrue( + StringUtils.contains( + illegalArgumentException.getMessage(), + "Index does not support complex fields in PostgreSQL")); + + illegalArgumentException = + assertThrows( + IllegalArgumentException.class, + () -> { + tableCatalog.createTable( + NameIdentifier.of(metalakeName, catalogName, schemaName, "test_failed"), + newColumns, + table_comment, + properties, + Transforms.EMPTY_TRANSFORM, + Distributions.NONE, + new SortOrder[0], + new Index[] {Indexes.unique("u1_key", new String[][] {{"col_2", "col_3"}})}); + }); + Assertions.assertTrue( + StringUtils.contains( + illegalArgumentException.getMessage(), + "Index does not support complex fields in PostgreSQL")); + + // Test create index with empty name success. + table = + tableCatalog.createTable( + NameIdentifier.of(metalakeName, catalogName, schemaName, "test_null_key"), + newColumns, + table_comment, + properties, + Transforms.EMPTY_TRANSFORM, + Distributions.NONE, + new SortOrder[0], + new Index[] { + Indexes.of( + Index.IndexType.UNIQUE_KEY, + null, + new String[][] {{"col_1"}, {"col_3"}, {"col_4"}}), + Indexes.of(Index.IndexType.UNIQUE_KEY, null, new String[][] {{"col_4"}}), + }); + Assertions.assertEquals(2, table.index().length); + Assertions.assertNotNull(table.index()[0].name()); + Assertions.assertNotNull(table.index()[1].name()); + + // Test create index with same col success. + table = + tableCatalog.createTable( + NameIdentifier.of(metalakeName, catalogName, schemaName, "many_index"), + newColumns, + table_comment, + properties, + Transforms.EMPTY_TRANSFORM, + Distributions.NONE, + new SortOrder[0], + new Index[] { + Indexes.unique("u4_key_2", new String[][] {{"col_2"}, {"col_3"}, {"col_4"}}), + Indexes.unique("u5_key_3", new String[][] {{"col_2"}, {"col_3"}, {"col_4"}}), + }); + Assertions.assertEquals(1, table.index().length); + Assertions.assertEquals("u4_key_2", table.index()[0].name()); + } } diff --git a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/TestPostgreSqlTableOperations.java b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/PostgreSqlTableOperationsIT.java similarity index 85% rename from integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/TestPostgreSqlTableOperations.java rename to integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/PostgreSqlTableOperationsIT.java index 035fad28a4d..5bd75416ad8 100644 --- a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/TestPostgreSqlTableOperations.java +++ b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/PostgreSqlTableOperationsIT.java @@ -12,9 +12,11 @@ import com.datastrato.gravitino.catalog.postgresql.converter.PostgreSqlTypeConverter; import com.datastrato.gravitino.catalog.postgresql.operation.PostgreSqlSchemaOperations; import com.datastrato.gravitino.catalog.postgresql.operation.PostgreSqlTableOperations; +import com.datastrato.gravitino.exceptions.GravitinoRuntimeException; import com.datastrato.gravitino.exceptions.NoSuchTableException; import com.datastrato.gravitino.integration.test.util.GravitinoITUtils; 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.rel.types.Types; @@ -33,7 +35,7 @@ import org.testcontainers.shaded.com.google.common.collect.Maps; @Tag("gravitino-docker-it") -public class TestPostgreSqlTableOperations extends TestPostgreSqlAbstractIT { +public class PostgreSqlTableOperationsIT extends TestPostgreSqlAbstractIT { private static Type VARCHAR = Types.VarCharType.of(255); private static Type INT = Types.IntegerType.get(); @@ -479,4 +481,85 @@ public void testCreateAutoIncrementTable() { Assertions.assertTrue( StringUtils.contains(illegalArgumentException.getMessage(), "Unsupported auto-increment")); } + + @Test + public void testCreateIndexTable() { + String tableName = GravitinoITUtils.genRandomName("index_table_"); + String tableComment = "test_comment"; + List columns = new ArrayList<>(); + columns.add( + new JdbcColumn.Builder() + .withName("col_1") + .withType(Types.LongType.get()) + .withComment("increment key") + .withNullable(false) + .withAutoIncrement(true) + .build()); + columns.add( + new JdbcColumn.Builder() + .withName("col_2") + .withType(INT) + .withNullable(false) + .withComment("id-1") + .build()); + columns.add( + new JdbcColumn.Builder() + .withName("col_3") + .withType(VARCHAR) + .withNullable(false) + .withComment("name") + .build()); + columns.add( + new JdbcColumn.Builder() + .withName("col_4") + .withType(VARCHAR) + .withNullable(false) + .withComment("city") + .build()); + Map properties = new HashMap<>(); + + Index[] indexes = + new Index[] { + Indexes.primary("test_pk", new String[][] {{"col_1"}, {"col_2"}}), + Indexes.unique("u1_key", new String[][] {{"col_2"}, {"col_3"}}), + Indexes.unique("u2_key", new String[][] {{"col_3"}, {"col_4"}}) + }; + + // Test create table index success. + TABLE_OPERATIONS.create( + TEST_DB_NAME, + tableName, + columns.toArray(new JdbcColumn[0]), + tableComment, + properties, + null, + indexes); + + JdbcTable load = TABLE_OPERATIONS.load(TEST_DB_NAME, tableName); + assertionsTableInfo(tableName, tableComment, columns, properties, indexes, load); + + TABLE_OPERATIONS.drop(TEST_DB_NAME, tableName); + + // Test create table index failed. + GravitinoRuntimeException gravitinoRuntimeException = + Assertions.assertThrows( + GravitinoRuntimeException.class, + () -> + TABLE_OPERATIONS.create( + TEST_DB_NAME, + tableName, + columns.toArray(new JdbcColumn[0]), + tableComment, + properties, + null, + new Index[] { + Indexes.primary("no_exist_pk", new String[][] {{"no_exist_1"}}), + Indexes.unique( + "no_exist_key", new String[][] {{"no_exist_2"}, {"no_exist_3"}}) + })); + Assertions.assertTrue( + StringUtils.contains( + gravitinoRuntimeException.getMessage(), + "column \"no_exist_1\" named in key does not exist")); + } } diff --git a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/util/AbstractIT.java b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/util/AbstractIT.java index cc9d088d212..146c20b25f7 100644 --- a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/util/AbstractIT.java +++ b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/util/AbstractIT.java @@ -33,7 +33,6 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Set; import java.util.stream.Collectors; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.ArrayUtils; @@ -247,10 +246,11 @@ protected static void assertionsTableInfo( Assertions.assertEquals( indexByName.get(table.index()[i].name()).type(), table.index()[i].type()); for (int j = 0; j < table.index()[i].fieldNames().length; j++) { - Set colNames = - Arrays.stream(indexByName.get(table.index()[i].name()).fieldNames()[j]) - .collect(Collectors.toSet()); - colNames.containsAll(Arrays.asList(table.index()[i].fieldNames()[j])); + for (int k = 0; k < table.index()[i].fieldNames()[j].length; k++) { + Assertions.assertEquals( + indexByName.get(table.index()[i].name()).fieldNames()[j][k], + table.index()[i].fieldNames()[j][k]); + } } } }