com.google.cloud.spanner.Mutation 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();
}
 

@Test
public void streamingWritesNoGrouping() throws Exception {

  // verify that grouping/sorting does not occur - batches should be created in received order.
  TestStream<Mutation> testStream =
      TestStream.create(SerializableCoder.of(Mutation.class))
          .addElements(m(1L), m(5L), m(2L), m(4L), m(3L), m(6L))
          .advanceWatermarkToInfinity();

  // verify that grouping/sorting does not occur when notset.
  pipeline
      .apply(testStream)
      .apply(
          SpannerIO.write()
              .withProjectId("test-project")
              .withInstanceId("test-instance")
              .withDatabaseId("test-database")
              .withServiceFactory(serviceFactory)
              .withMaxNumRows(2));
  pipeline.run();

  verifyBatches(batch(m(1L), m(5L)), batch(m(2L), m(4L)), batch(m(3L), m(6L)));
}
 

@Test
public void streamingWritesWithGrouping() throws Exception {

  // verify that grouping/sorting occurs when set.
  TestStream<Mutation> testStream =
      TestStream.create(SerializableCoder.of(Mutation.class))
          .addElements(m(1L), m(5L), m(2L), m(4L), m(3L), m(6L))
          .advanceWatermarkToInfinity();
  pipeline
      .apply(testStream)
      .apply(
          SpannerIO.write()
              .withProjectId("test-project")
              .withInstanceId("test-instance")
              .withDatabaseId("test-database")
              .withServiceFactory(serviceFactory)
              .withGroupingFactor(40)
              .withMaxNumRows(2));
  pipeline.run();

  // Output should be batches of sorted mutations.
  verifyBatches(batch(m(1L), m(2L)), batch(m(3L), m(4L)), batch(m(5L), m(6L)));
}
 

/** Example of a read write transaction. */
// [TARGET readWriteTransaction()]
// [VARIABLE my_singer_id]
public void readWriteTransaction(final long singerId) {
  // [START readWriteTransaction]
  TransactionRunner runner = dbClient.readWriteTransaction();
  runner.run(
      new TransactionCallable<Void>() {

        @Override
        public Void run(TransactionContext transaction) throws Exception {
          String column = "FirstName";
          Struct row =
              transaction.readRow("Singers", Key.of(singerId), Collections.singleton(column));
          String name = row.getString(column);
          transaction.buffer(
              Mutation.newUpdateBuilder("Singers").set(column).to(name.toUpperCase()).build());
          return null;
        }
      });
  // [END readWriteTransaction]
}
 

@Test
public void nullPrimitiveArrays() throws Exception {
  Mutation int64 =
      Mutation.newInsertOrUpdateBuilder("test").set("one").toInt64Array((long[]) null).build();
  Mutation float64 =
      Mutation.newInsertOrUpdateBuilder("test")
          .set("one")
          .toFloat64Array((double[]) null)
          .build();
  Mutation bool =
      Mutation.newInsertOrUpdateBuilder("test").set("one").toBoolArray((boolean[]) null).build();

  assertThat(MutationSizeEstimator.sizeOf(int64), is(0L));
  assertThat(MutationSizeEstimator.sizeOf(float64), is(0L));
  assertThat(MutationSizeEstimator.sizeOf(bool), is(0L));
}
 

private void spannerWriteWithRetryIfSchemaChange(Iterable<Mutation> batch)
    throws SpannerException {
  for (int retry = 1; ; retry++) {
    try {
      spannerAccessor.getDatabaseClient().writeAtLeastOnce(batch);
      return;
    } catch (AbortedException e) {
      if (retry >= ABORTED_RETRY_ATTEMPTS) {
        throw e;
      }
      if (e.isRetryable() || e.getMessage().contains(errString)) {
        continue;
      }
      throw e;
    }
  }
}
 

/**
 * Sends a QueueMessage to the spanner queue table
 *
 * <p>Example of sending data to a queue.
 *
 * <pre>{@code
 * MyProto myProto = MyProto.newBuilder().setMessage("My-Message").build();
 *
 * try {
 *   Queue.send(dbClient, queueName, "myKey", ByteArray.copyFrom(myProto.toByteArray()));
 * } catch (SpannerException e) {
 *   log.error("Could not write message to Queue", e);
 * }
 * }</pre>
 *
 * @param dbClient the Spanner database client
 * @param queueName the name of the queue to be polled
 * @param key the name used to partition the passed value for storage and lookup
 * @param value the message payload
 * @return Timestamp the timestamp that the message was written
 */
