@@ -0,0 +1,3300 @@
+/*%0A * Licensed to the Apache Software Foundation (ASF) under one or more%0A * contributor license agreements.  See the NOTICE file distributed with%0A * this work for additional information regarding copyright ownership.%0A * The ASF licenses this file to You under the Apache License, Version 2.0%0A * (the %22License%22); you may not use this file except in compliance with%0A * the License.  You may obtain a copy of the License at%0A *%0A *     http://www.apache.org/licenses/LICENSE-2.0%0A *%0A * Unless required by applicable law or agreed to in writing, software%0A * distributed under the License is distributed on an %22AS IS%22 BASIS,%0A * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.%0A * See the License for the specific language governing permissions and%0A * limitations under the License.%0A */%0A%0Apackage org.apache.shardingsphere.core.rewrite.token.generator;%0A%0Aimport com.google.common.base.Optional;%0Aimport org.apache.shardingsphere.core.parse.sql.segment.dml.assignment.AssignmentSegment;%0Aimport org.apache.shardingsphere.core.parse.sql.segment.dml.assignment.SetAssignmentsSegment;%0Aimport org.apache.shardingsphere.core.parse.sql.statement.SQLStatement;%0Aimport org.apache.shardingsphere.core.parse.sql.statement.dml.InsertStatement;%0Aimport org.apache.shardingsphere.core.parse.sql.token.impl.InsertSetAddGeneratedKeyToken;%0Aimport org.apache.shardingsphere.core.rule.ShardingRule;%0A%0Aimport java.util.ArrayList;%0Aimport java.util.List;%0A%0A/**%0A * Insert set add items token generator.%0A *%0A * @author panjuan%0A */%0Apublic final class InsertSetAddGeneratedKeyTokenGenerator implements OptionalSQLTokenGenerator%3CShardingRule%3E %7B%0A    %0A    @Override%0A    public Optional%3CInsertSetAddGeneratedKeyToken%3E generateSQLToken(final SQLStatement sqlStatement, final ShardingRule shardingRule) %7B%0A        Optional%3CSetAssignmentsSegment%3E setAssignmentsSegment = sqlStatement.findSQLSegment(SetAssignmentsSegment.class);%0A        if (!(sqlStatement instanceof InsertStatement && setAssignmentsSegment.isPresent())) %7B%0A            return Optional.absent();%0A        %7D%0A        return createInsertSetAddGeneratedKeyToken((InsertStatement) sqlStatement, shardingRule, setAssignmentsSegment.get());%0A    %7D%0A    %0A    private Optional%3CInsertSetAddGeneratedKeyToken%3E createInsertSetAddGeneratedKeyToken(final InsertStatement insertStatement, final ShardingRule shardingRule, final SetAssignmentsSegment segment) %7B%0A        Optional%3CString%3E generatedKeyColumn = getGeneratedKeyColumn(insertStatement, shardingRule);%0A        if (generatedKeyColumn.isPresent()) %7B%0A            List%3CAssignmentSegment%3E assignments = new ArrayList%3C%3E(segment.getAssignments());%0A            return Optional.of(new InsertSetAddGeneratedKeyToken(assignments.get(assignments.size() - 1).getStopIndex() + 1, generatedKeyColumn.get()));%0A        %7D%0A        return Optional.absent();%0A    %7D%0A    %0A    private Optional%3CString%3E getGeneratedKeyColumn(final InsertStatement insertStatement, final ShardingRule shardingRule) %7B%0A        String tableName = insertStatement.getTables().getSingleTableName();%0A        Optional%3CString%3E generateKeyColumn = shardingRule.findGenerateKeyColumnName(tableName);%0A        return generateKeyColumn.isPresent() && !insertStatement.getColumnNames().contains(generateKeyColumn.get()) ? generateKeyColumn : Optional.%3CString%3Eabsent();%0A    %7D%0A%7D%0A

add InsertSetAddGeneratedKeyTokenGenerator