-
Notifications
You must be signed in to change notification settings - Fork 6.6k
/
EncryptInsertOnUpdateTokenGenerator.java
128 lines (113 loc) · 8 KB
/
EncryptInsertOnUpdateTokenGenerator.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
/*
* 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.encrypt.rewrite.token.generator.impl;
import com.google.common.base.Preconditions;
import org.apache.shardingsphere.encrypt.rewrite.token.generator.BaseEncryptSQLTokenGenerator;
import org.apache.shardingsphere.encrypt.rewrite.token.pojo.EncryptAssignmentToken;
import org.apache.shardingsphere.encrypt.rewrite.token.pojo.EncryptLiteralAssignmentToken;
import org.apache.shardingsphere.encrypt.rewrite.token.pojo.EncryptParameterAssignmentToken;
import org.apache.shardingsphere.infra.rewrite.sql.token.generator.CollectionSQLTokenGenerator;
import org.apache.shardingsphere.sql.parser.binder.statement.SQLStatementContext;
import org.apache.shardingsphere.sql.parser.binder.statement.dml.InsertStatementContext;
import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.assignment.AssignmentSegment;
import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.column.OnDuplicateKeyColumnsSegment;
import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.expr.simple.LiteralExpressionSegment;
import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.expr.simple.ParameterMarkerExpressionSegment;
import org.apache.shardingsphere.sql.parser.sql.common.statement.dml.InsertStatement;
import org.apache.shardingsphere.sql.parser.sql.dialect.handler.dml.InsertStatementHandler;
import java.util.Collection;
import java.util.Collections;
import java.util.LinkedList;
import java.util.Optional;
/**
* Insert on update values token generator for encrypt.
*/
public final class EncryptInsertOnUpdateTokenGenerator extends BaseEncryptSQLTokenGenerator implements CollectionSQLTokenGenerator<InsertStatementContext> {
@Override
protected boolean isGenerateSQLTokenForEncrypt(final SQLStatementContext sqlStatementContext) {
return sqlStatementContext instanceof InsertStatementContext
&& InsertStatementHandler.getOnDuplicateKeyColumnsSegment(((InsertStatementContext) sqlStatementContext).getSqlStatement()).isPresent();
}
@Override
public Collection<EncryptAssignmentToken> generateSQLTokens(final InsertStatementContext insertStatementContext) {
Collection<EncryptAssignmentToken> result = new LinkedList<>();
InsertStatement insertStatement = insertStatementContext.getSqlStatement();
String tableName = insertStatement.getTable().getTableName().getIdentifier().getValue();
Preconditions.checkState(InsertStatementHandler.getOnDuplicateKeyColumnsSegment(insertStatement).isPresent());
OnDuplicateKeyColumnsSegment onDuplicateKeyColumnsSegment = InsertStatementHandler.getOnDuplicateKeyColumnsSegment(insertStatement).get();
Collection<AssignmentSegment> onDuplicateKeyColumnsSegments = onDuplicateKeyColumnsSegment.getColumns();
if (onDuplicateKeyColumnsSegments.isEmpty()) {
return result;
}
for (AssignmentSegment each : onDuplicateKeyColumnsSegments) {
if (getEncryptRule().findEncryptor(tableName, each.getColumn().getIdentifier().getValue()).isPresent()) {
generateSQLToken(tableName, each).ifPresent(result::add);
}
}
return result;
}
private Optional<EncryptAssignmentToken> generateSQLToken(final String tableName, final AssignmentSegment assignmentSegment) {
if (assignmentSegment.getValue() instanceof ParameterMarkerExpressionSegment) {
return Optional.of(generateParameterSQLToken(tableName, assignmentSegment));
}
if (assignmentSegment.getValue() instanceof LiteralExpressionSegment) {
return Optional.of(generateLiteralSQLToken(tableName, assignmentSegment));
}
return Optional.empty();
}
private EncryptAssignmentToken generateParameterSQLToken(final String tableName, final AssignmentSegment assignmentSegment) {
EncryptParameterAssignmentToken result = new EncryptParameterAssignmentToken(assignmentSegment.getColumn().getStartIndex(), assignmentSegment.getStopIndex());
String columnName = assignmentSegment.getColumn().getIdentifier().getValue();
addCipherColumn(tableName, columnName, result);
addAssistedQueryColumn(tableName, columnName, result);
addPlainColumn(tableName, columnName, result);
return result;
}
private EncryptAssignmentToken generateLiteralSQLToken(final String tableName, final AssignmentSegment assignmentSegment) {
EncryptLiteralAssignmentToken result = new EncryptLiteralAssignmentToken(assignmentSegment.getColumn().getStartIndex(), assignmentSegment.getStopIndex());
addCipherAssignment(tableName, assignmentSegment, result);
addAssistedQueryAssignment(tableName, assignmentSegment, result);
addPlainAssignment(tableName, assignmentSegment, result);
return result;
}
private void addCipherColumn(final String tableName, final String columnName, final EncryptParameterAssignmentToken token) {
token.addColumnName(getEncryptRule().getCipherColumn(tableName, columnName));
}
private void addAssistedQueryColumn(final String tableName, final String columnName, final EncryptParameterAssignmentToken token) {
getEncryptRule().findAssistedQueryColumn(tableName, columnName).ifPresent(token::addColumnName);
}
private void addPlainColumn(final String tableName, final String columnName, final EncryptParameterAssignmentToken token) {
getEncryptRule().findPlainColumn(tableName, columnName).ifPresent(token::addColumnName);
}
private void addCipherAssignment(final String tableName, final AssignmentSegment assignmentSegment, final EncryptLiteralAssignmentToken token) {
Object originalValue = ((LiteralExpressionSegment) assignmentSegment.getValue()).getLiterals();
Object cipherValue = getEncryptRule().getEncryptValues(tableName, assignmentSegment.getColumn().getIdentifier().getValue(), Collections.singletonList(originalValue)).iterator().next();
token.addAssignment(getEncryptRule().getCipherColumn(tableName, assignmentSegment.getColumn().getIdentifier().getValue()), cipherValue);
}
private void addAssistedQueryAssignment(final String tableName, final AssignmentSegment assignmentSegment, final EncryptLiteralAssignmentToken token) {
getEncryptRule().findAssistedQueryColumn(tableName, assignmentSegment.getColumn().getIdentifier().getValue()).ifPresent(assistedQueryColumn -> {
Object originalValue = ((LiteralExpressionSegment) assignmentSegment.getValue()).getLiterals();
Object assistedQueryValue = getEncryptRule().getEncryptAssistedQueryValues(tableName, assignmentSegment.getColumn().getIdentifier().getValue(), Collections.singletonList(originalValue))
.iterator().next();
token.addAssignment(assistedQueryColumn, assistedQueryValue);
});
}
private void addPlainAssignment(final String tableName, final AssignmentSegment assignmentSegment, final EncryptLiteralAssignmentToken token) {
Object originalValue = ((LiteralExpressionSegment) assignmentSegment.getValue()).getLiterals();
getEncryptRule().findPlainColumn(tableName, assignmentSegment.getColumn().getIdentifier().getValue()).ifPresent(plainColumn -> token.addAssignment(plainColumn, originalValue));
}
}