public static Timestamp send(DatabaseClient dbClient, String queueName, String key, byte[] value)
    throws SpannerException {
  Preconditions.checkNotNull(dbClient);
  Preconditions.checkNotNull(queueName);
  Preconditions.checkNotNull(key);
  Preconditions.checkNotNull(value);

  final List<Mutation> mutations = new ArrayList<>();
  mutations.add(
      Mutation.newInsertBuilder(queueName)
          .set("MessageId")
          .to(UUID.randomUUID().toString())
          .set("Key")
          .to(key)
          .set("Payload")
          .to(ByteArray.copyFrom(value))
          .set("Ack")
          .to(false)
          .set("Timestamp")
          .to(Value.COMMIT_TIMESTAMP)
          .build());

  return dbClient.write(mutations);
}
 

@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 streamingWrites() throws Exception {
  TestStream<Mutation> testStream =
      TestStream.create(SerializableCoder.of(Mutation.class))
          .addElements(m(1L), m(2L))
          .advanceProcessingTime(Duration.standardMinutes(1))
          .addElements(m(3L), m(4L))
          .advanceProcessingTime(Duration.standardMinutes(1))
          .addElements(m(5L), m(6L))
          .advanceWatermarkToInfinity();
  pipeline
      .apply(testStream)
      .apply(
          SpannerIO.write()
              .withProjectId("test-project")
              .withInstanceId("test-instance")
              .withDatabaseId("test-database")
              .withServiceFactory(serviceFactory));
  pipeline.run();

  verifyBatches(batch(m(1L), m(2L)), batch(m(3L), m(4L)), batch(m(5L), m(6L)));
}
 

@Test
public void updateColumnsArrayTest() {
	Mutation mutation = Mutation.newInsertOrUpdateBuilder("custom_test_table")
			.build();
	List<Mutation> mutations = Collections.singletonList(mutation);
	TestEntity entity = new TestEntity();
	Set<String> cols = new HashSet<>(Arrays.asList("a", "b"));
	when(this.mutationFactory.update(same(entity),
			eq(cols)))
					.thenReturn(mutations);

	verifyBeforeAndAfterEvents(new BeforeSaveEvent(Collections.singletonList(entity), cols),
			new AfterSaveEvent(mutations, Collections.singletonList(entity), cols),
			() -> this.spannerTemplate.update(entity, "a", "b"), x -> x.verify(this.databaseClient, times(1))
					.write(eq(mutations)));
}
 

private void makeTestData() {
  DatabaseClient databaseClient = getDatabaseClient();

  List<Mutation> mutations = new ArrayList<>();
  for (int i = 0; i < 5L; i++) {
    mutations.add(
        Mutation.newInsertOrUpdateBuilder(options.getTable())
            .set("key")
            .to((long) i)
            .set("value")
            .to(RandomUtils.randomAlphaNumeric(100))
            .build());
  }

  databaseClient.writeAtLeastOnce(mutations);
}
 

@Test
public void primitiveArrays() throws Exception {
  Mutation int64 =
      Mutation.newInsertOrUpdateBuilder("test")
          .set("one")
          .toInt64Array(new long[] {1L, 2L, 3L})
          .build();
  Mutation float64 =
      Mutation.newInsertOrUpdateBuilder("test")
          .set("one")
          .toFloat64Array(new double[] {1., 2.})
          .build();
  Mutation bool =
      Mutation.newInsertOrUpdateBuilder("test")
          .set("one")
          .toBoolArray(new boolean[] {true, true, false, true})
          .build();

  assertThat(MutationSizeEstimator.sizeOf(int64), is(24L));
  assertThat(MutationSizeEstimator.sizeOf(float64), is(16L));
  assertThat(MutationSizeEstimator.sizeOf(bool), is(4L));
}
 

