Skip to content

Commit 3e66135

Browse files
NicholasDColeNicholasDCole
committed
Port ES Scripted Upsert to 11.1.0 from confluentinc#759
Signed-off-by: Nicholas Cole <[email protected]>
1 parent cc26466 commit 3e66135

File tree

7 files changed

+337
-20
lines changed

7 files changed

+337
-20
lines changed

‎pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
<groupId>io.confluent</groupId>
1212
<artifactId>kafka-connect-elasticsearch</artifactId>
13-
<version>11.1.0</version>
13+
<version>11.1.0-LOCAL</version>
1414
<packaging>jar</packaging>
1515
<name>kafka-connect-elasticsearch</name>
1616
<organization>

‎src/main/java/io/confluent/connect/elasticsearch/DataConverter.java‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import com.fasterxml.jackson.databind.ObjectMapper;
2121
import com.fasterxml.jackson.databind.node.ObjectNode;
2222
import io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.BehaviorOnNullValues;
23+
import io.confluent.connect.elasticsearch.util.ScriptParser;
2324
import org.apache.kafka.connect.data.ConnectSchema;
2425
import org.apache.kafka.connect.data.Date;
2526
import org.apache.kafka.connect.data.Decimal;
@@ -40,6 +41,7 @@
4041
import org.elasticsearch.action.update.UpdateRequest;
4142
import org.elasticsearch.common.xcontent.XContentType;
4243
import org.elasticsearch.index.VersionType;
44+
import org.elasticsearch.script.Script;
4345
import org.slf4j.Logger;
4446
import org.slf4j.LoggerFactory;
4547

@@ -178,11 +180,44 @@ public DocWriteRequest<?> convertRecord(SinkRecord record, String index) {
178180
new IndexRequest(index).id(id).source(payload, XContentType.JSON).opType(opType),
179181
record
180182
);
183+
case SCRIPTED_UPSERT:
184+
try {
185+
186+
if (config.getIsPayloadAsParams()) {
187+
return buildUpdateRequestWithParams(index, payload, id);
188+
}
189+
190+
Script script = ScriptParser.parseScript(config.getScript());
191+
192+
return new UpdateRequest(index, id)
193+
.doc(payload, XContentType.JSON)
194+
.upsert(payload, XContentType.JSON)
195+
.retryOnConflict(Math.min(config.maxInFlightRequests(), 5))
196+
.script(script)
197+
.scriptedUpsert(true);
198+
199+
} catch (JsonProcessingException jsonProcessingException) {
200+
throw new RuntimeException(jsonProcessingException);
201+
}
181202
default:
182203
return null; // shouldn't happen
183204
}
184205
}
185206

