Java Code Examples for com.google.cloud.spanner.Mutation#newInsertOrUpdateBuilder()

private WriteBuilder writeBuilder(Op op, String tableName) {
	Mutation.WriteBuilder builder = null;
	switch (op) {
	case INSERT:
		builder = Mutation.newInsertBuilder(tableName);
		break;
	case INSERT_OR_UPDATE:
		builder = Mutation.newInsertOrUpdateBuilder(tableName);
		break;
	case UPDATE:
		builder = Mutation.newUpdateBuilder(tableName);
		break;
	}
	if (builder == null) {
		throw new IllegalArgumentException(
				"Unsupported save-mutation operation: " + op);
	}
	return builder;
}
 

@ProcessElement
public void processElement(ProcessContext c) throws IOException {
  /**
   * Input string is one line but Apache CSVParser process multiple lines, so we only take the
   * first item in the result list
   */
  KV<String, String> kv = c.element();
  String tableName = kv.getKey();
  Ddl ddl = c.sideInput(ddlView);
  Map<String, List<TableManifest.Column>> tableColumnsMap = c.sideInput(tableColumnsView);
  Table table = ddl.table(tableName);
  Reader in = new StringReader(kv.getValue());
  CSVFormat csvFormat =
      CSVFormat.newFormat(columnDelimiter.get())
          .withQuote(fieldQualifier.get())
          .withIgnoreEmptyLines(true)
          .withTrailingDelimiter(trailingDelimiter.get())
          .withEscape(escape.get())
          .withNullString(nullString.get());
  CSVParser parser = new CSVParser(in, csvFormat);
  List<CSVRecord> list = parser.getRecords();
  if (list.isEmpty()) {
    return;
  }
  if (list.size() > 1) {
    throw new RuntimeException("Unable to parse this row: " + c.element());
  }
  CSVRecord row = list.get(0);
  writeBuilder = Mutation.newInsertOrUpdateBuilder(table.name());
  try {
    c.output(parseRow(writeBuilder, row, table, tableColumnsMap.get(tableName)));
  } catch (IllegalArgumentException e) {
    throw new RuntimeException(
        String.format("Error to parseRow. row: %s, table: %s", row, table), e);
  }
}
 

@ProcessElement
public void processElement(ProcessContext c) {
  Mutation.WriteBuilder builder = Mutation.newInsertOrUpdateBuilder(table);
  Long key = c.element();
  builder.set("Key").to(key);
  String value = injectError.apply(key) ? null : RandomUtils.randomAlphaNumeric(valueSize);
  builder.set("Value").to(value);
  Mutation mutation = builder.build();
  c.output(mutation);
}
 

public Mutation generateMutation(Map<String, Value> overrides) {
  Mutation.WriteBuilder builder = Mutation.newInsertOrUpdateBuilder(randomCase(table.name()));
  for (Map.Entry<String, Iterator<Value>> values : valueGenerators.entrySet()) {
    String columnName = values.getKey();
    Value value = overrides.get(columnName);
    if (value == null) {
      value = values.getValue().next();
    }
    switch (value.getType().getCode()) {
      case BOOL:
        Boolean bool = value.isNull() ? null : value.getBool();
        builder.set(columnName).to(bool);
        break;
      case INT64:
        Long l = value.isNull() ? null : value.getInt64();
        builder.set(columnName).to(l);
        break;
      case FLOAT64:
        Double f = value.isNull() ? null : value.getFloat64();
        builder.set(columnName).to(f);
        break;
      case BYTES:
        ByteArray bytes = value.isNull() ? null : value.getBytes();
        builder.set(columnName).to(bytes);
        break;
      case STRING:
        String string = value.isNull() ? null : value.getString();
        builder.set(columnName).to(string);
        break;
      case TIMESTAMP:
        Timestamp timestamp = value.isNull() ? null : value.getTimestamp();
        builder.set(columnName).to(timestamp);
        break;
      case DATE:
        Date date = value.isNull() ? null : value.getDate();
        builder.set(columnName).to(date);
        break;
      case ARRAY:
        switch (value.getType().getArrayElementType().getCode()) {
          case BOOL:
            List<Boolean> bools = value.isNull() ? null : value.getBoolArray();
            builder.set(columnName).toBoolArray(bools);
            break;
          case INT64:
            List<Long> longs = value.isNull() ? null : value.getInt64Array();
            builder.set(columnName).toInt64Array(longs);
            break;
          case FLOAT64:
            List<Double> doubles = value.isNull() ? null : value.getFloat64Array();
            builder.set(columnName).toFloat64Array(doubles);
            break;
          case BYTES:
            List<ByteArray> bytesArray = value.isNull() ? null : value.getBytesArray();
            builder.set(columnName).toBytesArray(bytesArray);
            break;
          case STRING:
            List<String> strings = value.isNull() ? null : value.getStringArray();
            builder.set(columnName).toStringArray(strings);
            break;
          case TIMESTAMP:
            List<Timestamp> timestamps = value.isNull() ? null : value.getTimestampArray();
            builder.set(columnName).toTimestampArray(timestamps);
            break;
          case DATE:
            List<Date> dates = value.isNull() ? null : value.getDateArray();
            builder.set(columnName).toDateArray(dates);
            break;
        }
        break;
      default:
        throw new IllegalArgumentException("Unknown toValue " + value);
    }
  }
  return builder.build();
}
 

private Mutation createInsertMutation(Insert insert, boolean generateParameterMetaData)
    throws SQLException {
  ItemsList items = insert.getItemsList();
  if (generateParameterMetaData && items == null && insert.getSelect() != null) {
    // Just initialize the parameter meta data of the select statement
    createSelectBuilder(insert.getSelect(), insert.getSelect().toString());
    return null;
  }
  if (!(items instanceof ExpressionList)) {
    throw new CloudSpannerSQLException("Insert statement must specify a list of values",
        Code.INVALID_ARGUMENT);
  }
  if (insert.getColumns() == null || insert.getColumns().isEmpty()) {
    throw new CloudSpannerSQLException("Insert statement must specify a list of column names",
        Code.INVALID_ARGUMENT);
  }
  List<Expression> expressions = ((ExpressionList) items).getExpressions();
  String table = unquoteIdentifier(insert.getTable().getFullyQualifiedName());
  getParameterStore().setTable(table);
  WriteBuilder builder;
  if (insert.isUseDuplicate()) {
    /**
     * Do an insert-or-update. BUT: Cloud Spanner does not support supplying different values for
     * the insert and update statements, meaning that only the values specified in the INSERT part
     * of the statement will be considered. Anything specified in the 'ON DUPLICATE KEY UPDATE
     * ...' statement will be ignored.
     */
    if (this.forceUpdate)
      builder = Mutation.newUpdateBuilder(table);
    else
      builder = Mutation.newInsertOrUpdateBuilder(table);
  } else {
    /**
     * Just do an insert and throw an error if a row with the specified key alread exists.
     */
    builder = Mutation.newInsertBuilder(table);
  }
  int index = 0;
  for (Column col : insert.getColumns()) {
    String columnName = unquoteIdentifier(col.getFullyQualifiedName());
    expressions.get(index).accept(new ValueBinderExpressionVisitorAdapter<>(getParameterStore(),
        builder.set(columnName), columnName));
    index++;
  }
  return builder.build();
}