com.google.cloud.spanner.Mutation.WriteBuilder Java Examples

private Mutation createUpdateMutation(Update update, boolean generateParameterMetaData)
    throws SQLException {
  if (update.getTables().isEmpty())
    throw new CloudSpannerSQLException("No table found in update statement",
        Code.INVALID_ARGUMENT);
  if (update.getTables().size() > 1)
    throw new CloudSpannerSQLException(
        "Update statements for multiple tables at once are not supported", Code.INVALID_ARGUMENT);
  String table = unquoteIdentifier(update.getTables().get(0).getFullyQualifiedName());
  getParameterStore().setTable(table);
  List<Expression> expressions = update.getExpressions();
  WriteBuilder builder = Mutation.newUpdateBuilder(table);
  int index = 0;
  for (Column col : update.getColumns()) {
    String columnName = unquoteIdentifier(col.getFullyQualifiedName());
    expressions.get(index).accept(new ValueBinderExpressionVisitorAdapter<>(getParameterStore(),
        builder.set(columnName), columnName));
    index++;
  }
  visitUpdateWhereClause(update.getWhere(), builder, generateParameterMetaData);

  return builder.build();
}
 

private void visitUpdateWhereClause(Expression where, WriteBuilder builder,
    boolean generateParameterMetaData) throws SQLException {
  if (where != null) {
    DMLWhereClauseVisitor whereClauseVisitor = new DMLWhereClauseVisitor(getParameterStore()) {

      @Override
      protected void visitExpression(Column col, Expression expression) {
        String columnName = unquoteIdentifier(col.getFullyQualifiedName());
        expression.accept(new ValueBinderExpressionVisitorAdapter<>(getParameterStore(),
            builder.set(columnName), columnName));
      }

    };
    where.accept(whereClauseVisitor);
    if (!generateParameterMetaData && !whereClauseVisitor.isValid()) {
      throw new CloudSpannerSQLException(INVALID_WHERE_CLAUSE_UPDATE_MESSAGE,
          Code.INVALID_ARGUMENT);
    }
  } else {
    throw new SQLException(INVALID_WHERE_CLAUSE_UPDATE_MESSAGE);
  }
}
 

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;
}
 

private static boolean attemptSetIterableValue(Iterable<Object> value,
		ValueBinder<WriteBuilder> valueBinder,
		SpannerPersistentProperty spannerPersistentProperty, SpannerCustomConverter writeConverter) {

	Class innerType = ConversionUtils.boxIfNeeded(spannerPersistentProperty.getColumnInnerType());
	if (innerType == null) {
		return false;
	}

	boolean valueSet = false;

	// use the annotated column type if possible.
	if (spannerPersistentProperty.getAnnotatedColumnItemType() != null) {
		valueSet = attemptSetIterablePropertyWithTypeConversion(value, valueBinder, innerType,
				SpannerTypeMapper.getSimpleJavaClassFor(
						spannerPersistentProperty.getAnnotatedColumnItemType()),
				writeConverter);
	}
	else {
		if (!valueSet) {
			valueSet = attemptSetIterableValueOnBinder(value, valueBinder, writeConverter, innerType);
		}
	}
	return valueSet;
}
 

@SuppressWarnings("unchecked")
private static <T> boolean attemptSetSingleItemValue(Object value, Class<?> sourceType,
		ValueBinder<WriteBuilder> valueBinder, Class<T> targetType, SpannerCustomConverter writeConverter) {
	if (!writeConverter.canConvert(sourceType, targetType)) {
		return false;
	}
	Class innerType = ConversionUtils.boxIfNeeded(targetType);
	BiFunction<ValueBinder, T, ?> toMethod = (BiFunction<ValueBinder, T, ?>) singleItemTypeValueBinderMethodMap
			.get(innerType);
	if (toMethod == null) {
		return false;
	}
	// We're just checking for the bind to have succeeded, we don't need to chain the result.
	// Spanner allows binding of null values.
	Object ignored = toMethod.apply(valueBinder,
			(value != null) ? writeConverter.convert(value, targetType) : null);
	return true;
}
 