/** Example of unprotected blind write. */
// [TARGET writeAtLeastOnce(Iterable)]
// [VARIABLE my_singer_id]
public void writeAtLeastOnce(long singerId) {
  // [START writeAtLeastOnce]
  Mutation mutation =
      Mutation.newInsertBuilder("Singers")
          .set("SingerId")
          .to(singerId)
          .set("FirstName")
          .to("Billy")
          .set("LastName")
          .to("Joel")
          .build();
  dbClient.writeAtLeastOnce(Collections.singletonList(mutation));
  // [END writeAtLeastOnce]
}
 

@Test()
public void testDeleteStatementWithNullValueInKey()
    throws SQLException, NoSuchMethodException, SecurityException, IllegalAccessException,
    IllegalArgumentException, InvocationTargetException {
  CloudSpannerPreparedStatement ps =
      CloudSpannerTestObjects.createPreparedStatement("DELETE FROM FOO WHERE ID=?");
  ps.setNull(1, Types.BIGINT);
  Mutations mutations;
  Method createMutations = ps.getClass().getDeclaredMethod("createMutations");
  createMutations.setAccessible(true);
  mutations = (Mutations) createMutations.invoke(ps);

  Mutation deleteMutation = mutations.getMutations().get(0);
  Assert.assertNotNull(deleteMutation);
  Assert.assertEquals(Op.DELETE, deleteMutation.getOperation());
  List<Key> keys = Lists.newArrayList(deleteMutation.getKeySet().getKeys());
  Assert.assertEquals(1, keys.size());
  Assert.assertNull(keys.get(0).getParts().iterator().next());
}
 

private static void writeExampleData(PrintWriter pw) {
  List<Mutation> mutations = new ArrayList<>();
  for (Singer singer : SINGERS) {
    mutations.add(
        Mutation.newInsertBuilder("Singers")
            .set("SingerId")
            .to(singer.singerId)
            .set("FirstName")
            .to(singer.firstName)
            .set("LastName")
            .to(singer.lastName)
            .build());
  }
  for (Album album : ALBUMS) {
    mutations.add(
        Mutation.newInsertBuilder("Albums")
            .set("SingerId")
            .to(album.singerId)
            .set("AlbumId")
            .to(album.albumId)
            .set("AlbumTitle")
            .to(album.albumTitle)
            .build());
  }
  SpannerClient.getDatabaseClient().write(mutations);
}
 

/** Example of using {@link TransactionManager}. */
// [TARGET transactionManager()]
// [VARIABLE my_singer_id]
public void transactionManager(final long singerId) throws InterruptedException {
  // [START transactionManager]
  try (TransactionManager manager = dbClient.transactionManager()) {
    TransactionContext txn = manager.begin();
    while (true) {
      try {
        String column = "FirstName";
        Struct row = txn.readRow("Singers", Key.of(singerId), Collections.singleton(column));
        String name = row.getString(column);
        txn.buffer(
            Mutation.newUpdateBuilder("Singers").set(column).to(name.toUpperCase()).build());
        manager.commit();
        break;
      } catch (AbortedException e) {
        Thread.sleep(e.getRetryDelayInMillis() / 1000);
        txn = manager.resetForRetry();
      }
    }
  }
  // [END transactionManager]
}
 

/** Estimates a size of mutation in bytes. */
static long sizeOf(Mutation m) {
  if (m.getOperation() == Mutation.Op.DELETE) {
    return sizeOf(m.getKeySet());
  }
  long result = 0;
  for (Value v : m.getValues()) {
    switch (v.getType().getCode()) {
      case ARRAY:
        result += estimateArrayValue(v);
        break;
      case STRUCT:
        throw new IllegalArgumentException("Structs are not supported in mutation.");
      default:
        result += estimatePrimitiveValue(v);
    }
  }
  return result;
}
 

@Test
public void singleMutationPipeline() throws Exception {
  Mutation mutation = m(2L);
  PCollection<Mutation> mutations = pipeline.apply(Create.of(mutation));

  mutations.apply(
      SpannerIO.write()
          .withProjectId("test-project")
          .withInstanceId("test-instance")
          .withDatabaseId("test-database")
          .withServiceFactory(serviceFactory));
  pipeline.run();

  verifyBatches(batch(m(2L)));
}
 

