From bbd52f3c44a76d9078dca2deb259d13e65d6ea9e Mon Sep 17 00:00:00 2001 From: zhangliang Date: Sat, 25 Jul 2026 20:16:40 +0800 Subject: [PATCH 1/8] Fix PostgreSQL composite column type OID resolution --- .../query/PostgreSQLColumnDescription.java | 2 + .../query/PostgreSQLColumnTypeOIDLoader.java | 119 +++++++++++++ .../command/query/extended/Portal.java | 18 +- .../PostgreSQLComDescribeExecutor.java | 98 +++++++---- .../simple/PostgreSQLComQueryExecutor.java | 22 ++- .../PostgreSQLColumnTypeOIDLoaderTest.java | 166 ++++++++++++++++++ .../command/query/extended/PortalTest.java | 24 ++- .../PostgreSQLComDescribeExecutorTest.java | 41 ++++- .../PostgreSQLComQueryExecutorTest.java | 17 +- 9 files changed, 455 insertions(+), 52 deletions(-) create mode 100644 proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java create mode 100644 proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java diff --git a/database/protocol/dialect/postgresql/src/main/java/org/apache/shardingsphere/database/protocol/postgresql/packet/command/query/PostgreSQLColumnDescription.java b/database/protocol/dialect/postgresql/src/main/java/org/apache/shardingsphere/database/protocol/postgresql/packet/command/query/PostgreSQLColumnDescription.java index 2c6d16729a062..10597e2ada5a0 100644 --- a/database/protocol/dialect/postgresql/src/main/java/org/apache/shardingsphere/database/protocol/postgresql/packet/command/query/PostgreSQLColumnDescription.java +++ b/database/protocol/dialect/postgresql/src/main/java/org/apache/shardingsphere/database/protocol/postgresql/packet/command/query/PostgreSQLColumnDescription.java @@ -18,6 +18,7 @@ package org.apache.shardingsphere.database.protocol.postgresql.packet.command.query; import lombok.Getter; +import lombok.RequiredArgsConstructor; import org.apache.shardingsphere.database.protocol.postgresql.constant.PostgreSQLArrayColumnType; import org.apache.shardingsphere.database.protocol.postgresql.constant.PostgreSQLValueFormat; import org.apache.shardingsphere.database.protocol.postgresql.packet.command.query.extended.PostgreSQLBinaryColumnType; @@ -27,6 +28,7 @@ /** * Column description for PostgreSQL. */ +@RequiredArgsConstructor @Getter public final class PostgreSQLColumnDescription { diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java new file mode 100644 index 0000000000000..28de32e1d2f96 --- /dev/null +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java @@ -0,0 +1,119 @@ +/* + * 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.proxy.frontend.postgresql.command.query; + +import lombok.AccessLevel; +import lombok.NoArgsConstructor; +import org.apache.shardingsphere.infra.connection.kernel.KernelProcessor; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; +import org.apache.shardingsphere.infra.executor.sql.execute.engine.ConnectionMode; +import org.apache.shardingsphere.infra.session.query.QueryContext; +import org.apache.shardingsphere.proxy.backend.response.header.query.QueryHeader; +import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; +import org.postgresql.core.BaseConnection; +import org.postgresql.core.Oid; + +import java.sql.Connection; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Types; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Loader for PostgreSQL column type OIDs. + */ +@NoArgsConstructor(access = AccessLevel.PRIVATE) +public final class PostgreSQLColumnTypeOIDLoader { + + /** + * Load composite column type OIDs from the routed backend connection. + * + * @param connectionSession connection session + * @param queryContext query context + * @param queryHeaders query headers + * @return column indexes to type OIDs, or an empty map if no composite column type can be resolved + * @throws SQLException SQL exception + */ + public static Map load(final ConnectionSession connectionSession, final QueryContext queryContext, final List queryHeaders) throws SQLException { + if (queryHeaders.stream().noneMatch(each -> Types.STRUCT == each.getColumnType())) { + return Collections.emptyMap(); + } + ExecutionContext executionContext = new KernelProcessor().generateExecutionContext(queryContext, queryContext.getMetaData().getGlobalRuleMetaData(), queryContext.getMetaData().getProps()); + if (1 != executionContext.getExecutionUnits().size()) { + return Collections.emptyMap(); + } + ExecutionUnit executionUnit = executionContext.getExecutionUnits().iterator().next(); + List connections = connectionSession.getDatabaseConnectionManager().getConnections( + connectionSession.getUsedDatabaseName(), executionUnit.getDataSourceName(), 0, 1, ConnectionMode.CONNECTION_STRICTLY); + return load(connections.get(0), queryHeaders); + } + + /** + * Load composite column type OIDs from result set metadata. + * + * @param connection database connection + * @param metaData result set metadata + * @return column indexes to type OIDs, or an empty map if no composite column type can be resolved + * @throws SQLException SQL exception + */ + public static Map load(final Connection connection, final ResultSetMetaData metaData) throws SQLException { + if (!connection.isWrapperFor(BaseConnection.class)) { + return Collections.emptyMap(); + } + return getTypeOIDs(connection.unwrap(BaseConnection.class), metaData); + } + + private static Map load(final Connection connection, final List queryHeaders) throws SQLException { + if (!connection.isWrapperFor(BaseConnection.class)) { + return Collections.emptyMap(); + } + return getTypeOIDs(connection.unwrap(BaseConnection.class), queryHeaders); + } + + private static Map getTypeOIDs(final BaseConnection connection, final ResultSetMetaData metaData) throws SQLException { + int columnCount = metaData.getColumnCount(); + Map result = new HashMap<>(); + for (int columnIndex = 1; columnIndex <= columnCount; columnIndex++) { + if (Types.STRUCT == metaData.getColumnType(columnIndex)) { + int typeOID = connection.getTypeInfo().getPGType(metaData.getColumnTypeName(columnIndex)); + if (Oid.UNSPECIFIED != typeOID) { + result.put(columnIndex, typeOID); + } + } + } + return result; + } + + private static Map getTypeOIDs(final BaseConnection connection, final List queryHeaders) throws SQLException { + Map result = new HashMap<>(); + for (int i = 0; i < queryHeaders.size(); i++) { + QueryHeader each = queryHeaders.get(i); + if (Types.STRUCT == each.getColumnType()) { + int typeOID = connection.getTypeInfo().getPGType(each.getColumnTypeName()); + if (Oid.UNSPECIFIED != typeOID) { + result.put(i + 1, typeOID); + } + } + } + return result; + } +} diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java index 504c97bb980af..654444c152fc9 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java @@ -49,6 +49,7 @@ import org.apache.shardingsphere.proxy.backend.response.header.query.QueryHeader; import org.apache.shardingsphere.proxy.backend.response.header.query.QueryResponseHeader; import org.apache.shardingsphere.proxy.backend.response.header.update.UpdateResponseHeader; +import org.apache.shardingsphere.proxy.frontend.postgresql.command.query.PostgreSQLColumnTypeOIDLoader; import org.apache.shardingsphere.proxy.frontend.postgresql.command.query.PostgreSQLCommand; import org.apache.shardingsphere.sql.parser.statement.core.segment.dal.VariableAssignSegment; import org.apache.shardingsphere.sql.parser.statement.core.statement.SQLStatement; @@ -61,6 +62,7 @@ import java.util.Collection; import java.util.LinkedList; import java.util.List; +import java.util.Map; /** * PostgreSQL portal. @@ -81,8 +83,12 @@ public final class Portal { private final ProxyBackendHandler proxyBackendHandler; + private final QueryContext queryContext; + private final ProxyDatabaseConnectionManager databaseConnectionManager; + private Map columnTypeOIDs; + private ResponseHeader responseHeader; public Portal(final String name, final PostgreSQLServerPreparedStatement preparedStatement, final List params, final List resultFormats, @@ -97,7 +103,7 @@ public Portal(final String name, final PostgreSQLServerPreparedStatement prepare ((ParameterAware) sqlStatementContext).bindParameters(params); } DatabaseType protocolType = ProxyContext.getInstance().getContextManager().getDatabase(databaseName).getProtocolType(); - QueryContext queryContext = new QueryContext(sqlStatementContext, preparedStatement.getSql(), params, preparedStatement.getHintValueContext(), + queryContext = new QueryContext(sqlStatementContext, preparedStatement.getSql(), params, preparedStatement.getHintValueContext(), databaseConnectionManager.getConnectionSession().getConnectionContext(), ProxyContext.getInstance().getContextManager().getMetaDataContexts().getMetaData(), true); proxyBackendHandler = ProxyBackendHandlerFactory.newInstance(protocolType, queryContext, databaseConnectionManager.getConnectionSession(), true); } @@ -109,6 +115,10 @@ public Portal(final String name, final PostgreSQLServerPreparedStatement prepare */ public void bind() throws SQLException { responseHeader = proxyBackendHandler.execute(); + if (responseHeader instanceof QueryResponseHeader) { + QueryResponseHeader queryResponseHeader = (QueryResponseHeader) responseHeader; + columnTypeOIDs = PostgreSQLColumnTypeOIDLoader.load(databaseConnectionManager.getConnectionSession(), queryContext, queryResponseHeader.getQueryHeaders()); + } } /** @@ -136,7 +146,11 @@ private Collection createColumnDescriptions(final Q int columnIndex = 0; for (QueryHeader each : queryResponseHeader.getQueryHeaders()) { PostgreSQLValueFormat valueFormat = determineValueFormat(columnIndex); - result.add(new PostgreSQLColumnDescription(each.getColumnLabel(), ++columnIndex, each.getColumnType(), each.getColumnLength(), each.getColumnTypeName(), valueFormat.getCode())); + int currentColumnIndex = ++columnIndex; + Integer typeOID = columnTypeOIDs.get(currentColumnIndex); + result.add(PostgreSQLValueFormat.TEXT == valueFormat && null != typeOID + ? new PostgreSQLColumnDescription(each.getColumnLabel(), currentColumnIndex, typeOID, each.getColumnLength(), valueFormat.getCode()) + : new PostgreSQLColumnDescription(each.getColumnLabel(), currentColumnIndex, each.getColumnType(), each.getColumnLength(), each.getColumnTypeName(), valueFormat.getCode())); } return result; } diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutor.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutor.java index 986c476a9d650..c075108bf1546 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutor.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutor.java @@ -21,7 +21,7 @@ import org.apache.shardingsphere.database.connector.core.type.DatabaseTypeRegistry; import org.apache.shardingsphere.database.exception.core.exception.syntax.column.ColumnNotFoundException; import org.apache.shardingsphere.database.protocol.packet.DatabasePacket; -import org.apache.shardingsphere.database.protocol.postgresql.packet.PostgreSQLPacket; +import org.apache.shardingsphere.database.protocol.postgresql.constant.PostgreSQLValueFormat; import org.apache.shardingsphere.database.protocol.postgresql.packet.command.query.PostgreSQLColumnDescription; import org.apache.shardingsphere.database.protocol.postgresql.packet.command.query.PostgreSQLNoDataPacket; import org.apache.shardingsphere.database.protocol.postgresql.packet.command.query.PostgreSQLRowDescriptionPacket; @@ -37,6 +37,7 @@ import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; import org.apache.shardingsphere.proxy.frontend.command.executor.CommandExecutor; import org.apache.shardingsphere.proxy.frontend.postgresql.command.PortalContext; +import org.apache.shardingsphere.proxy.frontend.postgresql.command.query.PostgreSQLColumnTypeOIDLoader; import org.apache.shardingsphere.proxy.frontend.postgresql.command.query.extended.PostgreSQLPreparedStatementMetadataFactory; import org.apache.shardingsphere.proxy.frontend.postgresql.command.query.extended.PostgreSQLPreparedStatementParameterTypeResolver; import org.apache.shardingsphere.proxy.frontend.postgresql.command.query.extended.PostgreSQLServerPreparedStatement; @@ -63,6 +64,7 @@ import java.util.LinkedList; import java.util.List; import java.util.ListIterator; +import java.util.Map; import java.util.Optional; import java.util.stream.Collectors; @@ -111,7 +113,7 @@ private void tryDescribePreparedStatement(final PostgreSQLServerPreparedStatemen } } - private void describeInsertStatementByDatabaseMetaData(final PostgreSQLServerPreparedStatement preparedStatement) { + private void describeInsertStatementByDatabaseMetaData(final PostgreSQLServerPreparedStatement preparedStatement) throws SQLException { InsertStatement insertStatement = (InsertStatement) preparedStatement.getSqlStatementContext().getSqlStatement(); Collection unspecifiedTypeParameterIndexes = getUnspecifiedTypeParameterIndexes(preparedStatement); Optional returningSegment = insertStatement.getReturning(); @@ -121,7 +123,6 @@ private void describeInsertStatementByDatabaseMetaData(final PostgreSQLServerPre String logicTableName = insertStatement.getTable().map(optional -> optional.getTableName().getIdentifier().getValue()).orElse(""); ShardingSphereTable table = getTableFromMetaData(connectionSession.getUsedDatabaseName(), insertStatement, logicTableName); List columnNamesOfInsert = getColumnNamesOfInsertStatement(insertStatement, table); - preparedStatement.setRowDescription(returningSegment.map(returning -> describeReturning(returning, table)).orElseGet(PostgreSQLNoDataPacket::getInstance)); int parameterMarkerIndex = 0; for (InsertValuesSegment each : insertStatement.getValues()) { for (int i = 0; i < each.getValues().size(); i++) { @@ -138,38 +139,21 @@ private void describeInsertStatementByDatabaseMetaData(final PostgreSQLServerPre preparedStatement.getParameterTypes().set(parameterMarkerIndex++, PostgreSQLBinaryColumnType.valueOfJDBCType(table.getColumn(columnName).getDataType())); } } - } - - private Collection getUnspecifiedTypeParameterIndexes(final PostgreSQLServerPreparedStatement preparedStatement) { - Collection result = new HashSet<>(); - ListIterator parameterTypesListIterator = preparedStatement.getParameterTypes().listIterator(); - for (int index = parameterTypesListIterator.nextIndex(); parameterTypesListIterator.hasNext(); index = parameterTypesListIterator.nextIndex()) { - if (PostgreSQLBinaryColumnType.UNSPECIFIED == parameterTypesListIterator.next()) { - result.add(index); - } + if (returningSegment.isPresent()) { + Collection columnDescriptions = createReturningColumnDescriptions(returningSegment.get(), table); + Map columnTypeOIDs = loadColumnTypeOIDs(preparedStatement, columnDescriptions.size()); + preparedStatement.setRowDescription(new PostgreSQLRowDescriptionPacket(applyColumnTypeOIDs(columnDescriptions, columnTypeOIDs))); + } else { + preparedStatement.setRowDescription(PostgreSQLNoDataPacket.getInstance()); } - return result; } - private ShardingSphereTable getTableFromMetaData(final String databaseName, final InsertStatement insertStatement, final String logicTableName) { - ShardingSphereDatabase database = ProxyContext.getInstance().getContextManager().getDatabase(databaseName); - String schemaName = insertStatement.getTable().flatMap(SimpleTableSegment::getOwner).map(optional -> optional.getIdentifier() - .getValue()).orElseGet(() -> new DatabaseTypeRegistry(database.getProtocolType()).getDefaultSchemaName(databaseName)); - return database.getSchema(schemaName).getTable(logicTableName); - } - - private List getColumnNamesOfInsertStatement(final InsertStatement insertStatement, final ShardingSphereTable table) { - return insertStatement.getColumns().isEmpty() - ? table.getColumnNames() - : insertStatement.getColumns().stream().map(each -> new ShardingSphereIdentifier(each.getIdentifier().getValue())).collect(Collectors.toList()); - } - - private PostgreSQLRowDescriptionPacket describeReturning(final ReturningSegment returningSegment, final ShardingSphereTable table) { + private Collection createReturningColumnDescriptions(final ReturningSegment returningSegment, final ShardingSphereTable table) { Collection result = new LinkedList<>(); for (ProjectionSegment each : returningSegment.getProjections().getProjections()) { if (each instanceof ShorthandProjectionSegment) { - table.getAllColumns().stream() - .map(column -> new PostgreSQLColumnDescription(column.getName(), 0, column.getDataType(), estimateColumnLength(column.getDataType()), "")).forEach(result::add); + table.getAllColumns().stream().map(column -> new PostgreSQLColumnDescription(column.getName(), 0, column.getDataType(), estimateColumnLength(column.getDataType()), "")) + .forEach(result::add); } if (each instanceof ColumnProjectionSegment) { ColumnProjectionSegment segment = (ColumnProjectionSegment) each; @@ -182,7 +166,7 @@ private PostgreSQLRowDescriptionPacket describeReturning(final ReturningSegment result.add(convertExpressionToDescription((ExpressionProjectionSegment) each)); } } - return new PostgreSQLRowDescriptionPacket(result); + return result; } private ShardingSphereColumn generateDefaultColumn(final ColumnProjectionSegment segment) { @@ -223,6 +207,39 @@ private int estimateColumnLength(final int jdbcType) { } } + private Map loadColumnTypeOIDs(final PostgreSQLServerPreparedStatement preparedStatement, final int expectedColumnCount) throws SQLException { + try (PreparedStatement actualPreparedStatement = PostgreSQLPreparedStatementMetadataFactory.load(connectionSession, preparedStatement, Collections.emptyList())) { + ResultSetMetaData metaData = actualPreparedStatement.getMetaData(); + return null == metaData || expectedColumnCount != metaData.getColumnCount() + ? Collections.emptyMap() + : PostgreSQLColumnTypeOIDLoader.load(actualPreparedStatement.getConnection(), metaData); + } + } + + private Collection getUnspecifiedTypeParameterIndexes(final PostgreSQLServerPreparedStatement preparedStatement) { + Collection result = new HashSet<>(); + ListIterator parameterTypesListIterator = preparedStatement.getParameterTypes().listIterator(); + for (int index = parameterTypesListIterator.nextIndex(); parameterTypesListIterator.hasNext(); index = parameterTypesListIterator.nextIndex()) { + if (PostgreSQLBinaryColumnType.UNSPECIFIED == parameterTypesListIterator.next()) { + result.add(index); + } + } + return result; + } + + private ShardingSphereTable getTableFromMetaData(final String databaseName, final InsertStatement insertStatement, final String logicTableName) { + ShardingSphereDatabase database = ProxyContext.getInstance().getContextManager().getDatabase(databaseName); + String schemaName = insertStatement.getTable().flatMap(SimpleTableSegment::getOwner).map(optional -> optional.getIdentifier() + .getValue()).orElseGet(() -> new DatabaseTypeRegistry(database.getProtocolType()).getDefaultSchemaName(databaseName)); + return database.getSchema(schemaName).getTable(logicTableName); + } + + private List getColumnNamesOfInsertStatement(final InsertStatement insertStatement, final ShardingSphereTable table) { + return insertStatement.getColumns().isEmpty() + ? table.getColumnNames() + : insertStatement.getColumns().stream().map(each -> new ShardingSphereIdentifier(each.getIdentifier().getValue())).collect(Collectors.toList()); + } + private void tryDescribePreparedStatementByJDBC(final PostgreSQLServerPreparedStatement logicPreparedStatement) throws SQLException { try (PreparedStatement actualPreparedStatement = PostgreSQLPreparedStatementMetadataFactory.load(connectionSession, logicPreparedStatement, Collections.emptyList())) { PostgreSQLPreparedStatementParameterTypeResolver.resolveParameterTypes(logicPreparedStatement, actualPreparedStatement); @@ -239,14 +256,33 @@ private void populateColumnTypes(final PostgreSQLServerPreparedStatement logicPr logicPreparedStatement.setRowDescription(PostgreSQLNoDataPacket.getInstance()); return; } + Map columnTypeOIDs = PostgreSQLColumnTypeOIDLoader.load(actualPreparedStatement.getConnection(), resultSetMetaData); List columnDescriptions = new ArrayList<>(resultSetMetaData.getColumnCount()); for (int columnIndex = 1; columnIndex <= resultSetMetaData.getColumnCount(); columnIndex++) { String columnName = resultSetMetaData.getColumnName(columnIndex); int columnType = resultSetMetaData.getColumnType(columnIndex); int columnLength = resultSetMetaData.getColumnDisplaySize(columnIndex); String columnTypeName = resultSetMetaData.getColumnTypeName(columnIndex); - columnDescriptions.add(new PostgreSQLColumnDescription(columnName, columnIndex, columnType, columnLength, columnTypeName)); + Integer typeOID = columnTypeOIDs.get(columnIndex); + columnDescriptions.add(null != typeOID + ? new PostgreSQLColumnDescription(columnName, columnIndex, typeOID, columnLength, PostgreSQLValueFormat.TEXT.getCode()) + : new PostgreSQLColumnDescription(columnName, columnIndex, columnType, columnLength, columnTypeName)); } logicPreparedStatement.setRowDescription(new PostgreSQLRowDescriptionPacket(columnDescriptions)); } + + private Collection applyColumnTypeOIDs(final Collection columnDescriptions, final Map columnTypeOIDs) { + if (columnTypeOIDs.isEmpty()) { + return columnDescriptions; + } + Collection result = new ArrayList<>(columnDescriptions.size()); + int columnIndex = 1; + for (PostgreSQLColumnDescription each : columnDescriptions) { + Integer typeOID = columnTypeOIDs.get(columnIndex++); + result.add(null != typeOID + ? new PostgreSQLColumnDescription(each.getColumnName(), each.getColumnIndex(), typeOID, each.getColumnLength(), each.getDataFormat()) + : each); + } + return result; + } } diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutor.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutor.java index 90eb68e68f154..b344a6396983c 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutor.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutor.java @@ -21,6 +21,7 @@ import org.apache.shardingsphere.database.connector.core.metadata.database.enums.QuoteCharacter; import org.apache.shardingsphere.database.connector.core.type.DatabaseType; import org.apache.shardingsphere.database.protocol.packet.DatabasePacket; +import org.apache.shardingsphere.database.protocol.postgresql.constant.PostgreSQLValueFormat; import org.apache.shardingsphere.database.protocol.postgresql.packet.PostgreSQLPacket; import org.apache.shardingsphere.database.protocol.postgresql.packet.command.query.PostgreSQLColumnDescription; import org.apache.shardingsphere.database.protocol.postgresql.packet.command.query.PostgreSQLDataRowPacket; @@ -41,6 +42,7 @@ import org.apache.shardingsphere.proxy.frontend.command.executor.QueryCommandExecutor; import org.apache.shardingsphere.proxy.frontend.command.executor.ResponseType; import org.apache.shardingsphere.proxy.frontend.postgresql.command.PortalContext; +import org.apache.shardingsphere.proxy.frontend.postgresql.command.query.PostgreSQLColumnTypeOIDLoader; import org.apache.shardingsphere.proxy.frontend.postgresql.command.query.PostgreSQLCommand; import org.apache.shardingsphere.sql.parser.statement.core.segment.dal.VariableAssignSegment; import org.apache.shardingsphere.sql.parser.statement.core.statement.SQLStatement; @@ -55,6 +57,7 @@ import java.util.Collection; import java.util.Collections; import java.util.LinkedList; +import java.util.Map; /** * Command query executor for PostgreSQL. @@ -67,6 +70,8 @@ public final class PostgreSQLComQueryExecutor implements QueryCommandExecutor { private final PortalContext portalContext; + private final ConnectionSession connectionSession; + private final ProxyBackendHandler proxyBackendHandler; @Getter @@ -74,6 +79,7 @@ public final class PostgreSQLComQueryExecutor implements QueryCommandExecutor { public PostgreSQLComQueryExecutor(final PortalContext portalContext, final PostgreSQLComQueryPacket packet, final ConnectionSession connectionSession) throws SQLException { this.portalContext = portalContext; + this.connectionSession = connectionSession; DatabaseType databaseType = TypedSPILoader.getService(DatabaseType.class, "PostgreSQL"); SQLStatement sqlStatement = ProxySQLComQueryParser.parse(packet.getSQL(), databaseType, connectionSession); proxyBackendHandler = ProxyBackendHandlerFactory.newInstance(databaseType, packet.getSQL(), sqlStatement, connectionSession, packet.getHintValueContext()); @@ -83,22 +89,28 @@ public PostgreSQLComQueryExecutor(final PortalContext portalContext, final Postg public Collection execute() throws SQLException { ResponseHeader responseHeader = proxyBackendHandler.execute(); if (responseHeader instanceof QueryResponseHeader) { - return Collections.singleton(createRowDescriptionPacket((QueryResponseHeader) responseHeader)); + QueryResponseHeader queryResponseHeader = (QueryResponseHeader) responseHeader; + Map columnTypeOIDs = PostgreSQLColumnTypeOIDLoader.load(connectionSession, connectionSession.getQueryContext(), queryResponseHeader.getQueryHeaders()); + return Collections.singleton(createRowDescriptionPacket(queryResponseHeader, columnTypeOIDs)); } responseType = ResponseType.UPDATE; return createUpdatePacket((UpdateResponseHeader) responseHeader); } - private PostgreSQLRowDescriptionPacket createRowDescriptionPacket(final QueryResponseHeader queryResponseHeader) { + private PostgreSQLRowDescriptionPacket createRowDescriptionPacket(final QueryResponseHeader queryResponseHeader, final Map columnTypeOIDs) { responseType = ResponseType.QUERY; - return new PostgreSQLRowDescriptionPacket(createColumnDescriptions(queryResponseHeader)); + return new PostgreSQLRowDescriptionPacket(createColumnDescriptions(queryResponseHeader, columnTypeOIDs)); } - private Collection createColumnDescriptions(final QueryResponseHeader queryResponseHeader) { + private Collection createColumnDescriptions(final QueryResponseHeader queryResponseHeader, final Map columnTypeOIDs) { Collection result = new LinkedList<>(); int columnIndex = 0; for (QueryHeader each : queryResponseHeader.getQueryHeaders()) { - result.add(new PostgreSQLColumnDescription(each.getColumnLabel(), ++columnIndex, each.getColumnType(), each.getColumnLength(), each.getColumnTypeName())); + int currentColumnIndex = ++columnIndex; + Integer typeOID = columnTypeOIDs.get(currentColumnIndex); + result.add(null == typeOID + ? new PostgreSQLColumnDescription(each.getColumnLabel(), currentColumnIndex, each.getColumnType(), each.getColumnLength(), each.getColumnTypeName()) + : new PostgreSQLColumnDescription(each.getColumnLabel(), currentColumnIndex, typeOID, each.getColumnLength(), PostgreSQLValueFormat.TEXT.getCode())); } return result; } diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java new file mode 100644 index 0000000000000..ab8306463f287 --- /dev/null +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java @@ -0,0 +1,166 @@ +/* + * 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.proxy.frontend.postgresql.command.query; + +import org.apache.shardingsphere.infra.connection.kernel.KernelProcessor; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; +import org.apache.shardingsphere.infra.executor.sql.context.SQLUnit; +import org.apache.shardingsphere.infra.session.query.QueryContext; +import org.apache.shardingsphere.proxy.backend.connector.ProxyDatabaseConnectionManager; +import org.apache.shardingsphere.proxy.backend.response.header.query.QueryHeader; +import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; +import org.apache.shardingsphere.test.infra.framework.extension.mock.AutoMockExtension; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.MockedConstruction; +import org.postgresql.core.BaseConnection; +import org.postgresql.core.Oid; +import org.postgresql.core.TypeInfo; + +import java.sql.Connection; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Types; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.is; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockConstruction; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.RETURNS_DEEP_STUBS; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(AutoMockExtension.class) +class PostgreSQLColumnTypeOIDLoaderTest { + + @Test + void assertLoad() throws SQLException { + QueryContext queryContext = mock(QueryContext.class, RETURNS_DEEP_STUBS); + ExecutionContext executionContext = createExecutionContext( + Collections.singletonList(new ExecutionUnit("ds_0", new SQLUnit("SELECT record_value", Collections.emptyList())))); + BaseConnection connection = mock(BaseConnection.class); + mockTypeOID(connection, "record_type", 2249); + ConnectionSession connectionSession = createConnectionSession(connection); + try ( + MockedConstruction ignored = mockConstruction(KernelProcessor.class, + (mock, context) -> when(mock.generateExecutionContext(any(), any(), any())).thenReturn(executionContext))) { + Map actual = PostgreSQLColumnTypeOIDLoader.load( + connectionSession, queryContext, Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type"))); + assertThat(actual.get(1), is(2249)); + } + } + + @Test + void assertLoadWithoutCompositeColumn() throws SQLException { + assertTrue(PostgreSQLColumnTypeOIDLoader.load( + mock(ConnectionSession.class), mock(QueryContext.class), Collections.singletonList(createQueryHeader(Types.VARCHAR, "varchar"))).isEmpty()); + } + + @Test + void assertLoadWithMultipleExecutionUnits() throws SQLException { + ExecutionContext executionContext = createExecutionContext(Arrays.asList( + new ExecutionUnit("ds_0", new SQLUnit("SELECT record_value", Collections.emptyList())), + new ExecutionUnit("ds_1", new SQLUnit("SELECT record_value", Collections.emptyList())))); + try ( + MockedConstruction ignored = mockConstruction(KernelProcessor.class, + (mock, context) -> when(mock.generateExecutionContext(any(), any(), any())).thenReturn(executionContext))) { + assertTrue(PostgreSQLColumnTypeOIDLoader.load( + mock(ConnectionSession.class), mock(QueryContext.class, RETURNS_DEEP_STUBS), + Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type"))).isEmpty()); + } + } + + @Test + void assertLoadWithoutPgjdbcConnection() throws SQLException { + ExecutionContext executionContext = createExecutionContext( + Collections.singletonList(new ExecutionUnit("ds_0", new SQLUnit("SELECT record_value", Collections.emptyList())))); + Connection connection = mock(Connection.class); + ConnectionSession connectionSession = createConnectionSession(connection); + try ( + MockedConstruction ignored = mockConstruction(KernelProcessor.class, + (mock, context) -> when(mock.generateExecutionContext(any(), any(), any())).thenReturn(executionContext))) { + assertTrue(PostgreSQLColumnTypeOIDLoader.load( + connectionSession, mock(QueryContext.class, RETURNS_DEEP_STUBS), + Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type"))).isEmpty()); + verify(connection, never()).unwrap(BaseConnection.class); + } + } + + @Test + void assertLoadFromResultSetMetaData() throws SQLException { + BaseConnection connection = mock(BaseConnection.class); + mockTypeOID(connection, "record_type", 2249); + ResultSetMetaData metaData = mock(ResultSetMetaData.class); + when(metaData.getColumnCount()).thenReturn(2); + when(metaData.getColumnType(1)).thenReturn(Types.STRUCT); + when(metaData.getColumnType(2)).thenReturn(Types.VARCHAR); + when(metaData.getColumnTypeName(1)).thenReturn("record_type"); + Map actual = PostgreSQLColumnTypeOIDLoader.load(connection, metaData); + assertThat(actual.size(), is(1)); + assertThat(actual.get(1), is(2249)); + verify(metaData, never()).getColumnTypeName(2); + } + + @Test + void assertLoadUnspecifiedTypeOID() throws SQLException { + BaseConnection connection = mock(BaseConnection.class); + mockTypeOID(connection, "unknown_type", Oid.UNSPECIFIED); + ResultSetMetaData metaData = mock(ResultSetMetaData.class); + when(metaData.getColumnCount()).thenReturn(1); + when(metaData.getColumnType(1)).thenReturn(Types.STRUCT); + when(metaData.getColumnTypeName(1)).thenReturn("unknown_type"); + assertTrue(PostgreSQLColumnTypeOIDLoader.load(connection, metaData).isEmpty()); + } + + private QueryHeader createQueryHeader(final int columnType, final String columnTypeName) { + return new QueryHeader("", "", "record_value", "record_value", columnType, columnTypeName, -1, 0, false, false, false, false); + } + + private ExecutionContext createExecutionContext(final List executionUnits) { + ExecutionContext result = mock(ExecutionContext.class); + when(result.getExecutionUnits()).thenReturn(executionUnits); + return result; + } + + private ConnectionSession createConnectionSession(final Connection connection) throws SQLException { + ConnectionSession result = mock(ConnectionSession.class); + ProxyDatabaseConnectionManager connectionManager = mock(ProxyDatabaseConnectionManager.class); + when(connectionManager.getConnections(any(), anyString(), anyInt(), anyInt(), any())).thenReturn(Collections.singletonList(connection)); + when(result.getDatabaseConnectionManager()).thenReturn(connectionManager); + when(result.getUsedDatabaseName()).thenReturn("postgres"); + return result; + } + + private void mockTypeOID(final BaseConnection connection, final String columnTypeName, final int typeOID) throws SQLException { + TypeInfo typeInfo = mock(TypeInfo.class); + when(typeInfo.getPGType(columnTypeName)).thenReturn(typeOID); + when(connection.getTypeInfo()).thenReturn(typeInfo); + when(connection.isWrapperFor(BaseConnection.class)).thenReturn(true); + when(connection.unwrap(BaseConnection.class)).thenReturn(connection); + } +} diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PortalTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PortalTest.java index 1aa91f84c6ac1..1d4acacb35d79 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PortalTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PortalTest.java @@ -52,6 +52,7 @@ import org.apache.shardingsphere.proxy.backend.response.header.query.QueryResponseHeader; import org.apache.shardingsphere.proxy.backend.response.header.update.UpdateResponseHeader; import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; +import org.apache.shardingsphere.proxy.frontend.postgresql.command.query.PostgreSQLColumnTypeOIDLoader; import org.apache.shardingsphere.sql.parser.statement.core.segment.dal.VariableAssignSegment; import org.apache.shardingsphere.sql.parser.statement.core.segment.dal.VariableSegment; import org.apache.shardingsphere.sql.parser.statement.core.statement.SQLStatement; @@ -77,8 +78,10 @@ import java.util.Arrays; import java.util.Collection; import java.util.Collections; +import java.util.HashMap; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.Optional; import static org.hamcrest.Matchers.is; @@ -87,6 +90,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyList; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.RETURNS_DEEP_STUBS; @@ -95,7 +99,7 @@ import static org.mockito.Mockito.when; @ExtendWith(AutoMockExtension.class) -@StaticMockSettings({ProxyContext.class, ProxyBackendHandlerFactory.class}) +@StaticMockSettings({ProxyContext.class, ProxyBackendHandlerFactory.class, PostgreSQLColumnTypeOIDLoader.class}) @MockitoSettings(strictness = Strictness.LENIENT) class PortalTest { @@ -123,6 +127,7 @@ void setup() throws SQLException { when(ProxyBackendHandlerFactory.newInstance(eq(databaseType), anyString(), any(SQLStatement.class), eq(connectionSession), any(HintValueContext.class))).thenReturn(proxyBackendHandler); when(ProxyBackendHandlerFactory.newInstance(eq(databaseType), any(QueryContext.class), eq(connectionSession), anyBoolean())).thenReturn(proxyBackendHandler); when(databaseConnectionManager.getConnectionSession()).thenReturn(connectionSession); + when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), any(QueryContext.class), anyList())).thenReturn(Collections.emptyMap()); } @Test @@ -136,9 +141,9 @@ void assertGetName() throws SQLException { @Test void assertExecuteSelectStatementAndReturnAllRows() throws SQLException, ReflectiveOperationException { QueryResponseHeader responseHeader = mock(QueryResponseHeader.class); - QueryHeader queryHeader = new QueryHeader("schema", "table", "columnLabel", "columnName", Types.VARCHAR, "columnTypeName", 0, 0, false, false, false, false); - QueryHeader intColumnQueryHeader = new QueryHeader("schema", "table", "columnLabel", "columnName", Types.INTEGER, "columnTypeName", 0, 0, false, false, false, false); - when(responseHeader.getQueryHeaders()).thenReturn(Arrays.asList(queryHeader, intColumnQueryHeader)); + QueryHeader queryHeader = new QueryHeader("schema", "table", "columnLabel", "columnName", Types.STRUCT, "record_type", 0, 0, false, false, false, false); + QueryHeader binaryQueryHeader = new QueryHeader("schema", "table", "columnLabel", "columnName", Types.STRUCT, "record_type", 0, 0, false, false, false, false); + when(responseHeader.getQueryHeaders()).thenReturn(Arrays.asList(queryHeader, binaryQueryHeader)); when(proxyBackendHandler.execute()).thenReturn(responseHeader); when(proxyBackendHandler.next()).thenReturn(true, true, false); when(proxyBackendHandler.getRowData()).thenReturn(new QueryResponseRow(Collections.singletonList(new QueryResponseCell(Types.INTEGER, 0))), @@ -146,6 +151,11 @@ void assertExecuteSelectStatementAndReturnAllRows() throws SQLException, Reflect SelectStatementContext sqlStatementContext = mock(SelectStatementContext.class, RETURNS_DEEP_STUBS); when(sqlStatementContext.getSqlStatement()).thenReturn(SelectStatement.builder().databaseType(databaseType).build()); when(sqlStatementContext.getTablesContext().getDatabaseName()).thenReturn(Optional.empty()); + Map columnTypeOIDs = new HashMap<>(2, 1F); + columnTypeOIDs.put(1, 2249); + columnTypeOIDs.put(2, 2249); + when(PostgreSQLColumnTypeOIDLoader.load(eq(databaseConnectionManager.getConnectionSession()), any(QueryContext.class), anyList())) + .thenReturn(columnTypeOIDs); Portal portal = createPortal(sqlStatementContext, Arrays.asList(PostgreSQLValueFormat.TEXT, PostgreSQLValueFormat.BINARY)); PostgreSQLPacket portalDescription = portal.describe(); assertThat(portalDescription, isA(PostgreSQLRowDescriptionPacket.class)); @@ -153,9 +163,11 @@ void assertExecuteSelectStatementAndReturnAllRows() throws SQLException, Reflect .get(PostgreSQLRowDescriptionPacket.class.getDeclaredField("columnDescriptions"), portalDescription); Iterator columnDescriptionIterator = columnDescriptions.iterator(); PostgreSQLColumnDescription textColumnDescription = columnDescriptionIterator.next(); - PostgreSQLColumnDescription intColumnDescription = columnDescriptionIterator.next(); + PostgreSQLColumnDescription binaryColumnDescription = columnDescriptionIterator.next(); + assertThat(textColumnDescription.getTypeOID(), is(2249)); assertThat(textColumnDescription.getDataFormat(), is(PostgreSQLValueFormat.TEXT.getCode())); - assertThat(intColumnDescription.getDataFormat(), is(PostgreSQLValueFormat.BINARY.getCode())); + assertThat(binaryColumnDescription.getTypeOID(), is(PostgreSQLBinaryColumnType.VARCHAR.getValue())); + assertThat(binaryColumnDescription.getDataFormat(), is(PostgreSQLValueFormat.BINARY.getCode())); List actualPackets = portal.execute(0); assertThat(actualPackets.size(), is(3)); assertThat(actualPackets.get(0), isA(PostgreSQLDataRowPacket.class)); diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutorTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutorTest.java index a97d31e2b0262..1cd042fefde30 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutorTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutorTest.java @@ -21,6 +21,7 @@ import org.apache.shardingsphere.database.connector.core.type.DatabaseType; import org.apache.shardingsphere.database.exception.core.exception.syntax.column.ColumnNotFoundException; import org.apache.shardingsphere.database.protocol.packet.DatabasePacket; +import org.apache.shardingsphere.database.protocol.postgresql.constant.PostgreSQLValueFormat; import org.apache.shardingsphere.database.protocol.postgresql.packet.command.query.PostgreSQLColumnDescription; import org.apache.shardingsphere.database.protocol.postgresql.packet.command.query.PostgreSQLNoDataPacket; import org.apache.shardingsphere.database.protocol.postgresql.packet.command.query.PostgreSQLParameterDescriptionPacket; @@ -51,6 +52,7 @@ import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; import org.apache.shardingsphere.proxy.backend.session.ServerPreparedStatementRegistry; import org.apache.shardingsphere.proxy.frontend.postgresql.command.PortalContext; +import org.apache.shardingsphere.proxy.frontend.postgresql.command.query.PostgreSQLColumnTypeOIDLoader; import org.apache.shardingsphere.proxy.frontend.postgresql.command.query.extended.Portal; import org.apache.shardingsphere.proxy.frontend.postgresql.command.query.extended.PostgreSQLServerPreparedStatement; import org.apache.shardingsphere.sql.parser.engine.api.CacheOption; @@ -60,6 +62,7 @@ import org.apache.shardingsphere.sqltranslator.rule.builder.DefaultSQLTranslatorRuleConfigurationBuilder; import org.apache.shardingsphere.test.infra.framework.extension.mock.AutoMockExtension; import org.apache.shardingsphere.test.infra.framework.extension.mock.StaticMockSettings; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.junit.jupiter.params.ParameterizedTest; @@ -81,6 +84,7 @@ import java.util.Arrays; import java.util.Collection; import java.util.Collections; +import java.util.HashMap; import java.util.Iterator; import java.util.List; import java.util.Map; @@ -106,7 +110,7 @@ import static org.mockito.Mockito.when; @ExtendWith(AutoMockExtension.class) -@StaticMockSettings(ProxyContext.class) +@StaticMockSettings({ProxyContext.class, PostgreSQLColumnTypeOIDLoader.class}) @MockitoSettings(strictness = Strictness.LENIENT) class PostgreSQLComDescribeExecutorTest { @@ -130,6 +134,11 @@ class PostgreSQLComDescribeExecutorTest { @InjectMocks private PostgreSQLComDescribeExecutor executor; + @BeforeEach + void setUp() throws SQLException { + when(PostgreSQLColumnTypeOIDLoader.load(any(Connection.class), any(ResultSetMetaData.class))).thenReturn(Collections.emptyMap()); + } + @Test void assertDescribePortal() throws SQLException { when(packet.getType()).thenReturn('P'); @@ -254,7 +263,6 @@ void assertDescribePreparedStatementInsertWithSchemaAndMixedParameterTypes() thr when(packet.getName()).thenReturn(statementId); String sql = "INSERT INTO public.t_small (col1, col2) VALUES (?, ?) RETURNING *, col1 + col2 expr_sum"; SQLStatement sqlStatement = SQL_PARSER_ENGINE.parse(sql, false); - List parameterTypes = Arrays.asList(PostgreSQLBinaryColumnType.INT4, PostgreSQLBinaryColumnType.UNSPECIFIED); SQLStatementContext sqlStatementContext = mock(InsertStatementContext.class); when(sqlStatementContext.getSqlStatement()).thenReturn(sqlStatement); ShardingSphereTable table = new ShardingSphereTable("t_small", @@ -264,6 +272,8 @@ void assertDescribePreparedStatementInsertWithSchemaAndMixedParameterTypes() thr Collections.emptyList(), Collections.emptyList()); ContextManager contextManager = mockContextManager(table); when(ProxyContext.getInstance().getContextManager()).thenReturn(contextManager); + prepareJDBCBackendConnectionForResultMetaData(sql, 3); + List parameterTypes = Arrays.asList(PostgreSQLBinaryColumnType.INT4, PostgreSQLBinaryColumnType.UNSPECIFIED); List parameterIndexes = IntStream.range(0, sqlStatement.getParameterCount()).boxed().collect(Collectors.toList()); connectionSession.getServerPreparedStatementRegistry().addPreparedStatement( statementId, new PostgreSQLServerPreparedStatement(sql, sqlStatementContext, new HintValueContext(), parameterTypes, parameterIndexes)); @@ -377,6 +387,13 @@ void assertDescribePreparedStatementInsertWithReturning(final String testName, f when(sqlStatementContext.getSqlStatement()).thenReturn(sqlStatement); ContextManager contextManager = mockContextManager(); when(ProxyContext.getInstance().getContextManager()).thenReturn(contextManager); + prepareJDBCBackendConnectionForResultMetaData(sql, expectedColumns.size()); + if ("returning complex columns".equals(testName)) { + Map columnTypeOIDs = new HashMap<>(2, 1F); + columnTypeOIDs.put(12, 2249); + columnTypeOIDs.put(13, 2249); + when(PostgreSQLColumnTypeOIDLoader.load(any(Connection.class), any(ResultSetMetaData.class))).thenReturn(columnTypeOIDs); + } List parameterIndexes = IntStream.range(0, sqlStatement.getParameterCount()).boxed().collect(Collectors.toList()); connectionSession.getServerPreparedStatementRegistry().addPreparedStatement( statementId, new PostgreSQLServerPreparedStatement(sql, sqlStatementContext, new HintValueContext(), parameterTypes, parameterIndexes)); @@ -420,6 +437,7 @@ void assertDescribeSelectPreparedStatement() throws SQLException { ConnectionContext connectionContext = mock(ConnectionContext.class); when(connectionContext.getCurrentDatabaseName()).thenReturn(Optional.of(DATABASE_NAME)); when(connectionSession.getConnectionContext()).thenReturn(connectionContext); + when(PostgreSQLColumnTypeOIDLoader.load(any(Connection.class), any(ResultSetMetaData.class))).thenReturn(Collections.singletonMap(1, 2249)); connectionSession.getServerPreparedStatementRegistry().addPreparedStatement( statementId, new PostgreSQLServerPreparedStatement(sql, sqlStatementContext, new HintValueContext(), parameterTypes, parameterIndexes)); Collection actual = executor.execute(); @@ -434,7 +452,7 @@ void assertDescribeSelectPreparedStatement() throws SQLException { PostgreSQLRowDescriptionPacket actualRowDescription = (PostgreSQLRowDescriptionPacket) actualPacketsIterator.next(); List actualColumnDescriptions = getColumnDescriptions(actualRowDescription); List expectedColumnDescriptions = Arrays.asList( - new PostgreSQLColumnDescription("id", 1, Types.INTEGER, 11, "int4"), + new PostgreSQLColumnDescription("id", 1, 2249, 11, PostgreSQLValueFormat.TEXT.getCode()), new PostgreSQLColumnDescription("k", 2, Types.INTEGER, 11, "int4"), new PostgreSQLColumnDescription("c", 3, Types.CHAR, 60, "int4"), new PostgreSQLColumnDescription("pad", 4, Types.CHAR, 120, "int4")); @@ -644,6 +662,19 @@ private void prepareJDBCBackendConnectionWithPreparedStatement(final String sql) when(connectionSession.getDatabaseConnectionManager()).thenReturn(databaseConnectionManager); } + private void prepareJDBCBackendConnectionForResultMetaData(final String sql, final int columnCount) throws SQLException { + ProxyDatabaseConnectionManager databaseConnectionManager = mock(ProxyDatabaseConnectionManager.class); + Connection connection = mock(Connection.class); + PreparedStatement preparedStatement = mock(PreparedStatement.class); + ResultSetMetaData metaData = mock(ResultSetMetaData.class); + when(metaData.getColumnCount()).thenReturn(columnCount); + when(preparedStatement.getMetaData()).thenReturn(metaData); + when(preparedStatement.getConnection()).thenReturn(connection); + when(connection.prepareStatement(sql)).thenReturn(preparedStatement); + when(databaseConnectionManager.getConnections(any(), nullable(String.class), anyInt(), anyInt(), any(ConnectionMode.class))).thenReturn(Collections.singletonList(connection)); + when(connectionSession.getDatabaseConnectionManager()).thenReturn(databaseConnectionManager); + } + private void prepareJDBCBackendConnectionWithParamTypes(final String sql, final int[] paramTypes, final String[] paramTypeNames) throws SQLException { ParameterMetaData parameterMetaData = mock(ParameterMetaData.class); for (int i = 0; i < paramTypes.length; i++) { @@ -738,7 +769,7 @@ private static List getExpectedReturningColumns() { expectedColumn("k", Types.INTEGER, 4, "int4"), expectedColumn("c", Types.CHAR, -1, "char"), expectedColumn("pad", Types.CHAR, -1, "char"), - expectedColumn("t_order", Types.VARCHAR, -1, "varchar"), - expectedColumn("alias_t_order", Types.VARCHAR, -1, "varchar")); + new PostgreSQLColumnDescription("t_order", 0, 2249, -1, PostgreSQLValueFormat.TEXT.getCode()), + new PostgreSQLColumnDescription("alias_t_order", 0, 2249, -1, PostgreSQLValueFormat.TEXT.getCode())); } } diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutorTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutorTest.java index 00cf2bb59842d..5123ab125b3cd 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutorTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutorTest.java @@ -29,6 +29,7 @@ import org.apache.shardingsphere.database.protocol.postgresql.packet.handshake.PostgreSQLParameterStatusPacket; import org.apache.shardingsphere.database.exception.core.exception.data.InvalidParameterValueException; import org.apache.shardingsphere.infra.hint.HintValueContext; +import org.apache.shardingsphere.infra.session.query.QueryContext; import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader; import org.apache.shardingsphere.proxy.backend.handler.ProxyBackendHandler; import org.apache.shardingsphere.proxy.backend.handler.ProxyBackendHandlerFactory; @@ -40,6 +41,7 @@ import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; import org.apache.shardingsphere.proxy.frontend.command.executor.ResponseType; import org.apache.shardingsphere.proxy.frontend.postgresql.command.PortalContext; +import org.apache.shardingsphere.proxy.frontend.postgresql.command.query.PostgreSQLColumnTypeOIDLoader; import org.apache.shardingsphere.sql.parser.statement.core.segment.dal.VariableAssignSegment; import org.apache.shardingsphere.sql.parser.statement.core.segment.dal.VariableSegment; import org.apache.shardingsphere.sql.parser.statement.core.statement.SQLStatement; @@ -63,6 +65,7 @@ import org.mockito.quality.Strictness; import java.sql.SQLException; +import java.sql.Types; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; @@ -76,13 +79,15 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @ExtendWith(AutoMockExtension.class) -@StaticMockSettings({ProxySQLComQueryParser.class, ProxyBackendHandlerFactory.class}) +@StaticMockSettings({ProxySQLComQueryParser.class, ProxyBackendHandlerFactory.class, PostgreSQLColumnTypeOIDLoader.class}) @MockitoSettings(strictness = Strictness.LENIENT) class PostgreSQLComQueryExecutorTest { @@ -109,13 +114,19 @@ void setUp() throws SQLException { SQLStatement sqlStatement = new SQLStatement(DATABASE_TYPE); when(ProxySQLComQueryParser.parse(queryPacket.getSQL(), DATABASE_TYPE, connectionSession)).thenReturn(sqlStatement); when(ProxyBackendHandlerFactory.newInstance(DATABASE_TYPE, queryPacket.getSQL(), sqlStatement, connectionSession, queryPacket.getHintValueContext())).thenReturn(proxyBackendHandler); + QueryContext queryContext = mock(QueryContext.class); + when(connectionSession.getQueryContext()).thenReturn(queryContext); + when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), eq(queryContext), anyList())).thenReturn(Collections.emptyMap()); queryExecutor = new PostgreSQLComQueryExecutor(portalContext, queryPacket, connectionSession); } @Test void assertExecuteQueryWithColumnDescription() throws SQLException, ReflectiveOperationException { + QueryContext queryContext = connectionSession.getQueryContext(); + when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), eq(queryContext), anyList())).thenReturn(Collections.singletonMap(1, 2249)); QueryResponseHeader queryResponseHeader = mock(QueryResponseHeader.class); - when(queryResponseHeader.getQueryHeaders()).thenReturn(Collections.singletonList(new QueryHeader("schema", "table", "label", "column", 1, "type", 2, 3, true, true, true, true))); + when(queryResponseHeader.getQueryHeaders()).thenReturn( + Collections.singletonList(new QueryHeader("schema", "table", "label", "column", Types.STRUCT, "record_type", 2, 3, true, true, true, true))); when(proxyBackendHandler.execute()).thenReturn(queryResponseHeader); Collection actual = queryExecutor.execute(); PostgreSQLRowDescriptionPacket rowDescriptionPacket = (PostgreSQLRowDescriptionPacket) actual.iterator().next(); @@ -126,7 +137,7 @@ void assertExecuteQueryWithColumnDescription() throws SQLException, ReflectiveOp assertThat(columnDescription.getColumnName(), is("label")); assertThat(columnDescription.getColumnIndex(), is(1)); assertThat(columnDescription.getColumnLength(), is(2)); - assertThat(columnDescription.getTypeOID(), is(new PostgreSQLColumnDescription("column", 1, 1, 2, "type").getTypeOID())); + assertThat(columnDescription.getTypeOID(), is(2249)); assertThat(queryExecutor.getResponseType(), is(ResponseType.QUERY)); } From ca8527dd5dd5c515c79c94ef3ee5f77245b88840 Mon Sep 17 00:00:00 2001 From: zhangliang Date: Sat, 25 Jul 2026 20:42:38 +0800 Subject: [PATCH 2/8] Fix PostgreSQL composite column type OID resolution --- RELEASE-NOTES.md | 1 + .../query/PostgreSQLColumnTypeOIDLoader.java | 47 ++++++++++--------- 2 files changed, 25 insertions(+), 23 deletions(-) diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index adfc108040725..4a78208c24d8b 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -36,6 +36,7 @@ 1. Proxy: Add MySQL exception mapping for ColumnNotFoundException - [#39126](https://github.com/apache/shardingsphere/pull/39126) 1. Proxy: Fix MySQL prepared statement parameter signedness decoding - [#39204](https://github.com/apache/shardingsphere/pull/39204) 1. Proxy: Fix Proxy Native Docker image failing to start due to unexpanded LOCAL_PATH in ENTRYPOINT - [#39146](https://github.com/apache/shardingsphere/pull/39146) +1. Proxy: Fix incorrect PostgreSQL composite column type OIDs in simple and extended query row descriptions - [#39241](https://github.com/apache/shardingsphere/pull/39241) 1. JDBC & Proxy: Remove default MySQL prepared statement query properties when creating data sources - [#38593](https://github.com/apache/shardingsphere/pull/38593) 1. Mode: Fix rule metadata not removed from memory after dropping rules in Etcd cluster mode - [#38561](https://github.com/apache/shardingsphere/pull/38561) 1. Agent: Fix wrong target class name in StaticMethodAdviceExecutor error logs - [#39077](https://github.com/apache/shardingsphere/pull/39077) diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java index 28de32e1d2f96..e011a31ee7fb8 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java @@ -37,6 +37,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Optional; /** * Loader for PostgreSQL column type OIDs. @@ -67,6 +68,10 @@ public static Map load(final ConnectionSession connectionSessi return load(connections.get(0), queryHeaders); } + private static Map load(final Connection connection, final List queryHeaders) throws SQLException { + return connection.isWrapperFor(BaseConnection.class) ? getTypeOIDs(connection.unwrap(BaseConnection.class), queryHeaders) : Collections.emptyMap(); + } + /** * Load composite column type OIDs from result set metadata. * @@ -76,44 +81,40 @@ public static Map load(final ConnectionSession connectionSessi * @throws SQLException SQL exception */ public static Map load(final Connection connection, final ResultSetMetaData metaData) throws SQLException { - if (!connection.isWrapperFor(BaseConnection.class)) { - return Collections.emptyMap(); - } - return getTypeOIDs(connection.unwrap(BaseConnection.class), metaData); + return connection.isWrapperFor(BaseConnection.class) ? getTypeOIDs(connection.unwrap(BaseConnection.class), metaData) : Collections.emptyMap(); } - private static Map load(final Connection connection, final List queryHeaders) throws SQLException { - if (!connection.isWrapperFor(BaseConnection.class)) { - return Collections.emptyMap(); + private static Map getTypeOIDs(final BaseConnection connection, final List queryHeaders) throws SQLException { + Map result = new HashMap<>(); + for (int i = 0; i < queryHeaders.size(); i++) { + QueryHeader each = queryHeaders.get(i); + Optional typeOID = findTypeOID(connection, each.getColumnType(), each.getColumnTypeName()); + if (typeOID.isPresent()) { + result.put(i + 1, typeOID.get()); + } } - return getTypeOIDs(connection.unwrap(BaseConnection.class), queryHeaders); + return result; } private static Map getTypeOIDs(final BaseConnection connection, final ResultSetMetaData metaData) throws SQLException { int columnCount = metaData.getColumnCount(); Map result = new HashMap<>(); for (int columnIndex = 1; columnIndex <= columnCount; columnIndex++) { - if (Types.STRUCT == metaData.getColumnType(columnIndex)) { - int typeOID = connection.getTypeInfo().getPGType(metaData.getColumnTypeName(columnIndex)); - if (Oid.UNSPECIFIED != typeOID) { - result.put(columnIndex, typeOID); - } + Optional typeOID = findTypeOID(connection, metaData.getColumnType(columnIndex), metaData.getColumnTypeName(columnIndex)); + if (typeOID.isPresent()) { + result.put(columnIndex, typeOID.get()); } } return result; } - private static Map getTypeOIDs(final BaseConnection connection, final List queryHeaders) throws SQLException { - Map result = new HashMap<>(); - for (int i = 0; i < queryHeaders.size(); i++) { - QueryHeader each = queryHeaders.get(i); - if (Types.STRUCT == each.getColumnType()) { - int typeOID = connection.getTypeInfo().getPGType(each.getColumnTypeName()); - if (Oid.UNSPECIFIED != typeOID) { - result.put(i + 1, typeOID); - } + private static Optional findTypeOID(final BaseConnection connection, final int columnType, final String columnTypeName) throws SQLException { + if (Types.STRUCT == columnType) { + int typeOID = connection.getTypeInfo().getPGType(columnTypeName); + if (Oid.UNSPECIFIED != typeOID) { + return Optional.of(typeOID); } } - return result; + return Optional.empty(); } } From 2c78deac1f9bdd1beb122f2b0375fda885f0bb5d Mon Sep 17 00:00:00 2001 From: zhangliang Date: Sun, 26 Jul 2026 14:04:16 +0800 Subject: [PATCH 3/8] Fix PostgreSQL composite type OID routing - resolve composite type OIDs from the actual routed data source - reject ambiguous composite OIDs across multiple data sources - document error 30031 and fix column metadata argument order --- .../query/PostgreSQLColumnDescription.java | 4 +- .../PostgreSQLColumnDescriptionTest.java | 7 + .../error-code/sql-error-code.cn.md | 23 ++-- .../error-code/sql-error-code.en.md | 23 ++-- .../backend/connector/ProxySQLExecutor.java | 7 + .../DialectJDBCResultMetadataChecker.java | 42 ++++++ .../connector/ProxySQLExecutorTest.java | 32 +++++ .../PostgreSQLJDBCResultMetadataChecker.java | 80 +++++++++++ ...mpositeTypeAcrossDataSourcesException.java | 33 +++++ ....executor.DialectJDBCResultMetadataChecker | 18 +++ ...stgreSQLJDBCResultMetadataCheckerTest.java | 129 ++++++++++++++++++ ...iteTypeAcrossDataSourcesExceptionTest.java | 37 +++++ .../query/PostgreSQLColumnTypeOIDLoader.java | 47 ++++--- ...reSQLPreparedStatementMetadataFactory.java | 17 ++- .../PostgreSQLColumnTypeOIDLoaderTest.java | 18 ++- ...LPreparedStatementMetadataFactoryTest.java | 26 ++++ .../PostgreSQLComDescribeExecutorTest.java | 1 + 17 files changed, 493 insertions(+), 51 deletions(-) create mode 100644 proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectJDBCResultMetadataChecker.java create mode 100644 proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataChecker.java create mode 100644 proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesException.java create mode 100644 proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker create mode 100644 proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataCheckerTest.java create mode 100644 proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesExceptionTest.java diff --git a/database/protocol/dialect/postgresql/src/main/java/org/apache/shardingsphere/database/protocol/postgresql/packet/command/query/PostgreSQLColumnDescription.java b/database/protocol/dialect/postgresql/src/main/java/org/apache/shardingsphere/database/protocol/postgresql/packet/command/query/PostgreSQLColumnDescription.java index 10597e2ada5a0..43ed351fc4314 100644 --- a/database/protocol/dialect/postgresql/src/main/java/org/apache/shardingsphere/database/protocol/postgresql/packet/command/query/PostgreSQLColumnDescription.java +++ b/database/protocol/dialect/postgresql/src/main/java/org/apache/shardingsphere/database/protocol/postgresql/packet/command/query/PostgreSQLColumnDescription.java @@ -38,10 +38,10 @@ public final class PostgreSQLColumnDescription { private final int columnIndex; - private final int columnLength; - private final int typeOID; + private final int columnLength; + private final int typeModifier = -1; private final int dataFormat; diff --git a/database/protocol/dialect/postgresql/src/test/java/org/apache/shardingsphere/database/protocol/postgresql/packet/command/query/PostgreSQLColumnDescriptionTest.java b/database/protocol/dialect/postgresql/src/test/java/org/apache/shardingsphere/database/protocol/postgresql/packet/command/query/PostgreSQLColumnDescriptionTest.java index a819e62d28348..c986e5250637a 100644 --- a/database/protocol/dialect/postgresql/src/test/java/org/apache/shardingsphere/database/protocol/postgresql/packet/command/query/PostgreSQLColumnDescriptionTest.java +++ b/database/protocol/dialect/postgresql/src/test/java/org/apache/shardingsphere/database/protocol/postgresql/packet/command/query/PostgreSQLColumnDescriptionTest.java @@ -26,6 +26,13 @@ class PostgreSQLColumnDescriptionTest { + @Test + void assertNativeTypeOid() { + PostgreSQLColumnDescription description = new PostgreSQLColumnDescription("record_value", 1, 2249, 32, 0); + assertThat(description.getTypeOID(), is(2249)); + assertThat(description.getColumnLength(), is(32)); + } + @Test void assertIntegerTypeOid() { PostgreSQLColumnDescription description = new PostgreSQLColumnDescription("age", 1, Types.INTEGER, 4, null); diff --git a/docs/document/content/user-manual/error-code/sql-error-code.cn.md b/docs/document/content/user-manual/error-code/sql-error-code.cn.md index bb0cf35bb0843..d572274e3de9b 100644 --- a/docs/document/content/user-manual/error-code/sql-error-code.cn.md +++ b/docs/document/content/user-manual/error-code/sql-error-code.cn.md @@ -247,14 +247,15 @@ SQL 错误码以标准的 SQL State，Vendor Code 和详细错误信息提供， ## 其他异常 -| Vendor Code | SQL State | 错误信息 | -|-------------|-----------|------------------------------------------------------| -| 30000 | HY000 | Unknown exception: %s | -| 30001 | 0A000 | Unsupported SQL operation: %s | -| 30002 | HY000 | Database protocol exception: %s | -| 30003 | 0A000 | Unsupported command: %s | -| 30004 | HY000 | Server exception: %s | -| 30005 | HY000 | Underlying SQL state: %s, underlying error code: %s. | -| 30010 | HY000 | Can not find plugin class '%s'. | -| 30020 | HY000 | File access failed, file is: %s | -| 30030 | 0A000 | No tableless route info found. | +| Vendor Code | SQL State | 错误信息 | +|-------------|-----------|------------------------------------------------------------------------------------------------| +| 30000 | HY000 | Unknown exception: %s | +| 30001 | 0A000 | Unsupported SQL operation: %s | +| 30002 | HY000 | Database protocol exception: %s | +| 30003 | 0A000 | Unsupported command: %s | +| 30004 | HY000 | Server exception: %s | +| 30005 | HY000 | Underlying SQL state: %s, underlying error code: %s. | +| 30010 | HY000 | Can not find plugin class '%s'. | +| 30020 | HY000 | File access failed, file is: %s | +| 30030 | 0A000 | No tableless route info found. | +| 30031 | 0A000 | PostgreSQL composite result columns are not supported for SQL routed to multiple data sources. | diff --git a/docs/document/content/user-manual/error-code/sql-error-code.en.md b/docs/document/content/user-manual/error-code/sql-error-code.en.md index 7d35fa946573c..3b04c87a22d88 100644 --- a/docs/document/content/user-manual/error-code/sql-error-code.en.md +++ b/docs/document/content/user-manual/error-code/sql-error-code.en.md @@ -248,14 +248,15 @@ SQL error codes provide by standard `SQL State`, `Vendor Code` and `Reason`, whi ## Other Exception -| Vendor Code | SQL State | Reason | -|-------------|-----------|------------------------------------------------------| -| 30000 | HY000 | Unknown exception: %s | -| 30001 | 0A000 | Unsupported SQL operation: %s | -| 30002 | HY000 | Database protocol exception: %s | -| 30003 | 0A000 | Unsupported command: %s | -| 30004 | HY000 | Server exception: %s | -| 30005 | HY000 | Underlying SQL state: %s, underlying error code: %s. | -| 30010 | HY000 | Can not find plugin class '%s'. | -| 30020 | HY000 | File access failed, file is: %s | -| 30030 | 0A000 | No tableless route info found. | +| Vendor Code | SQL State | Reason | +|-------------|-----------|------------------------------------------------------------------------------------------------| +| 30000 | HY000 | Unknown exception: %s | +| 30001 | 0A000 | Unsupported SQL operation: %s | +| 30002 | HY000 | Database protocol exception: %s | +| 30003 | 0A000 | Unsupported command: %s | +| 30004 | HY000 | Server exception: %s | +| 30005 | HY000 | Underlying SQL state: %s, underlying error code: %s. | +| 30010 | HY000 | Can not find plugin class '%s'. | +| 30020 | HY000 | File access failed, file is: %s | +| 30030 | 0A000 | No tableless route info found. | +| 30031 | 0A000 | PostgreSQL composite result columns are not supported for SQL routed to multiple data sources. | diff --git a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutor.java b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutor.java index 4e9a3b9c4e206..28009ad61b193 100644 --- a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutor.java +++ b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutor.java @@ -55,6 +55,7 @@ import org.apache.shardingsphere.infra.session.connection.transaction.TransactionConnectionContext; import org.apache.shardingsphere.infra.spi.type.ordered.OrderedSPILoader; import org.apache.shardingsphere.mode.metadata.MetaDataContexts; +import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker; import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.ProxyJDBCExecutor; import org.apache.shardingsphere.proxy.backend.connector.jdbc.statement.JDBCBackendStatement; import org.apache.shardingsphere.proxy.backend.connector.sane.DialectSaneQueryResultEngine; @@ -238,6 +239,12 @@ type, maxConnectionsSizePerQuery, databaseConnectionManager, statementManager, n } catch (final SQLException ex) { return getSaneExecuteResults(executionContext, ex); } + DatabaseType databaseType = ProxyContext.getInstance().getContextManager().getDatabase(databaseName).getProtocolType(); + Optional checker = DatabaseTypedSPILoader.findService(DialectJDBCResultMetadataChecker.class, databaseType); + if (checker.isPresent()) { + JDBCExecutionUnit sample = executionGroupContext.getInputGroups().iterator().next().getInputs().get(0); + checker.get().check(executionContext, sample.getStorageResource(), sample.getExecutionUnit().getSqlUnit().getSql()); + } executeTransactionHooksBeforeExecuteSQL(databaseConnectionManager.getConnectionSession()); return regularExecutor.execute(executionContext.getQueryContext(), executionGroupContext, isReturnGeneratedKeys, isExceptionThrown); } diff --git a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectJDBCResultMetadataChecker.java b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectJDBCResultMetadataChecker.java new file mode 100644 index 0000000000000..5081dd47f1a31 --- /dev/null +++ b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectJDBCResultMetadataChecker.java @@ -0,0 +1,42 @@ +/* + * 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.proxy.backend.connector.jdbc.executor; + +import org.apache.shardingsphere.database.connector.core.spi.DatabaseTypedSPI; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; +import org.apache.shardingsphere.infra.spi.annotation.SingletonSPI; + +import java.sql.SQLException; +import java.sql.Statement; + +/** + * Dialect JDBC result metadata checker. + */ +@SingletonSPI +public interface DialectJDBCResultMetadataChecker extends DatabaseTypedSPI { + + /** + * Check JDBC result metadata after routed resources are prepared. + * + * @param executionContext execution context + * @param statement JDBC statement + * @param sql routed SQL corresponding to the statement + * @throws SQLException SQL exception + */ + void check(ExecutionContext executionContext, Statement statement, String sql) throws SQLException; +} diff --git a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutorTest.java b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutorTest.java index 9d9a64e68e746..4e563f0def4cd 100644 --- a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutorTest.java +++ b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutorTest.java @@ -32,9 +32,13 @@ import org.apache.shardingsphere.infra.binder.context.statement.type.dml.InsertStatementContext; import org.apache.shardingsphere.infra.config.props.ConfigurationPropertyKey; import org.apache.shardingsphere.infra.config.rule.RuleConfiguration; +import org.apache.shardingsphere.infra.executor.kernel.model.ExecutionGroup; import org.apache.shardingsphere.infra.executor.kernel.model.ExecutionGroupContext; import org.apache.shardingsphere.infra.executor.kernel.model.ExecutionGroupReportContext; import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; +import org.apache.shardingsphere.infra.executor.sql.context.SQLUnit; +import org.apache.shardingsphere.infra.executor.sql.execute.engine.ConnectionMode; import org.apache.shardingsphere.infra.executor.sql.execute.engine.driver.jdbc.JDBCExecutionUnit; import org.apache.shardingsphere.infra.executor.sql.execute.engine.raw.RawExecutor; import org.apache.shardingsphere.infra.executor.sql.execute.engine.raw.RawSQLExecutionUnit; @@ -54,6 +58,7 @@ import org.apache.shardingsphere.infra.session.query.QueryContext; import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader; import org.apache.shardingsphere.mode.manager.ContextManager; +import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker; import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.ProxyJDBCExecutor; import org.apache.shardingsphere.proxy.backend.connector.jdbc.statement.JDBCBackendStatement; import org.apache.shardingsphere.proxy.backend.connector.sane.DialectSaneQueryResultEngine; @@ -98,6 +103,7 @@ import org.mockito.quality.Strictness; import java.sql.SQLException; +import java.sql.Statement; import java.util.Arrays; import java.util.Collection; import java.util.Collections; @@ -116,6 +122,7 @@ import static org.mockito.Mockito.anyBoolean; import static org.mockito.Mockito.anyCollection; import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mockConstruction; @@ -315,6 +322,31 @@ private Stream executeScenarios() { Arguments.of("execute-with-driver-and-no-transaction", false, createInsertStatement(postgresqlDatabaseType), false, false, false)); } + @Test + void assertExecuteWithDialectMetadataCheckerException() throws SQLException { + when(connectionSession.getUsedDatabaseName()).thenReturn("foo_db"); + when(database.getRuleMetaData().getRules()).thenReturn(createRules(false)); + ProxySQLExecutor proxySQLExecutor = createProxySQLExecutor("foo_schema", true); + setExecutorField(proxySQLExecutor, "regularExecutor", regularExecutor); + ExecutionContext executionContext = createExecutionContext("execute-with-dialect-metadata-checker", new SQLStatement(fixtureDatabaseType), false); + Statement statement = mock(Statement.class); + ExecutionUnit executionUnit = new ExecutionUnit("foo_ds", new SQLUnit("SELECT 1", Collections.emptyList())); + JDBCExecutionUnit jdbcExecutionUnit = new JDBCExecutionUnit(executionUnit, ConnectionMode.CONNECTION_STRICTLY, statement); + ExecutionGroupContext executionGroupContext = new ExecutionGroupContext<>( + Collections.singletonList(new ExecutionGroup<>(Collections.singletonList(jdbcExecutionUnit))), mock(ExecutionGroupReportContext.class)); + DialectJDBCResultMetadataChecker checker = mock(DialectJDBCResultMetadataChecker.class); + SQLException expected = new SQLException("expected"); + doThrow(expected).when(checker).check(executionContext, statement, "SELECT 1"); + try ( + MockedConstruction ignored = mockConstruction(DriverExecutionPrepareEngine.class, + (mock, context) -> when(mock.prepare(anyString(), eq(executionContext), anyCollection(), any(ExecutionGroupReportContext.class))).thenReturn(executionGroupContext)); + MockedStatic spiLoader = mockStatic(DatabaseTypedSPILoader.class, CALLS_REAL_METHODS)) { + spiLoader.when(() -> DatabaseTypedSPILoader.findService(DialectJDBCResultMetadataChecker.class, fixtureDatabaseType)).thenReturn(Optional.of(checker)); + assertThat(assertThrows(SQLException.class, () -> proxySQLExecutor.execute(executionContext)), is(expected)); + verify(regularExecutor, never()).execute(any(), any(), anyBoolean(), anyBoolean()); + } + } + @ParameterizedTest(name = "{0}") @MethodSource("executeFallbackScenarios") void assertExecuteFallback(final String name, final boolean hasRawExecutionRule, final SQLStatement sqlStatement, final boolean hasSaneResult) throws SQLException { diff --git a/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataChecker.java b/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataChecker.java new file mode 100644 index 0000000000000..e5a2e9e7e6d4a --- /dev/null +++ b/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataChecker.java @@ -0,0 +1,80 @@ +/* + * 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.proxy.backend.postgresql.connector.jdbc; + +import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext; +import org.apache.shardingsphere.infra.binder.context.statement.type.dml.SelectStatementContext; +import org.apache.shardingsphere.infra.exception.ShardingSpherePreconditions; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; +import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker; +import org.apache.shardingsphere.proxy.backend.postgresql.exception.PostgreSQLCompositeTypeAcrossDataSourcesException; + +import java.sql.PreparedStatement; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; +import java.sql.Types; +import java.util.Collection; + +/** + * JDBC result metadata checker for PostgreSQL. + */ +public final class PostgreSQLJDBCResultMetadataChecker implements DialectJDBCResultMetadataChecker { + + @Override + public void check(final ExecutionContext executionContext, final Statement statement, final String sql) throws SQLException { + if (!isRoutedToMultipleDataSources(executionContext.getExecutionUnits())) { + return; + } + if (statement instanceof PreparedStatement) { + checkCompositeType(((PreparedStatement) statement).getMetaData(), executionContext.getSqlStatementContext()); + return; + } + try (PreparedStatement preparedStatement = statement.getConnection().prepareStatement(sql)) { + checkCompositeType(preparedStatement.getMetaData(), executionContext.getSqlStatementContext()); + } + } + + private boolean isRoutedToMultipleDataSources(final Collection executionUnits) { + String dataSourceName = executionUnits.iterator().next().getDataSourceName(); + for (ExecutionUnit each : executionUnits) { + if (!dataSourceName.equals(each.getDataSourceName())) { + return true; + } + } + return false; + } + + private void checkCompositeType(final ResultSetMetaData metaData, final SQLStatementContext sqlStatementContext) throws SQLException { + if (null == metaData) { + return; + } + int columnCount = sqlStatementContext instanceof SelectStatementContext && ((SelectStatementContext) sqlStatementContext).containsDerivedProjections() + ? ((SelectStatementContext) sqlStatementContext).getProjectionsContext().getExpandProjections().size() + : metaData.getColumnCount(); + for (int columnIndex = 1; columnIndex <= columnCount; columnIndex++) { + ShardingSpherePreconditions.checkState(Types.STRUCT != metaData.getColumnType(columnIndex), PostgreSQLCompositeTypeAcrossDataSourcesException::new); + } + } + + @Override + public String getDatabaseType() { + return "PostgreSQL"; + } +} diff --git a/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesException.java b/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesException.java new file mode 100644 index 0000000000000..d77ac86973198 --- /dev/null +++ b/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesException.java @@ -0,0 +1,33 @@ +/* + * 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.proxy.backend.postgresql.exception; + +import org.apache.shardingsphere.infra.exception.external.sql.sqlstate.XOpenSQLState; +import org.apache.shardingsphere.infra.exception.external.sql.type.generic.GenericSQLException; + +/** + * PostgreSQL composite type across data sources exception. + */ +public final class PostgreSQLCompositeTypeAcrossDataSourcesException extends GenericSQLException { + + private static final long serialVersionUID = 2994368346256039324L; + + public PostgreSQLCompositeTypeAcrossDataSourcesException() { + super(XOpenSQLState.FEATURE_NOT_SUPPORTED, 31, "PostgreSQL composite result columns are not supported for SQL routed to multiple data sources."); + } +} diff --git a/proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker b/proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker new file mode 100644 index 0000000000000..cf750c54ea713 --- /dev/null +++ b/proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker @@ -0,0 +1,18 @@ +# +# 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. +# + +org.apache.shardingsphere.proxy.backend.postgresql.connector.jdbc.PostgreSQLJDBCResultMetadataChecker diff --git a/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataCheckerTest.java b/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataCheckerTest.java new file mode 100644 index 0000000000000..3e0626f79911d --- /dev/null +++ b/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataCheckerTest.java @@ -0,0 +1,129 @@ +/* + * 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.proxy.backend.postgresql.connector.jdbc; + +import org.apache.shardingsphere.database.connector.core.spi.DatabaseTypedSPILoader; +import org.apache.shardingsphere.database.connector.core.type.DatabaseType; +import org.apache.shardingsphere.infra.binder.context.segment.select.projection.Projection; +import org.apache.shardingsphere.infra.binder.context.statement.type.dml.SelectStatementContext; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; +import org.apache.shardingsphere.infra.executor.sql.context.SQLUnit; +import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader; +import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker; +import org.apache.shardingsphere.proxy.backend.postgresql.exception.PostgreSQLCompositeTypeAcrossDataSourcesException; +import org.junit.jupiter.api.Test; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; +import java.sql.Types; +import java.util.Arrays; +import java.util.Collections; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.RETURNS_DEEP_STUBS; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class PostgreSQLJDBCResultMetadataCheckerTest { + + private final DatabaseType databaseType = TypedSPILoader.getService(DatabaseType.class, "PostgreSQL"); + + private final DialectJDBCResultMetadataChecker checker = DatabaseTypedSPILoader.getService(DialectJDBCResultMetadataChecker.class, databaseType); + + @Test + void assertCheckWithSingleDataSource() throws SQLException { + Statement statement = mock(Statement.class); + checker.check(createExecutionContext("ds_0", "ds_0"), statement, "SELECT record_value"); + verify(statement, never()).getConnection(); + } + + @Test + void assertCheckPreparedStatement() throws SQLException { + PreparedStatement preparedStatement = mock(PreparedStatement.class); + ResultSetMetaData metaData = createMetaData(Types.INTEGER); + when(preparedStatement.getMetaData()).thenReturn(metaData); + assertDoesNotThrow(() -> checker.check(createExecutionContext("ds_0", "ds_1"), preparedStatement, "SELECT record_value")); + verify(preparedStatement).getMetaData(); + } + + @Test + void assertCheckCompositeType() throws SQLException { + PreparedStatement preparedStatement = mock(PreparedStatement.class); + ResultSetMetaData metaData = createMetaData(Types.STRUCT); + when(preparedStatement.getMetaData()).thenReturn(metaData); + assertThrows(PostgreSQLCompositeTypeAcrossDataSourcesException.class, + () -> checker.check(createExecutionContext("ds_0", "ds_1"), preparedStatement, "SELECT record_value")); + } + + @Test + void assertCheckStatement() throws SQLException { + PreparedStatement preparedStatement = mock(PreparedStatement.class); + ResultSetMetaData metaData = createMetaData(Types.INTEGER); + when(preparedStatement.getMetaData()).thenReturn(metaData); + Connection connection = mock(Connection.class); + when(connection.prepareStatement("SELECT record_value")).thenReturn(preparedStatement); + Statement statement = mock(Statement.class); + when(statement.getConnection()).thenReturn(connection); + checker.check(createExecutionContext("ds_0", "ds_1"), statement, "SELECT record_value"); + verify(preparedStatement).close(); + } + + @Test + void assertCheckCompositeTypeWithoutMetaData() { + assertDoesNotThrow(() -> checker.check(createExecutionContext("ds_0", "ds_1"), mock(PreparedStatement.class), "SELECT record_value")); + } + + @Test + void assertCheckCompositeTypeWithDerivedProjection() throws SQLException { + ResultSetMetaData metaData = mock(ResultSetMetaData.class); + when(metaData.getColumnType(1)).thenReturn(Types.INTEGER); + when(metaData.getColumnType(2)).thenReturn(Types.STRUCT); + SelectStatementContext sqlStatementContext = mock(SelectStatementContext.class, RETURNS_DEEP_STUBS); + when(sqlStatementContext.containsDerivedProjections()).thenReturn(true); + when(sqlStatementContext.getProjectionsContext().getExpandProjections()).thenReturn(Collections.singletonList(mock(Projection.class))); + ExecutionContext executionContext = createExecutionContext("ds_0", "ds_1"); + when(executionContext.getSqlStatementContext()).thenReturn(sqlStatementContext); + PreparedStatement preparedStatement = mock(PreparedStatement.class); + when(preparedStatement.getMetaData()).thenReturn(metaData); + assertDoesNotThrow(() -> checker.check(executionContext, preparedStatement, "SELECT record_value")); + } + + private ExecutionContext createExecutionContext(final String... dataSourceNames) { + ExecutionContext result = mock(ExecutionContext.class); + ExecutionUnit[] executionUnits = new ExecutionUnit[dataSourceNames.length]; + for (int i = 0; i < dataSourceNames.length; i++) { + executionUnits[i] = new ExecutionUnit(dataSourceNames[i], new SQLUnit("SELECT record_value", Collections.emptyList())); + } + when(result.getExecutionUnits()).thenReturn(Arrays.asList(executionUnits)); + return result; + } + + private ResultSetMetaData createMetaData(final int columnType) throws SQLException { + ResultSetMetaData result = mock(ResultSetMetaData.class); + when(result.getColumnCount()).thenReturn(1); + when(result.getColumnType(1)).thenReturn(columnType); + return result; + } +} diff --git a/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesExceptionTest.java b/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesExceptionTest.java new file mode 100644 index 0000000000000..5d8fd77f39318 --- /dev/null +++ b/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesExceptionTest.java @@ -0,0 +1,37 @@ +/* + * 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.proxy.backend.postgresql.exception; + +import org.apache.shardingsphere.infra.exception.external.sql.sqlstate.XOpenSQLState; +import org.junit.jupiter.api.Test; + +import java.sql.SQLException; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.is; + +class PostgreSQLCompositeTypeAcrossDataSourcesExceptionTest { + + @Test + void assertToSQLException() { + SQLException actual = new PostgreSQLCompositeTypeAcrossDataSourcesException().toSQLException(); + assertThat(actual.getSQLState(), is(XOpenSQLState.FEATURE_NOT_SUPPORTED.getValue())); + assertThat(actual.getErrorCode(), is(30031)); + assertThat(actual.getMessage(), is("PostgreSQL composite result columns are not supported for SQL routed to multiple data sources.")); + } +} diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java index e011a31ee7fb8..bf940626c000f 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java @@ -19,6 +19,7 @@ import lombok.AccessLevel; import lombok.NoArgsConstructor; +import org.apache.shardingsphere.infra.annotation.HighFrequencyInvocation; import org.apache.shardingsphere.infra.connection.kernel.KernelProcessor; import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; @@ -42,6 +43,7 @@ /** * Loader for PostgreSQL column type OIDs. */ +@HighFrequencyInvocation @NoArgsConstructor(access = AccessLevel.PRIVATE) public final class PostgreSQLColumnTypeOIDLoader { @@ -55,14 +57,19 @@ public final class PostgreSQLColumnTypeOIDLoader { * @throws SQLException SQL exception */ public static Map load(final ConnectionSession connectionSession, final QueryContext queryContext, final List queryHeaders) throws SQLException { - if (queryHeaders.stream().noneMatch(each -> Types.STRUCT == each.getColumnType())) { + if (!containsCompositeType(queryHeaders)) { return Collections.emptyMap(); } ExecutionContext executionContext = new KernelProcessor().generateExecutionContext(queryContext, queryContext.getMetaData().getGlobalRuleMetaData(), queryContext.getMetaData().getProps()); - if (1 != executionContext.getExecutionUnits().size()) { + if (executionContext.getExecutionUnits().isEmpty()) { return Collections.emptyMap(); } ExecutionUnit executionUnit = executionContext.getExecutionUnits().iterator().next(); + for (ExecutionUnit each : executionContext.getExecutionUnits()) { + if (!executionUnit.getDataSourceName().equals(each.getDataSourceName())) { + return Collections.emptyMap(); + } + } List connections = connectionSession.getDatabaseConnectionManager().getConnections( connectionSession.getUsedDatabaseName(), executionUnit.getDataSourceName(), 0, 1, ConnectionMode.CONNECTION_STRICTLY); return load(connections.get(0), queryHeaders); @@ -84,13 +91,24 @@ public static Map load(final Connection connection, final Resu return connection.isWrapperFor(BaseConnection.class) ? getTypeOIDs(connection.unwrap(BaseConnection.class), metaData) : Collections.emptyMap(); } + private static boolean containsCompositeType(final List queryHeaders) { + for (QueryHeader each : queryHeaders) { + if (Types.STRUCT == each.getColumnType()) { + return true; + } + } + return false; + } + private static Map getTypeOIDs(final BaseConnection connection, final List queryHeaders) throws SQLException { Map result = new HashMap<>(); for (int i = 0; i < queryHeaders.size(); i++) { QueryHeader each = queryHeaders.get(i); - Optional typeOID = findTypeOID(connection, each.getColumnType(), each.getColumnTypeName()); - if (typeOID.isPresent()) { - result.put(i + 1, typeOID.get()); + if (Types.STRUCT == each.getColumnType()) { + Optional typeOID = findTypeOID(connection, each.getColumnTypeName()); + if (typeOID.isPresent()) { + result.put(i + 1, typeOID.get()); + } } } return result; @@ -100,21 +118,18 @@ private static Map getTypeOIDs(final BaseConnection connection int columnCount = metaData.getColumnCount(); Map result = new HashMap<>(); for (int columnIndex = 1; columnIndex <= columnCount; columnIndex++) { - Optional typeOID = findTypeOID(connection, metaData.getColumnType(columnIndex), metaData.getColumnTypeName(columnIndex)); - if (typeOID.isPresent()) { - result.put(columnIndex, typeOID.get()); + if (Types.STRUCT == metaData.getColumnType(columnIndex)) { + Optional typeOID = findTypeOID(connection, metaData.getColumnTypeName(columnIndex)); + if (typeOID.isPresent()) { + result.put(columnIndex, typeOID.get()); + } } } return result; } - private static Optional findTypeOID(final BaseConnection connection, final int columnType, final String columnTypeName) throws SQLException { - if (Types.STRUCT == columnType) { - int typeOID = connection.getTypeInfo().getPGType(columnTypeName); - if (Oid.UNSPECIFIED != typeOID) { - return Optional.of(typeOID); - } - } - return Optional.empty(); + private static Optional findTypeOID(final BaseConnection connection, final String columnTypeName) throws SQLException { + int typeOID = connection.getTypeInfo().getPGType(columnTypeName); + return Oid.UNSPECIFIED == typeOID ? Optional.empty() : Optional.of(typeOID); } } diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java index c0709b3319b66..3f90d950d5b5b 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java @@ -19,9 +19,12 @@ import lombok.AccessLevel; import lombok.NoArgsConstructor; +import org.apache.shardingsphere.database.connector.core.spi.DatabaseTypedSPILoader; +import org.apache.shardingsphere.infra.annotation.HighFrequencyInvocation; import org.apache.shardingsphere.infra.binder.context.aware.ParameterAware; import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext; import org.apache.shardingsphere.infra.binder.engine.SQLBindEngine; +import org.apache.shardingsphere.infra.exception.external.sql.ShardingSphereSQLException; import org.apache.shardingsphere.infra.connection.kernel.KernelProcessor; import org.apache.shardingsphere.infra.exception.ShardingSpherePreconditions; import org.apache.shardingsphere.infra.exception.kernel.metadata.PreparedStatementMetadataResolutionException; @@ -30,7 +33,9 @@ import org.apache.shardingsphere.infra.executor.sql.execute.engine.ConnectionMode; import org.apache.shardingsphere.infra.metadata.ShardingSphereMetaData; import org.apache.shardingsphere.infra.session.query.QueryContext; +import org.apache.shardingsphere.infra.util.close.QuietlyCloser; import org.apache.shardingsphere.proxy.backend.connector.ProxyDatabaseConnectionManager; +import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker; import org.apache.shardingsphere.proxy.backend.context.ProxyContext; import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; @@ -42,6 +47,7 @@ /** * Metadata factory for PostgreSQL prepared statements. */ +@HighFrequencyInvocation @NoArgsConstructor(access = AccessLevel.PRIVATE) public final class PostgreSQLPreparedStatementMetadataFactory { @@ -72,6 +78,15 @@ public static PreparedStatement load(final ConnectionSession connectionSession, databaseConnectionManager.getConnections(connectionSession.getUsedDatabaseName(), executionUnit.getDataSourceName(), 0, 1, ConnectionMode.CONNECTION_STRICTLY); ShardingSpherePreconditions.checkNotEmpty(connections, () -> new PreparedStatementMetadataResolutionException("no backend connection was acquired")); - return connections.iterator().next().prepareStatement(executionUnit.getSqlUnit().getSql()); + String sql = executionUnit.getSqlUnit().getSql(); + PreparedStatement result = connections.iterator().next().prepareStatement(sql); + try { + DatabaseTypedSPILoader.getService(DialectJDBCResultMetadataChecker.class, connectionSession.getProtocolType()) + .check(executionContext, result, sql); + } catch (final SQLException | ShardingSphereSQLException ex) { + QuietlyCloser.close(result); + throw ex; + } + return result; } } diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java index ab8306463f287..9bfdc4347f7ee 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java @@ -25,9 +25,7 @@ import org.apache.shardingsphere.proxy.backend.connector.ProxyDatabaseConnectionManager; import org.apache.shardingsphere.proxy.backend.response.header.query.QueryHeader; import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; -import org.apache.shardingsphere.test.infra.framework.extension.mock.AutoMockExtension; import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.MockedConstruction; import org.postgresql.core.BaseConnection; import org.postgresql.core.Oid; @@ -38,8 +36,8 @@ import java.sql.SQLException; import java.sql.Types; import java.util.Arrays; +import java.util.Collection; import java.util.Collections; -import java.util.List; import java.util.Map; import static org.hamcrest.MatcherAssert.assertThat; @@ -48,21 +46,21 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.RETURNS_DEEP_STUBS; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mockConstruction; import static org.mockito.Mockito.never; -import static org.mockito.Mockito.RETURNS_DEEP_STUBS; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -@ExtendWith(AutoMockExtension.class) class PostgreSQLColumnTypeOIDLoaderTest { @Test - void assertLoad() throws SQLException { + void assertLoadWithSameDataSource() throws SQLException { QueryContext queryContext = mock(QueryContext.class, RETURNS_DEEP_STUBS); - ExecutionContext executionContext = createExecutionContext( - Collections.singletonList(new ExecutionUnit("ds_0", new SQLUnit("SELECT record_value", Collections.emptyList())))); + ExecutionContext executionContext = createExecutionContext(Arrays.asList( + new ExecutionUnit("ds_0", new SQLUnit("SELECT record_value", Collections.emptyList())), + new ExecutionUnit("ds_0", new SQLUnit("SELECT record_value", Collections.emptyList())))); BaseConnection connection = mock(BaseConnection.class); mockTypeOID(connection, "record_type", 2249); ConnectionSession connectionSession = createConnectionSession(connection); @@ -82,7 +80,7 @@ void assertLoadWithoutCompositeColumn() throws SQLException { } @Test - void assertLoadWithMultipleExecutionUnits() throws SQLException { + void assertLoadWithMultipleDataSources() throws SQLException { ExecutionContext executionContext = createExecutionContext(Arrays.asList( new ExecutionUnit("ds_0", new SQLUnit("SELECT record_value", Collections.emptyList())), new ExecutionUnit("ds_1", new SQLUnit("SELECT record_value", Collections.emptyList())))); @@ -141,7 +139,7 @@ private QueryHeader createQueryHeader(final int columnType, final String columnT return new QueryHeader("", "", "record_value", "record_value", columnType, columnTypeName, -1, 0, false, false, false, false); } - private ExecutionContext createExecutionContext(final List executionUnits) { + private ExecutionContext createExecutionContext(final Collection executionUnits) { ExecutionContext result = mock(ExecutionContext.class); when(result.getExecutionUnits()).thenReturn(executionUnits); return result; diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java index 6881558cdbc69..5878fa4bbbfab 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java @@ -17,6 +17,7 @@ package org.apache.shardingsphere.proxy.frontend.postgresql.command.query.extended; +import org.apache.shardingsphere.database.connector.core.spi.DatabaseTypedSPILoader; import org.apache.shardingsphere.database.connector.core.type.DatabaseType; import org.apache.shardingsphere.database.protocol.postgresql.packet.command.query.extended.PostgreSQLBinaryColumnType; import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext; @@ -38,7 +39,9 @@ import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader; import org.apache.shardingsphere.mode.manager.ContextManager; import org.apache.shardingsphere.proxy.backend.connector.ProxyDatabaseConnectionManager; +import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker; import org.apache.shardingsphere.proxy.backend.context.ProxyContext; +import org.apache.shardingsphere.proxy.backend.postgresql.exception.PostgreSQLCompositeTypeAcrossDataSourcesException; import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; import org.apache.shardingsphere.sql.parser.engine.api.CacheOption; import org.apache.shardingsphere.sql.parser.statement.core.statement.SQLStatement; @@ -51,6 +54,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.MockedConstruction; +import org.mockito.MockedStatic; import java.sql.Connection; import java.sql.PreparedStatement; @@ -66,11 +70,16 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.nullable; +import static org.mockito.Mockito.CALLS_REAL_METHODS; import static org.mockito.Mockito.RETURNS_DEEP_STUBS; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mockConstruction; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @ExtendWith(AutoMockExtension.class) @@ -89,6 +98,7 @@ class PostgreSQLPreparedStatementMetadataFactoryTest { @Test void assertLoad() throws SQLException { PreparedStatement expected = prepareJDBCBackendConnection(null); + when(connectionSession.getProtocolType()).thenReturn(databaseType); assertThat(PostgreSQLPreparedStatementMetadataFactory.load(connectionSession, createPreparedStatement(true), PARAMETERS), is(expected)); } @@ -115,6 +125,22 @@ void assertLoadWithEmptyExecutionUnits() { } } + @Test + void assertClosePreparedStatementOnCheckerException() throws SQLException { + PreparedStatement expected = prepareJDBCBackendConnection(null); + PostgreSQLServerPreparedStatement preparedStatement = createPreparedStatement(true); + when(connectionSession.getProtocolType()).thenReturn(databaseType); + PostgreSQLCompositeTypeAcrossDataSourcesException exception = new PostgreSQLCompositeTypeAcrossDataSourcesException(); + DialectJDBCResultMetadataChecker checker = mock(DialectJDBCResultMetadataChecker.class); + doThrow(exception).when(checker).check(any(ExecutionContext.class), eq(expected), anyString()); + try (MockedStatic spiLoader = mockStatic(DatabaseTypedSPILoader.class, CALLS_REAL_METHODS)) { + spiLoader.when(() -> DatabaseTypedSPILoader.getService(DialectJDBCResultMetadataChecker.class, databaseType)).thenReturn(checker); + assertThat(assertThrows(PostgreSQLCompositeTypeAcrossDataSourcesException.class, + () -> PostgreSQLPreparedStatementMetadataFactory.load(connectionSession, preparedStatement, PARAMETERS)), is(exception)); + verify(expected).close(); + } + } + private PostgreSQLServerPreparedStatement createPreparedStatement(final boolean withUsedDatabaseName) { SQLStatement sqlStatement = sqlParserEngine.parse("SELECT id FROM foo_tbl WHERE id=?", false); SQLStatementContext sqlStatementContext = mock(SelectStatementContext.class); diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutorTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutorTest.java index 1cd042fefde30..a689fdaad8d8a 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutorTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutorTest.java @@ -136,6 +136,7 @@ class PostgreSQLComDescribeExecutorTest { @BeforeEach void setUp() throws SQLException { + when(connectionSession.getProtocolType()).thenReturn(DATABASE_TYPE); when(PostgreSQLColumnTypeOIDLoader.load(any(Connection.class), any(ResultSetMetaData.class))).thenReturn(Collections.emptyMap()); } From 72fccbbbb4f268f758ac28951dcace20470b1a69 Mon Sep 17 00:00:00 2001 From: zhangliang Date: Sun, 26 Jul 2026 14:12:28 +0800 Subject: [PATCH 4/8] Fix PostgreSQL composite type OID routing - resolve composite type OIDs from the actual routed data source - reject ambiguous composite OIDs across multiple data sources - document error 30031 and fix column metadata argument order --- .../proxy/backend/connector/ProxySQLExecutor.java | 4 ++-- ...ker.java => DialectResultSetMetadataChecker.java} | 4 ++-- .../backend/connector/ProxySQLExecutorTest.java | 6 +++--- ....java => PostgreSQLResultSetMetadataChecker.java} | 6 +++--- ...or.jdbc.executor.DialectResultSetMetadataChecker} | 2 +- .../PostgreSQLJDBCResultMetadataCheckerTest.java | 4 ++-- .../PostgreSQLPreparedStatementMetadataFactory.java | 12 +++++++----- ...stgreSQLPreparedStatementMetadataFactoryTest.java | 6 +++--- 8 files changed, 23 insertions(+), 21 deletions(-) rename proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/{DialectJDBCResultMetadataChecker.java => DialectResultSetMetadataChecker.java} (93%) rename proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/{PostgreSQLJDBCResultMetadataChecker.java => PostgreSQLResultSetMetadataChecker.java} (95%) rename proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/{org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker => org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker} (95%) diff --git a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutor.java b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutor.java index 28009ad61b193..a150cea3c8905 100644 --- a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutor.java +++ b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutor.java @@ -55,7 +55,7 @@ import org.apache.shardingsphere.infra.session.connection.transaction.TransactionConnectionContext; import org.apache.shardingsphere.infra.spi.type.ordered.OrderedSPILoader; import org.apache.shardingsphere.mode.metadata.MetaDataContexts; -import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker; +import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker; import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.ProxyJDBCExecutor; import org.apache.shardingsphere.proxy.backend.connector.jdbc.statement.JDBCBackendStatement; import org.apache.shardingsphere.proxy.backend.connector.sane.DialectSaneQueryResultEngine; @@ -240,7 +240,7 @@ type, maxConnectionsSizePerQuery, databaseConnectionManager, statementManager, n return getSaneExecuteResults(executionContext, ex); } DatabaseType databaseType = ProxyContext.getInstance().getContextManager().getDatabase(databaseName).getProtocolType(); - Optional checker = DatabaseTypedSPILoader.findService(DialectJDBCResultMetadataChecker.class, databaseType); + Optional checker = DatabaseTypedSPILoader.findService(DialectResultSetMetadataChecker.class, databaseType); if (checker.isPresent()) { JDBCExecutionUnit sample = executionGroupContext.getInputGroups().iterator().next().getInputs().get(0); checker.get().check(executionContext, sample.getStorageResource(), sample.getExecutionUnit().getSqlUnit().getSql()); diff --git a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectJDBCResultMetadataChecker.java b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectResultSetMetadataChecker.java similarity index 93% rename from proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectJDBCResultMetadataChecker.java rename to proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectResultSetMetadataChecker.java index 5081dd47f1a31..30c1f2386c9df 100644 --- a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectJDBCResultMetadataChecker.java +++ b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectResultSetMetadataChecker.java @@ -25,10 +25,10 @@ import java.sql.Statement; /** - * Dialect JDBC result metadata checker. + * Dialect ResultSet metadata checker. */ @SingletonSPI -public interface DialectJDBCResultMetadataChecker extends DatabaseTypedSPI { +public interface DialectResultSetMetadataChecker extends DatabaseTypedSPI { /** * Check JDBC result metadata after routed resources are prepared. diff --git a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutorTest.java b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutorTest.java index 4e563f0def4cd..552c4488db50d 100644 --- a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutorTest.java +++ b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutorTest.java @@ -58,7 +58,7 @@ import org.apache.shardingsphere.infra.session.query.QueryContext; import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader; import org.apache.shardingsphere.mode.manager.ContextManager; -import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker; +import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker; import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.ProxyJDBCExecutor; import org.apache.shardingsphere.proxy.backend.connector.jdbc.statement.JDBCBackendStatement; import org.apache.shardingsphere.proxy.backend.connector.sane.DialectSaneQueryResultEngine; @@ -334,14 +334,14 @@ void assertExecuteWithDialectMetadataCheckerException() throws SQLException { JDBCExecutionUnit jdbcExecutionUnit = new JDBCExecutionUnit(executionUnit, ConnectionMode.CONNECTION_STRICTLY, statement); ExecutionGroupContext executionGroupContext = new ExecutionGroupContext<>( Collections.singletonList(new ExecutionGroup<>(Collections.singletonList(jdbcExecutionUnit))), mock(ExecutionGroupReportContext.class)); - DialectJDBCResultMetadataChecker checker = mock(DialectJDBCResultMetadataChecker.class); + DialectResultSetMetadataChecker checker = mock(DialectResultSetMetadataChecker.class); SQLException expected = new SQLException("expected"); doThrow(expected).when(checker).check(executionContext, statement, "SELECT 1"); try ( MockedConstruction ignored = mockConstruction(DriverExecutionPrepareEngine.class, (mock, context) -> when(mock.prepare(anyString(), eq(executionContext), anyCollection(), any(ExecutionGroupReportContext.class))).thenReturn(executionGroupContext)); MockedStatic spiLoader = mockStatic(DatabaseTypedSPILoader.class, CALLS_REAL_METHODS)) { - spiLoader.when(() -> DatabaseTypedSPILoader.findService(DialectJDBCResultMetadataChecker.class, fixtureDatabaseType)).thenReturn(Optional.of(checker)); + spiLoader.when(() -> DatabaseTypedSPILoader.findService(DialectResultSetMetadataChecker.class, fixtureDatabaseType)).thenReturn(Optional.of(checker)); assertThat(assertThrows(SQLException.class, () -> proxySQLExecutor.execute(executionContext)), is(expected)); verify(regularExecutor, never()).execute(any(), any(), anyBoolean(), anyBoolean()); } diff --git a/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataChecker.java b/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLResultSetMetadataChecker.java similarity index 95% rename from proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataChecker.java rename to proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLResultSetMetadataChecker.java index e5a2e9e7e6d4a..3056ac34ee8a4 100644 --- a/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataChecker.java +++ b/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLResultSetMetadataChecker.java @@ -22,7 +22,7 @@ import org.apache.shardingsphere.infra.exception.ShardingSpherePreconditions; import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; -import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker; +import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker; import org.apache.shardingsphere.proxy.backend.postgresql.exception.PostgreSQLCompositeTypeAcrossDataSourcesException; import java.sql.PreparedStatement; @@ -33,9 +33,9 @@ import java.util.Collection; /** - * JDBC result metadata checker for PostgreSQL. + * ResultSet metadata checker for PostgreSQL. */ -public final class PostgreSQLJDBCResultMetadataChecker implements DialectJDBCResultMetadataChecker { +public final class PostgreSQLResultSetMetadataChecker implements DialectResultSetMetadataChecker { @Override public void check(final ExecutionContext executionContext, final Statement statement, final String sql) throws SQLException { diff --git a/proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker b/proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker similarity index 95% rename from proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker rename to proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker index cf750c54ea713..4e6c3292c9072 100644 --- a/proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker +++ b/proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker @@ -15,4 +15,4 @@ # limitations under the License. # -org.apache.shardingsphere.proxy.backend.postgresql.connector.jdbc.PostgreSQLJDBCResultMetadataChecker +org.apache.shardingsphere.proxy.backend.postgresql.connector.jdbc.PostgreSQLResultSetMetadataChecker diff --git a/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataCheckerTest.java b/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataCheckerTest.java index 3e0626f79911d..0cecc9e7191c8 100644 --- a/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataCheckerTest.java +++ b/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataCheckerTest.java @@ -25,7 +25,7 @@ import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; import org.apache.shardingsphere.infra.executor.sql.context.SQLUnit; import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader; -import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker; +import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker; import org.apache.shardingsphere.proxy.backend.postgresql.exception.PostgreSQLCompositeTypeAcrossDataSourcesException; import org.junit.jupiter.api.Test; @@ -50,7 +50,7 @@ class PostgreSQLJDBCResultMetadataCheckerTest { private final DatabaseType databaseType = TypedSPILoader.getService(DatabaseType.class, "PostgreSQL"); - private final DialectJDBCResultMetadataChecker checker = DatabaseTypedSPILoader.getService(DialectJDBCResultMetadataChecker.class, databaseType); + private final DialectResultSetMetadataChecker checker = DatabaseTypedSPILoader.getService(DialectResultSetMetadataChecker.class, databaseType); @Test void assertCheckWithSingleDataSource() throws SQLException { diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java index 3f90d950d5b5b..fd5ec39297b24 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java @@ -24,7 +24,6 @@ import org.apache.shardingsphere.infra.binder.context.aware.ParameterAware; import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext; import org.apache.shardingsphere.infra.binder.engine.SQLBindEngine; -import org.apache.shardingsphere.infra.exception.external.sql.ShardingSphereSQLException; import org.apache.shardingsphere.infra.connection.kernel.KernelProcessor; import org.apache.shardingsphere.infra.exception.ShardingSpherePreconditions; import org.apache.shardingsphere.infra.exception.kernel.metadata.PreparedStatementMetadataResolutionException; @@ -35,7 +34,7 @@ import org.apache.shardingsphere.infra.session.query.QueryContext; import org.apache.shardingsphere.infra.util.close.QuietlyCloser; import org.apache.shardingsphere.proxy.backend.connector.ProxyDatabaseConnectionManager; -import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker; +import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker; import org.apache.shardingsphere.proxy.backend.context.ProxyContext; import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; @@ -43,6 +42,7 @@ import java.sql.PreparedStatement; import java.sql.SQLException; import java.util.List; +import java.util.Optional; /** * Metadata factory for PostgreSQL prepared statements. @@ -81,9 +81,11 @@ public static PreparedStatement load(final ConnectionSession connectionSession, String sql = executionUnit.getSqlUnit().getSql(); PreparedStatement result = connections.iterator().next().prepareStatement(sql); try { - DatabaseTypedSPILoader.getService(DialectJDBCResultMetadataChecker.class, connectionSession.getProtocolType()) - .check(executionContext, result, sql); - } catch (final SQLException | ShardingSphereSQLException ex) { + Optional resultSetMetadataChecker = DatabaseTypedSPILoader.findService(DialectResultSetMetadataChecker.class, connectionSession.getProtocolType()); + if (resultSetMetadataChecker.isPresent()) { + resultSetMetadataChecker.get().check(executionContext, result, sql); + } + } catch (final SQLException ex) { QuietlyCloser.close(result); throw ex; } diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java index 5878fa4bbbfab..e7f01cde613f9 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java @@ -39,7 +39,7 @@ import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader; import org.apache.shardingsphere.mode.manager.ContextManager; import org.apache.shardingsphere.proxy.backend.connector.ProxyDatabaseConnectionManager; -import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectJDBCResultMetadataChecker; +import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker; import org.apache.shardingsphere.proxy.backend.context.ProxyContext; import org.apache.shardingsphere.proxy.backend.postgresql.exception.PostgreSQLCompositeTypeAcrossDataSourcesException; import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; @@ -131,10 +131,10 @@ void assertClosePreparedStatementOnCheckerException() throws SQLException { PostgreSQLServerPreparedStatement preparedStatement = createPreparedStatement(true); when(connectionSession.getProtocolType()).thenReturn(databaseType); PostgreSQLCompositeTypeAcrossDataSourcesException exception = new PostgreSQLCompositeTypeAcrossDataSourcesException(); - DialectJDBCResultMetadataChecker checker = mock(DialectJDBCResultMetadataChecker.class); + DialectResultSetMetadataChecker checker = mock(DialectResultSetMetadataChecker.class); doThrow(exception).when(checker).check(any(ExecutionContext.class), eq(expected), anyString()); try (MockedStatic spiLoader = mockStatic(DatabaseTypedSPILoader.class, CALLS_REAL_METHODS)) { - spiLoader.when(() -> DatabaseTypedSPILoader.getService(DialectJDBCResultMetadataChecker.class, databaseType)).thenReturn(checker); + spiLoader.when(() -> DatabaseTypedSPILoader.getService(DialectResultSetMetadataChecker.class, databaseType)).thenReturn(checker); assertThat(assertThrows(PostgreSQLCompositeTypeAcrossDataSourcesException.class, () -> PostgreSQLPreparedStatementMetadataFactory.load(connectionSession, preparedStatement, PARAMETERS)), is(exception)); verify(expected).close(); From 0e9cd80a7d541227f96bfa6601dff87af864f673 Mon Sep 17 00:00:00 2001 From: zhangliang Date: Sun, 26 Jul 2026 14:14:08 +0800 Subject: [PATCH 5/8] Fix PostgreSQL composite type OID routing - resolve composite type OIDs from the actual routed data source - reject ambiguous composite OIDs across multiple data sources - document error 30031 and fix column metadata argument order --- .../frontend/postgresql/command/query/extended/Portal.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java index 654444c152fc9..2ab80658f7d9f 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java @@ -116,8 +116,7 @@ public Portal(final String name, final PostgreSQLServerPreparedStatement prepare public void bind() throws SQLException { responseHeader = proxyBackendHandler.execute(); if (responseHeader instanceof QueryResponseHeader) { - QueryResponseHeader queryResponseHeader = (QueryResponseHeader) responseHeader; - columnTypeOIDs = PostgreSQLColumnTypeOIDLoader.load(databaseConnectionManager.getConnectionSession(), queryContext, queryResponseHeader.getQueryHeaders()); + columnTypeOIDs = PostgreSQLColumnTypeOIDLoader.load(databaseConnectionManager.getConnectionSession(), queryContext, ((QueryResponseHeader) responseHeader).getQueryHeaders()); } } From ea2a12241ae2dfbbbcb49951c66c2fc8a0621096 Mon Sep 17 00:00:00 2001 From: zhangliang Date: Sun, 26 Jul 2026 14:23:47 +0800 Subject: [PATCH 6/8] Fix PostgreSQL composite type OID routing - resolve composite type OIDs from the actual routed data source - reject ambiguous composite OIDs across multiple data sources - document error 30031 and fix column metadata argument order --- .../describe/PostgreSQLComDescribeExecutor.java | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutor.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutor.java index c075108bf1546..0346e9a2af080 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutor.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutor.java @@ -152,8 +152,8 @@ private Collection createReturningColumnDescription Collection result = new LinkedList<>(); for (ProjectionSegment each : returningSegment.getProjections().getProjections()) { if (each instanceof ShorthandProjectionSegment) { - table.getAllColumns().stream().map(column -> new PostgreSQLColumnDescription(column.getName(), 0, column.getDataType(), estimateColumnLength(column.getDataType()), "")) - .forEach(result::add); + table.getAllColumns().stream() + .map(column -> new PostgreSQLColumnDescription(column.getName(), 0, column.getDataType(), estimateColumnLength(column.getDataType()), "")).forEach(result::add); } if (each instanceof ColumnProjectionSegment) { ColumnProjectionSegment segment = (ColumnProjectionSegment) each; @@ -275,13 +275,11 @@ private Collection applyColumnTypeOIDs(final Collec if (columnTypeOIDs.isEmpty()) { return columnDescriptions; } - Collection result = new ArrayList<>(columnDescriptions.size()); + Collection result = new LinkedList<>(); int columnIndex = 1; for (PostgreSQLColumnDescription each : columnDescriptions) { Integer typeOID = columnTypeOIDs.get(columnIndex++); - result.add(null != typeOID - ? new PostgreSQLColumnDescription(each.getColumnName(), each.getColumnIndex(), typeOID, each.getColumnLength(), each.getDataFormat()) - : each); + result.add(null == typeOID ? each : new PostgreSQLColumnDescription(each.getColumnName(), each.getColumnIndex(), typeOID, each.getColumnLength(), each.getDataFormat())); } return result; } From 1addf6db66bc2c241608e619b89ea461ba8819e2 Mon Sep 17 00:00:00 2001 From: zhangliang Date: Sun, 26 Jul 2026 16:13:38 +0800 Subject: [PATCH 7/8] Fix PostgreSQL composite type OID routing - resolve composite type OIDs from the actual routed data source - reject ambiguous composite OIDs across multiple data sources - document error 30031 and fix column metadata argument order --- .../StandardDatabaseProxyConnector.java | 8 +-- .../header/query/QueryResponseHeader.java | 11 +++ .../StandardDatabaseProxyConnectorTest.java | 7 +- .../query/PostgreSQLColumnTypeOIDLoader.java | 23 ++----- .../command/query/extended/Portal.java | 6 +- ...reSQLPreparedStatementMetadataFactory.java | 6 +- .../simple/PostgreSQLComQueryExecutor.java | 2 +- .../PostgreSQLColumnTypeOIDLoaderTest.java | 67 +++++-------------- .../command/query/extended/PortalTest.java | 5 +- ...LPreparedStatementMetadataFactoryTest.java | 5 +- .../PostgreSQLComQueryExecutorTest.java | 10 +-- 11 files changed, 58 insertions(+), 92 deletions(-) diff --git a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnector.java b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnector.java index 812d30fb76b71..fdc83d2d89530 100644 --- a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnector.java +++ b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnector.java @@ -249,7 +249,7 @@ private ResponseHeader doExecute(final ExecutionContext executionContext) throws } Object executeResultSample = executeResults.iterator().next(); return executeResultSample instanceof QueryResult - ? processExecuteQuery(queryContext.getSqlStatementContext(), executeResults.stream().map(QueryResult.class::cast).collect(Collectors.toList()), (QueryResult) executeResultSample) + ? processExecuteQuery(executionContext, executeResults.stream().map(QueryResult.class::cast).collect(Collectors.toList()), (QueryResult) executeResultSample) : processExecuteUpdate(executeResults.stream().map(UpdateResult.class::cast).collect(Collectors.toList())); } @@ -291,10 +291,10 @@ private ResponseHeader processExecuteQueryFederation(final ResultSet resultSet) return new QueryResponseHeader(queryHeaders); } - private QueryResponseHeader processExecuteQuery(final SQLStatementContext sqlStatementContext, final List queryResults, final QueryResult queryResultSample) throws SQLException { - queryHeaders = createQueryHeaders(sqlStatementContext, queryResultSample); + private QueryResponseHeader processExecuteQuery(final ExecutionContext executionContext, final List queryResults, final QueryResult queryResultSample) throws SQLException { + queryHeaders = createQueryHeaders(executionContext.getSqlStatementContext(), queryResultSample); mergedResult = mergeQuery(queryResults); - return new QueryResponseHeader(queryHeaders); + return new QueryResponseHeader(queryHeaders, executionContext.getExecutionUnits().iterator().next().getDataSourceName()); } private List createQueryHeaders(final SQLStatementContext sqlStatementContext, final QueryResult queryResultSample) throws SQLException { diff --git a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/response/header/query/QueryResponseHeader.java b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/response/header/query/QueryResponseHeader.java index aacbaba30bcef..d67edf9d5688f 100644 --- a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/response/header/query/QueryResponseHeader.java +++ b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/response/header/query/QueryResponseHeader.java @@ -31,4 +31,15 @@ public final class QueryResponseHeader implements ResponseHeader { private final List queryHeaders; + + private final String dataSourceName; + + /** + * Constructs a query response header without routed data sources. + * + * @param queryHeaders query headers + */ + public QueryResponseHeader(final List queryHeaders) { + this(queryHeaders, null); + } } diff --git a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnectorTest.java b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnectorTest.java index fd947f30c34aa..abb9ce0050a9b 100644 --- a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnectorTest.java +++ b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnectorTest.java @@ -489,7 +489,9 @@ void assertExecuteWithQueryResult() throws SQLException { DatabaseProxyConnector engine = createDatabaseProxyConnector(JDBCDriverType.STATEMENT, queryContext); setField(engine, "proxySQLExecutor", proxySQLExecutor); ExecutionContext executionContext = mock(ExecutionContext.class, RETURNS_DEEP_STUBS); - when(executionContext.getExecutionUnits()).thenReturn(Collections.singletonList(mock(ExecutionUnit.class))); + ExecutionUnit executionUnit = mock(ExecutionUnit.class); + when(executionUnit.getDataSourceName()).thenReturn("ds_0"); + when(executionContext.getExecutionUnits()).thenReturn(Collections.singletonList(executionUnit)); when(executionContext.getSqlStatementContext()).thenReturn(sqlStatementContext); QueryResult queryResult = mock(QueryResult.class); QueryResultMetaData queryResultMetaData = mock(QueryResultMetaData.class); @@ -517,7 +519,8 @@ void assertExecuteWithQueryResult() throws SQLException { MockedStatic serviceLoader = mockStatic(ShardingSphereServiceLoader.class)) { spiLoader.when(() -> DatabaseTypedSPILoader.getService(eq(QueryHeaderBuilder.class), any(DatabaseType.class))).thenReturn(new QueryHeaderBuilderFixture()); serviceLoader.when(() -> ShardingSphereServiceLoader.getServiceInstances(AdvancedProxySQLExecutor.class)).thenReturn(Collections.emptyList()); - assertThat(engine.execute(), isA(QueryResponseHeader.class)); + QueryResponseHeader actual = (QueryResponseHeader) engine.execute(); + assertThat(actual.getDataSourceName(), is("ds_0")); assertThat(mockedKernelProcessor.constructed().size(), is(1)); assertThat(mockedDatabaseTypeRegistry.constructed().size(), is(1)); assertThat(mockedMergeEngine.constructed().size(), is(1)); diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java index bf940626c000f..c9e3aac78c18a 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java @@ -20,12 +20,9 @@ import lombok.AccessLevel; import lombok.NoArgsConstructor; import org.apache.shardingsphere.infra.annotation.HighFrequencyInvocation; -import org.apache.shardingsphere.infra.connection.kernel.KernelProcessor; -import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; -import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; import org.apache.shardingsphere.infra.executor.sql.execute.engine.ConnectionMode; -import org.apache.shardingsphere.infra.session.query.QueryContext; import org.apache.shardingsphere.proxy.backend.response.header.query.QueryHeader; +import org.apache.shardingsphere.proxy.backend.response.header.query.QueryResponseHeader; import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; import org.postgresql.core.BaseConnection; import org.postgresql.core.Oid; @@ -51,27 +48,21 @@ public final class PostgreSQLColumnTypeOIDLoader { * Load composite column type OIDs from the routed backend connection. * * @param connectionSession connection session - * @param queryContext query context - * @param queryHeaders query headers + * @param queryResponseHeader query response header * @return column indexes to type OIDs, or an empty map if no composite column type can be resolved * @throws SQLException SQL exception */ - public static Map load(final ConnectionSession connectionSession, final QueryContext queryContext, final List queryHeaders) throws SQLException { + public static Map load(final ConnectionSession connectionSession, final QueryResponseHeader queryResponseHeader) throws SQLException { + List queryHeaders = queryResponseHeader.getQueryHeaders(); if (!containsCompositeType(queryHeaders)) { return Collections.emptyMap(); } - ExecutionContext executionContext = new KernelProcessor().generateExecutionContext(queryContext, queryContext.getMetaData().getGlobalRuleMetaData(), queryContext.getMetaData().getProps()); - if (executionContext.getExecutionUnits().isEmpty()) { + String dataSourceName = queryResponseHeader.getDataSourceName(); + if (null == dataSourceName) { return Collections.emptyMap(); } - ExecutionUnit executionUnit = executionContext.getExecutionUnits().iterator().next(); - for (ExecutionUnit each : executionContext.getExecutionUnits()) { - if (!executionUnit.getDataSourceName().equals(each.getDataSourceName())) { - return Collections.emptyMap(); - } - } List connections = connectionSession.getDatabaseConnectionManager().getConnections( - connectionSession.getUsedDatabaseName(), executionUnit.getDataSourceName(), 0, 1, ConnectionMode.CONNECTION_STRICTLY); + connectionSession.getUsedDatabaseName(), dataSourceName, 0, 1, ConnectionMode.CONNECTION_STRICTLY); return load(connections.get(0), queryHeaders); } diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java index 2ab80658f7d9f..d06c8de1eecb4 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java @@ -83,8 +83,6 @@ public final class Portal { private final ProxyBackendHandler proxyBackendHandler; - private final QueryContext queryContext; - private final ProxyDatabaseConnectionManager databaseConnectionManager; private Map columnTypeOIDs; @@ -103,7 +101,7 @@ public Portal(final String name, final PostgreSQLServerPreparedStatement prepare ((ParameterAware) sqlStatementContext).bindParameters(params); } DatabaseType protocolType = ProxyContext.getInstance().getContextManager().getDatabase(databaseName).getProtocolType(); - queryContext = new QueryContext(sqlStatementContext, preparedStatement.getSql(), params, preparedStatement.getHintValueContext(), + QueryContext queryContext = new QueryContext(sqlStatementContext, preparedStatement.getSql(), params, preparedStatement.getHintValueContext(), databaseConnectionManager.getConnectionSession().getConnectionContext(), ProxyContext.getInstance().getContextManager().getMetaDataContexts().getMetaData(), true); proxyBackendHandler = ProxyBackendHandlerFactory.newInstance(protocolType, queryContext, databaseConnectionManager.getConnectionSession(), true); } @@ -116,7 +114,7 @@ public Portal(final String name, final PostgreSQLServerPreparedStatement prepare public void bind() throws SQLException { responseHeader = proxyBackendHandler.execute(); if (responseHeader instanceof QueryResponseHeader) { - columnTypeOIDs = PostgreSQLColumnTypeOIDLoader.load(databaseConnectionManager.getConnectionSession(), queryContext, ((QueryResponseHeader) responseHeader).getQueryHeaders()); + columnTypeOIDs = PostgreSQLColumnTypeOIDLoader.load(databaseConnectionManager.getConnectionSession(), (QueryResponseHeader) responseHeader); } } diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java index fd5ec39297b24..e8410c318d0a5 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java @@ -26,6 +26,7 @@ import org.apache.shardingsphere.infra.binder.engine.SQLBindEngine; import org.apache.shardingsphere.infra.connection.kernel.KernelProcessor; import org.apache.shardingsphere.infra.exception.ShardingSpherePreconditions; +import org.apache.shardingsphere.infra.exception.external.sql.ShardingSphereSQLException; import org.apache.shardingsphere.infra.exception.kernel.metadata.PreparedStatementMetadataResolutionException; import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; @@ -81,11 +82,12 @@ public static PreparedStatement load(final ConnectionSession connectionSession, String sql = executionUnit.getSqlUnit().getSql(); PreparedStatement result = connections.iterator().next().prepareStatement(sql); try { - Optional resultSetMetadataChecker = DatabaseTypedSPILoader.findService(DialectResultSetMetadataChecker.class, connectionSession.getProtocolType()); + Optional resultSetMetadataChecker = + DatabaseTypedSPILoader.findService(DialectResultSetMetadataChecker.class, sqlStatementContext.getSqlStatement().getDatabaseType()); if (resultSetMetadataChecker.isPresent()) { resultSetMetadataChecker.get().check(executionContext, result, sql); } - } catch (final SQLException ex) { + } catch (final SQLException | ShardingSphereSQLException ex) { QuietlyCloser.close(result); throw ex; } diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutor.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutor.java index b344a6396983c..7fa2ac6430423 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutor.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutor.java @@ -90,7 +90,7 @@ public Collection execute() throws SQLException { ResponseHeader responseHeader = proxyBackendHandler.execute(); if (responseHeader instanceof QueryResponseHeader) { QueryResponseHeader queryResponseHeader = (QueryResponseHeader) responseHeader; - Map columnTypeOIDs = PostgreSQLColumnTypeOIDLoader.load(connectionSession, connectionSession.getQueryContext(), queryResponseHeader.getQueryHeaders()); + Map columnTypeOIDs = PostgreSQLColumnTypeOIDLoader.load(connectionSession, queryResponseHeader); return Collections.singleton(createRowDescriptionPacket(queryResponseHeader, columnTypeOIDs)); } responseType = ResponseType.UPDATE; diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java index 9bfdc4347f7ee..5512c9d7faa08 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java @@ -17,16 +17,12 @@ package org.apache.shardingsphere.proxy.frontend.postgresql.command.query; -import org.apache.shardingsphere.infra.connection.kernel.KernelProcessor; -import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; -import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; -import org.apache.shardingsphere.infra.executor.sql.context.SQLUnit; -import org.apache.shardingsphere.infra.session.query.QueryContext; import org.apache.shardingsphere.proxy.backend.connector.ProxyDatabaseConnectionManager; import org.apache.shardingsphere.proxy.backend.response.header.query.QueryHeader; +import org.apache.shardingsphere.proxy.backend.response.header.query.QueryResponseHeader; import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; +import org.apache.shardingsphere.infra.executor.sql.execute.engine.ConnectionMode; import org.junit.jupiter.api.Test; -import org.mockito.MockedConstruction; import org.postgresql.core.BaseConnection; import org.postgresql.core.Oid; import org.postgresql.core.TypeInfo; @@ -35,8 +31,6 @@ import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Types; -import java.util.Arrays; -import java.util.Collection; import java.util.Collections; import java.util.Map; @@ -46,9 +40,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.Mockito.RETURNS_DEEP_STUBS; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.mockConstruction; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -57,56 +49,37 @@ class PostgreSQLColumnTypeOIDLoaderTest { @Test void assertLoadWithSameDataSource() throws SQLException { - QueryContext queryContext = mock(QueryContext.class, RETURNS_DEEP_STUBS); - ExecutionContext executionContext = createExecutionContext(Arrays.asList( - new ExecutionUnit("ds_0", new SQLUnit("SELECT record_value", Collections.emptyList())), - new ExecutionUnit("ds_0", new SQLUnit("SELECT record_value", Collections.emptyList())))); BaseConnection connection = mock(BaseConnection.class); mockTypeOID(connection, "record_type", 2249); ConnectionSession connectionSession = createConnectionSession(connection); - try ( - MockedConstruction ignored = mockConstruction(KernelProcessor.class, - (mock, context) -> when(mock.generateExecutionContext(any(), any(), any())).thenReturn(executionContext))) { - Map actual = PostgreSQLColumnTypeOIDLoader.load( - connectionSession, queryContext, Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type"))); - assertThat(actual.get(1), is(2249)); - } + QueryResponseHeader queryResponseHeader = new QueryResponseHeader( + Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type")), "ds_0"); + Map actual = PostgreSQLColumnTypeOIDLoader.load(connectionSession, queryResponseHeader); + assertThat(actual.get(1), is(2249)); + verify(connectionSession.getDatabaseConnectionManager()).getConnections("postgres", "ds_0", 0, 1, ConnectionMode.CONNECTION_STRICTLY); } @Test void assertLoadWithoutCompositeColumn() throws SQLException { - assertTrue(PostgreSQLColumnTypeOIDLoader.load( - mock(ConnectionSession.class), mock(QueryContext.class), Collections.singletonList(createQueryHeader(Types.VARCHAR, "varchar"))).isEmpty()); + QueryResponseHeader queryResponseHeader = new QueryResponseHeader( + Collections.singletonList(createQueryHeader(Types.VARCHAR, "varchar")), "ds_0"); + assertTrue(PostgreSQLColumnTypeOIDLoader.load(mock(ConnectionSession.class), queryResponseHeader).isEmpty()); } @Test - void assertLoadWithMultipleDataSources() throws SQLException { - ExecutionContext executionContext = createExecutionContext(Arrays.asList( - new ExecutionUnit("ds_0", new SQLUnit("SELECT record_value", Collections.emptyList())), - new ExecutionUnit("ds_1", new SQLUnit("SELECT record_value", Collections.emptyList())))); - try ( - MockedConstruction ignored = mockConstruction(KernelProcessor.class, - (mock, context) -> when(mock.generateExecutionContext(any(), any(), any())).thenReturn(executionContext))) { - assertTrue(PostgreSQLColumnTypeOIDLoader.load( - mock(ConnectionSession.class), mock(QueryContext.class, RETURNS_DEEP_STUBS), - Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type"))).isEmpty()); - } + void assertLoadWithoutDataSource() throws SQLException { + QueryResponseHeader queryResponseHeader = new QueryResponseHeader(Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type"))); + assertTrue(PostgreSQLColumnTypeOIDLoader.load(mock(ConnectionSession.class), queryResponseHeader).isEmpty()); } @Test void assertLoadWithoutPgjdbcConnection() throws SQLException { - ExecutionContext executionContext = createExecutionContext( - Collections.singletonList(new ExecutionUnit("ds_0", new SQLUnit("SELECT record_value", Collections.emptyList())))); Connection connection = mock(Connection.class); ConnectionSession connectionSession = createConnectionSession(connection); - try ( - MockedConstruction ignored = mockConstruction(KernelProcessor.class, - (mock, context) -> when(mock.generateExecutionContext(any(), any(), any())).thenReturn(executionContext))) { - assertTrue(PostgreSQLColumnTypeOIDLoader.load( - connectionSession, mock(QueryContext.class, RETURNS_DEEP_STUBS), - Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type"))).isEmpty()); - verify(connection, never()).unwrap(BaseConnection.class); - } + QueryResponseHeader queryResponseHeader = new QueryResponseHeader( + Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type")), "ds_0"); + assertTrue(PostgreSQLColumnTypeOIDLoader.load(connectionSession, queryResponseHeader).isEmpty()); + verify(connection, never()).unwrap(BaseConnection.class); } @Test @@ -139,12 +112,6 @@ private QueryHeader createQueryHeader(final int columnType, final String columnT return new QueryHeader("", "", "record_value", "record_value", columnType, columnTypeName, -1, 0, false, false, false, false); } - private ExecutionContext createExecutionContext(final Collection executionUnits) { - ExecutionContext result = mock(ExecutionContext.class); - when(result.getExecutionUnits()).thenReturn(executionUnits); - return result; - } - private ConnectionSession createConnectionSession(final Connection connection) throws SQLException { ConnectionSession result = mock(ConnectionSession.class); ProxyDatabaseConnectionManager connectionManager = mock(ProxyDatabaseConnectionManager.class); diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PortalTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PortalTest.java index 1d4acacb35d79..4cda63c32f03f 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PortalTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PortalTest.java @@ -90,7 +90,6 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyBoolean; -import static org.mockito.ArgumentMatchers.anyList; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.RETURNS_DEEP_STUBS; @@ -127,7 +126,7 @@ void setup() throws SQLException { when(ProxyBackendHandlerFactory.newInstance(eq(databaseType), anyString(), any(SQLStatement.class), eq(connectionSession), any(HintValueContext.class))).thenReturn(proxyBackendHandler); when(ProxyBackendHandlerFactory.newInstance(eq(databaseType), any(QueryContext.class), eq(connectionSession), anyBoolean())).thenReturn(proxyBackendHandler); when(databaseConnectionManager.getConnectionSession()).thenReturn(connectionSession); - when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), any(QueryContext.class), anyList())).thenReturn(Collections.emptyMap()); + when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), any(QueryResponseHeader.class))).thenReturn(Collections.emptyMap()); } @Test @@ -154,7 +153,7 @@ void assertExecuteSelectStatementAndReturnAllRows() throws SQLException, Reflect Map columnTypeOIDs = new HashMap<>(2, 1F); columnTypeOIDs.put(1, 2249); columnTypeOIDs.put(2, 2249); - when(PostgreSQLColumnTypeOIDLoader.load(eq(databaseConnectionManager.getConnectionSession()), any(QueryContext.class), anyList())) + when(PostgreSQLColumnTypeOIDLoader.load(eq(databaseConnectionManager.getConnectionSession()), any(QueryResponseHeader.class))) .thenReturn(columnTypeOIDs); Portal portal = createPortal(sqlStatementContext, Arrays.asList(PostgreSQLValueFormat.TEXT, PostgreSQLValueFormat.BINARY)); PostgreSQLPacket portalDescription = portal.describe(); diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java index e7f01cde613f9..2fc19d00c3a57 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java @@ -62,6 +62,7 @@ import java.sql.Types; import java.util.Collections; import java.util.List; +import java.util.Optional; import java.util.Properties; import static org.hamcrest.MatcherAssert.assertThat; @@ -98,7 +99,6 @@ class PostgreSQLPreparedStatementMetadataFactoryTest { @Test void assertLoad() throws SQLException { PreparedStatement expected = prepareJDBCBackendConnection(null); - when(connectionSession.getProtocolType()).thenReturn(databaseType); assertThat(PostgreSQLPreparedStatementMetadataFactory.load(connectionSession, createPreparedStatement(true), PARAMETERS), is(expected)); } @@ -129,12 +129,11 @@ void assertLoadWithEmptyExecutionUnits() { void assertClosePreparedStatementOnCheckerException() throws SQLException { PreparedStatement expected = prepareJDBCBackendConnection(null); PostgreSQLServerPreparedStatement preparedStatement = createPreparedStatement(true); - when(connectionSession.getProtocolType()).thenReturn(databaseType); PostgreSQLCompositeTypeAcrossDataSourcesException exception = new PostgreSQLCompositeTypeAcrossDataSourcesException(); DialectResultSetMetadataChecker checker = mock(DialectResultSetMetadataChecker.class); doThrow(exception).when(checker).check(any(ExecutionContext.class), eq(expected), anyString()); try (MockedStatic spiLoader = mockStatic(DatabaseTypedSPILoader.class, CALLS_REAL_METHODS)) { - spiLoader.when(() -> DatabaseTypedSPILoader.getService(DialectResultSetMetadataChecker.class, databaseType)).thenReturn(checker); + spiLoader.when(() -> DatabaseTypedSPILoader.findService(DialectResultSetMetadataChecker.class, databaseType)).thenReturn(Optional.of(checker)); assertThat(assertThrows(PostgreSQLCompositeTypeAcrossDataSourcesException.class, () -> PostgreSQLPreparedStatementMetadataFactory.load(connectionSession, preparedStatement, PARAMETERS)), is(exception)); verify(expected).close(); diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutorTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutorTest.java index 5123ab125b3cd..eb228f4478cb5 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutorTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutorTest.java @@ -29,7 +29,6 @@ import org.apache.shardingsphere.database.protocol.postgresql.packet.handshake.PostgreSQLParameterStatusPacket; import org.apache.shardingsphere.database.exception.core.exception.data.InvalidParameterValueException; import org.apache.shardingsphere.infra.hint.HintValueContext; -import org.apache.shardingsphere.infra.session.query.QueryContext; import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader; import org.apache.shardingsphere.proxy.backend.handler.ProxyBackendHandler; import org.apache.shardingsphere.proxy.backend.handler.ProxyBackendHandlerFactory; @@ -79,7 +78,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -114,16 +113,13 @@ void setUp() throws SQLException { SQLStatement sqlStatement = new SQLStatement(DATABASE_TYPE); when(ProxySQLComQueryParser.parse(queryPacket.getSQL(), DATABASE_TYPE, connectionSession)).thenReturn(sqlStatement); when(ProxyBackendHandlerFactory.newInstance(DATABASE_TYPE, queryPacket.getSQL(), sqlStatement, connectionSession, queryPacket.getHintValueContext())).thenReturn(proxyBackendHandler); - QueryContext queryContext = mock(QueryContext.class); - when(connectionSession.getQueryContext()).thenReturn(queryContext); - when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), eq(queryContext), anyList())).thenReturn(Collections.emptyMap()); + when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), any(QueryResponseHeader.class))).thenReturn(Collections.emptyMap()); queryExecutor = new PostgreSQLComQueryExecutor(portalContext, queryPacket, connectionSession); } @Test void assertExecuteQueryWithColumnDescription() throws SQLException, ReflectiveOperationException { - QueryContext queryContext = connectionSession.getQueryContext(); - when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), eq(queryContext), anyList())).thenReturn(Collections.singletonMap(1, 2249)); + when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), any(QueryResponseHeader.class))).thenReturn(Collections.singletonMap(1, 2249)); QueryResponseHeader queryResponseHeader = mock(QueryResponseHeader.class); when(queryResponseHeader.getQueryHeaders()).thenReturn( Collections.singletonList(new QueryHeader("schema", "table", "label", "column", Types.STRUCT, "record_type", 2, 3, true, true, true, true))); From fc14554cc7851b5ee1a003ac4c5d48edc70c52aa Mon Sep 17 00:00:00 2001 From: zhangliang Date: Sun, 26 Jul 2026 20:29:44 +0800 Subject: [PATCH 8/8] Fix PostgreSQL composite type OID routing - resolve composite type OIDs from the actual routed data source - reject ambiguous composite OIDs across multiple data sources - document error 30031 and fix column metadata argument order --- .../error-code/sql-error-code.cn.md | 23 ++-- .../error-code/sql-error-code.en.md | 23 ++-- .../backend/connector/ProxySQLExecutor.java | 7 - .../StandardDatabaseProxyConnector.java | 14 +- .../DialectResultSetMetadataChecker.java | 42 ------ .../UnicastDatabaseProxyBackendHandler.java | 14 ++ .../header/query/QueryResponseHeader.java | 11 -- .../connector/ProxySQLExecutorTest.java | 32 ----- .../StandardDatabaseProxyConnectorTest.java | 7 +- .../PostgreSQLResultSetMetadataChecker.java | 80 ----------- ...mpositeTypeAcrossDataSourcesException.java | 33 ----- ...c.executor.DialectResultSetMetadataChecker | 18 --- ...stgreSQLJDBCResultMetadataCheckerTest.java | 129 ------------------ ...iteTypeAcrossDataSourcesExceptionTest.java | 37 ----- .../query/PostgreSQLColumnTypeOIDLoader.java | 40 ++++-- .../command/query/extended/Portal.java | 3 +- ...reSQLPreparedStatementMetadataFactory.java | 30 ++-- .../PostgreSQLComDescribeExecutor.java | 61 ++++----- .../simple/PostgreSQLComQueryExecutor.java | 2 +- .../PostgreSQLColumnTypeOIDLoaderTest.java | 97 ++++++------- .../command/query/extended/PortalTest.java | 6 +- ...LPreparedStatementMetadataFactoryTest.java | 33 +++-- .../PostgreSQLComDescribeExecutorTest.java | 14 +- .../PostgreSQLComQueryExecutorTest.java | 6 +- 24 files changed, 199 insertions(+), 563 deletions(-) delete mode 100644 proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectResultSetMetadataChecker.java delete mode 100644 proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLResultSetMetadataChecker.java delete mode 100644 proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesException.java delete mode 100644 proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker delete mode 100644 proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataCheckerTest.java delete mode 100644 proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesExceptionTest.java diff --git a/docs/document/content/user-manual/error-code/sql-error-code.cn.md b/docs/document/content/user-manual/error-code/sql-error-code.cn.md index d572274e3de9b..bb0cf35bb0843 100644 --- a/docs/document/content/user-manual/error-code/sql-error-code.cn.md +++ b/docs/document/content/user-manual/error-code/sql-error-code.cn.md @@ -247,15 +247,14 @@ SQL 错误码以标准的 SQL State，Vendor Code 和详细错误信息提供， ## 其他异常 -| Vendor Code | SQL State | 错误信息 | -|-------------|-----------|------------------------------------------------------------------------------------------------| -| 30000 | HY000 | Unknown exception: %s | -| 30001 | 0A000 | Unsupported SQL operation: %s | -| 30002 | HY000 | Database protocol exception: %s | -| 30003 | 0A000 | Unsupported command: %s | -| 30004 | HY000 | Server exception: %s | -| 30005 | HY000 | Underlying SQL state: %s, underlying error code: %s. | -| 30010 | HY000 | Can not find plugin class '%s'. | -| 30020 | HY000 | File access failed, file is: %s | -| 30030 | 0A000 | No tableless route info found. | -| 30031 | 0A000 | PostgreSQL composite result columns are not supported for SQL routed to multiple data sources. | +| Vendor Code | SQL State | 错误信息 | +|-------------|-----------|------------------------------------------------------| +| 30000 | HY000 | Unknown exception: %s | +| 30001 | 0A000 | Unsupported SQL operation: %s | +| 30002 | HY000 | Database protocol exception: %s | +| 30003 | 0A000 | Unsupported command: %s | +| 30004 | HY000 | Server exception: %s | +| 30005 | HY000 | Underlying SQL state: %s, underlying error code: %s. | +| 30010 | HY000 | Can not find plugin class '%s'. | +| 30020 | HY000 | File access failed, file is: %s | +| 30030 | 0A000 | No tableless route info found. | diff --git a/docs/document/content/user-manual/error-code/sql-error-code.en.md b/docs/document/content/user-manual/error-code/sql-error-code.en.md index 3b04c87a22d88..7d35fa946573c 100644 --- a/docs/document/content/user-manual/error-code/sql-error-code.en.md +++ b/docs/document/content/user-manual/error-code/sql-error-code.en.md @@ -248,15 +248,14 @@ SQL error codes provide by standard `SQL State`, `Vendor Code` and `Reason`, whi ## Other Exception -| Vendor Code | SQL State | Reason | -|-------------|-----------|------------------------------------------------------------------------------------------------| -| 30000 | HY000 | Unknown exception: %s | -| 30001 | 0A000 | Unsupported SQL operation: %s | -| 30002 | HY000 | Database protocol exception: %s | -| 30003 | 0A000 | Unsupported command: %s | -| 30004 | HY000 | Server exception: %s | -| 30005 | HY000 | Underlying SQL state: %s, underlying error code: %s. | -| 30010 | HY000 | Can not find plugin class '%s'. | -| 30020 | HY000 | File access failed, file is: %s | -| 30030 | 0A000 | No tableless route info found. | -| 30031 | 0A000 | PostgreSQL composite result columns are not supported for SQL routed to multiple data sources. | +| Vendor Code | SQL State | Reason | +|-------------|-----------|------------------------------------------------------| +| 30000 | HY000 | Unknown exception: %s | +| 30001 | 0A000 | Unsupported SQL operation: %s | +| 30002 | HY000 | Database protocol exception: %s | +| 30003 | 0A000 | Unsupported command: %s | +| 30004 | HY000 | Server exception: %s | +| 30005 | HY000 | Underlying SQL state: %s, underlying error code: %s. | +| 30010 | HY000 | Can not find plugin class '%s'. | +| 30020 | HY000 | File access failed, file is: %s | +| 30030 | 0A000 | No tableless route info found. | diff --git a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutor.java b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutor.java index a150cea3c8905..4e9a3b9c4e206 100644 --- a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutor.java +++ b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutor.java @@ -55,7 +55,6 @@ import org.apache.shardingsphere.infra.session.connection.transaction.TransactionConnectionContext; import org.apache.shardingsphere.infra.spi.type.ordered.OrderedSPILoader; import org.apache.shardingsphere.mode.metadata.MetaDataContexts; -import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker; import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.ProxyJDBCExecutor; import org.apache.shardingsphere.proxy.backend.connector.jdbc.statement.JDBCBackendStatement; import org.apache.shardingsphere.proxy.backend.connector.sane.DialectSaneQueryResultEngine; @@ -239,12 +238,6 @@ type, maxConnectionsSizePerQuery, databaseConnectionManager, statementManager, n } catch (final SQLException ex) { return getSaneExecuteResults(executionContext, ex); } - DatabaseType databaseType = ProxyContext.getInstance().getContextManager().getDatabase(databaseName).getProtocolType(); - Optional checker = DatabaseTypedSPILoader.findService(DialectResultSetMetadataChecker.class, databaseType); - if (checker.isPresent()) { - JDBCExecutionUnit sample = executionGroupContext.getInputGroups().iterator().next().getInputs().get(0); - checker.get().check(executionContext, sample.getStorageResource(), sample.getExecutionUnit().getSqlUnit().getSql()); - } executeTransactionHooksBeforeExecuteSQL(databaseConnectionManager.getConnectionSession()); return regularExecutor.execute(executionContext.getQueryContext(), executionGroupContext, isReturnGeneratedKeys, isExceptionThrown); } diff --git a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnector.java b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnector.java index fdc83d2d89530..6ff072951b1a7 100644 --- a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnector.java +++ b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnector.java @@ -17,6 +17,7 @@ package org.apache.shardingsphere.proxy.backend.connector; +import lombok.Getter; import org.apache.shardingsphere.database.connector.core.metadata.database.metadata.DialectDatabaseMetaData; import org.apache.shardingsphere.database.connector.core.metadata.database.metadata.option.transaction.DDLCommitPolicy; import org.apache.shardingsphere.database.connector.core.metadata.database.metadata.option.transaction.DialectTransactionOption; @@ -36,6 +37,7 @@ import org.apache.shardingsphere.infra.exception.kernel.metadata.resource.storageunit.EmptyStorageUnitException; import org.apache.shardingsphere.infra.exception.kernel.metadata.rule.EmptyRuleException; import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; import org.apache.shardingsphere.infra.executor.sql.execute.engine.SQLExecutorExceptionHandler; import org.apache.shardingsphere.infra.executor.sql.execute.engine.driver.jdbc.JDBCExecutionUnit; import org.apache.shardingsphere.infra.executor.sql.execute.result.ExecuteResult; @@ -125,6 +127,9 @@ public final class StandardDatabaseProxyConnector implements DatabaseProxyConnec private MergedResult mergedResult; + @Getter + private Collection executionUnits = Collections.emptyList(); + public StandardDatabaseProxyConnector(final JDBCDriverType driverType, final QueryContext queryContext, final ProxyDatabaseConnectionManager databaseConnectionManager) { this.driverType = driverType; this.queryContext = queryContext; @@ -190,6 +195,7 @@ public ResponseHeader execute() throws SQLException { return new UpdateResponseHeader(queryContext.getSqlStatementContext().getSqlStatement()); } ExecutionContext executionContext = generateExecutionContext(); + executionUnits = executionContext.getExecutionUnits(); return isNeedImplicitCommitTransaction(queryContext.getSqlStatementContext().getSqlStatement(), executionContext.getExecutionUnits().size() > 1) ? doExecuteWithImplicitCommitTransaction(() -> doExecute(executionContext)) : doExecute(executionContext); @@ -249,7 +255,7 @@ private ResponseHeader doExecute(final ExecutionContext executionContext) throws } Object executeResultSample = executeResults.iterator().next(); return executeResultSample instanceof QueryResult - ? processExecuteQuery(executionContext, executeResults.stream().map(QueryResult.class::cast).collect(Collectors.toList()), (QueryResult) executeResultSample) + ? processExecuteQuery(queryContext.getSqlStatementContext(), executeResults.stream().map(QueryResult.class::cast).collect(Collectors.toList()), (QueryResult) executeResultSample) : processExecuteUpdate(executeResults.stream().map(UpdateResult.class::cast).collect(Collectors.toList())); } @@ -291,10 +297,10 @@ private ResponseHeader processExecuteQueryFederation(final ResultSet resultSet) return new QueryResponseHeader(queryHeaders); } - private QueryResponseHeader processExecuteQuery(final ExecutionContext executionContext, final List queryResults, final QueryResult queryResultSample) throws SQLException { - queryHeaders = createQueryHeaders(executionContext.getSqlStatementContext(), queryResultSample); + private QueryResponseHeader processExecuteQuery(final SQLStatementContext sqlStatementContext, final List queryResults, final QueryResult queryResultSample) throws SQLException { + queryHeaders = createQueryHeaders(sqlStatementContext, queryResultSample); mergedResult = mergeQuery(queryResults); - return new QueryResponseHeader(queryHeaders, executionContext.getExecutionUnits().iterator().next().getDataSourceName()); + return new QueryResponseHeader(queryHeaders); } private List createQueryHeaders(final SQLStatementContext sqlStatementContext, final QueryResult queryResultSample) throws SQLException { diff --git a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectResultSetMetadataChecker.java b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectResultSetMetadataChecker.java deleted file mode 100644 index 30c1f2386c9df..0000000000000 --- a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/jdbc/executor/DialectResultSetMetadataChecker.java +++ /dev/null @@ -1,42 +0,0 @@ -/* - * 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.proxy.backend.connector.jdbc.executor; - -import org.apache.shardingsphere.database.connector.core.spi.DatabaseTypedSPI; -import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; -import org.apache.shardingsphere.infra.spi.annotation.SingletonSPI; - -import java.sql.SQLException; -import java.sql.Statement; - -/** - * Dialect ResultSet metadata checker. - */ -@SingletonSPI -public interface DialectResultSetMetadataChecker extends DatabaseTypedSPI { - - /** - * Check JDBC result metadata after routed resources are prepared. - * - * @param executionContext execution context - * @param statement JDBC statement - * @param sql routed SQL corresponding to the statement - * @throws SQLException SQL exception - */ - void check(ExecutionContext executionContext, Statement statement, String sql) throws SQLException; -} diff --git a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/handler/data/type/UnicastDatabaseProxyBackendHandler.java b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/handler/data/type/UnicastDatabaseProxyBackendHandler.java index e54d9285804da..1512934eea4aa 100644 --- a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/handler/data/type/UnicastDatabaseProxyBackendHandler.java +++ b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/handler/data/type/UnicastDatabaseProxyBackendHandler.java @@ -23,10 +23,12 @@ import org.apache.shardingsphere.database.exception.core.exception.syntax.database.NoDatabaseSelectedException; import org.apache.shardingsphere.infra.exception.ShardingSpherePreconditions; import org.apache.shardingsphere.infra.exception.kernel.metadata.resource.storageunit.EmptyStorageUnitException; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; import org.apache.shardingsphere.infra.session.query.QueryContext; import org.apache.shardingsphere.mode.manager.ContextManager; import org.apache.shardingsphere.proxy.backend.connector.DatabaseProxyConnector; import org.apache.shardingsphere.proxy.backend.connector.DatabaseProxyConnectorFactory; +import org.apache.shardingsphere.proxy.backend.connector.StandardDatabaseProxyConnector; import org.apache.shardingsphere.proxy.backend.context.ProxyContext; import org.apache.shardingsphere.proxy.backend.handler.data.DatabaseProxyBackendHandler; import org.apache.shardingsphere.proxy.backend.response.data.QueryResponseRow; @@ -35,6 +37,7 @@ import java.sql.SQLException; import java.util.Collection; +import java.util.Collections; import java.util.Optional; import java.util.stream.Collectors; @@ -90,6 +93,17 @@ public QueryResponseRow getRowData() throws SQLException { return databaseProxyConnector.getRowData(); } + /** + * Get execution units. + * + * @return execution units + */ + public Collection getExecutionUnits() { + return databaseProxyConnector instanceof StandardDatabaseProxyConnector + ? ((StandardDatabaseProxyConnector) databaseProxyConnector).getExecutionUnits() + : Collections.emptyList(); + } + @Override public void close() throws SQLException { if (null != databaseProxyConnector) { diff --git a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/response/header/query/QueryResponseHeader.java b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/response/header/query/QueryResponseHeader.java index d67edf9d5688f..aacbaba30bcef 100644 --- a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/response/header/query/QueryResponseHeader.java +++ b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/response/header/query/QueryResponseHeader.java @@ -31,15 +31,4 @@ public final class QueryResponseHeader implements ResponseHeader { private final List queryHeaders; - - private final String dataSourceName; - - /** - * Constructs a query response header without routed data sources. - * - * @param queryHeaders query headers - */ - public QueryResponseHeader(final List queryHeaders) { - this(queryHeaders, null); - } } diff --git a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutorTest.java b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutorTest.java index 552c4488db50d..9d9a64e68e746 100644 --- a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutorTest.java +++ b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/ProxySQLExecutorTest.java @@ -32,13 +32,9 @@ import org.apache.shardingsphere.infra.binder.context.statement.type.dml.InsertStatementContext; import org.apache.shardingsphere.infra.config.props.ConfigurationPropertyKey; import org.apache.shardingsphere.infra.config.rule.RuleConfiguration; -import org.apache.shardingsphere.infra.executor.kernel.model.ExecutionGroup; import org.apache.shardingsphere.infra.executor.kernel.model.ExecutionGroupContext; import org.apache.shardingsphere.infra.executor.kernel.model.ExecutionGroupReportContext; import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; -import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; -import org.apache.shardingsphere.infra.executor.sql.context.SQLUnit; -import org.apache.shardingsphere.infra.executor.sql.execute.engine.ConnectionMode; import org.apache.shardingsphere.infra.executor.sql.execute.engine.driver.jdbc.JDBCExecutionUnit; import org.apache.shardingsphere.infra.executor.sql.execute.engine.raw.RawExecutor; import org.apache.shardingsphere.infra.executor.sql.execute.engine.raw.RawSQLExecutionUnit; @@ -58,7 +54,6 @@ import org.apache.shardingsphere.infra.session.query.QueryContext; import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader; import org.apache.shardingsphere.mode.manager.ContextManager; -import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker; import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.ProxyJDBCExecutor; import org.apache.shardingsphere.proxy.backend.connector.jdbc.statement.JDBCBackendStatement; import org.apache.shardingsphere.proxy.backend.connector.sane.DialectSaneQueryResultEngine; @@ -103,7 +98,6 @@ import org.mockito.quality.Strictness; import java.sql.SQLException; -import java.sql.Statement; import java.util.Arrays; import java.util.Collection; import java.util.Collections; @@ -122,7 +116,6 @@ import static org.mockito.Mockito.anyBoolean; import static org.mockito.Mockito.anyCollection; import static org.mockito.Mockito.anyString; -import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mockConstruction; @@ -322,31 +315,6 @@ private Stream executeScenarios() { Arguments.of("execute-with-driver-and-no-transaction", false, createInsertStatement(postgresqlDatabaseType), false, false, false)); } - @Test - void assertExecuteWithDialectMetadataCheckerException() throws SQLException { - when(connectionSession.getUsedDatabaseName()).thenReturn("foo_db"); - when(database.getRuleMetaData().getRules()).thenReturn(createRules(false)); - ProxySQLExecutor proxySQLExecutor = createProxySQLExecutor("foo_schema", true); - setExecutorField(proxySQLExecutor, "regularExecutor", regularExecutor); - ExecutionContext executionContext = createExecutionContext("execute-with-dialect-metadata-checker", new SQLStatement(fixtureDatabaseType), false); - Statement statement = mock(Statement.class); - ExecutionUnit executionUnit = new ExecutionUnit("foo_ds", new SQLUnit("SELECT 1", Collections.emptyList())); - JDBCExecutionUnit jdbcExecutionUnit = new JDBCExecutionUnit(executionUnit, ConnectionMode.CONNECTION_STRICTLY, statement); - ExecutionGroupContext executionGroupContext = new ExecutionGroupContext<>( - Collections.singletonList(new ExecutionGroup<>(Collections.singletonList(jdbcExecutionUnit))), mock(ExecutionGroupReportContext.class)); - DialectResultSetMetadataChecker checker = mock(DialectResultSetMetadataChecker.class); - SQLException expected = new SQLException("expected"); - doThrow(expected).when(checker).check(executionContext, statement, "SELECT 1"); - try ( - MockedConstruction ignored = mockConstruction(DriverExecutionPrepareEngine.class, - (mock, context) -> when(mock.prepare(anyString(), eq(executionContext), anyCollection(), any(ExecutionGroupReportContext.class))).thenReturn(executionGroupContext)); - MockedStatic spiLoader = mockStatic(DatabaseTypedSPILoader.class, CALLS_REAL_METHODS)) { - spiLoader.when(() -> DatabaseTypedSPILoader.findService(DialectResultSetMetadataChecker.class, fixtureDatabaseType)).thenReturn(Optional.of(checker)); - assertThat(assertThrows(SQLException.class, () -> proxySQLExecutor.execute(executionContext)), is(expected)); - verify(regularExecutor, never()).execute(any(), any(), anyBoolean(), anyBoolean()); - } - } - @ParameterizedTest(name = "{0}") @MethodSource("executeFallbackScenarios") void assertExecuteFallback(final String name, final boolean hasRawExecutionRule, final SQLStatement sqlStatement, final boolean hasSaneResult) throws SQLException { diff --git a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnectorTest.java b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnectorTest.java index abb9ce0050a9b..fd947f30c34aa 100644 --- a/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnectorTest.java +++ b/proxy/backend/core/src/test/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseProxyConnectorTest.java @@ -489,9 +489,7 @@ void assertExecuteWithQueryResult() throws SQLException { DatabaseProxyConnector engine = createDatabaseProxyConnector(JDBCDriverType.STATEMENT, queryContext); setField(engine, "proxySQLExecutor", proxySQLExecutor); ExecutionContext executionContext = mock(ExecutionContext.class, RETURNS_DEEP_STUBS); - ExecutionUnit executionUnit = mock(ExecutionUnit.class); - when(executionUnit.getDataSourceName()).thenReturn("ds_0"); - when(executionContext.getExecutionUnits()).thenReturn(Collections.singletonList(executionUnit)); + when(executionContext.getExecutionUnits()).thenReturn(Collections.singletonList(mock(ExecutionUnit.class))); when(executionContext.getSqlStatementContext()).thenReturn(sqlStatementContext); QueryResult queryResult = mock(QueryResult.class); QueryResultMetaData queryResultMetaData = mock(QueryResultMetaData.class); @@ -519,8 +517,7 @@ void assertExecuteWithQueryResult() throws SQLException { MockedStatic serviceLoader = mockStatic(ShardingSphereServiceLoader.class)) { spiLoader.when(() -> DatabaseTypedSPILoader.getService(eq(QueryHeaderBuilder.class), any(DatabaseType.class))).thenReturn(new QueryHeaderBuilderFixture()); serviceLoader.when(() -> ShardingSphereServiceLoader.getServiceInstances(AdvancedProxySQLExecutor.class)).thenReturn(Collections.emptyList()); - QueryResponseHeader actual = (QueryResponseHeader) engine.execute(); - assertThat(actual.getDataSourceName(), is("ds_0")); + assertThat(engine.execute(), isA(QueryResponseHeader.class)); assertThat(mockedKernelProcessor.constructed().size(), is(1)); assertThat(mockedDatabaseTypeRegistry.constructed().size(), is(1)); assertThat(mockedMergeEngine.constructed().size(), is(1)); diff --git a/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLResultSetMetadataChecker.java b/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLResultSetMetadataChecker.java deleted file mode 100644 index 3056ac34ee8a4..0000000000000 --- a/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLResultSetMetadataChecker.java +++ /dev/null @@ -1,80 +0,0 @@ -/* - * 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.proxy.backend.postgresql.connector.jdbc; - -import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext; -import org.apache.shardingsphere.infra.binder.context.statement.type.dml.SelectStatementContext; -import org.apache.shardingsphere.infra.exception.ShardingSpherePreconditions; -import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; -import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; -import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker; -import org.apache.shardingsphere.proxy.backend.postgresql.exception.PostgreSQLCompositeTypeAcrossDataSourcesException; - -import java.sql.PreparedStatement; -import java.sql.ResultSetMetaData; -import java.sql.SQLException; -import java.sql.Statement; -import java.sql.Types; -import java.util.Collection; - -/** - * ResultSet metadata checker for PostgreSQL. - */ -public final class PostgreSQLResultSetMetadataChecker implements DialectResultSetMetadataChecker { - - @Override - public void check(final ExecutionContext executionContext, final Statement statement, final String sql) throws SQLException { - if (!isRoutedToMultipleDataSources(executionContext.getExecutionUnits())) { - return; - } - if (statement instanceof PreparedStatement) { - checkCompositeType(((PreparedStatement) statement).getMetaData(), executionContext.getSqlStatementContext()); - return; - } - try (PreparedStatement preparedStatement = statement.getConnection().prepareStatement(sql)) { - checkCompositeType(preparedStatement.getMetaData(), executionContext.getSqlStatementContext()); - } - } - - private boolean isRoutedToMultipleDataSources(final Collection executionUnits) { - String dataSourceName = executionUnits.iterator().next().getDataSourceName(); - for (ExecutionUnit each : executionUnits) { - if (!dataSourceName.equals(each.getDataSourceName())) { - return true; - } - } - return false; - } - - private void checkCompositeType(final ResultSetMetaData metaData, final SQLStatementContext sqlStatementContext) throws SQLException { - if (null == metaData) { - return; - } - int columnCount = sqlStatementContext instanceof SelectStatementContext && ((SelectStatementContext) sqlStatementContext).containsDerivedProjections() - ? ((SelectStatementContext) sqlStatementContext).getProjectionsContext().getExpandProjections().size() - : metaData.getColumnCount(); - for (int columnIndex = 1; columnIndex <= columnCount; columnIndex++) { - ShardingSpherePreconditions.checkState(Types.STRUCT != metaData.getColumnType(columnIndex), PostgreSQLCompositeTypeAcrossDataSourcesException::new); - } - } - - @Override - public String getDatabaseType() { - return "PostgreSQL"; - } -} diff --git a/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesException.java b/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesException.java deleted file mode 100644 index d77ac86973198..0000000000000 --- a/proxy/backend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesException.java +++ /dev/null @@ -1,33 +0,0 @@ -/* - * 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.proxy.backend.postgresql.exception; - -import org.apache.shardingsphere.infra.exception.external.sql.sqlstate.XOpenSQLState; -import org.apache.shardingsphere.infra.exception.external.sql.type.generic.GenericSQLException; - -/** - * PostgreSQL composite type across data sources exception. - */ -public final class PostgreSQLCompositeTypeAcrossDataSourcesException extends GenericSQLException { - - private static final long serialVersionUID = 2994368346256039324L; - - public PostgreSQLCompositeTypeAcrossDataSourcesException() { - super(XOpenSQLState.FEATURE_NOT_SUPPORTED, 31, "PostgreSQL composite result columns are not supported for SQL routed to multiple data sources."); - } -} diff --git a/proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker b/proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker deleted file mode 100644 index 4e6c3292c9072..0000000000000 --- a/proxy/backend/dialect/postgresql/src/main/resources/META-INF/services/org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker +++ /dev/null @@ -1,18 +0,0 @@ -# -# 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. -# - -org.apache.shardingsphere.proxy.backend.postgresql.connector.jdbc.PostgreSQLResultSetMetadataChecker diff --git a/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataCheckerTest.java b/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataCheckerTest.java deleted file mode 100644 index 0cecc9e7191c8..0000000000000 --- a/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/connector/jdbc/PostgreSQLJDBCResultMetadataCheckerTest.java +++ /dev/null @@ -1,129 +0,0 @@ -/* - * 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.proxy.backend.postgresql.connector.jdbc; - -import org.apache.shardingsphere.database.connector.core.spi.DatabaseTypedSPILoader; -import org.apache.shardingsphere.database.connector.core.type.DatabaseType; -import org.apache.shardingsphere.infra.binder.context.segment.select.projection.Projection; -import org.apache.shardingsphere.infra.binder.context.statement.type.dml.SelectStatementContext; -import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; -import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; -import org.apache.shardingsphere.infra.executor.sql.context.SQLUnit; -import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader; -import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker; -import org.apache.shardingsphere.proxy.backend.postgresql.exception.PostgreSQLCompositeTypeAcrossDataSourcesException; -import org.junit.jupiter.api.Test; - -import java.sql.Connection; -import java.sql.PreparedStatement; -import java.sql.ResultSetMetaData; -import java.sql.SQLException; -import java.sql.Statement; -import java.sql.Types; -import java.util.Arrays; -import java.util.Collections; - -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.RETURNS_DEEP_STUBS; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -class PostgreSQLJDBCResultMetadataCheckerTest { - - private final DatabaseType databaseType = TypedSPILoader.getService(DatabaseType.class, "PostgreSQL"); - - private final DialectResultSetMetadataChecker checker = DatabaseTypedSPILoader.getService(DialectResultSetMetadataChecker.class, databaseType); - - @Test - void assertCheckWithSingleDataSource() throws SQLException { - Statement statement = mock(Statement.class); - checker.check(createExecutionContext("ds_0", "ds_0"), statement, "SELECT record_value"); - verify(statement, never()).getConnection(); - } - - @Test - void assertCheckPreparedStatement() throws SQLException { - PreparedStatement preparedStatement = mock(PreparedStatement.class); - ResultSetMetaData metaData = createMetaData(Types.INTEGER); - when(preparedStatement.getMetaData()).thenReturn(metaData); - assertDoesNotThrow(() -> checker.check(createExecutionContext("ds_0", "ds_1"), preparedStatement, "SELECT record_value")); - verify(preparedStatement).getMetaData(); - } - - @Test - void assertCheckCompositeType() throws SQLException { - PreparedStatement preparedStatement = mock(PreparedStatement.class); - ResultSetMetaData metaData = createMetaData(Types.STRUCT); - when(preparedStatement.getMetaData()).thenReturn(metaData); - assertThrows(PostgreSQLCompositeTypeAcrossDataSourcesException.class, - () -> checker.check(createExecutionContext("ds_0", "ds_1"), preparedStatement, "SELECT record_value")); - } - - @Test - void assertCheckStatement() throws SQLException { - PreparedStatement preparedStatement = mock(PreparedStatement.class); - ResultSetMetaData metaData = createMetaData(Types.INTEGER); - when(preparedStatement.getMetaData()).thenReturn(metaData); - Connection connection = mock(Connection.class); - when(connection.prepareStatement("SELECT record_value")).thenReturn(preparedStatement); - Statement statement = mock(Statement.class); - when(statement.getConnection()).thenReturn(connection); - checker.check(createExecutionContext("ds_0", "ds_1"), statement, "SELECT record_value"); - verify(preparedStatement).close(); - } - - @Test - void assertCheckCompositeTypeWithoutMetaData() { - assertDoesNotThrow(() -> checker.check(createExecutionContext("ds_0", "ds_1"), mock(PreparedStatement.class), "SELECT record_value")); - } - - @Test - void assertCheckCompositeTypeWithDerivedProjection() throws SQLException { - ResultSetMetaData metaData = mock(ResultSetMetaData.class); - when(metaData.getColumnType(1)).thenReturn(Types.INTEGER); - when(metaData.getColumnType(2)).thenReturn(Types.STRUCT); - SelectStatementContext sqlStatementContext = mock(SelectStatementContext.class, RETURNS_DEEP_STUBS); - when(sqlStatementContext.containsDerivedProjections()).thenReturn(true); - when(sqlStatementContext.getProjectionsContext().getExpandProjections()).thenReturn(Collections.singletonList(mock(Projection.class))); - ExecutionContext executionContext = createExecutionContext("ds_0", "ds_1"); - when(executionContext.getSqlStatementContext()).thenReturn(sqlStatementContext); - PreparedStatement preparedStatement = mock(PreparedStatement.class); - when(preparedStatement.getMetaData()).thenReturn(metaData); - assertDoesNotThrow(() -> checker.check(executionContext, preparedStatement, "SELECT record_value")); - } - - private ExecutionContext createExecutionContext(final String... dataSourceNames) { - ExecutionContext result = mock(ExecutionContext.class); - ExecutionUnit[] executionUnits = new ExecutionUnit[dataSourceNames.length]; - for (int i = 0; i < dataSourceNames.length; i++) { - executionUnits[i] = new ExecutionUnit(dataSourceNames[i], new SQLUnit("SELECT record_value", Collections.emptyList())); - } - when(result.getExecutionUnits()).thenReturn(Arrays.asList(executionUnits)); - return result; - } - - private ResultSetMetaData createMetaData(final int columnType) throws SQLException { - ResultSetMetaData result = mock(ResultSetMetaData.class); - when(result.getColumnCount()).thenReturn(1); - when(result.getColumnType(1)).thenReturn(columnType); - return result; - } -} diff --git a/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesExceptionTest.java b/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesExceptionTest.java deleted file mode 100644 index 5d8fd77f39318..0000000000000 --- a/proxy/backend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/backend/postgresql/exception/PostgreSQLCompositeTypeAcrossDataSourcesExceptionTest.java +++ /dev/null @@ -1,37 +0,0 @@ -/* - * 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.proxy.backend.postgresql.exception; - -import org.apache.shardingsphere.infra.exception.external.sql.sqlstate.XOpenSQLState; -import org.junit.jupiter.api.Test; - -import java.sql.SQLException; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.is; - -class PostgreSQLCompositeTypeAcrossDataSourcesExceptionTest { - - @Test - void assertToSQLException() { - SQLException actual = new PostgreSQLCompositeTypeAcrossDataSourcesException().toSQLException(); - assertThat(actual.getSQLState(), is(XOpenSQLState.FEATURE_NOT_SUPPORTED.getValue())); - assertThat(actual.getErrorCode(), is(30031)); - assertThat(actual.getMessage(), is("PostgreSQL composite result columns are not supported for SQL routed to multiple data sources.")); - } -} diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java index c9e3aac78c18a..d60d7972c2e0a 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoader.java @@ -19,10 +19,14 @@ import lombok.AccessLevel; import lombok.NoArgsConstructor; -import org.apache.shardingsphere.infra.annotation.HighFrequencyInvocation; +import org.apache.shardingsphere.infra.exception.ShardingSpherePreconditions; +import org.apache.shardingsphere.infra.exception.generic.UnsupportedSQLOperationException; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; import org.apache.shardingsphere.infra.executor.sql.execute.engine.ConnectionMode; +import org.apache.shardingsphere.proxy.backend.connector.StandardDatabaseProxyConnector; +import org.apache.shardingsphere.proxy.backend.handler.ProxyBackendHandler; +import org.apache.shardingsphere.proxy.backend.handler.data.type.UnicastDatabaseProxyBackendHandler; import org.apache.shardingsphere.proxy.backend.response.header.query.QueryHeader; -import org.apache.shardingsphere.proxy.backend.response.header.query.QueryResponseHeader; import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; import org.postgresql.core.BaseConnection; import org.postgresql.core.Oid; @@ -31,6 +35,7 @@ import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Types; +import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -40,33 +45,40 @@ /** * Loader for PostgreSQL column type OIDs. */ -@HighFrequencyInvocation @NoArgsConstructor(access = AccessLevel.PRIVATE) public final class PostgreSQLColumnTypeOIDLoader { /** - * Load composite column type OIDs from the routed backend connection. + * Load composite column type OIDs from the executed backend route. * * @param connectionSession connection session - * @param queryResponseHeader query response header + * @param backendHandler backend handler + * @param queryHeaders query headers * @return column indexes to type OIDs, or an empty map if no composite column type can be resolved * @throws SQLException SQL exception */ - public static Map load(final ConnectionSession connectionSession, final QueryResponseHeader queryResponseHeader) throws SQLException { - List queryHeaders = queryResponseHeader.getQueryHeaders(); + public static Map load(final ConnectionSession connectionSession, final ProxyBackendHandler backendHandler, + final List queryHeaders) throws SQLException { if (!containsCompositeType(queryHeaders)) { return Collections.emptyMap(); } - String dataSourceName = queryResponseHeader.getDataSourceName(); - if (null == dataSourceName) { + Collection executionUnits; + if (backendHandler instanceof StandardDatabaseProxyConnector) { + executionUnits = ((StandardDatabaseProxyConnector) backendHandler).getExecutionUnits(); + } else if (backendHandler instanceof UnicastDatabaseProxyBackendHandler) { + executionUnits = ((UnicastDatabaseProxyBackendHandler) backendHandler).getExecutionUnits(); + } else { return Collections.emptyMap(); } + if (executionUnits.isEmpty()) { + return Collections.emptyMap(); + } + ShardingSpherePreconditions.checkState(1 == executionUnits.size(), + () -> new UnsupportedSQLOperationException("PostgreSQL composite result columns routed to multiple execution units")); + ExecutionUnit executionUnit = executionUnits.iterator().next(); List connections = connectionSession.getDatabaseConnectionManager().getConnections( - connectionSession.getUsedDatabaseName(), dataSourceName, 0, 1, ConnectionMode.CONNECTION_STRICTLY); - return load(connections.get(0), queryHeaders); - } - - private static Map load(final Connection connection, final List queryHeaders) throws SQLException { + connectionSession.getUsedDatabaseName(), executionUnit.getDataSourceName(), 0, 1, ConnectionMode.CONNECTION_STRICTLY); + Connection connection = connections.get(0); return connection.isWrapperFor(BaseConnection.class) ? getTypeOIDs(connection.unwrap(BaseConnection.class), queryHeaders) : Collections.emptyMap(); } diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java index d06c8de1eecb4..8bb82240cf8b7 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/Portal.java @@ -114,7 +114,8 @@ public Portal(final String name, final PostgreSQLServerPreparedStatement prepare public void bind() throws SQLException { responseHeader = proxyBackendHandler.execute(); if (responseHeader instanceof QueryResponseHeader) { - columnTypeOIDs = PostgreSQLColumnTypeOIDLoader.load(databaseConnectionManager.getConnectionSession(), (QueryResponseHeader) responseHeader); + columnTypeOIDs = PostgreSQLColumnTypeOIDLoader.load( + databaseConnectionManager.getConnectionSession(), proxyBackendHandler, ((QueryResponseHeader) responseHeader).getQueryHeaders()); } } diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java index e8410c318d0a5..3dd1ef10529d4 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactory.java @@ -19,14 +19,12 @@ import lombok.AccessLevel; import lombok.NoArgsConstructor; -import org.apache.shardingsphere.database.connector.core.spi.DatabaseTypedSPILoader; -import org.apache.shardingsphere.infra.annotation.HighFrequencyInvocation; import org.apache.shardingsphere.infra.binder.context.aware.ParameterAware; import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext; import org.apache.shardingsphere.infra.binder.engine.SQLBindEngine; import org.apache.shardingsphere.infra.connection.kernel.KernelProcessor; import org.apache.shardingsphere.infra.exception.ShardingSpherePreconditions; -import org.apache.shardingsphere.infra.exception.external.sql.ShardingSphereSQLException; +import org.apache.shardingsphere.infra.exception.generic.UnsupportedSQLOperationException; import org.apache.shardingsphere.infra.exception.kernel.metadata.PreparedStatementMetadataResolutionException; import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; @@ -35,20 +33,19 @@ import org.apache.shardingsphere.infra.session.query.QueryContext; import org.apache.shardingsphere.infra.util.close.QuietlyCloser; import org.apache.shardingsphere.proxy.backend.connector.ProxyDatabaseConnectionManager; -import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker; import org.apache.shardingsphere.proxy.backend.context.ProxyContext; import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; import java.sql.Connection; import java.sql.PreparedStatement; +import java.sql.ResultSetMetaData; import java.sql.SQLException; +import java.sql.Types; import java.util.List; -import java.util.Optional; /** * Metadata factory for PostgreSQL prepared statements. */ -@HighFrequencyInvocation @NoArgsConstructor(access = AccessLevel.PRIVATE) public final class PostgreSQLPreparedStatementMetadataFactory { @@ -82,15 +79,24 @@ public static PreparedStatement load(final ConnectionSession connectionSession, String sql = executionUnit.getSqlUnit().getSql(); PreparedStatement result = connections.iterator().next().prepareStatement(sql); try { - Optional resultSetMetadataChecker = - DatabaseTypedSPILoader.findService(DialectResultSetMetadataChecker.class, sqlStatementContext.getSqlStatement().getDatabaseType()); - if (resultSetMetadataChecker.isPresent()) { - resultSetMetadataChecker.get().check(executionContext, result, sql); - } - } catch (final SQLException | ShardingSphereSQLException ex) { + ShardingSpherePreconditions.checkState(1 == executionContext.getExecutionUnits().size() || !containsCompositeType(result.getMetaData()), + () -> new UnsupportedSQLOperationException("PostgreSQL composite result columns routed to multiple execution units")); + } catch (final SQLException | UnsupportedSQLOperationException ex) { QuietlyCloser.close(result); throw ex; } return result; } + + private static boolean containsCompositeType(final ResultSetMetaData metaData) throws SQLException { + if (null == metaData) { + return false; + } + for (int columnIndex = 1; columnIndex <= metaData.getColumnCount(); columnIndex++) { + if (Types.STRUCT == metaData.getColumnType(columnIndex)) { + return true; + } + } + return false; + } } diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutor.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutor.java index 0346e9a2af080..14447e677b4d7 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutor.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutor.java @@ -123,6 +123,9 @@ private void describeInsertStatementByDatabaseMetaData(final PostgreSQLServerPre String logicTableName = insertStatement.getTable().map(optional -> optional.getTableName().getIdentifier().getValue()).orElse(""); ShardingSphereTable table = getTableFromMetaData(connectionSession.getUsedDatabaseName(), insertStatement, logicTableName); List columnNamesOfInsert = getColumnNamesOfInsertStatement(insertStatement, table); + preparedStatement.setRowDescription(returningSegment.isPresent() + ? describeReturning(returningSegment.get(), table, preparedStatement) + : PostgreSQLNoDataPacket.getInstance()); int parameterMarkerIndex = 0; for (InsertValuesSegment each : insertStatement.getValues()) { for (int i = 0; i < each.getValues().size(); i++) { @@ -139,16 +142,34 @@ private void describeInsertStatementByDatabaseMetaData(final PostgreSQLServerPre preparedStatement.getParameterTypes().set(parameterMarkerIndex++, PostgreSQLBinaryColumnType.valueOfJDBCType(table.getColumn(columnName).getDataType())); } } - if (returningSegment.isPresent()) { - Collection columnDescriptions = createReturningColumnDescriptions(returningSegment.get(), table); - Map columnTypeOIDs = loadColumnTypeOIDs(preparedStatement, columnDescriptions.size()); - preparedStatement.setRowDescription(new PostgreSQLRowDescriptionPacket(applyColumnTypeOIDs(columnDescriptions, columnTypeOIDs))); - } else { - preparedStatement.setRowDescription(PostgreSQLNoDataPacket.getInstance()); + } + + private Collection getUnspecifiedTypeParameterIndexes(final PostgreSQLServerPreparedStatement preparedStatement) { + Collection result = new HashSet<>(); + ListIterator parameterTypesListIterator = preparedStatement.getParameterTypes().listIterator(); + for (int index = parameterTypesListIterator.nextIndex(); parameterTypesListIterator.hasNext(); index = parameterTypesListIterator.nextIndex()) { + if (PostgreSQLBinaryColumnType.UNSPECIFIED == parameterTypesListIterator.next()) { + result.add(index); + } } + return result; } - private Collection createReturningColumnDescriptions(final ReturningSegment returningSegment, final ShardingSphereTable table) { + private ShardingSphereTable getTableFromMetaData(final String databaseName, final InsertStatement insertStatement, final String logicTableName) { + ShardingSphereDatabase database = ProxyContext.getInstance().getContextManager().getDatabase(databaseName); + String schemaName = insertStatement.getTable().flatMap(SimpleTableSegment::getOwner).map(optional -> optional.getIdentifier() + .getValue()).orElseGet(() -> new DatabaseTypeRegistry(database.getProtocolType()).getDefaultSchemaName(databaseName)); + return database.getSchema(schemaName).getTable(logicTableName); + } + + private List getColumnNamesOfInsertStatement(final InsertStatement insertStatement, final ShardingSphereTable table) { + return insertStatement.getColumns().isEmpty() + ? table.getColumnNames() + : insertStatement.getColumns().stream().map(each -> new ShardingSphereIdentifier(each.getIdentifier().getValue())).collect(Collectors.toList()); + } + + private PostgreSQLRowDescriptionPacket describeReturning(final ReturningSegment returningSegment, final ShardingSphereTable table, + final PostgreSQLServerPreparedStatement preparedStatement) throws SQLException { Collection result = new LinkedList<>(); for (ProjectionSegment each : returningSegment.getProjections().getProjections()) { if (each instanceof ShorthandProjectionSegment) { @@ -166,7 +187,7 @@ private Collection createReturningColumnDescription result.add(convertExpressionToDescription((ExpressionProjectionSegment) each)); } } - return result; + return new PostgreSQLRowDescriptionPacket(applyColumnTypeOIDs(result, loadColumnTypeOIDs(preparedStatement, result.size()))); } private ShardingSphereColumn generateDefaultColumn(final ColumnProjectionSegment segment) { @@ -216,30 +237,6 @@ private Map loadColumnTypeOIDs(final PostgreSQLServerPreparedS } } - private Collection getUnspecifiedTypeParameterIndexes(final PostgreSQLServerPreparedStatement preparedStatement) { - Collection result = new HashSet<>(); - ListIterator parameterTypesListIterator = preparedStatement.getParameterTypes().listIterator(); - for (int index = parameterTypesListIterator.nextIndex(); parameterTypesListIterator.hasNext(); index = parameterTypesListIterator.nextIndex()) { - if (PostgreSQLBinaryColumnType.UNSPECIFIED == parameterTypesListIterator.next()) { - result.add(index); - } - } - return result; - } - - private ShardingSphereTable getTableFromMetaData(final String databaseName, final InsertStatement insertStatement, final String logicTableName) { - ShardingSphereDatabase database = ProxyContext.getInstance().getContextManager().getDatabase(databaseName); - String schemaName = insertStatement.getTable().flatMap(SimpleTableSegment::getOwner).map(optional -> optional.getIdentifier() - .getValue()).orElseGet(() -> new DatabaseTypeRegistry(database.getProtocolType()).getDefaultSchemaName(databaseName)); - return database.getSchema(schemaName).getTable(logicTableName); - } - - private List getColumnNamesOfInsertStatement(final InsertStatement insertStatement, final ShardingSphereTable table) { - return insertStatement.getColumns().isEmpty() - ? table.getColumnNames() - : insertStatement.getColumns().stream().map(each -> new ShardingSphereIdentifier(each.getIdentifier().getValue())).collect(Collectors.toList()); - } - private void tryDescribePreparedStatementByJDBC(final PostgreSQLServerPreparedStatement logicPreparedStatement) throws SQLException { try (PreparedStatement actualPreparedStatement = PostgreSQLPreparedStatementMetadataFactory.load(connectionSession, logicPreparedStatement, Collections.emptyList())) { PostgreSQLPreparedStatementParameterTypeResolver.resolveParameterTypes(logicPreparedStatement, actualPreparedStatement); diff --git a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutor.java b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutor.java index 7fa2ac6430423..faf6b88edc434 100644 --- a/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutor.java +++ b/proxy/frontend/dialect/postgresql/src/main/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutor.java @@ -90,7 +90,7 @@ public Collection execute() throws SQLException { ResponseHeader responseHeader = proxyBackendHandler.execute(); if (responseHeader instanceof QueryResponseHeader) { QueryResponseHeader queryResponseHeader = (QueryResponseHeader) responseHeader; - Map columnTypeOIDs = PostgreSQLColumnTypeOIDLoader.load(connectionSession, queryResponseHeader); + Map columnTypeOIDs = PostgreSQLColumnTypeOIDLoader.load(connectionSession, proxyBackendHandler, queryResponseHeader.getQueryHeaders()); return Collections.singleton(createRowDescriptionPacket(queryResponseHeader, columnTypeOIDs)); } responseType = ResponseType.UPDATE; diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java index 5512c9d7faa08..7896810e5ab0e 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/PostgreSQLColumnTypeOIDLoaderTest.java @@ -17,29 +17,29 @@ package org.apache.shardingsphere.proxy.frontend.postgresql.command.query; +import org.apache.shardingsphere.infra.exception.generic.UnsupportedSQLOperationException; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; +import org.apache.shardingsphere.infra.executor.sql.execute.engine.ConnectionMode; import org.apache.shardingsphere.proxy.backend.connector.ProxyDatabaseConnectionManager; +import org.apache.shardingsphere.proxy.backend.connector.StandardDatabaseProxyConnector; import org.apache.shardingsphere.proxy.backend.response.header.query.QueryHeader; -import org.apache.shardingsphere.proxy.backend.response.header.query.QueryResponseHeader; import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; -import org.apache.shardingsphere.infra.executor.sql.execute.engine.ConnectionMode; import org.junit.jupiter.api.Test; import org.postgresql.core.BaseConnection; import org.postgresql.core.Oid; import org.postgresql.core.TypeInfo; -import java.sql.Connection; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Types; +import java.util.Collection; import java.util.Collections; import java.util.Map; import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.is; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyInt; -import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -48,84 +48,77 @@ class PostgreSQLColumnTypeOIDLoaderTest { @Test - void assertLoadWithSameDataSource() throws SQLException { - BaseConnection connection = mock(BaseConnection.class); - mockTypeOID(connection, "record_type", 2249); - ConnectionSession connectionSession = createConnectionSession(connection); - QueryResponseHeader queryResponseHeader = new QueryResponseHeader( - Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type")), "ds_0"); - Map actual = PostgreSQLColumnTypeOIDLoader.load(connectionSession, queryResponseHeader); - assertThat(actual.get(1), is(2249)); - verify(connectionSession.getDatabaseConnectionManager()).getConnections("postgres", "ds_0", 0, 1, ConnectionMode.CONNECTION_STRICTLY); - } - - @Test - void assertLoadWithoutCompositeColumn() throws SQLException { - QueryResponseHeader queryResponseHeader = new QueryResponseHeader( - Collections.singletonList(createQueryHeader(Types.VARCHAR, "varchar")), "ds_0"); - assertTrue(PostgreSQLColumnTypeOIDLoader.load(mock(ConnectionSession.class), queryResponseHeader).isEmpty()); + void assertLoadFromExecutedRoute() throws SQLException { + ConnectionSession connectionSession = mock(ConnectionSession.class); + StandardDatabaseProxyConnector backendHandler = mock(StandardDatabaseProxyConnector.class); + Collection executionUnits = mockExecutionUnits(1); + when(backendHandler.getExecutionUnits()).thenReturn(executionUnits); + BaseConnection connection = mockConnection("record_type", 2249); + ProxyDatabaseConnectionManager connectionManager = mock(ProxyDatabaseConnectionManager.class); + when(connectionSession.getDatabaseConnectionManager()).thenReturn(connectionManager); + when(connectionSession.getUsedDatabaseName()).thenReturn("postgres"); + when(connectionManager.getConnections("postgres", "ds_0", 0, 1, ConnectionMode.CONNECTION_STRICTLY)).thenReturn(Collections.singletonList(connection)); + Map actual = PostgreSQLColumnTypeOIDLoader.load(connectionSession, backendHandler, Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type"))); + assertThat(actual, is(Collections.singletonMap(1, 2249))); } @Test - void assertLoadWithoutDataSource() throws SQLException { - QueryResponseHeader queryResponseHeader = new QueryResponseHeader(Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type"))); - assertTrue(PostgreSQLColumnTypeOIDLoader.load(mock(ConnectionSession.class), queryResponseHeader).isEmpty()); + void assertLoadFromMultipleExecutionUnits() { + StandardDatabaseProxyConnector backendHandler = mock(StandardDatabaseProxyConnector.class); + Collection executionUnits = mockExecutionUnits(2); + when(backendHandler.getExecutionUnits()).thenReturn(executionUnits); + assertThrows(UnsupportedSQLOperationException.class, () -> PostgreSQLColumnTypeOIDLoader.load( + mock(ConnectionSession.class), backendHandler, Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type")))); } @Test - void assertLoadWithoutPgjdbcConnection() throws SQLException { - Connection connection = mock(Connection.class); - ConnectionSession connectionSession = createConnectionSession(connection); - QueryResponseHeader queryResponseHeader = new QueryResponseHeader( - Collections.singletonList(createQueryHeader(Types.STRUCT, "record_type")), "ds_0"); - assertTrue(PostgreSQLColumnTypeOIDLoader.load(connectionSession, queryResponseHeader).isEmpty()); - verify(connection, never()).unwrap(BaseConnection.class); + void assertLoadWithoutCompositeType() throws SQLException { + StandardDatabaseProxyConnector backendHandler = mock(StandardDatabaseProxyConnector.class); + assertTrue(PostgreSQLColumnTypeOIDLoader.load( + mock(ConnectionSession.class), backendHandler, Collections.singletonList(createQueryHeader(Types.VARCHAR, "varchar"))).isEmpty()); + verify(backendHandler, never()).getExecutionUnits(); } @Test void assertLoadFromResultSetMetaData() throws SQLException { - BaseConnection connection = mock(BaseConnection.class); - mockTypeOID(connection, "record_type", 2249); + BaseConnection connection = mockConnection("record_type", 2249); ResultSetMetaData metaData = mock(ResultSetMetaData.class); when(metaData.getColumnCount()).thenReturn(2); when(metaData.getColumnType(1)).thenReturn(Types.STRUCT); when(metaData.getColumnType(2)).thenReturn(Types.VARCHAR); when(metaData.getColumnTypeName(1)).thenReturn("record_type"); Map actual = PostgreSQLColumnTypeOIDLoader.load(connection, metaData); - assertThat(actual.size(), is(1)); - assertThat(actual.get(1), is(2249)); + assertThat(actual, is(Collections.singletonMap(1, 2249))); verify(metaData, never()).getColumnTypeName(2); } @Test void assertLoadUnspecifiedTypeOID() throws SQLException { - BaseConnection connection = mock(BaseConnection.class); - mockTypeOID(connection, "unknown_type", Oid.UNSPECIFIED); ResultSetMetaData metaData = mock(ResultSetMetaData.class); when(metaData.getColumnCount()).thenReturn(1); when(metaData.getColumnType(1)).thenReturn(Types.STRUCT); when(metaData.getColumnTypeName(1)).thenReturn("unknown_type"); + BaseConnection connection = mockConnection("unknown_type", Oid.UNSPECIFIED); assertTrue(PostgreSQLColumnTypeOIDLoader.load(connection, metaData).isEmpty()); } - private QueryHeader createQueryHeader(final int columnType, final String columnTypeName) { - return new QueryHeader("", "", "record_value", "record_value", columnType, columnTypeName, -1, 0, false, false, false, false); + private Collection mockExecutionUnits(final int executionUnitCount) { + ExecutionUnit executionUnit = mock(ExecutionUnit.class); + when(executionUnit.getDataSourceName()).thenReturn("ds_0"); + return Collections.nCopies(executionUnitCount, executionUnit); } - private ConnectionSession createConnectionSession(final Connection connection) throws SQLException { - ConnectionSession result = mock(ConnectionSession.class); - ProxyDatabaseConnectionManager connectionManager = mock(ProxyDatabaseConnectionManager.class); - when(connectionManager.getConnections(any(), anyString(), anyInt(), anyInt(), any())).thenReturn(Collections.singletonList(connection)); - when(result.getDatabaseConnectionManager()).thenReturn(connectionManager); - when(result.getUsedDatabaseName()).thenReturn("postgres"); + private BaseConnection mockConnection(final String columnTypeName, final int typeOID) throws SQLException { + TypeInfo typeInfo = mock(TypeInfo.class); + when(typeInfo.getPGType(columnTypeName)).thenReturn(typeOID); + BaseConnection result = mock(BaseConnection.class); + when(result.getTypeInfo()).thenReturn(typeInfo); + when(result.isWrapperFor(BaseConnection.class)).thenReturn(true); + when(result.unwrap(BaseConnection.class)).thenReturn(result); return result; } - private void mockTypeOID(final BaseConnection connection, final String columnTypeName, final int typeOID) throws SQLException { - TypeInfo typeInfo = mock(TypeInfo.class); - when(typeInfo.getPGType(columnTypeName)).thenReturn(typeOID); - when(connection.getTypeInfo()).thenReturn(typeInfo); - when(connection.isWrapperFor(BaseConnection.class)).thenReturn(true); - when(connection.unwrap(BaseConnection.class)).thenReturn(connection); + private QueryHeader createQueryHeader(final int columnType, final String columnTypeName) { + return new QueryHeader("", "", "", "", columnType, columnTypeName, 0, 0, false, false, false, false); } } diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PortalTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PortalTest.java index 4cda63c32f03f..582d1c01ca5c5 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PortalTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PortalTest.java @@ -90,6 +90,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyList; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.RETURNS_DEEP_STUBS; @@ -126,7 +127,7 @@ void setup() throws SQLException { when(ProxyBackendHandlerFactory.newInstance(eq(databaseType), anyString(), any(SQLStatement.class), eq(connectionSession), any(HintValueContext.class))).thenReturn(proxyBackendHandler); when(ProxyBackendHandlerFactory.newInstance(eq(databaseType), any(QueryContext.class), eq(connectionSession), anyBoolean())).thenReturn(proxyBackendHandler); when(databaseConnectionManager.getConnectionSession()).thenReturn(connectionSession); - when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), any(QueryResponseHeader.class))).thenReturn(Collections.emptyMap()); + when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), eq(proxyBackendHandler), anyList())).thenReturn(Collections.emptyMap()); } @Test @@ -153,8 +154,7 @@ void assertExecuteSelectStatementAndReturnAllRows() throws SQLException, Reflect Map columnTypeOIDs = new HashMap<>(2, 1F); columnTypeOIDs.put(1, 2249); columnTypeOIDs.put(2, 2249); - when(PostgreSQLColumnTypeOIDLoader.load(eq(databaseConnectionManager.getConnectionSession()), any(QueryResponseHeader.class))) - .thenReturn(columnTypeOIDs); + when(PostgreSQLColumnTypeOIDLoader.load(any(ConnectionSession.class), eq(proxyBackendHandler), anyList())).thenReturn(columnTypeOIDs); Portal portal = createPortal(sqlStatementContext, Arrays.asList(PostgreSQLValueFormat.TEXT, PostgreSQLValueFormat.BINARY)); PostgreSQLPacket portalDescription = portal.describe(); assertThat(portalDescription, isA(PostgreSQLRowDescriptionPacket.class)); diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java index 2fc19d00c3a57..e0f37acbe6278 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/PostgreSQLPreparedStatementMetadataFactoryTest.java @@ -17,7 +17,6 @@ package org.apache.shardingsphere.proxy.frontend.postgresql.command.query.extended; -import org.apache.shardingsphere.database.connector.core.spi.DatabaseTypedSPILoader; import org.apache.shardingsphere.database.connector.core.type.DatabaseType; import org.apache.shardingsphere.database.protocol.postgresql.packet.command.query.extended.PostgreSQLBinaryColumnType; import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext; @@ -25,8 +24,11 @@ import org.apache.shardingsphere.infra.config.props.ConfigurationProperties; import org.apache.shardingsphere.infra.connection.kernel.KernelProcessor; import org.apache.shardingsphere.infra.exception.external.sql.ShardingSphereSQLException; +import org.apache.shardingsphere.infra.exception.generic.UnsupportedSQLOperationException; import org.apache.shardingsphere.infra.exception.kernel.metadata.PreparedStatementMetadataResolutionException; import org.apache.shardingsphere.infra.executor.sql.context.ExecutionContext; +import org.apache.shardingsphere.infra.executor.sql.context.ExecutionUnit; +import org.apache.shardingsphere.infra.executor.sql.context.SQLUnit; import org.apache.shardingsphere.infra.hint.HintValueContext; import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase; import org.apache.shardingsphere.infra.metadata.database.resource.unit.StorageUnit; @@ -39,9 +41,7 @@ import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader; import org.apache.shardingsphere.mode.manager.ContextManager; import org.apache.shardingsphere.proxy.backend.connector.ProxyDatabaseConnectionManager; -import org.apache.shardingsphere.proxy.backend.connector.jdbc.executor.DialectResultSetMetadataChecker; import org.apache.shardingsphere.proxy.backend.context.ProxyContext; -import org.apache.shardingsphere.proxy.backend.postgresql.exception.PostgreSQLCompositeTypeAcrossDataSourcesException; import org.apache.shardingsphere.proxy.backend.session.ConnectionSession; import org.apache.shardingsphere.sql.parser.engine.api.CacheOption; import org.apache.shardingsphere.sql.parser.statement.core.statement.SQLStatement; @@ -54,15 +54,14 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.MockedConstruction; -import org.mockito.MockedStatic; import java.sql.Connection; import java.sql.PreparedStatement; +import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Types; import java.util.Collections; import java.util.List; -import java.util.Optional; import java.util.Properties; import static org.hamcrest.MatcherAssert.assertThat; @@ -71,15 +70,11 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.nullable; -import static org.mockito.Mockito.CALLS_REAL_METHODS; import static org.mockito.Mockito.RETURNS_DEEP_STUBS; -import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mockConstruction; -import static org.mockito.Mockito.mockStatic; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -126,16 +121,20 @@ void assertLoadWithEmptyExecutionUnits() { } @Test - void assertClosePreparedStatementOnCheckerException() throws SQLException { + void assertLoadCompositeTypeWithMultipleExecutionUnits() throws SQLException { PreparedStatement expected = prepareJDBCBackendConnection(null); PostgreSQLServerPreparedStatement preparedStatement = createPreparedStatement(true); - PostgreSQLCompositeTypeAcrossDataSourcesException exception = new PostgreSQLCompositeTypeAcrossDataSourcesException(); - DialectResultSetMetadataChecker checker = mock(DialectResultSetMetadataChecker.class); - doThrow(exception).when(checker).check(any(ExecutionContext.class), eq(expected), anyString()); - try (MockedStatic spiLoader = mockStatic(DatabaseTypedSPILoader.class, CALLS_REAL_METHODS)) { - spiLoader.when(() -> DatabaseTypedSPILoader.findService(DialectResultSetMetadataChecker.class, databaseType)).thenReturn(Optional.of(checker)); - assertThat(assertThrows(PostgreSQLCompositeTypeAcrossDataSourcesException.class, - () -> PostgreSQLPreparedStatementMetadataFactory.load(connectionSession, preparedStatement, PARAMETERS)), is(exception)); + ResultSetMetaData metaData = mock(ResultSetMetaData.class); + when(metaData.getColumnCount()).thenReturn(1); + when(metaData.getColumnType(1)).thenReturn(Types.STRUCT); + when(expected.getMetaData()).thenReturn(metaData); + ExecutionUnit executionUnit = new ExecutionUnit("ds_0", new SQLUnit("SELECT id FROM foo_tbl WHERE id=?", PARAMETERS)); + ExecutionContext executionContext = mock(ExecutionContext.class); + when(executionContext.getExecutionUnits()).thenReturn(Collections.nCopies(2, executionUnit)); + try ( + MockedConstruction ignored = mockConstruction(KernelProcessor.class, + (mock, context) -> when(mock.generateExecutionContext(any(), any(), any())).thenReturn(executionContext))) { + assertThrows(UnsupportedSQLOperationException.class, () -> PostgreSQLPreparedStatementMetadataFactory.load(connectionSession, preparedStatement, PARAMETERS)); verify(expected).close(); } } diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutorTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutorTest.java index a689fdaad8d8a..7e6f48b9cea00 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutorTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/extended/describe/PostgreSQLComDescribeExecutorTest.java @@ -131,13 +131,16 @@ class PostgreSQLComDescribeExecutorTest { @Mock private ConnectionSession connectionSession; + private Map columnTypeOIDs; + @InjectMocks private PostgreSQLComDescribeExecutor executor; @BeforeEach void setUp() throws SQLException { when(connectionSession.getProtocolType()).thenReturn(DATABASE_TYPE); - when(PostgreSQLColumnTypeOIDLoader.load(any(Connection.class), any(ResultSetMetaData.class))).thenReturn(Collections.emptyMap()); + columnTypeOIDs = Collections.emptyMap(); + when(PostgreSQLColumnTypeOIDLoader.load(any(Connection.class), any(ResultSetMetaData.class))).thenAnswer(ignored -> columnTypeOIDs); } @Test @@ -390,10 +393,9 @@ void assertDescribePreparedStatementInsertWithReturning(final String testName, f when(ProxyContext.getInstance().getContextManager()).thenReturn(contextManager); prepareJDBCBackendConnectionForResultMetaData(sql, expectedColumns.size()); if ("returning complex columns".equals(testName)) { - Map columnTypeOIDs = new HashMap<>(2, 1F); + columnTypeOIDs = new HashMap<>(2, 1F); columnTypeOIDs.put(12, 2249); columnTypeOIDs.put(13, 2249); - when(PostgreSQLColumnTypeOIDLoader.load(any(Connection.class), any(ResultSetMetaData.class))).thenReturn(columnTypeOIDs); } List parameterIndexes = IntStream.range(0, sqlStatement.getParameterCount()).boxed().collect(Collectors.toList()); connectionSession.getServerPreparedStatementRegistry().addPreparedStatement( @@ -431,14 +433,14 @@ void assertDescribeSelectPreparedStatement() throws SQLException { SQLStatementContext sqlStatementContext = mock(SelectStatementContext.class); when(sqlStatementContext.getSqlStatement()).thenReturn(sqlStatement); prepareJDBCBackendConnection(sql); - List parameterTypes = new ArrayList<>(Collections.singleton(PostgreSQLBinaryColumnType.UNSPECIFIED)); ContextManager contextManager = mockContextManager(); when(ProxyContext.getInstance().getContextManager()).thenReturn(contextManager); - List parameterIndexes = IntStream.range(0, sqlStatement.getParameterCount()).boxed().collect(Collectors.toList()); ConnectionContext connectionContext = mock(ConnectionContext.class); when(connectionContext.getCurrentDatabaseName()).thenReturn(Optional.of(DATABASE_NAME)); when(connectionSession.getConnectionContext()).thenReturn(connectionContext); - when(PostgreSQLColumnTypeOIDLoader.load(any(Connection.class), any(ResultSetMetaData.class))).thenReturn(Collections.singletonMap(1, 2249)); + columnTypeOIDs = Collections.singletonMap(1, 2249); + List parameterTypes = new ArrayList<>(Collections.singleton(PostgreSQLBinaryColumnType.UNSPECIFIED)); + List parameterIndexes = IntStream.range(0, sqlStatement.getParameterCount()).boxed().collect(Collectors.toList()); connectionSession.getServerPreparedStatementRegistry().addPreparedStatement( statementId, new PostgreSQLServerPreparedStatement(sql, sqlStatementContext, new HintValueContext(), parameterTypes, parameterIndexes)); Collection actual = executor.execute(); diff --git a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutorTest.java b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutorTest.java index eb228f4478cb5..88fb2754505ed 100644 --- a/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutorTest.java +++ b/proxy/frontend/dialect/postgresql/src/test/java/org/apache/shardingsphere/proxy/frontend/postgresql/command/query/simple/PostgreSQLComQueryExecutorTest.java @@ -78,7 +78,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyList; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -113,16 +113,16 @@ void setUp() throws SQLException { SQLStatement sqlStatement = new SQLStatement(DATABASE_TYPE); when(ProxySQLComQueryParser.parse(queryPacket.getSQL(), DATABASE_TYPE, connectionSession)).thenReturn(sqlStatement); when(ProxyBackendHandlerFactory.newInstance(DATABASE_TYPE, queryPacket.getSQL(), sqlStatement, connectionSession, queryPacket.getHintValueContext())).thenReturn(proxyBackendHandler); - when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), any(QueryResponseHeader.class))).thenReturn(Collections.emptyMap()); + when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), eq(proxyBackendHandler), anyList())).thenReturn(Collections.emptyMap()); queryExecutor = new PostgreSQLComQueryExecutor(portalContext, queryPacket, connectionSession); } @Test void assertExecuteQueryWithColumnDescription() throws SQLException, ReflectiveOperationException { - when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), any(QueryResponseHeader.class))).thenReturn(Collections.singletonMap(1, 2249)); QueryResponseHeader queryResponseHeader = mock(QueryResponseHeader.class); when(queryResponseHeader.getQueryHeaders()).thenReturn( Collections.singletonList(new QueryHeader("schema", "table", "label", "column", Types.STRUCT, "record_type", 2, 3, true, true, true, true))); + when(PostgreSQLColumnTypeOIDLoader.load(eq(connectionSession), eq(proxyBackendHandler), anyList())).thenReturn(Collections.singletonMap(1, 2249)); when(proxyBackendHandler.execute()).thenReturn(queryResponseHeader); Collection actual = queryExecutor.execute(); PostgreSQLRowDescriptionPacket rowDescriptionPacket = (PostgreSQLRowDescriptionPacket) actual.iterator().next();