@Test
public void testCommitTimestampsType() {
	CommitTimestamps entity = new CommitTimestamps();

	doWithFields(CommitTimestamps.class,
			f -> setField(f, entity, CommitTimestamp.of(f.getType())),
			ff -> !ff.isSynthetic() && Objects.isNull(ff.getAnnotation(PrimaryKey.class)));

	WriteBuilder writeBuilder = Mutation.newInsertBuilder("commit_timestamps_table");
	this.spannerEntityWriter.write(entity, writeBuilder::set);
	Mutation mutation = writeBuilder.build();
	assertThat(mutation.asMap().entrySet().stream()
			.filter(e -> !"id".equals(e.getKey()))
			.map(Map.Entry::getValue)
			.collect(Collectors.toList())).allMatch(Value::isCommitTimestamp);
}
 

@Test
public void writeNullColumnsTest() {
	TestEntity t = new TestEntity();

	t.dateField = null;
	t.doubleList = null;

	WriteBuilder writeBuilder = mock(WriteBuilder.class);

	ValueBinder<WriteBuilder> dateFieldBinder = mock(ValueBinder.class);
	when(dateFieldBinder.to((Date) any())).thenReturn(null);
	when(writeBuilder.set(eq("dateField"))).thenReturn(dateFieldBinder);

	ValueBinder<WriteBuilder> doubleListFieldBinder = mock(ValueBinder.class);
	when(doubleListFieldBinder.toFloat64Array((Iterable<Double>) any()))
			.thenReturn(null);
	when(writeBuilder.set(eq("doubleList"))).thenReturn(doubleListFieldBinder);

	this.spannerEntityWriter.write(t, writeBuilder::set,
			Collections.unmodifiableSet(new HashSet<String>(Arrays.asList("dateField", "doubleList"))));
	verify(dateFieldBinder, times(1)).to((Date) isNull());
	verify(doubleListFieldBinder, times(1))
			.toFloat64Array((Iterable<Double>) isNull());
}
 

static void prepareMutations(TransactionContext transaction, String xid, List<Mutation> mutations)
    throws SQLException {
  int index = 0;
  for (Mutation mutation : mutations) {
    WriteBuilder prepared =
        Mutation.newInsertBuilder(CloudSpannerXAConnection.XA_PREPARED_MUTATIONS_TABLE);
    prepared.set(CloudSpannerXAConnection.XA_XID_COLUMN).to(xid);
    prepared.set(CloudSpannerXAConnection.XA_NUMBER_COLUMN).to(index);
    prepared.set(CloudSpannerXAConnection.XA_MUTATION_COLUMN).to(serializeMutation(mutation));
    transaction.buffer(prepared.build());
    index++;
  }
}
 

private ValueBinderExpressionVisitorAdapter<WriteBuilder> create() {
  ParameterStore parameterStore = new ParameterStore();
  WriteBuilder builder =
      Mutation.newInsertBuilder("INSERT INTO FOO (ID, COL1, COL2) VALUES (?, ?, ?)");
  ValueBinder<WriteBuilder> binder = builder.set("COL1");
  ValueBinderExpressionVisitorAdapter<WriteBuilder> res =
      new ValueBinderExpressionVisitorAdapter<>(parameterStore, binder, "COL1");
  return res;
}
 

private List<Mutation> saveObject(Op op, Object object,
		Set<String> includeProperties) {
	SpannerPersistentEntity<?> persistentEntity = this.spannerMappingContext
			.getPersistentEntity(object.getClass());
	List<Mutation> mutations = new ArrayList<>();
	Mutation.WriteBuilder writeBuilder = writeBuilder(op,
			persistentEntity.tableName());
	this.spannerEntityProcessor.write(object, writeBuilder::set, includeProperties);
	mutations.add(writeBuilder.build());

	persistentEntity.doWithInterleavedProperties((spannerPersistentProperty) -> {
		if (includeProperties == null
				|| includeProperties.contains(spannerPersistentProperty.getName())) {

			Iterable kids = (Iterable) persistentEntity.getPropertyAccessor(object)
					.getProperty(spannerPersistentProperty);

			if (kids != null && !ConversionUtils.ignoreForWriteLazyProxy(kids)) {
				for (Object child : kids) {
					verifyChildHasParentId(persistentEntity, object,
							this.spannerMappingContext.getPersistentEntity(
									spannerPersistentProperty.getColumnInnerType()),
							child);
					mutations.addAll(saveObject(op, child, includeProperties));
				}
			}
		}
	});
	return mutations;
}
 