@Test
public void deleteAllObjectTest() {
	Mutation mutation = Mutation.delete("custom_test_table", Key.of("key"));
	TestEntity entity = new TestEntity();
	List entities = Arrays.asList(entity, entity, entity);
	List<Mutation> mutations = Arrays.asList(mutation, mutation, mutation);
	when(this.mutationFactory.delete(entity)).thenReturn(mutation);

	verifyBeforeAndAfterEvents(new BeforeDeleteEvent(mutations, entities, null, null),
			new AfterDeleteEvent(mutations, entities, null, null),
			() -> this.spannerTemplate.deleteAll(Arrays.asList(entity, entity, entity)),
			x -> x.verify(this.databaseClient, times(1))
					.write(eq(mutations)));
}
 

/** Estimates a size of the mutation group in bytes. */
public static long sizeOf(MutationGroup group) {
  long result = 0;
  for (Mutation m : group) {
    result += sizeOf(m);
  }
  return result;
}
 

private void applySaveMutations(Supplier<List<Mutation>> mutationsSupplier, Iterable<?> entities,
		Set<String> includeProperties) {
	maybeEmitEvent(new BeforeSaveEvent(entities, includeProperties));
	List<Mutation> mutations = mutationsSupplier.get();
	applyMutations(mutations);
	maybeEmitEvent(new AfterSaveEvent(mutations, entities, includeProperties));
}
 

@Test
public void lazyWriteTest() {
	TestEntity t = new TestEntity();
	t.id = "a";
	ChildEntity c1 = new ChildEntity();
	c1.keyComponents = new EmbeddedKeyComponents();
	c1.keyComponents.id = t.id;
	c1.keyComponents.id2 = "c2";
	ChildEntity c2 = new ChildEntity();
	c2.keyComponents = new EmbeddedKeyComponents();
	c2.keyComponents.id = t.id;
	c2.keyComponents.id2 = "c3";

	// intentionally setting it as an untouched-proxy
	t.childEntities = ConversionUtils.wrapSimpleLazyProxy(() -> Arrays.asList(c1, c2), List.class);

	List<Mutation> mutations = this.spannerMutationFactory.upsert(t, null);
	Mutation parentMutation = mutations.get(0);

	// The size is 1 because the child is an un-evaluated proxy.
	assertThat(mutations).hasSize(1);
	assertThat(parentMutation.getTable()).isEqualTo("custom_test_table");
	assertThat(parentMutation.getOperation()).isEqualTo(Op.INSERT_OR_UPDATE);

	t.childEntities.size();

	// mutations should now be generated since the child proxy has been touched.
	mutations = this.spannerMutationFactory.upsert(t, null);
	parentMutation = mutations.get(0);
	assertThat(mutations).hasSize(3);
	List<Mutation> childMutations = mutations.subList(1, mutations.size());
	assertThat(parentMutation.getTable()).isEqualTo("custom_test_table");
	assertThat(parentMutation.getOperation()).isEqualTo(Op.INSERT_OR_UPDATE);
	for (Mutation childMutation : childMutations) {
		assertThat(childMutation.getTable()).isEqualTo("child_test_table");
		assertThat(childMutation.getOperation()).isEqualTo(Op.INSERT_OR_UPDATE);
	}
}
 

/**
 * Executes a batch of inserts to a table.
 */
public void batchInsert(int batchSize) {
  ArrayList<Mutation> mutations = new ArrayList<>();
  for (int i = 0; i < batchSize; i++) {
    mutations.add(buildSingleMutationInsert());
  }
  this.databaseClient.write(mutations);
}
 

static void writeDatatypesData(DatabaseClient dbClient) {
  List<Mutation> mutations = new ArrayList<>();
  for (Venue venue : VENUES) {
    mutations.add(
        Mutation.newInsertBuilder("Venues")
            .set("VenueId")
            .to(venue.venueId)
            .set("VenueName")
            .to(venue.venueName)
            .set("VenueInfo")
            .to(venue.venueInfo)
            .set("Capacity")
            .to(venue.capacity)
            .set("AvailableDates")
            .to(venue.availableDates)
            .set("LastContactDate")
            .to(venue.lastContactDate)
            .set("OutdoorVenue")
            .to(venue.outdoorVenue)
            .set("PopularityScore")
            .to(venue.popularityScore)
            .set("LastUpdateTime")
            .to(Value.COMMIT_TIMESTAMP)
            .build());
  }
  dbClient.write(mutations);
}
 

