This is an automated email from the ASF dual-hosted git repository.

zhaojinchao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shardingsphere.git


The following commit(s) were added to refs/heads/master by this push:
     new 7e50a41fd58 Support get unique key for og (#28459)
7e50a41fd58 is described below

commit 7e50a41fd58f2db3e2f6868020fba3e5fbca5da4
Author: Chuxin Chen <[email protected]>
AuthorDate: Wed Sep 20 11:36:35 2023 +0800

    Support get unique key for og (#28459)
---
 .../data/loader/OpenGaussMetaDataLoader.java       |  28 ++++
 .../data/loader/OpenGaussMetaDataLoaderTest.java   | 183 +++++++++++++++++++++
 2 files changed, 211 insertions(+)

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 58159925aac..f3501c46473 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.HashSet;
 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<SchemaMetaData> load(final MetaDataLoaderMaterial 
material) throws SQLException {
         try (Connection connection = material.getDataSource().getConnection()) 
{
@@ -87,6 +93,24 @@ public final class OpenGaussMetaDataLoader implements 
DialectMetaDataLoader {
                 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<IndexMetaData> indexMetaDatas = 
result.getOrDefault(schemaName, LinkedHashMultimap.create()).get(tableName);
+                if (null == indexMetaDatas || indexMetaDatas.isEmpty()) {
+                    continue;
+                }
+                Optional<IndexMetaData> 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 @@ public final class OpenGaussMetaDataLoader implements 
DialectMetaDataLoader {
         return String.format(BASIC_INDEX_META_DATA_SQL, 
schemaNames.stream().map(each -> String.format("'%s'", 
each)).collect(Collectors.joining(",")));
     }
     
+    private String getAdvanceIndexMetaDataSQL(final Collection<String> 
schemaNames) {
+        return String.format(ADVANCE_INDEX_META_DATA_SQL, 
schemaNames.stream().map(each -> String.format("'%s'", 
each)).collect(Collectors.joining(",")));
+    }
+    
     private Map<String, Multimap<String, ColumnMetaData>> 
loadColumnMetaDataMap(final Connection connection, final Collection<String> 
tables,
                                                                                
 final Collection<String> schemaNames) throws SQLException {
         Map<String, Multimap<String, ColumnMetaData>> 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 00000000000..9b0ecaf786b
--- /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<DialectMetaDataLoader> result = 
DatabaseTypedSPILoader.findService(DialectMetaDataLoader.class, 
TypedSPILoader.getService(DatabaseType.class, "openGauss"));
+        assertTrue(result.isPresent());
+        return result.get();
+    }
+    
+    private void assertTableMetaDataMap(final Collection<SchemaMetaData> 
schemaMetaDataList) {
+        assertThat(schemaMetaDataList.size(), is(1));
+        TableMetaData actualTableMetaData = 
schemaMetaDataList.iterator().next().getTables().iterator().next();
+        assertThat(actualTableMetaData.getColumns().size(), is(2));
+        Iterator<ColumnMetaData> 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<IndexMetaData> indexesIterator = 
actualTableMetaData.getIndexes().iterator();
+        IndexMetaData indexMetaData = new IndexMetaData("id");
+        indexMetaData.setUnique(true);
+        indexMetaData.getColumns().add("id");
+        assertThat(indexesIterator.next(), is(indexMetaData));
+    }
+}

Reply via email to