com.couchbase.connect.kafka.handler.sink.N1qlSinkHandler Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of kafka-connect-couchbase Show documentation
Show all versions of kafka-connect-couchbase Show documentation
A Kafka Connect Couchbase connector for copying data between Kafka and Couchbase Server.
The newest version!
/*
* Copyright 2021 Couchbase, Inc.
*
* Licensed 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 com.couchbase.connect.kafka.handler.sink;
import com.couchbase.client.java.ReactiveCollection;
import com.couchbase.client.java.json.JsonObject;
import com.couchbase.client.java.query.ReactiveQueryResult;
import com.couchbase.connect.kafka.config.sink.CouchbaseSinkConfig;
import com.couchbase.connect.kafka.config.sink.N1qlSinkHandlerConfig;
import com.couchbase.connect.kafka.util.config.ConfigHelper;
import org.apache.kafka.common.config.ConfigException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import static com.couchbase.client.java.query.QueryOptions.queryOptions;
public class N1qlSinkHandler implements SinkHandler {
private static final Logger log = LoggerFactory.getLogger(N1qlSinkHandler.class);
private static final String ID_FIELD = "__id__";
private N1qlSinkHandlerConfig.Operation mode;
protected String conditions;
protected boolean createDocuments;
protected String bucketName;
/**
* Returns a modified copy of the given map, using the alias property
* as the default for the canonical property.
*
* @param properties the property map to modify.
* @return a copy of the given map
*/
static Map withAlias(Map properties, String canonicalPropertyName, String aliasPropertyName) {
Map copy = new HashMap<>(properties);
String aliasValue = copy.get(aliasPropertyName);
String canonicalValue = copy.get(canonicalPropertyName);
if (canonicalValue == null && aliasValue != null) {
copy.put(canonicalPropertyName, aliasValue);
}
return copy;
}
@Override
public void init(SinkHandlerContext context) {
// Accept the legacy property "couchbase.create.document" as an alias for "couchbase.n1ql.create.document".
Map properties = withAlias(context.configProperties(),
ConfigHelper.keyName(N1qlSinkHandlerConfig.class, N1qlSinkHandlerConfig::n1qlCreateDocument),
"couchbase.create.document");
CouchbaseSinkConfig config = ConfigHelper.parse(CouchbaseSinkConfig.class, properties);
this.mode = config.n1qlOperation();
this.conditions = conditions(config.n1qlWhereFields());
this.createDocuments = config.n1qlCreateDocument();
this.bucketName = config.bucket();
if (mode == N1qlSinkHandlerConfig.Operation.UPDATE_WHERE && conditions.isEmpty()) {
String keyName = ConfigHelper.keyName(N1qlSinkHandlerConfig.class, N1qlSinkHandlerConfig::n1qlWhereFields);
throw new ConfigException("Missing required configuration \"" + keyName + "\".");
}
}
@Override
public SinkAction handle(SinkHandlerParams params) {
String documentId = getDocumentId(params);
SinkDocument doc = params.document().orElse(null);
String keySpace = keyspace(params);
if (doc == null) {
return SinkAction.remove(params, params.collection(), documentId);
}
final JsonObject node;
try {
node = JsonObject.fromJson(doc.content());
} catch (Exception e) {
log.warn("could not generate n1ql statement from node (not json)", e);
return SinkAction.ignore();
}
if (node.isEmpty()) {
log.warn("could not generate n1ql statement from empty node");
return SinkAction.ignore();
}
for (String name : node.getNames()) {
if (name.contains("`")) {
// todo figure out how to escape backticks when generating N1QL statements.
// For now, bail out to avoid N1QL injection.
log.warn("could not generate n1ql statement from node with backtick (`) in field name");
return SinkAction.ignore();
}
}
String statement = getStatement(keySpace, node);
node.put(ID_FIELD, documentId);
// ReactiveCluster.query is an unholy blend of hot and cold.
// Make it truly cold by wrapping it with Mono.defer().
Mono action = Mono.defer(() ->
params.cluster()
.query(statement, queryOptions().parameters(node))
.map(ReactiveQueryResult::metaData)); // metadata arrival signals query completion
ConcurrencyHint concurrencyHint = mode == N1qlSinkHandlerConfig.Operation.UPDATE
? ConcurrencyHint.of(documentId) // UPDATE affects only this document
: ConcurrencyHint.neverConcurrent(); // UPDATE_WHERE affects unknown documents (usually the sames ones)
return new SinkAction(action, concurrencyHint);
}
/**
* Returns the target keyspace (bucket + scope + collection)
* pre-escaped for inclusion in a query statement.
*/
protected static String keyspace(SinkHandlerParams params) {
ReactiveCollection c = params.collection();
String bucket = c.bucketName();
String scope = defaultIfEmpty(c.scopeName(), "_default");
String collection = defaultIfEmpty(c.name(), "_default");
List components = new ArrayList<>();
components.add(bucket); // always include bucket
// For compatibility with pre-7.0 servers, omit scope and collection
// when the keyspace is the default collection.
boolean defaultCollection =
scope.equals("_default") && collection.equals("_default");
if (!defaultCollection) {
components.add(scope);
components.add(collection);
}
// Escape each component by enclosing in backticks (`),
// then put a dot (.) in between each component.
return components.stream()
.map(it -> "`" + it + "`")
.collect(Collectors.joining("."));
}
private static String defaultIfEmpty(String s, String defaultValue) {
return s.isEmpty() ? defaultValue : s;
}
private String getStatement(String keySpace, JsonObject kafkaMessage) {
switch (this.mode) {
case UPDATE_WHERE:
return updateWithConditionStatement(keySpace, kafkaMessage);
case UPDATE:
return createDocuments
? mergeStatement(keySpace, kafkaMessage)
: updateStatement(keySpace, kafkaMessage);
default:
throw new AssertionError("unrecognized n1ql mode");
}
}
private String updateStatement(String keySpace, JsonObject values) {
return "UPDATE " + keySpace +
" USE KEYS $" + ID_FIELD +
" SET " + assignments(values) +
" RETURNING meta().id;";
}
private String updateWithConditionStatement(String keySpace, JsonObject values) {
return "UPDATE " + keySpace +
" SET " + assignments(values) +
" WHERE " + conditions +
" RETURNING meta().id;";
}
private String mergeStatement(String keyspace, JsonObject values) {
return "MERGE INTO " + keyspace + " AS doc" +
" USING 1 AS o" + // dummy to satisfy the MERGE INTO syntax?
" ON KEY $" + ID_FIELD +
" WHEN MATCHED THEN UPDATE SET " + assignments(values, "doc.") +
" WHEN NOT MATCHED THEN INSERT " + values;
}
private static String assignments(JsonObject values) {
return assignments(values, "");
}
private static String assignments(JsonObject values, String prefix) {
List assignments = new ArrayList<>();
for (String name : values.getNames()) {
assignments.add(prefix + "`" + name + "` = $" + name);
}
return String.join(", ", assignments);
}
private static String conditions(List fields) {
List conditions = new ArrayList<>();
for (String name : fields) {
final String value;
int colonIndex = name.indexOf(':');
if (colonIndex != -1) {
// compare against a string constant (whatever's after the colon)
value = "'" + name.substring(colonIndex + 1) + "'";
name = name.substring(0, colonIndex);
} else {
value = "$" + name;
}
conditions.add("`" + name + "` = " + value);
}
return String.join(" AND ", conditions);
}
@Override
public String toString() {
return "N1qlSinkHandler{" +
"mode=" + mode +
", conditions='" + conditions + '\'' +
", createDocuments=" + createDocuments +
", bucketName='" + bucketName + '\'' +
'}';
}
}