diff --git a/infra/database/type/opengauss/src/main/java/org/apache/shardingsphere/infra/database/opengauss/metadata/data/loader/OpenGaussMetaDataLoader.java b/infra/database/type/opengauss/src/main/java/org/apache/shardingsphere/infra/database/opengauss/metadata/data/loader/OpenGaussMetaDataLoader.java index 56d586b019e37..7bd2e3fbf87f6 100644 --- a/infra/database/type/opengauss/src/main/java/org/apache/shardingsphere/infra/database/opengauss/metadata/data/loader/OpenGaussMetaDataLoader.java +++ b/infra/database/type/opengauss/src/main/java/org/apache/shardingsphere/infra/database/opengauss/metadata/data/loader/OpenGaussMetaDataLoader.java @@ -40,6 +40,7 @@ import java.util.LinkedHashMap; import java.util.LinkedList; import java.util.Map; +import java.util.Optional; import java.util.stream.Collectors; /** @@ -60,6 +61,11 @@ public final class OpenGaussMetaDataLoader implements DialectMetaDataLoader { private static final String BASIC_INDEX_META_DATA_SQL = "SELECT tablename, indexname, schemaname FROM pg_indexes WHERE schemaname IN (%s)"; + private static final String ADVANCE_INDEX_META_DATA_SQL = + "SELECT idx.relname as index_name, insp.nspname as index_schema, tbl.relname as table_name, att.attname AS column_name, pgi.indisunique as is_unique" + + " FROM pg_index pgi JOIN pg_class idx ON idx.oid = pgi.indexrelid JOIN pg_namespace insp ON insp.oid = idx.relnamespace JOIN pg_class tbl ON tbl.oid = pgi.indrelid" + + " JOIN pg_namespace tnsp ON tnsp.oid = tbl.relnamespace JOIN pg_attribute att ON att.attrelid = tbl.oid AND att.attnum = ANY(pgi.indkey) WHERE tnsp.nspname IN (%s)"; + @Override public Collection load(final DataSource dataSource, final Collection tables, final String defaultSchemaName) throws SQLException { try (Connection connection = dataSource.getConnection()) { @@ -87,6 +93,24 @@ private Map> loadIndexMetaDataMap(final indexMetaDataMap.put(tableName, new IndexMetaData(indexName)); } } + try (PreparedStatement preparedStatement = connection.prepareStatement(getAdvanceIndexMetaDataSQL(schemaNames)); ResultSet resultSet = preparedStatement.executeQuery()) { + while (resultSet.next()) { + String schemaName = resultSet.getString("index_schema"); + String tableName = resultSet.getString("table_name"); + String columnName = resultSet.getString("column_name"); + String indexName = resultSet.getString("index_name"); + boolean isUnique = resultSet.getBoolean("is_unique"); + Collection indexMetaDatas = result.getOrDefault(schemaName, LinkedHashMultimap.create()).get(tableName); + if (null == indexMetaDatas || indexMetaDatas.isEmpty()) { + continue; + } + Optional indexMetaData = indexMetaDatas.stream().filter(each -> each.getName().equals(indexName)).findFirst(); + if (indexMetaData.isPresent()) { + indexMetaData.get().setUnique(isUnique); + indexMetaData.get().getColumns().add(columnName); + } + } + } return result; } @@ -94,6 +118,10 @@ private String getIndexMetaDataSQL(final Collection schemaNames) { return String.format(BASIC_INDEX_META_DATA_SQL, schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(","))); } + private String getAdvanceIndexMetaDataSQL(final Collection schemaNames) { + return String.format(ADVANCE_INDEX_META_DATA_SQL, schemaNames.stream().map(each -> String.format("'%s'", each)).collect(Collectors.joining(","))); + } + private Map> loadColumnMetaDataMap(final Connection connection, final Collection tables, final Collection schemaNames) throws SQLException { Map> result = new LinkedHashMap<>(); diff --git a/infra/database/type/opengauss/src/test/java/org/apache/shardingsphere/infra/database/opengauss/metadata/data/loader/OpenGaussMetaDataLoaderTest.java b/infra/database/type/opengauss/src/test/java/org/apache/shardingsphere/infra/database/opengauss/metadata/data/loader/OpenGaussMetaDataLoaderTest.java new file mode 100644 index 0000000000000..9b0ecaf786ba4 --- /dev/null +++ b/infra/database/type/opengauss/src/test/java/org/apache/shardingsphere/infra/database/opengauss/metadata/data/loader/OpenGaussMetaDataLoaderTest.java @@ -0,0 +1,183 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.shardingsphere.infra.database.opengauss.metadata.data.loader; + +import org.apache.shardingsphere.infra.database.core.metadata.data.loader.DialectMetaDataLoader; +import org.apache.shardingsphere.infra.database.core.metadata.data.model.ColumnMetaData; +import org.apache.shardingsphere.infra.database.core.metadata.data.model.IndexMetaData; +import org.apache.shardingsphere.infra.database.core.metadata.data.model.SchemaMetaData; +import org.apache.shardingsphere.infra.database.core.metadata.data.model.TableMetaData; +import org.apache.shardingsphere.infra.database.core.spi.DatabaseTypedSPILoader; +import org.apache.shardingsphere.infra.database.core.type.DatabaseType; +import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader; +import org.junit.jupiter.api.Test; + +import javax.sql.DataSource; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Types; +import java.util.Collection; +import java.util.Collections; +import java.util.Iterator; +import java.util.Optional; + +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.RETURNS_DEEP_STUBS; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +class OpenGaussMetaDataLoaderTest { + + private static final String BASIC_TABLE_META_DATA_SQL = "SELECT table_name, column_name, ordinal_position, data_type, udt_name, column_default, table_schema, is_nullable" + + " FROM information_schema.columns WHERE table_schema IN ('public')"; + + private static final String TABLE_META_DATA_SQL_WITHOUT_TABLES = BASIC_TABLE_META_DATA_SQL + " ORDER BY ordinal_position"; + + private static final String TABLE_META_DATA_SQL_WITH_TABLES = BASIC_TABLE_META_DATA_SQL + " AND table_name IN ('tbl') ORDER BY ordinal_position"; + + private static final String PRIMARY_KEY_META_DATA_SQL = "SELECT tc.table_name, kc.column_name, kc.table_schema FROM information_schema.table_constraints tc" + + " JOIN information_schema.key_column_usage kc ON kc.table_schema = tc.table_schema AND kc.table_name = tc.table_name AND kc.constraint_name = tc.constraint_name" + + " WHERE tc.constraint_type = 'PRIMARY KEY' AND kc.ordinal_position IS NOT NULL AND kc.table_schema IN ('public')"; + + private static final String BASIC_INDEX_META_DATA_SQL = "SELECT tablename, indexname, schemaname FROM pg_indexes WHERE schemaname IN ('public')"; + + private static final String ADVANCE_INDEX_META_DATA_SQL = + "SELECT idx.relname as index_name, insp.nspname as index_schema, tbl.relname as table_name, att.attname AS column_name, pgi.indisunique as is_unique" + + " FROM pg_index pgi JOIN pg_class idx ON idx.oid = pgi.indexrelid JOIN pg_namespace insp ON insp.oid = idx.relnamespace JOIN pg_class tbl ON tbl.oid = pgi.indrelid" + + " JOIN pg_namespace tnsp ON tnsp.oid = tbl.relnamespace JOIN pg_attribute att ON att.attrelid = tbl.oid AND att.attnum = ANY(pgi.indkey) WHERE tnsp.nspname IN ('public')"; + + @Test + void assertLoadWithoutTables() throws SQLException { + DataSource dataSource = mockDataSource(); + ResultSet schemaResultSet = mockSchemaMetaDataResultSet(); + when(dataSource.getConnection().getMetaData().getSchemas()).thenReturn(schemaResultSet); + ResultSet tableResultSet = mockTableMetaDataResultSet(); + when(dataSource.getConnection().prepareStatement(TABLE_META_DATA_SQL_WITHOUT_TABLES).executeQuery()).thenReturn(tableResultSet); + ResultSet primaryKeyResultSet = mockPrimaryKeyMetaDataResultSet(); + when(dataSource.getConnection().prepareStatement(PRIMARY_KEY_META_DATA_SQL).executeQuery()).thenReturn(primaryKeyResultSet); + ResultSet indexResultSet = mockIndexMetaDataResultSet(); + when(dataSource.getConnection().prepareStatement(BASIC_INDEX_META_DATA_SQL).executeQuery()).thenReturn(indexResultSet); + ResultSet advanceIndexResultSet = mockAdvanceIndexMetaDataResultSet(); + when(dataSource.getConnection().prepareStatement(ADVANCE_INDEX_META_DATA_SQL).executeQuery()).thenReturn(advanceIndexResultSet); + assertTableMetaDataMap(getDialectTableMetaDataLoader().load(dataSource, Collections.emptyList(), "sharding_db")); + } + + private ResultSet mockSchemaMetaDataResultSet() throws SQLException { + ResultSet result = mock(ResultSet.class); + when(result.next()).thenReturn(true, false); + when(result.getString("TABLE_SCHEM")).thenReturn("public"); + return result; + } + + @Test + void assertLoadWithTables() throws SQLException { + DataSource dataSource = mockDataSource(); + ResultSet schemaResultSet = mockSchemaMetaDataResultSet(); + when(dataSource.getConnection().getMetaData().getSchemas()).thenReturn(schemaResultSet); + ResultSet tableResultSet = mockTableMetaDataResultSet(); + when(dataSource.getConnection().prepareStatement(TABLE_META_DATA_SQL_WITH_TABLES).executeQuery()).thenReturn(tableResultSet); + ResultSet primaryKeyResultSet = mockPrimaryKeyMetaDataResultSet(); + when(dataSource.getConnection().prepareStatement(PRIMARY_KEY_META_DATA_SQL).executeQuery()).thenReturn(primaryKeyResultSet); + ResultSet indexResultSet = mockIndexMetaDataResultSet(); + when(dataSource.getConnection().prepareStatement(BASIC_INDEX_META_DATA_SQL).executeQuery()).thenReturn(indexResultSet); + ResultSet advanceIndexResultSet = mockAdvanceIndexMetaDataResultSet(); + when(dataSource.getConnection().prepareStatement(ADVANCE_INDEX_META_DATA_SQL).executeQuery()).thenReturn(advanceIndexResultSet); + assertTableMetaDataMap(getDialectTableMetaDataLoader().load(dataSource, Collections.singletonList("tbl"), "sharding_db")); + } + + private DataSource mockDataSource() throws SQLException { + DataSource result = mock(DataSource.class, RETURNS_DEEP_STUBS); + ResultSet typeInfoResultSet = mockTypeInfoResultSet(); + when(result.getConnection().getMetaData().getTypeInfo()).thenReturn(typeInfoResultSet); + return result; + } + + private ResultSet mockTypeInfoResultSet() throws SQLException { + ResultSet result = mock(ResultSet.class); + when(result.next()).thenReturn(true, true, false); + when(result.getString("TYPE_NAME")).thenReturn("int4", "varchar"); + when(result.getInt("DATA_TYPE")).thenReturn(Types.INTEGER, Types.VARCHAR); + return result; + } + + private ResultSet mockTableMetaDataResultSet() throws SQLException { + ResultSet result = mock(ResultSet.class); + when(result.next()).thenReturn(true, true, false); + when(result.getString("table_name")).thenReturn("tbl"); + when(result.getString("column_name")).thenReturn("id", "name"); + when(result.getInt("ordinal_position")).thenReturn(1, 2); + when(result.getString("data_type")).thenReturn("integer", "character varying"); + when(result.getString("udt_name")).thenReturn("int4", "varchar"); + when(result.getString("column_default")).thenReturn("nextval('id_seq'::regclass)", ""); + when(result.getString("table_schema")).thenReturn("public", "public"); + when(result.getString("is_nullable")).thenReturn("NO", "YES"); + return result; + } + + private ResultSet mockPrimaryKeyMetaDataResultSet() throws SQLException { + ResultSet result = mock(ResultSet.class); + when(result.next()).thenReturn(true, false); + when(result.getString("table_name")).thenReturn("tbl"); + when(result.getString("column_name")).thenReturn("id"); + when(result.getString("table_schema")).thenReturn("public"); + return result; + } + + private ResultSet mockIndexMetaDataResultSet() throws SQLException { + ResultSet result = mock(ResultSet.class); + when(result.next()).thenReturn(true, false); + when(result.getString("tablename")).thenReturn("tbl"); + when(result.getString("indexname")).thenReturn("id"); + when(result.getString("schemaname")).thenReturn("public"); + return result; + } + + private ResultSet mockAdvanceIndexMetaDataResultSet() throws SQLException { + ResultSet result = mock(ResultSet.class); + when(result.next()).thenReturn(true, false); + when(result.getString("table_name")).thenReturn("tbl"); + when(result.getString("column_name")).thenReturn("id"); + when(result.getString("index_name")).thenReturn("id"); + when(result.getString("index_schema")).thenReturn("public"); + when(result.getBoolean("is_unique")).thenReturn(true); + return result; + } + + private DialectMetaDataLoader getDialectTableMetaDataLoader() { + Optional result = DatabaseTypedSPILoader.findService(DialectMetaDataLoader.class, TypedSPILoader.getService(DatabaseType.class, "openGauss")); + assertTrue(result.isPresent()); + return result.get(); + } + + private void assertTableMetaDataMap(final Collection schemaMetaDataList) { + assertThat(schemaMetaDataList.size(), is(1)); + TableMetaData actualTableMetaData = schemaMetaDataList.iterator().next().getTables().iterator().next(); + assertThat(actualTableMetaData.getColumns().size(), is(2)); + Iterator columnsIterator = actualTableMetaData.getColumns().iterator(); + assertThat(columnsIterator.next(), is(new ColumnMetaData("id", Types.INTEGER, true, true, true, true, false, false))); + assertThat(columnsIterator.next(), is(new ColumnMetaData("name", Types.VARCHAR, false, false, true, true, false, true))); + assertThat(actualTableMetaData.getIndexes().size(), is(1)); + Iterator indexesIterator = actualTableMetaData.getIndexes().iterator(); + IndexMetaData indexMetaData = new IndexMetaData("id"); + indexMetaData.setUnique(true); + indexMetaData.getColumns().add("id"); + assertThat(indexesIterator.next(), is(indexMetaData)); + } +}