207+
private UpdateRequest buildUpdateRequestWithParams(String index, String payload, String id)
208+
throws JsonProcessingException {
209+
210+
Script script = ScriptParser.parseScriptWithParams(config.getScript(), payload);
211+
212+
UpdateRequest updateRequest =
213+
new UpdateRequest(index, id)
214+
.retryOnConflict(Math.min(config.maxInFlightRequests(), 5))
215+
.script(script)
216+
.scriptedUpsert(true);
217+
218+
return updateRequest;
219+
}
220+
186221
private String getPayload(SinkRecord record) {
187222
if (record.value() == null) {
188223
return null;

‎src/main/java/io/confluent/connect/elasticsearch/ElasticsearchSinkConnectorConfig.java‎

Lines changed: 68 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424
import java.util.Set;
2525
import java.util.stream.Collectors;
2626
import java.util.concurrent.TimeUnit;
27+
28+
import io.confluent.connect.elasticsearch.validator.ScriptValidator;
2729
import org.apache.kafka.common.config.AbstractConfig;
2830
import org.apache.kafka.common.config.ConfigDef;
2931
import org.apache.kafka.common.config.ConfigException;
@@ -256,6 +258,22 @@ public class ElasticsearchSinkConnectorConfig extends AbstractConfig {
256258
private static final String WRITE_METHOD_DISPLAY = "Write Method";
257259
private static final String WRITE_METHOD_DEFAULT = WriteMethod.INSERT.name();
258260

261+
public static final String UPSERT_SCRIPT_CONFIG = "upsert.script";
262+
263+
private static final String UPSERT_SCRIPT_DOC = "Script used for"
264+
+ " upserting data to Elasticsearch. This script allows for"
265+
+ " customizable behavior upon upserting a document. Please refer to"
266+
+ " Elasticsearch scripted upsert documentation";
267+
268+
private static final String UPSERT_SCRIPT_DISPLAY = "Upsert Script";
269+
270+
public static final String PAYLOAD_AS_PARAMS_CONFIG = "payload.as.params";
271+
272+
private static final String PAYLOAD_AS_PARAMS_DOC = "Defines Payload to be injected"
273+
+ " into upsert.script script component as params object";
274+
275+
private static final String PAYLOAD_AS_PARAMS_DISPLAY = "Payload as Params";
276+
259277
// Proxy group
260278
public static final String PROXY_HOST_CONFIG = "proxy.host";
261279
private static final String PROXY_HOST_DISPLAY = "Proxy Host";
@@ -379,7 +397,8 @@ public enum SecurityProtocol {
379397

380398
public enum WriteMethod {
381399
INSERT,
382-
UPSERT
400+
UPSERT,
401+
SCRIPTED_UPSERT
383402
}
384403

385404
protected static ConfigDef baseConfigDef() {
@@ -573,8 +592,8 @@ private static void addConversionConfigs(ConfigDef configDef) {
573592
DATA_CONVERSION_GROUP,
574593
++order,
575594
Width.SHORT,
576-
IGNORE_KEY_DISPLAY
577-
).define(
595+
IGNORE_KEY_DISPLAY)
596+
.define(
578597
IGNORE_SCHEMA_CONFIG,
579598
Type.BOOLEAN,
580599
IGNORE_SCHEMA_DEFAULT,
@@ -583,8 +602,8 @@ private static void addConversionConfigs(ConfigDef configDef) {
583602
DATA_CONVERSION_GROUP,
584603
++order,
585604
Width.SHORT,
586-
IGNORE_SCHEMA_DISPLAY
587-
).define(
605+
IGNORE_SCHEMA_DISPLAY)
606+
.define(
588607
COMPACT_MAP_ENTRIES_CONFIG,
589608
Type.BOOLEAN,
590609
COMPACT_MAP_ENTRIES_DEFAULT,
@@ -593,8 +612,8 @@ private static void addConversionConfigs(ConfigDef configDef) {
593612
DATA_CONVERSION_GROUP,
594613
++order,
595614
Width.SHORT,
596-
COMPACT_MAP_ENTRIES_DISPLAY
597-
).define(
615+
COMPACT_MAP_ENTRIES_DISPLAY)
616+
.define(
598617
IGNORE_KEY_TOPICS_CONFIG,
599618
Type.LIST,
600619
IGNORE_KEY_TOPICS_DEFAULT,
@@ -603,8 +622,8 @@ private static void addConversionConfigs(ConfigDef configDef) {
603622
DATA_CONVERSION_GROUP,
604623
++order,
605624
Width.LONG,
606-
IGNORE_KEY_TOPICS_DISPLAY
607-
).define(
625+
IGNORE_KEY_TOPICS_DISPLAY)
626+
.define(
608627
IGNORE_SCHEMA_TOPICS_CONFIG,
609628
Type.LIST,
610629
IGNORE_SCHEMA_TOPICS_DEFAULT,
@@ -613,8 +632,8 @@ private static void addConversionConfigs(ConfigDef configDef) {
613632
DATA_CONVERSION_GROUP,
614633
++order,
615634
Width.LONG,
616-
IGNORE_SCHEMA_TOPICS_DISPLAY
617-
).define(
635+
IGNORE_SCHEMA_TOPICS_DISPLAY)
636+
.define(
618637
DROP_INVALID_MESSAGE_CONFIG,
619638
Type.BOOLEAN,
620639
DROP_INVALID_MESSAGE_DEFAULT,
@@ -623,8 +642,8 @@ private static void addConversionConfigs(ConfigDef configDef) {
623642
DATA_CONVERSION_GROUP,
624643
++order,
625644
Width.LONG,
626-
DROP_INVALID_MESSAGE_DISPLAY
627-
).define(
645+
DROP_INVALID_MESSAGE_DISPLAY)
646+
.define(
628647
BEHAVIOR_ON_NULL_VALUES_CONFIG,
629648
Type.STRING,
630649
BEHAVIOR_ON_NULL_VALUES_DEFAULT.name(),
@@ -635,8 +654,8 @@ private static void addConversionConfigs(ConfigDef configDef) {
635654
++order,
636655
Width.SHORT,
637656
BEHAVIOR_ON_NULL_VALUES_DISPLAY,
638-
new EnumRecommender<>(BehaviorOnNullValues.class)
639-
).define(
657+
new EnumRecommender<>(BehaviorOnNullValues.class))
658+
.define(
640659
BEHAVIOR_ON_MALFORMED_DOCS_CONFIG,
641660
Type.STRING,
642661
BEHAVIOR_ON_MALFORMED_DOCS_DEFAULT.name(),
@@ -647,8 +666,8 @@ private static void addConversionConfigs(ConfigDef configDef) {
647666
++order,
648667
Width.SHORT,
649668
BEHAVIOR_ON_MALFORMED_DOCS_DISPLAY,
650-
new EnumRecommender<>(BehaviorOnMalformedDoc.class)
651-
).define(
669+
new EnumRecommender<>(BehaviorOnMalformedDoc.class))
670+
.define(
652671
WRITE_METHOD_CONFIG,
653672
Type.STRING,
654673
WRITE_METHOD_DEFAULT,
@@ -659,8 +678,30 @@ private static void addConversionConfigs(ConfigDef configDef) {
659678
++order,
660679
Width.SHORT,
661680
WRITE_METHOD_DISPLAY,
662-
new EnumRecommender<>(WriteMethod.class)
663-
);
681+
new EnumRecommender<>(WriteMethod.class))
682+
.define(
683+
UPSERT_SCRIPT_CONFIG,
684+
Type.STRING,
685+
null,
686+
new ScriptValidator(),
687+
Importance.LOW,
688+
UPSERT_SCRIPT_DOC,
689+
DATA_CONVERSION_GROUP,
690+
++order,
691+
Width.SHORT,
692+
UPSERT_SCRIPT_DISPLAY,
693+
new ScriptValidator())
694+
.define(
695+
PAYLOAD_AS_PARAMS_CONFIG,
696+
Type.BOOLEAN,
697+
false,
698+
Importance.LOW,
699+
PAYLOAD_AS_PARAMS_DOC,
700+
DATA_CONVERSION_GROUP,
701+
++order,
702+
Width.SHORT,
703+
PAYLOAD_AS_PARAMS_DISPLAY);
704+
;
664705
}
665706

666707
private static void addProxyConfigs(ConfigDef configDef) {
@@ -989,6 +1030,14 @@ public WriteMethod writeMethod() {
9891030
return WriteMethod.valueOf(getString(WRITE_METHOD_CONFIG).toUpperCase());
9901031
}
9911032

1033+
public String getScript() {
1034+
return getString(UPSERT_SCRIPT_CONFIG);
1035+
}
1036+
1037+
public Boolean getIsPayloadAsParams() {
1038+
return getBoolean(PAYLOAD_AS_PARAMS_CONFIG);
1039+
}
1040+
9921041
private static class DataStreamDatasetValidator implements Validator {
9931042

9941043
@Override

‎src/main/java/io/confluent/connect/elasticsearch/util/ScriptParser.java‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
/*
2+
* Copyright 2018 Confluent Inc.
3+
*
4+
* Licensed under the Confluent Community License (the "License"); you may not use
5+
* this file except in compliance with the License. You may obtain a copy of the
6+
* License at
7+
*
8+
* http://www.confluent.io/confluent-community-license
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
12+
* WARRANTIES OF ANY KIND, either express or implied. See the License for the
13+
* specific language governing permissions and limitations under the License.
14+
*/
15+
16+
package io.confluent.connect.elasticsearch.util;
17+
18+
import com.fasterxml.jackson.core.JsonProcessingException;
19+
import com.fasterxml.jackson.core.type.TypeReference;
20+
import com.fasterxml.jackson.databind.ObjectMapper;
21+
import org.elasticsearch.script.Script;
22+
23+
import java.util.Map;
24+
25+
public class ScriptParser {
26+
27+
private static final ObjectMapper objectMapper = new ObjectMapper();
28+
29+
public static Script parseScript(String scriptJson) throws JsonProcessingException {
30+
31+
Map<String, Object> map = ScriptParser.parseSchemaStringAsJson(scriptJson);
32+
33+
return Script.parse(map);
34+
}
35+
36+
private static Map<String, Object> parseSchemaStringAsJson(String scriptJson)
37+
throws JsonProcessingException {
38+
39+
ObjectMapper objectMapper = new ObjectMapper();
40+
41+
Map<String, Object> scriptConverted;
42+
43+
scriptConverted = objectMapper.readValue(
44+
scriptJson, new TypeReference<Map<String, Object>>(){});
45+
46+
return scriptConverted;
47+
}
48+
49+
public static Script parseScriptWithParams(String scriptJson, String jsonPayload)
50+
throws JsonProcessingException {
51+
52+
Map<String, Object> map = ScriptParser.parseSchemaStringAsJson(scriptJson);
53+
54+
Map<String, Object> fields = objectMapper.readValue(jsonPayload,
55+
new TypeReference<Map<String, Object>>() {});
56+
57+
map.put("params", fields);
58+
59+
return Script.parse(map);
60+
}
61+
}

‎src/main/java/io/confluent/connect/elasticsearch/validator/ScriptValidator.java‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
/*
2+
* Copyright 2018 Confluent Inc.
3+
*
4+
* Licensed under the Confluent Community License (the "License"); you may not use
5+
* this file except in compliance with the License. You may obtain a copy of the
6+
* License at
7+
*
8+
* http://www.confluent.io/confluent-community-license
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
12+
* WARRANTIES OF ANY KIND, either express or implied. See the License for the
13+
* specific language governing permissions and limitations under the License.
14+
*/
15+
16+
package io.confluent.connect.elasticsearch.validator;
17+
18+
import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.WRITE_METHOD_CONFIG;
19+
import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.WriteMethod.SCRIPTED_UPSERT;
20+
21+
import com.fasterxml.jackson.core.JsonProcessingException;
22+
import io.confluent.connect.elasticsearch.util.ScriptParser;
23+
import java.util.ArrayList;
24+
import java.util.List;
25+
import java.util.Map;
26+
import org.apache.kafka.common.config.ConfigDef;
27+
import org.apache.kafka.common.config.ConfigException;
28+
import org.elasticsearch.script.Script;
29+
30+
public class ScriptValidator implements ConfigDef.Validator, ConfigDef.Recommender {
31+
32+
@Override
33+
@SuppressWarnings("unchecked")
34+
public void ensureValid(String name, Object value) {
35+
36+
if (value == null) {
37+
return;
38+
}
39+
40+
String script = (String) value;
41+
42+
try {
43+
Script parsedScript = ScriptParser.parseScript(script);
44+
45+
if (parsedScript.getIdOrCode() == null) {
46+
throw new ConfigException(name, script, "The specified script is missing code");
47+
} else if (parsedScript.getLang() == null) {
48+
throw new ConfigException(name, script, "The specified script is missing lang");
49+
}
50+
51+
} catch (JsonProcessingException jsonProcessingException) {
52+
throw new ConfigException(
53+
name, script, "The specified script is not a valid Elasticsearch painless script");
54+
}
55+
}
56+
57+
@Override
58+
public String toString() {
59+
return "A valid script that is able to be parsed";
60+
}
61+
62+
@Override
63+
public List<Object> validValues(String name, Map<String, Object> parsedConfig) {
64+
if (!parsedConfig.get(WRITE_METHOD_CONFIG).equals(SCRIPTED_UPSERT)) {
65+
return new ArrayList<>();
66+
}
67+
return null;
68+
}
69+
70+
@Override
71+
public boolean visible(String name, Map<String, Object> parsedConfig) {
72+
return parsedConfig.get(WRITE_METHOD_CONFIG).equals(SCRIPTED_UPSERT.name());
73+
}
74+
}

0 commit comments

Comments
 (0)