(shardingsphere) branch master updated: Add generated key support check for database dialects (#35417)

[zh...@apache.org]

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

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


The following commit(s) were added to refs/heads/master by this push:
     new eb8fd6fb2ab Add generated key support check for database dialects (#35417)
eb8fd6fb2ab is described below

commit eb8fd6fb2ab6f5e9d60cb6c34afbd6cbe9731a94
Author: Liang Zhang <zh...@apache.org>
AuthorDate: Thu May 15 10:47:48 2025 +0800

    Add generated key support check for database dialects (#35417)
    
    * Add generated key support check for database dialects
    
    - Introduce DialectGeneratedKeyOption to indicate generated key support
    - Add getGeneratedKeyOption method to DialectDatabaseMetaData
    - Update MySQLDatabaseMetaData to support generated keys
    - Modify ProxySQLExecutor and StandardDatabaseConnector to use the new generated key option
    
    * Add generated key support check for database dialects
    
    - Introduce DialectGeneratedKeyOption to indicate generated key support
    - Add getGeneratedKeyOption method to DialectDatabaseMetaData
    - Update MySQLDatabaseMetaData to support generated keys
    - Modify ProxySQLExecutor and StandardDatabaseConnector to use the new generated key option
---
 .../database/metadata/DialectDatabaseMetaData.java | 10 +++++++
 .../option/keygen/DialectGeneratedKeyOption.java   | 31 ++++++++++++++++++++++
 .../metadata/database/MySQLDatabaseMetaData.java   |  6 +++++
 .../proxy/backend/connector/ProxySQLExecutor.java  | 10 ++++---
 .../connector/StandardDatabaseConnector.java       | 10 ++++---
 5 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/infra/database/core/src/main/java/org/apache/shardingsphere/infra/database/core/metadata/database/metadata/DialectDatabaseMetaData.java b/infra/database/core/src/main/java/org/apache/shardingsphere/infra/database/core/metadata/database/metadata/DialectDatabaseMetaData.java
index 08df94ead55..c8119c6c963 100644
--- a/infra/database/core/src/main/java/org/apache/shardingsphere/infra/database/core/metadata/database/metadata/DialectDatabaseMetaData.java
+++ b/infra/database/core/src/main/java/org/apache/shardingsphere/infra/database/core/metadata/database/metadata/DialectDatabaseMetaData.java
@@ -27,6 +27,7 @@ import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.
 import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.datatype.DialectDataTypeOption;
 import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.index.DialectIndexOption;
 import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.join.DialectJoinOption;
+import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.keygen.DialectGeneratedKeyOption;
 import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.schema.DefaultSchemaOption;
 import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.schema.DialectSchemaOption;
 import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.table.DialectDriverQuerySystemCatalogOption;
@@ -137,6 +138,15 @@ public interface DialectDatabaseMetaData extends DatabaseTypedSPI {
         return new DialectJoinOption(false, false);
     }
     
+    /**
+     * Get generated key option.
+     *
+     * @return generated key option
+     */
+    default DialectGeneratedKeyOption getGeneratedKeyOption() {
+        return new DialectGeneratedKeyOption(false);
+    }
+    
     /**
      * Get alter table option.
      *
diff --git a/infra/database/core/src/main/java/org/apache/shardingsphere/infra/database/core/metadata/database/metadata/option/keygen/DialectGeneratedKeyOption.java b/infra/database/core/src/main/java/org/apache/shardingsphere/infra/database/core/metadata/database/metadata/option/keygen/DialectGeneratedKeyOption.java
new file mode 100644
index 00000000000..b86435f511e
--- /dev/null
+++ b/infra/database/core/src/main/java/org/apache/shardingsphere/infra/database/core/metadata/database/metadata/option/keygen/DialectGeneratedKeyOption.java
@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.keygen;
+
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+
+/**
+ * Dialect generated key option.
+ */
+@RequiredArgsConstructor
+@Getter
+public final class DialectGeneratedKeyOption {
+    
+    private final boolean isSupportReturnGeneratedKeys;
+}
diff --git a/infra/database/type/mysql/src/main/java/org/apache/shardingsphere/infra/database/mysql/metadata/database/MySQLDatabaseMetaData.java b/infra/database/type/mysql/src/main/java/org/apache/shardingsphere/infra/database/mysql/metadata/database/MySQLDatabaseMetaData.java
index 436a7dc0eb5..74a9e3691ae 100644
--- a/infra/database/type/mysql/src/main/java/org/apache/shardingsphere/infra/database/mysql/metadata/database/MySQLDatabaseMetaData.java
+++ b/infra/database/type/mysql/src/main/java/org/apache/shardingsphere/infra/database/mysql/metadata/database/MySQLDatabaseMetaData.java
@@ -25,6 +25,7 @@ import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.
 import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.connection.DialectConnectionOption;
 import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.datatype.DialectDataTypeOption;
 import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.join.DialectJoinOption;
+import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.keygen.DialectGeneratedKeyOption;
 import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.transaction.DialectTransactionOption;
 import org.apache.shardingsphere.infra.database.mysql.metadata.database.option.MySQLDataTypeOption;
 
@@ -75,6 +76,11 @@ public final class MySQLDatabaseMetaData implements DialectDatabaseMetaData {
         return new DialectJoinOption(true, true);
     }
     
+    @Override
+    public DialectGeneratedKeyOption getGeneratedKeyOption() {
+        return new DialectGeneratedKeyOption(true);
+    }
+    
     @Override
     public String getDatabaseType() {
         return "MySQL";
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 6e9abf88b79..c5d5efef4e4 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
@@ -22,6 +22,7 @@ import lombok.Getter;
 import org.apache.shardingsphere.infra.binder.context.statement.SQLStatementContext;
 import org.apache.shardingsphere.infra.binder.context.type.TableAvailable;
 import org.apache.shardingsphere.infra.config.props.ConfigurationPropertyKey;
+import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.DialectDatabaseMetaData;
 import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.option.transaction.DialectTransactionOption;
 import org.apache.shardingsphere.infra.database.core.type.DatabaseType;
 import org.apache.shardingsphere.infra.database.core.type.DatabaseTypeRegistry;
@@ -64,7 +65,7 @@ import org.apache.shardingsphere.sql.parser.statement.core.statement.ddl.DDLStat
 import org.apache.shardingsphere.sql.parser.statement.core.statement.ddl.FetchStatement;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.ddl.MoveStatement;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.ddl.TruncateStatement;
-import org.apache.shardingsphere.sql.parser.statement.mysql.dml.MySQLInsertStatement;
+import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.InsertStatement;
 import org.apache.shardingsphere.sql.parser.statement.opengauss.ddl.OpenGaussCursorStatement;
 import org.apache.shardingsphere.sqlfederation.engine.SQLFederationEngine;
 import org.apache.shardingsphere.transaction.api.TransactionType;
@@ -185,8 +186,11 @@ public final class ProxySQLExecutor {
         Collection<ShardingSphereRule> rules = ProxyContext.getInstance().getContextManager().getMetaDataContexts().getMetaData().getDatabase(databaseName).getRuleMetaData().getRules();
         int maxConnectionsSizePerQuery = ProxyContext.getInstance()
                 .getContextManager().getMetaDataContexts().getMetaData().getProps().<Integer>getValue(ConfigurationPropertyKey.MAX_CONNECTIONS_SIZE_PER_QUERY);
-        boolean isReturnGeneratedKeys = executionContext.getSqlStatementContext().getSqlStatement() instanceof MySQLInsertStatement;
-        return hasRawExecutionRule(rules) ? rawExecute(executionContext, rules, maxConnectionsSizePerQuery)
+        SQLStatement sqlStatement = executionContext.getSqlStatementContext().getSqlStatement();
+        DialectDatabaseMetaData dialectDatabaseMetaData = new DatabaseTypeRegistry(sqlStatement.getDatabaseType()).getDialectDatabaseMetaData();
+        boolean isReturnGeneratedKeys = sqlStatement instanceof InsertStatement && dialectDatabaseMetaData.getGeneratedKeyOption().isSupportReturnGeneratedKeys();
+        return hasRawExecutionRule(rules)
+                ? rawExecute(executionContext, rules, maxConnectionsSizePerQuery)
                 : useDriverToExecute(executionContext, rules, maxConnectionsSizePerQuery, isReturnGeneratedKeys, SQLExecutorExceptionHandler.isExceptionThrown());
     }
     
diff --git a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseConnector.java b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseConnector.java
index 5edd3c7256c..af3d626d02b 100644
--- a/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseConnector.java
+++ b/proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/connector/StandardDatabaseConnector.java
@@ -28,7 +28,9 @@ import org.apache.shardingsphere.infra.binder.context.type.CursorAvailable;
 import org.apache.shardingsphere.infra.binder.context.type.TableAvailable;
 import org.apache.shardingsphere.infra.config.props.ConfigurationPropertyKey;
 import org.apache.shardingsphere.infra.connection.kernel.KernelProcessor;
+import org.apache.shardingsphere.infra.database.core.metadata.database.metadata.DialectDatabaseMetaData;
 import org.apache.shardingsphere.infra.database.core.type.DatabaseType;
+import org.apache.shardingsphere.infra.database.core.type.DatabaseTypeRegistry;
 import org.apache.shardingsphere.infra.exception.core.ShardingSpherePreconditions;
 import org.apache.shardingsphere.infra.exception.dialect.SQLExceptionTransformEngine;
 import org.apache.shardingsphere.infra.exception.kernel.metadata.resource.storageunit.EmptyStorageUnitException;
@@ -73,8 +75,8 @@ import org.apache.shardingsphere.proxy.backend.util.TransactionUtils;
 import org.apache.shardingsphere.sharding.merge.common.IteratorStreamMergedResult;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.SQLStatement;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.DMLStatement;
+import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.InsertStatement;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.SelectStatement;
-import org.apache.shardingsphere.sql.parser.statement.mysql.dml.MySQLInsertStatement;
 import org.apache.shardingsphere.sqlfederation.executor.context.SQLFederationContext;
 import org.apache.shardingsphere.transaction.api.TransactionType;
 import org.apache.shardingsphere.transaction.implicit.ImplicitTransactionCallback;
@@ -248,10 +250,12 @@ public final class StandardDatabaseConnector implements DatabaseConnector {
     }
     
     private ResultSet doExecuteFederation() {
-        boolean isReturnGeneratedKeys = queryContext.getSqlStatementContext().getSqlStatement() instanceof MySQLInsertStatement;
+        SQLStatement sqlStatement = queryContext.getSqlStatementContext().getSqlStatement();
+        DialectDatabaseMetaData dialectDatabaseMetaData = new DatabaseTypeRegistry(sqlStatement.getDatabaseType()).getDialectDatabaseMetaData();
+        boolean isReturnGeneratedKeys = sqlStatement instanceof InsertStatement && dialectDatabaseMetaData.getGeneratedKeyOption().isSupportReturnGeneratedKeys();
         DatabaseType protocolType = database.getProtocolType();
         ProxyJDBCExecutorCallback callback = ProxyJDBCExecutorCallbackFactory.newInstance(driverType, protocolType, database.getResourceMetaData(),
-                queryContext.getSqlStatementContext().getSqlStatement(), this, isReturnGeneratedKeys, SQLExecutorExceptionHandler.isExceptionThrown(), true);
+                sqlStatement, this, isReturnGeneratedKeys, SQLExecutorExceptionHandler.isExceptionThrown(), true);
         DriverExecutionPrepareEngine<JDBCExecutionUnit, Connection> prepareEngine = createDriverExecutionPrepareEngine(isReturnGeneratedKeys, contextManager.getMetaDataContexts());
         SQLFederationContext context = new SQLFederationContext(
                 false, queryContext, contextManager.getMetaDataContexts().getMetaData(), databaseConnectionManager.getConnectionSession().getProcessId());