private void encodeKey(OrderedCode orderedCode, Mutation m) {
  Map<String, Value> mutationMap = mutationAsMap(m);
  for (SpannerSchema.KeyPart part : schema.getKeyParts(m.getTable())) {
    Value val = mutationMap.get(part.getField());
    if (val == null || val.isNull()) {
      if (part.isDesc()) {
        orderedCode.writeInfinityDecreasing();
      } else {
        orderedCode.writeInfinity();
      }
    } else {
      Type.Code code = val.getType().getCode();
      switch (code) {
        case BOOL:
          writeNumber(orderedCode, part, (long) (val.getBool() ? 0 : 1));
          break;
        case INT64:
          writeNumber(orderedCode, part, val.getInt64());
          break;
        case FLOAT64:
          writeNumber(orderedCode, part, Double.doubleToLongBits(val.getFloat64()));
          break;
        case STRING:
          writeString(orderedCode, part, val.getString());
          break;
        case BYTES:
          writeBytes(orderedCode, part, val.getBytes());
          break;
        case TIMESTAMP:
          writeTimestamp(orderedCode, part, val.getTimestamp());
          break;
        case DATE:
          writeNumber(orderedCode, part, encodeDate(val.getDate()));
          break;
        default:
          throw new IllegalArgumentException("Unknown type " + val.getType());
      }
    }
  }
}
 

@Test
public void testDeleteStatementWithoutWhereClause() throws SQLException {
  Mutation deleteMutation = getMutation("DELETE FROM FOO");
  Assert.assertNotNull(deleteMutation);
  Assert.assertEquals(Op.DELETE, deleteMutation.getOperation());
  Assert.assertTrue(deleteMutation.getKeySet().isAll());
  Assert.assertNotNull(deleteMutation);
  Assert.assertEquals(Op.DELETE, deleteMutation.getOperation());
  Assert.assertTrue(deleteMutation.getKeySet().isAll());
}
 

private Mutation serializeDeserialize(Mutation original)
    throws NoSuchMethodException, SecurityException, IllegalAccessException,
    IllegalArgumentException, InvocationTargetException {
  Method serialize = XATransaction.class.getDeclaredMethod("serializeMutation", Mutation.class);
  serialize.setAccessible(true);
  String serialized = (String) serialize.invoke(null, original);

  Method deserialize = XATransaction.class.getDeclaredMethod("deserializeMutation", String.class);
  deserialize.setAccessible(true);
  Mutation deserialized = (Mutation) deserialize.invoke(null, serialized);

  return deserialized;
}
 

@Test
public void updateAllTest() {
	Mutation mutation = Mutation.newInsertOrUpdateBuilder("custom_test_table")
			.build();
	TestEntity entity = new TestEntity();
	List<Mutation> mutations = Arrays.asList(mutation, mutation, mutation);
	List entities = Arrays.asList(entity, entity, entity);
	when(this.mutationFactory.update(same(entity), isNull()))
			.thenReturn(Collections.singletonList(mutation));

	verifyBeforeAndAfterEvents(new BeforeSaveEvent(entities, null),
			new AfterSaveEvent(mutations, entities, null),
			() -> this.spannerTemplate.updateAll(entities), x -> x.verify(this.databaseClient, times(1))
					.write(eq(mutations)));
}
 

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

@Test
public void upsertColumnsSetTest() {
	Mutation mutation = Mutation.newInsertOrUpdateBuilder("custom_test_table")
			.build();
	TestEntity entity = new TestEntity();
	List<Mutation> mutations = Collections.singletonList(mutation);
	Set<String> cols = new HashSet<>(Arrays.asList("a", "b"));
	when(this.mutationFactory.upsert(same(entity), eq(cols)))
			.thenReturn(Collections.singletonList(mutation));

	verifyBeforeAndAfterEvents(new BeforeSaveEvent(Collections.singletonList(entity), cols),
			new AfterSaveEvent(mutations, Collections.singletonList(entity), cols),
			() -> this.spannerTemplate.upsert(entity, cols), x -> x.verify(this.databaseClient, times(1))
					.write(eq(mutations)));
}