private static boolean attemptSetIterablePropertyWithTypeConversion(Iterable<Object> value,
		ValueBinder<WriteBuilder> valueBinder, Class innerType, Class<?> targetType,
		SpannerCustomConverter writeConverter) {
	if (writeConverter.canConvert(innerType, targetType)) {
		BiConsumer<ValueBinder<?>, Iterable> toMethod = iterablePropertyTypeToMethodMap
				.get(targetType);
		toMethod.accept(valueBinder,
				(value != null) ? ConversionUtils.convertIterable(value, targetType, writeConverter) : null);
		return true;
	}
	return false;
}
 

@Test
@SuppressWarnings("unchecked")
public void writeSomeColumnsTest() {
	TestEntity t = new TestEntity();
	t.id = "key1";
	t.enumField = TestEntity.Color.BLACK;

	WriteBuilder writeBuilder = mock(WriteBuilder.class);

	ValueBinder<WriteBuilder> idBinder = mock(ValueBinder.class);
	when(idBinder.to(anyString())).thenReturn(null);
	when(writeBuilder.set(eq("id"))).thenReturn(idBinder);

	ValueBinder<WriteBuilder> stringFieldBinder = mock(ValueBinder.class);
	when(stringFieldBinder.to(anyString())).thenReturn(null);
	when(writeBuilder.set(eq("custom_col"))).thenReturn(stringFieldBinder);

	ValueBinder<WriteBuilder> booleanFieldBinder = mock(ValueBinder.class);
	when(booleanFieldBinder.to((Boolean) any())).thenReturn(null);
	when(writeBuilder.set(eq("booleanField"))).thenReturn(booleanFieldBinder);

	this.spannerEntityWriter.write(t, writeBuilder::set,
			new HashSet<>(Arrays.asList("id", "custom_col")));

	verify(idBinder, times(1)).to(eq(t.id));
	verify(stringFieldBinder, times(1)).to(eq(t.enumField.toString()));
	verifyNoInteractions(booleanFieldBinder);
}
 

@Test
public void writeUnsupportedTypeIterableTest() {
	this.expectedEx.expect(SpannerDataException.class);
	this.expectedEx.expectMessage("Unsupported mapping for type: class java.util.ArrayList");
	FaultyTestEntity2 ft = new FaultyTestEntity2();
	ft.listWithUnsupportedInnerType = new ArrayList<>();
	WriteBuilder writeBuilder = Mutation.newInsertBuilder("faulty_test_table_2");
	this.spannerEntityWriter.write(ft, writeBuilder::set);
}
 

@Test
public void writeIncompatibleTypeTest() {
	this.expectedEx.expect(SpannerDataException.class);
	this.expectedEx.expectMessage("Unsupported mapping for type: " +
			"class org.springframework.cloud.gcp.data.spanner.core.convert.TestEntities$TestEntity");
	FaultyTestEntity ft = new FaultyTestEntity();
	ft.fieldWithUnsupportedType = new TestEntity();
	WriteBuilder writeBuilder = Mutation.newInsertBuilder("faulty_test_table");
	this.spannerEntityWriter.write(ft, writeBuilder::set);
}
 

@Test
public void testUserSetUnconvertableColumnType() {
	this.expectedEx.expect(SpannerDataException.class);
	this.expectedEx.expectMessage("Unsupported mapping for type: class java.lang.Boolean");
	UserSetUnconvertableColumnType userSetUnconvertableColumnType = new UserSetUnconvertableColumnType();
	WriteBuilder writeBuilder = Mutation.newInsertBuilder("faulty_test_table");
	this.spannerEntityWriter.write(userSetUnconvertableColumnType, writeBuilder::set);
}
 

