Added upsert by API

Author: t3t5u

README.md

@@ -17,7 +17,7 @@ kintone output plugin for Embulk stores app records from kintone.
 - **basic_auth_password**: kintone basic auth password (string, optional)
 - **guest_space_id**: kintone app belongs to guest space, guest space id is required. (integer, optional)
 - **mode**: kintone mode (string, required)
-- **update_key**: Column name to set update key (string, required if mode is update or insert_or_update)
+- **update_key**: Column name to set update key (string, required if mode is update, upsert or insert_or_update)
 - **reduce_key**: Key column name to reduce expanded SUBTABLE (string, optional)
 - **sort_columns**: List of columns for sorting input records (array of objects, optional)
     - **name**: Column name (string, required)
@@ -29,11 +29,13 @@ kintone output plugin for Embulk stores app records from kintone.
 - **skip_if_non_existing_id_or_update_key**: The skip policy if the record corresponding to the id or update key does not exist (string `auto`, `never` or `always`, default is `auto`). No effect for insert mode.
     - **auto**:
         - **update mode**: Skip the record if no id or update key value is specified.
+        - **upsert mode**: Skip the record if no update key value is specified. Insert the record if no id is specified.
         - **insert_or_update mode**: Skip the record if corresponds to the id does not exist or no update key value is specified.
     - **never**: Never skip the record even if corresponds to the id or update key does not exist.
         - **update mode**: Throw exception if no id or update key value is specified.
+        - **upsert mode**: Insert the record even if no id or update key value is specified.
         - **insert_or_update mode**: Insert the record if corresponds to the id or update key does not exist (also, if no id or update key value is specified).
-    - **always**: Always skip the record if corresponds to the id or update key does not exist (also, if no id or update key value is specified). update mode and insert_or_update mode will the same behavior (only updated, never inserted).
+    - **always**: Always skip the record if corresponds to the id or update key does not exist (also, if no id or update key value is specified). update mode, upsert mode and insert_or_update mode will the same behavior (only updated, never inserted).
 - **column_options** advanced: a key-value pairs where key is a column name and value is options for the column.
     - **field_code**: field code (string, required)
     - **type**: field type (string, required). See [this page](https://cybozu.dev/ja/kintone/docs/overview/field-types/#field-type-update) for list of available types. However, following types are not yet supported

src/main/java/org/embulk/output/kintone/KintoneMode.java

@@ -36,6 +36,22 @@ public void add(Page page, Skip skip, KintonePageOutput output) {
       consumer.accept(page);
     }
   },
