Skip to content

[feature][connector] JDBC sinks: support upsert and row deletion #16448

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nicoloboschi merged 1 commit into from
Jul 14, 2022
Merged

[feature][connector] JDBC sinks: support upsert and row deletion #16448

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ flexible messaging model and an intuitive client API.</description>
<joda.version>2.10.5</joda.version>
<jclouds.version>2.5.0</jclouds.version>
<guice.version>5.1.0</guice.version>
<sqlite-jdbc.version>3.8.11.2</sqlite-jdbc.version>
<sqlite-jdbc.version>3.36.0.3</sqlite-jdbc.version>
<mysql-jdbc.version>8.0.11</mysql-jdbc.version>
<postgresql-jdbc.version>42.3.3</postgresql-jdbc.version>
<clickhouse-jdbc.version>0.3.2</clickhouse-jdbc.version>
Expand Down
123 changes: 90 additions & 33 deletions pulsar-io/jdbc/core/src/main/java/org/apache/pulsar/io/jdbc/BaseJdbcAutoSchemaSink.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,8 @@

import com.fasterxml.jackson.databind.JsonNode;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Lists;
import java.sql.PreparedStatement;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
Expand All @@ -44,48 +44,39 @@
public abstract class BaseJdbcAutoSchemaSink extends JdbcAbstractSink<GenericObject> {

@Override
public void bindValue(PreparedStatement statement,
Record<GenericObject> message, String action) throws Exception {
final GenericObject record = message.getValue();
Function<String, Object> recordValueGetter;
if (message.getSchema() != null && message.getSchema() instanceof KeyValueSchema) {
KeyValueSchema<GenericObject, GenericObject> keyValueSchema = (KeyValueSchema) message.getSchema();

final org.apache.pulsar.client.api.Schema<GenericObject> keySchema = keyValueSchema.getKeySchema();
final org.apache.pulsar.client.api.Schema<GenericObject> valueSchema = keyValueSchema.getValueSchema();
KeyValue<GenericObject, GenericObject> keyValue =
(KeyValue<GenericObject, GenericObject>) record.getNativeObject();

final GenericObject key = keyValue.getKey();
final GenericObject value = keyValue.getValue();

Map<String, Object> data = new HashMap<>();
fillKeyValueSchemaData(keySchema, key, data);
fillKeyValueSchemaData(valueSchema, value, data);
recordValueGetter = (k) -> data.get(k);
} else {
recordValueGetter = (key) -> ((GenericRecord) record).getField(key);
}
public String generateUpsertQueryStatement() {
throw new IllegalStateException("UPSERT not supported");
}

List<ColumnId> columns = Lists.newArrayList();
if (action == null || action.equals(INSERT)) {
columns = tableDefinition.getColumns();
} else if (action.equals(DELETE)){
columns.addAll(tableDefinition.getKeyColumns());
} else if (action.equals(UPDATE)){
columns.addAll(tableDefinition.getNonKeyColumns());
columns.addAll(tableDefinition.getKeyColumns());
@Override
public void bindValue(PreparedStatement statement, Mutation mutation) throws Exception {
final List<ColumnId> columns = new ArrayList<>();
switch (mutation.getType()) {
case INSERT:
columns.addAll(tableDefinition.getColumns());
break;
case UPSERT:
columns.addAll(tableDefinition.getColumns());
columns.addAll(tableDefinition.getNonKeyColumns());
break;
case UPDATE:
columns.addAll(tableDefinition.getNonKeyColumns());
columns.addAll(tableDefinition.getKeyColumns());
break;
case DELETE:
columns.addAll(tableDefinition.getKeyColumns());
break;
}

int index = 1;
for (ColumnId columnId : columns) {
String colName = columnId.getName();
int colType = columnId.getType();
if (log.isDebugEnabled()) {
log.debug("colName: {} colType: {}", colName, colType);
log.debug("getting value for column: {} type: {}", colName, colType);
}
try {
Object obj = recordValueGetter.apply(colName);
Object obj = mutation.getValues().apply(colName);
if (obj != null) {
setColumnValue(statement, index++, obj);
} else {
Expand All @@ -105,6 +96,66 @@ public void bindValue(PreparedStatement statement,
}
}

@Override
public Mutation createMutation(Record<GenericObject> message) {
final GenericObject record = message.getValue();
Function<String, Object> recordValueGetter;
MutationType mutationType = null;
if (message.getSchema() != null && message.getSchema() instanceof KeyValueSchema) {
KeyValueSchema<GenericObject, GenericObject> keyValueSchema = (KeyValueSchema) message.getSchema();

final org.apache.pulsar.client.api.Schema<GenericObject> keySchema = keyValueSchema.getKeySchema();
final org.apache.pulsar.client.api.Schema<GenericObject> valueSchema = keyValueSchema.getValueSchema();
KeyValue<GenericObject, GenericObject> keyValue =
(KeyValue<GenericObject, GenericObject>) record.getNativeObject();

final GenericObject key = keyValue.getKey();
final GenericObject value = keyValue.getValue();

boolean isDelete = false;
if (value == null) {
switch (jdbcSinkConfig.getNullValueAction()) {
case DELETE:
isDelete = true;
break;
case FAIL:
throw new IllegalArgumentException("Got record with value NULL with nullValueAction=FAIL");
default:
break;
}
}
Map<String, Object> data = new HashMap<>();
fillKeyValueSchemaData(keySchema, key, data);
if (isDelete) {
mutationType = MutationType.DELETE;
} else {
fillKeyValueSchemaData(valueSchema, value, data);
}
recordValueGetter = (k) -> data.get(k);
} else {
recordValueGetter = (key) -> ((GenericRecord) record).getField(key);
}
String action = message.getProperties().get(ACTION_PROPERTY);
if (action != null) {
mutationType = MutationType.valueOf(action);
} else if (mutationType == null) {
switch (jdbcSinkConfig.getInsertMode()) {
case INSERT:
mutationType = MutationType.INSERT;
break;
case UPSERT:
mutationType = MutationType.UPSERT;
break;
case UPDATE:
mutationType = MutationType.UPDATE;
break;
default:
throw new IllegalArgumentException("Unknown insert mode: " + jdbcSinkConfig.getInsertMode());
}
}
return new Mutation(mutationType, recordValueGetter);
}

private static void setColumnNull(PreparedStatement statement, int index, int type) throws Exception {
if (log.isDebugEnabled()) {
log.debug("Setting column value to null, statement: {}, index: {}", statement.toString(), index);
Expand Down Expand Up @@ -163,6 +214,9 @@ private static Object getValueFromJsonNode(final JsonNode fn) {
private static void fillKeyValueSchemaData(org.apache.pulsar.client.api.Schema<GenericObject> schema,
GenericObject record,
Map<String, Object> data) {
if (record == null) {
return;
}
switch (schema.getSchemaInfo().getType()) {
case JSON:
final JsonNode jsonNode = (JsonNode) record.getNativeObject();
Expand Down Expand Up @@ -190,6 +244,9 @@ private static void fillKeyValueSchemaData(org.apache.pulsar.client.api.Schema<G

@VisibleForTesting
static Object convertAvroField(Object avroValue, Schema schema) {
if (avroValue == null) {
return null;
}
switch (schema.getType()) {
case NULL:
case INT:
Expand Down
88 changes: 68 additions & 20 deletions pulsar-io/jdbc/core/src/main/java/org/apache/pulsar/io/jdbc/JdbcAbstractSink.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,9 @@
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Function;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.pulsar.functions.api.Record;
Expand All @@ -43,7 +46,7 @@
@Slf4j
public abstract class JdbcAbstractSink<T> implements Sink<T> {
// ----- Runtime fields
private JdbcSinkConfig jdbcSinkConfig;
protected JdbcSinkConfig jdbcSinkConfig;
@Getter
private Connection connection;
private String jdbcUrl;
Expand All @@ -52,13 +55,11 @@ public abstract class JdbcAbstractSink<T> implements Sink<T> {
private JdbcUtils.TableId tableId;
private PreparedStatement insertStatement;
private PreparedStatement updateStatement;
private PreparedStatement upsertStatement;
private PreparedStatement deleteStatement;


protected static final String ACTION = "ACTION";
protected static final String INSERT = "INSERT";
protected static final String UPDATE = "UPDATE";
protected static final String DELETE = "DELETE";
protected static final String ACTION_PROPERTY = "ACTION";

protected JdbcUtils.TableDefinition tableDefinition;

Expand Down Expand Up @@ -122,12 +123,15 @@ private void initStatement() throws Exception {
}

tableDefinition = JdbcUtils.getTableDefinition(connection, tableId, keyList, nonKeyList);
insertStatement = JdbcUtils.buildInsertStatement(connection, JdbcUtils.buildInsertSql(tableDefinition));
insertStatement = JdbcUtils.buildInsertStatement(connection, generateInsertQueryStatement());
if (jdbcSinkConfig.getInsertMode() == JdbcSinkConfig.InsertMode.UPSERT) {
upsertStatement = JdbcUtils.buildInsertStatement(connection, generateUpsertQueryStatement());
}
if (!nonKeyList.isEmpty()) {
updateStatement = JdbcUtils.buildUpdateStatement(connection, JdbcUtils.buildUpdateSql(tableDefinition));
updateStatement = JdbcUtils.buildUpdateStatement(connection, generateUpdateQueryStatement());
}
if (!keyList.isEmpty()) {
deleteStatement = JdbcUtils.buildDeleteStatement(connection, JdbcUtils.buildDeleteSql(tableDefinition));
deleteStatement = JdbcUtils.buildDeleteStatement(connection, generateDeleteQueryStatement());
}
}

Expand All @@ -136,6 +140,18 @@ public void close() throws Exception {
if (connection != null && !connection.getAutoCommit()) {
connection.commit();
}
if (insertStatement != null) {
insertStatement.close();
}
if (updateStatement != null) {
updateStatement.close();
}
if (upsertStatement != null) {
upsertStatement.close();
}
if (deleteStatement != null) {
deleteStatement.close();
}
if (flushExecutor != null) {
flushExecutor.shutdown();
flushExecutor = null;
Expand All @@ -159,10 +175,40 @@ public void write(Record<T> record) throws Exception {
}
}

public String generateInsertQueryStatement() {
return JdbcUtils.buildInsertSql(tableDefinition);
}

public String generateUpdateQueryStatement() {
return JdbcUtils.buildUpdateSql(tableDefinition);
}

public abstract String generateUpsertQueryStatement();

public String generateDeleteQueryStatement() {
return JdbcUtils.buildDeleteSql(tableDefinition);
}

// bind value with a PreparedStetement
public abstract void bindValue(
PreparedStatement statement,
Record<T> message, String action) throws Exception;
Mutation mutation) throws Exception;

public abstract Mutation createMutation(Record<T> message);

@Data
@AllArgsConstructor
protected static class Mutation {
private MutationType type;
private Function<String, Object> values;
}
protected enum MutationType {
INSERT,
UPDATE,
UPSERT,
DELETE
}


private void flush() {
// if not in flushing state, do flush, else return;
Expand All @@ -187,42 +233,44 @@ private void flush() {
try {
// bind each record value
for (Record<T> record : swapList) {
String action = record.getProperties().get(ACTION);
if (action == null) {
action = INSERT;
}
switch (action) {
final Mutation mutation = createMutation(record);
switch (mutation.getType()) {
case DELETE:
bindValue(deleteStatement, record, action);
bindValue(deleteStatement, mutation);
count += 1;
deleteStatement.execute();
break;
case UPDATE:
bindValue(updateStatement, record, action);
bindValue(updateStatement, mutation);
count += 1;
updateStatement.execute();
break;
case INSERT:
bindValue(insertStatement, record, action);
bindValue(insertStatement, mutation);
count += 1;
insertStatement.execute();
break;
case UPSERT:
bindValue(upsertStatement, mutation);
count += 1;
upsertStatement.execute();
break;
default:
String msg = String.format(
"Unsupported action %s, can be one of %s, or not set which indicate %s",
action, Arrays.asList(INSERT, UPDATE, DELETE), INSERT);
mutation.getType(), Arrays.toString(MutationType.values()), MutationType.INSERT);
throw new IllegalArgumentException(msg);
}
}
connection.commit();
swapList.forEach(Record::ack);
} catch (Exception e) {
log.error("Got exception ", e);
log.error("Got exception ", e.getMessage(), e);
swapList.forEach(Record::fail);
}

if (swapList.size() != count) {
log.error("Update count {} not match total number of records {}", count, swapList.size());
log.error("Update count {} not match total number of records {}", count, swapList.size());
}

// finish flush
Expand Down
29 changes: 29 additions & 0 deletions pulsar-io/jdbc/core/src/main/java/org/apache/pulsar/io/jdbc/JdbcSinkConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,35 @@ public class JdbcSinkConfig implements Serializable {
)
private int batchSize = 200;

@FieldDoc(
required = false,
defaultValue = "INSERT",
help = "If it is configured as UPSERT, the sink will use upsert semantics rather than "
+ "plain INSERT/UPDATE statements. Upsert semantics refer to atomically adding a new row or "
+ "updating the existing row if there is a primary key constraint violation, "
+ "which provides idempotence."
)
private InsertMode insertMode = InsertMode.INSERT;

@FieldDoc(
required = false,
defaultValue = "FAIL",
help = "How to handle records with null values, possible options are DELETE or FAIL."
)
private NullValueAction nullValueAction = NullValueAction.FAIL;

public enum InsertMode {
INSERT,
UPSERT,
UPDATE;
}

public enum NullValueAction {
FAIL,
DELETE
}


public static JdbcSinkConfig load(String yamlFile) throws IOException {
ObjectMapper mapper = new ObjectMapper(new YAMLFactory());
return mapper.readValue(new File(yamlFile), JdbcSinkConfig.class);
Expand Down
Loading