public static void main(String[] args) {
  Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);
  Pipeline p = Pipeline.create(options);
  
  String instanceId = options.getInstanceId();
  String databaseId = options.getDatabaseId();
 
  // Read singers from a tab-delimited file
  p.apply("ReadDPH", TextIO.read().from(options.getDphFilename()))
      // Parse the tab-delimited lines into Singer objects
      .apply("ParseDPH", ParDo.of(new ParseDPH()))
      // Spanner expects a Mutation object, so create it using the Singer's data
      .apply("CreateDPHMutation", ParDo.of(new DoFn<DPH, Mutation>() {
        @ProcessElement
        public void processElement(ProcessContext c) {
          DPH dph = c.element();
          WriteBuilder b = Mutation.newInsertOrUpdateBuilder("DPH");
          b.set("subject").to(dph.subject);
          
          if (dph.col_0 != null) {
            b.set("col_0").toStringArray(dph.col_0);
          }
          
          if (dph.col_1 != null) {
            b.set("col_1").toStringArray(dph.col_1);
          }
          
          if (dph.col_2 != null) {
            b.set("col_2").to(dph.col_2);
          }
          
          if (dph.col_3 != null) {
            b.set("col_3").to(dph.col_3);
          }
          
          if (dph.col_4 != null) {
            b.set("col_4").to(dph.col_4);
          }
 
          if (dph.col_5 != null) {
            b.set("col_5").to(dph.col_5);
          }
          
          if (dph.col_6 != null) {
            b.set("col_6").to(dph.col_6);
          }
          
          if (dph.col_7 != null) {
            b.set("col_7").toStringArray(dph.col_7);
          }
          
          if (dph.col_8 != null) {
            b.set("col_8").to(dph.col_8);
          }
          
          if (dph.col_9 != null) {
            b.set("col_9").to(dph.col_9);
          }
          
          if (dph.col_10 != null) {
            b.set("col_10").to(dph.col_10);
          }
          
          if (dph.col_11 != null) {
            b.set("col_11").to(dph.col_11);
          }
          
          if (dph.col_12 != null) {
            b.set("col_12").to(dph.col_12);
          }
          
          if (dph.col_13 != null) {
            b.set("col_13").to(dph.col_13);
          }
          
          if (dph.col_14 != null) {
            b.set("col_14").toStringArray(dph.col_14);
          } 

          if (dph.col_15 != null) {
            b.set("col_15").to(dph.col_15);
          }
          
          if (dph.col_16 != null) {
            b.set("col_16").to(dph.col_16);
          }
          
          c.output(b.build());
        }
      }))
      // Finally write the Mutations to Spanner
      .apply("WriteDPH", SpannerIO.write()
          //.withBatchSizeBytes(1024) // 1KB
          .withInstanceId(instanceId)
          .withDatabaseId(databaseId));

 

  p.run().waitUntilFinish();
}
 

/**
 * <p>
 * For each property this method "set"s the column name and finds the corresponding "to"
 * method on the {@link ValueBinder} interface.
 * </p>
 * <pre>
 * {
 * 	&#64;code
 *
 * 	long singerId = my_singer_id;
 * 	Mutation.WriteBuilder writeBuilder = Mutation.newInsertBuilder("Singer")
 * 			.set("SingerId")
 * 			.to(singerId)
 * 			.set("FirstName")
 * 			.to("Billy")
 * 			.set("LastName")
 * 			.to("Joel");
 * }
 * </pre>
 *
 * @param accessor the accessor used to get the value to write
 * @param property the property that will be written
 * @param sink the object that will accept the value to be written
 */
// @formatter:on
@SuppressWarnings("unchecked")
private void writeProperty(MultipleValueBinder sink,
		PersistentPropertyAccessor accessor,
		SpannerPersistentProperty property) {
	Object propertyValue = accessor.getProperty(property);

	Class<?> propertyType = property.getType();
	ValueBinder<WriteBuilder> valueBinder = sink.set(property.getColumnName());

	boolean valueSet = false;

	/*
	 * Due to type erasure, binder methods for Iterable properties must be manually specified.
	 * ByteArray must be excluded since it implements Iterable, but is also explicitly
	 * supported by spanner.
	 */
	if (ConversionUtils.isIterableNonByteArrayType(propertyType)) {
		valueSet = attemptSetIterableValue((Iterable<Object>) propertyValue, valueBinder,
				property, this.writeConverter);
	}
	else {

		// if the property is a commit timestamp, then its Spanner column type is always TIMESTAMP
		// and only the dummy value needs to be written to trigger auto-population of the commit
		// time
		if (property.isCommitTimestamp()) {
			valueSet = attemptSetSingleItemValue(Value.COMMIT_TIMESTAMP, Timestamp.class, valueBinder,
					Timestamp.class, this.writeConverter);
		}
		// use the user's annotated column type if possible
		else if (property.getAnnotatedColumnItemType() != null) {
			valueSet = attemptSetSingleItemValue(propertyValue, propertyType,
					valueBinder,
					SpannerTypeMapper.getSimpleJavaClassFor(property.getAnnotatedColumnItemType()),
					this.writeConverter);
		}
		else {
			if (!valueSet) {
				valueSet = attemptBindSingleValue(propertyValue, propertyType, valueBinder, this.writeConverter);
			}
		}
	}

	if (!valueSet) {
		throw new SpannerDataException(String.format(
				"Unsupported mapping for type: %s", propertyValue.getClass()));
	}
}
 

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();
}