+  UPSERT("upsert") {
+    @Override
+    public void validate(PluginTask task, KintoneClient client) {
+      if (!task.getUpdateKeyName().isPresent() && client.getColumn(Id.FIELD) == null) {
+        throw new ConfigException("When mode is upsert, require update_key or id column.");
+      }
+      client.validateIdOrUpdateKey(task.getUpdateKeyName().orElse(Id.FIELD));
+    }
+
+    @Override
+    public void add(Page page, Skip skip, KintonePageOutput output) {
+      Consumer<Page> consumer =
+          skip == Skip.ALWAYS ? output::insertOrUpdatePage : output::upsertPage;
+      consumer.accept(page);
+    }
+  },
   INSERT_OR_UPDATE("insert_or_update") {
     @Override
     public void validate(PluginTask task, KintoneClient client) {

src/main/java/org/embulk/output/kintone/KintonePageOutput.java

@@ -163,6 +163,15 @@ private void update(List<RecordForUpdate> records, Map<Operation, Integer> resul
     Operation.UPDATE.compute(results, list.size());
   }
 
+  private void upsert(List<RecordForUpdate> records, Map<Operation, Integer> results) {
+    List<RecordRevision> list =
+        executeWithRetry(
+            () -> client.get().record().updateRecords(task.getAppId(), records, true), records);
+    list.stream()
+        .collect(Collectors.groupingBy(RecordRevision::getOperation))
+        .forEach((key, value) -> Operation.valueOf(key).compute(results, value.size()));
+  }
+
   private <T> T executeWithRetry(Supplier<T> operation) {
     return executeWithRetry(operation, null);
   }
@@ -296,6 +305,51 @@ public void updatePage(Page page) {
     Operation.UPDATE.info(results);
   }
 
+  public void upsertPage(Page page) {
+    Skip skip = task.getSkipIfNonExistingIdOrUpdateKey();
+    String columnName = task.getUpdateKeyName().orElse(Id.FIELD);
+    boolean isId = columnName.equals(Id.FIELD);
+    int increments = 0;
+    List<RecordForUpdate> records = new ArrayList<>();
+    reader.setPage(page);
+    KintoneColumnVisitor visitor =
+        new KintoneColumnVisitor(
+            reader,
+            task.getDerivedColumns(),
+            task.getColumnOptions(),
+            task.getPreferNulls(),
+            task.getIgnoreNulls(),
+            task.getReduceKeyName().orElse(null),
+            task.getUpdateKeyName().orElse(Id.FIELD));
+    Map<Operation, Integer> results = new TreeMap<>();
+    while (reader.nextRecord()) {
+      Record record = new Record();
+      IdOrUpdateKey idOrUpdateKey = new IdOrUpdateKey();
+      visitor.setRecord(record);
+      visitor.setIdOrUpdateKey(idOrUpdateKey);
+      reader.getSchema().visitColumns(visitor);
+      putWrongTypeFields(record);
+      if (isId && !idOrUpdateKey.isIdPresent()) {
+        // Record inserted because no id was specified
+        idOrUpdateKey.getId().setValue(Long.MAX_VALUE - increments++);
+      } else if (skip == Skip.AUTO && !isId && !idOrUpdateKey.isUpdateKeyPresent()) {
+        LOGGER.warn("Record skipped because no update key value was specified");
+        continue;
+      } else if (!isId && !idOrUpdateKey.isUpdateKeyPresent()) {
+        LOGGER.warn("Record inserted though no update key value was specified");
+      }
+      records.add(idOrUpdateKey.forUpdate(record));
+      if (records.size() == task.getChunkSize()) {
+        upsert(records, results);
+        records.clear();
+      }
+    }
+    if (!records.isEmpty()) {
+      upsert(records, results);
+    }
+    results.keySet().forEach(operation -> operation.info(results));
+  }
+
   public void insertOrUpdatePage(Page page) {
     List<Record> records = new ArrayList<>();
     List<IdOrUpdateKey> idOrUpdateKeys = new ArrayList<>();

src/test/java/org/embulk/output/kintone/KintoneClientTest.java

@@ -59,6 +59,27 @@ public void testUpdate() {
     assertConfigException("The id column must be 'long'.", id(Types.STRING));
   }
 
+  @Test
+  public void testUpsert() {
+    merge(config("mode: upsert"));
+    merge(config("update_key: null"));
+    assertConfigException("When mode is upsert, require update_key or id column.");
+    merge(config("update_key: non_existing_column"));
+    assertConfigException("The column 'non_existing_column' for update does not exist.");
+    merge(config("update_key: non_existing_field"));
+    assertConfigException("The field 'non_existing_field' for update does not exist.");
+    merge(config("update_key: invalid_type_field_multi_line_text"));
+    assertConfigException("The update_key must be 'SINGLE_LINE_TEXT' or 'NUMBER'.");
+    merge(config("update_key: long_number"));
+    runWithMockClient(Lazy::get);
+    merge(config("update_key: string_single_line_text"));
+    runWithMockClient(Lazy::get);
+    merge(config("update_key: $id"));
+    runWithMockClient(Lazy::get, id(Types.LONG));
+    merge(config("update_key: null"));
+    assertConfigException("The id column must be 'long'.", id(Types.STRING));
+  }
+
   @Test
   public void testInsertOrUpdate() {
     merge(config("mode: insert_or_update"));

src/test/java/org/embulk/output/kintone/KintonePageOutputVerifier.java

@@ -22,12 +22,15 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.function.BiConsumer;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import org.embulk.config.TaskReport;
 import org.embulk.exec.PooledBufferAllocator;
 import org.embulk.output.kintone.record.Id;
+import org.embulk.output.kintone.record.Skip;
 import org.embulk.spi.BufferAllocator;
 import org.embulk.spi.Exec;
 import org.embulk.spi.Page;
@@ -40,37 +43,49 @@
 public class KintonePageOutputVerifier implements TransactionalPageOutput {
   private final TransactionalPageOutput transactionalPageOutput;
   private final String domain;
+  private final KintoneMode mode;
   private final String field;
+  private final Skip skip;
   private final List<String> values;
   private final List<String> addValues;
   private final List<Record> addRecords;
   private final List<RecordForUpdate> updateRecords;
+  private final Predicate<Record> nonNull;
 
   public KintonePageOutputVerifier(
       String domain,
+      KintoneMode mode,
       String field,
+      Skip skip,
       List<String> values,
       List<String> addValues,
       List<Record> addRecords,
-      List<RecordForUpdate> updateRecords) {
-    this(null, domain, field, values, addValues, addRecords, updateRecords);
+      List<RecordForUpdate> updateRecords,
+      Predicate<Record> nonNull) {
+    this(null, domain, mode, field, skip, values, addValues, addRecords, updateRecords, nonNull);
   }
 
   public KintonePageOutputVerifier(
       TransactionalPageOutput transactionalPageOutput,
       String domain,
+      KintoneMode mode,
       String field,
+      Skip skip,
       List<String> values,
       List<String> addValues,
       List<Record> addRecords,
-      List<RecordForUpdate> updateRecords) {
+      List<RecordForUpdate> updateRecords,
+      Predicate<Record> nonNull) {
     this.transactionalPageOutput = transactionalPageOutput;
     this.domain = domain;
+    this.mode = mode;
     this.field = field;
+    this.skip = skip;
     this.values = values;
     this.addValues = addValues;
     this.addRecords = addRecords;
     this.updateRecords = updateRecords;
+    this.nonNull = nonNull;
   }
 
   @Override
@@ -124,7 +139,7 @@ private void runWithMockClient(MockClient.Runnable runnable) throws Exception {
     mockClient.run(runnable);
     @SuppressWarnings("unchecked")
     ArgumentCaptor<List<Record>> addRecordsArgumentCaptor = ArgumentCaptor.forClass(List.class);
-    verify(mockRecordClient, atLeast(0)).addRecords(eq(0L), addRecordsArgumentCaptor.capture());
+    addRecordsVerifier().accept(mockRecordClient, addRecordsArgumentCaptor);
     assertRecords(
         domain,
         addRecordsArgumentCaptor.getAllValues().stream()
@@ -134,8 +149,7 @@ private void runWithMockClient(MockClient.Runnable runnable) throws Exception {
     @SuppressWarnings("unchecked")
     ArgumentCaptor<List<RecordForUpdate>> updateRecordsArgumentCaptor =
         ArgumentCaptor.forClass(List.class);
-    verify(mockRecordClient, atLeast(0))
-        .updateRecords(eq(0L), updateRecordsArgumentCaptor.capture());
+    updateRecordsVerifier().accept(mockRecordClient, updateRecordsArgumentCaptor);
     assertRecordForUpdates(
         domain,
         updateRecordsArgumentCaptor.getAllValues().stream()
@@ -144,6 +158,37 @@ private void runWithMockClient(MockClient.Runnable runnable) throws Exception {
         updateRecords);
   }
 
+  private BiConsumer<RecordClient, ArgumentCaptor<List<Record>>> addRecordsVerifier() {
+    return mode == KintoneMode.UPSERT && skip != Skip.ALWAYS
+        ? (rc, ac) -> {}
+        : this::verifyAddRecords;
+  }
+
+  private BiConsumer<RecordClient, ArgumentCaptor<List<RecordForUpdate>>> updateRecordsVerifier() {
+    return mode == KintoneMode.UPSERT && skip != Skip.ALWAYS
+        ? this::verifyUpsertRecords
+        : this::verifyUpdateRecords;
+  }
+
+  private void verifyAddRecords(
+      RecordClient mockRecordClient, ArgumentCaptor<List<Record>> addRecordsArgumentCaptor) {
+    verify(mockRecordClient, atLeast(0)).addRecords(eq(0L), addRecordsArgumentCaptor.capture());
+  }
+
+  private void verifyUpdateRecords(
+      RecordClient mockRecordClient,
+      ArgumentCaptor<List<RecordForUpdate>> updateRecordsArgumentCaptor) {
+    verify(mockRecordClient, atLeast(0))
+        .updateRecords(eq(0L), updateRecordsArgumentCaptor.capture());
+  }
+
+  private void verifyUpsertRecords(
+      RecordClient mockRecordClient,
+      ArgumentCaptor<List<RecordForUpdate>> updateRecordsArgumentCaptor) {
+    verify(mockRecordClient, atLeast(0))
+        .updateRecords(eq(0L), updateRecordsArgumentCaptor.capture(), eq(true));
+  }
+
   private List<String> getValues() {
     if (values.isEmpty() || updateRecords.isEmpty()) {
       return values;
@@ -164,13 +209,13 @@ record ->
   }
 
   private List<Record> getRecords() {
-    return updateRecords.stream().map(this::getRecord).collect(Collectors.toList());
+    return updateRecords.stream().map(this::getRecord).filter(nonNull).collect(Collectors.toList());
   }
 
   private Record getRecord(RecordForUpdate updateRecord) {
     Long id = updateRecord.getId();
     UpdateKey key = updateRecord.getUpdateKey();
-    return id != null ? getRecord(id) : getRecord(key);
+    return id != null ? getRecord(id) : key.getField() != null ? getRecord(key) : null;
   }
 
   private Record getRecord(Long id) {

src/test/java/org/embulk/output/kintone/MockClient.java

@@ -92,6 +92,8 @@ public void run(Runnable runnable) throws Exception {
         .thenReturn(mockGetRecordsByCursorResponseBody);
     when(mockRecordClient.addRecords(eq(0L), anyList())).thenReturn(Collections.emptyList());
     when(mockRecordClient.updateRecords(eq(0L), anyList())).thenReturn(Collections.emptyList());
+    when(mockRecordClient.updateRecords(eq(0L), anyList(), eq(true)))
+        .thenReturn(Collections.emptyList());
     com.kintone.client.KintoneClient mockKintoneClient = mock(KintoneClient.class);
     when(mockKintoneClient.app()).thenReturn(mockAppClient);
     when(mockKintoneClient.record()).thenReturn(mockRecordClient);