com.google.cloud.spanner.KeySet Java Examples

@Test
@SuppressWarnings("unchecked")
public void getColumnsStringForSelectMultipleTest() {
	final SpannerPersistentEntity<TestEntity> entity = (SpannerPersistentEntity<TestEntity>)
			this.spannerMappingContext.getPersistentEntity(TestEntity.class);
	Statement childrenRowsQuery = SpannerStatementQueryExecutor.buildQuery(
			KeySet.newBuilder().addKey(Key.of("k1.1", "k1.2")).addKey(Key.of("k2.1", "k2.2")).build(),
			entity, new SpannerWriteConverter(),
			this.spannerMappingContext, entity.getWhere());

	assertThat(childrenRowsQuery.getSql())
			.isEqualTo(
					"SELECT other, deleted, id, custom_col, id_2, ARRAY (SELECT AS STRUCT deleted, id3, id, id_2 " +
							"FROM child_test_table WHERE (child_test_table.id = custom_test_table.id " +
							"AND child_test_table.id_2 = custom_test_table.id_2) AND (deleted = false)) AS childEntities " +
							"FROM custom_test_table WHERE ((id = @tag0 AND id_2 = @tag1) " +
							"OR (id = @tag2 AND id_2 = @tag3)) AND (deleted = false)");
}
 

@Test
public void deleteOrdering() throws Exception {
  SpannerSchema.Builder builder = SpannerSchema.builder();

  builder.addColumn("test1", "key", "INT64");
  builder.addKeyPart("test1", "key", false);

  builder.addColumn("test2", "key", "INT64");
  builder.addKeyPart("test2", "key", false);

  SpannerSchema schema = builder.build();

  // Verify that the encoded keys are ordered by table name then key
  List<Mutation> sortedMutations =
      Arrays.asList(
          Mutation.delete("test1", KeySet.all()), // non-point deletes come first
          Mutation.delete("test1", Key.of(1L)),
          Mutation.delete("test1", Key.of(2L)),
          Mutation.delete("test2", KeySet.prefixRange(Key.of(1L))),
          Mutation.delete("test2", Key.of(2L)));

  verifyEncodedOrdering(schema, sortedMutations);
}
 

@Test
public void testMutationSerialization()
    throws IOException, ClassNotFoundException, NoSuchMethodException, SecurityException,
    IllegalAccessException, IllegalArgumentException, InvocationTargetException {
  Mutation original = Mutation.newInsertBuilder("FOO").set("BAR").to("test").build();
  Mutation deserialized = serializeDeserialize(original);
  assertEquals(original, deserialized);

  original = Mutation.newUpdateBuilder("FOO").set("BAR").to("bla").build();
  deserialized = serializeDeserialize(original);
  assertEquals(original, deserialized);

  original = Mutation.delete("FOO", Key.of("bla"));
  deserialized = serializeDeserialize(original);
  assertEquals(original, deserialized);

  original = Mutation.delete("FOO", KeySet.all());
  deserialized = serializeDeserialize(original);
  assertEquals(original, deserialized);

  original =
      Mutation.delete("FOO", KeySet.range(KeyRange.closedClosed(Key.of("foo"), Key.of("bar"))));
  deserialized = serializeDeserialize(original);
  assertEquals(original, deserialized);
}
 

private Mutation createDeleteMutation(Delete delete, boolean generateParameterMetaData)
    throws SQLException {
  String table = unquoteIdentifier(delete.getTable().getFullyQualifiedName());
  getParameterStore().setTable(table);
  Expression where = delete.getWhere();
  if (where == null) {
    // Delete all
    return Mutation.delete(table, KeySet.all());
  } else {
    // Delete one
    DeleteKeyBuilder keyBuilder =
        new DeleteKeyBuilder(getConnection().getTable(table), generateParameterMetaData);
    visitDeleteWhereClause(where, keyBuilder, generateParameterMetaData);
    return Mutation.delete(table, keyBuilder.getKeyBuilder().build());
  }
}
 

private static void readStoringIndex(PrintWriter pw) {
  // We can read MarketingBudget also from the index since it stores a copy of MarketingBudget.
  ResultSet resultSet =
      SpannerClient.getDatabaseClient()
          .singleUse()
          .readUsingIndex(
              "Albums",
              "AlbumsByAlbumTitle2",
              KeySet.all(),
              Arrays.asList("AlbumId", "AlbumTitle", "MarketingBudget"));
  while (resultSet.next()) {
    pw.printf(
        "%d %s %s\n",
        resultSet.getLong(0),
        resultSet.getString(1),
        resultSet.isNull("MarketingBudget") ? "NULL" : resultSet.getLong("MarketingBudget"));
  }
}
 

private static void readOnlyTransaction(PrintWriter pw) {
  // ReadOnlyTransaction must be closed by calling close() on it to release resources held by it.
  // We use a try-with-resource block to automatically do so.
  try (ReadOnlyTransaction transaction =
      SpannerClient.getDatabaseClient().readOnlyTransaction()) {
    ResultSet queryResultSet =
        transaction.executeQuery(
            Statement.of("SELECT SingerId, AlbumId, AlbumTitle FROM Albums"));
    while (queryResultSet.next()) {
      pw.printf(
          "%d %d %s\n",
          queryResultSet.getLong(0), queryResultSet.getLong(1), queryResultSet.getString(2));
    }
    ResultSet readResultSet =
        transaction.read(
            "Albums", KeySet.all(), Arrays.asList("SingerId", "AlbumId", "AlbumTitle"));
    while (readResultSet.next()) {
      pw.printf(
          "%d %d %s\n",
          readResultSet.getLong(0), readResultSet.getLong(1), readResultSet.getString(2));
    }
  }
}
 

static void deleteExampleData(DatabaseClient dbClient) {
  List<Mutation> mutations = new ArrayList<>();

  // KeySet.Builder can be used to delete a specific set of rows.
  // Delete the Albums with the key values (2,1) and (2,3).
  mutations.add(
      Mutation.delete(
          "Albums", KeySet.newBuilder().addKey(Key.of(2, 1)).addKey(Key.of(2, 3)).build()));

  // KeyRange can be used to delete rows with a key in a specific range.
  // Delete a range of rows where the column key is >=3 and <5
  mutations.add(
      Mutation.delete("Singers", KeySet.range(KeyRange.closedOpen(Key.of(3), Key.of(5)))));

  // KeySet.all() can be used to delete all the rows in a table.
  // Delete remaining Singers rows, which will also delete the remaining Albums rows since it was
  // defined with ON DELETE CASCADE.
  mutations.add(Mutation.delete("Singers", KeySet.all()));

  dbClient.write(mutations);
  System.out.printf("Records deleted.\n");
}
 

private static void readOnlyTransaction(PrintWriter pw) {
  // ReadOnlyTransaction must be closed by calling close() on it to release resources held by it.
  // We use a try-with-resource block to automatically do so.
  try (ReadOnlyTransaction transaction =
      SpannerClient.getDatabaseClient().readOnlyTransaction()) {
    ResultSet queryResultSet =
        transaction.executeQuery(
            Statement.of("SELECT SingerId, AlbumId, AlbumTitle FROM Albums"));
    while (queryResultSet.next()) {
      pw.printf(
          "%d %d %s\n",
          queryResultSet.getLong(0), queryResultSet.getLong(1), queryResultSet.getString(2));
    }
    ResultSet readResultSet =
        transaction.read(
            "Albums", KeySet.all(), Arrays.asList("SingerId", "AlbumId", "AlbumTitle"));
    while (readResultSet.next()) {
      pw.printf(
          "%d %d %s\n",
          readResultSet.getLong(0), readResultSet.getLong(1), readResultSet.getString(2));
    }
  }
}
 

/** Count the number of cells modified by {@link MutationGroup}. */
public static long countOf(SpannerSchema spannerSchema, MutationGroup mutationGroup) {
  long mutatedCells = 0L;
  for (Mutation mutation : mutationGroup) {
    if (mutation.getOperation() == Op.DELETE) {
      // For single key deletes sum up all the columns in the schema.
      // There is no clear way to estimate range deletes, so they are ignored.
      if (isPointDelete(mutation)) {
        final KeySet keySet = mutation.getKeySet();

        final long rows = Iterables.size(keySet.getKeys());
        mutatedCells += rows * spannerSchema.getCellsMutatedPerRow(mutation.getTable());
      }
    } else {
      // sum the cells of the columns included in the mutation
      for (String column : mutation.getColumns()) {
        mutatedCells += spannerSchema.getCellsMutatedPerColumn(mutation.getTable(), column);
      }
    }
  }

  return mutatedCells;
}
 

static void readOnlyTransaction(DatabaseClient dbClient) {
  // ReadOnlyTransaction must be closed by calling close() on it to release resources held by it.
  // We use a try-with-resource block to automatically do so.
  try (ReadOnlyTransaction transaction = dbClient.readOnlyTransaction()) {
    ResultSet queryResultSet =
        transaction.executeQuery(
            Statement.of("SELECT SingerId, AlbumId, AlbumTitle FROM Albums"));
    while (queryResultSet.next()) {
      System.out.printf(
          "%d %d %s\n",
          queryResultSet.getLong(0), queryResultSet.getLong(1), queryResultSet.getString(2));
    }
    try (ResultSet readResultSet =
        transaction.read(
            "Albums", KeySet.all(), Arrays.asList("SingerId", "AlbumId", "AlbumTitle"))) {
      while (readResultSet.next()) {
        System.out.printf(
            "%d %d %s\n",
            readResultSet.getLong(0), readResultSet.getLong(1), readResultSet.getString(2));
      }
    }
  }
}
 

static void readStoringIndex(DatabaseClient dbClient) {
  // We can read MarketingBudget also from the index since it stores a copy of MarketingBudget.
  try (ResultSet resultSet =
      dbClient
          .singleUse()
          .readUsingIndex(
              "Albums",
              "AlbumsByAlbumTitle2",
              KeySet.all(),
              Arrays.asList("AlbumId", "AlbumTitle", "MarketingBudget"))) {
    while (resultSet.next()) {
      System.out.printf(
          "%d %s %s\n",
          resultSet.getLong(0),
          resultSet.getString(1),
          resultSet.isNull("MarketingBudget") ? "NULL" : resultSet.getLong("MarketingBudget"));
    }
  }
}
 

@Override
public <T> List<T> read(Class<T> entityClass, KeySet keys, SpannerReadOptions options) {
	SpannerPersistentEntity<T> persistentEntity =
			(SpannerPersistentEntity<T>) this.mappingContext.getPersistentEntity(entityClass);
	List<T> entities;
	if (persistentEntity.hasEagerlyLoadedProperties() || persistentEntity.hasWhere()) {
		entities = executeReadQueryAndResolveChildren(keys, persistentEntity,
				toQueryOption(keys, options), options != null ? options.getIndex() : null);
	}
	else {
		entities = mapToListAndResolveChildren(executeRead(persistentEntity.tableName(), keys,
				persistentEntity.columns(), options), entityClass,
				(options != null) ? options.getIncludeProperties() : null,
				options != null && options.isAllowPartialRead());
	}
	maybeEmitEvent(new AfterReadEvent(entities, keys, options));
	return entities;
}
 

@Test
public void existsByIdEmbeddedKeyTest() {
	ResultSet results = mock(ResultSet.class);
	when(results.next()).thenReturn(false);

	when(this.readContext.read(any(), any(), any(), any())).thenReturn(results);
	when(this.databaseClient.singleUse(any())).thenReturn(this.readContext);

	Key key = Key.of("key");
	KeySet keySet = KeySet.singleKey(key);
	assertThat(this.spannerTemplate.existsById(TestEntityEmbeddedPK.class, key)).isFalse();

	verify(this.databaseClient, times(1)).singleUse();
	verify(this.readContext, times(1))
			.read(eq("test_table_embedded_pk"), eq(keySet), eq(Collections.singleton("stringId")));
}
 

@Test
public void existsByIdTest() {
	ResultSet results = mock(ResultSet.class);
	when(results.next()).thenReturn(true);

	when(this.readContext.read(any(), any(), any(), any())).thenReturn(results);
	when(this.databaseClient.singleUse(any())).thenReturn(this.readContext);

	Key key = Key.of("key");
	KeySet keySet = KeySet.singleKey(key);
	assertThat(this.spannerTemplate.existsById(TestEntity.class, key)).isTrue();

	verify(this.databaseClient, times(1)).singleUse();
	verify(this.readContext, times(1))
			.read(eq("custom_test_table"), eq(keySet), eq(Collections.singleton("id")));
}
 

void partitionRead() {
  // [START partition_read]
  final BatchReadOnlyTransaction txn =
      batchClient.batchReadOnlyTransaction(TimestampBound.strong());
  List<Partition> partitions =
      txn.partitionRead(
          PartitionOptions.getDefaultInstance(),
          "Singers",
          KeySet.all(),
          Arrays.asList("SingerId", "FirstName", "LastName"));
  for (final Partition p : partitions) {
    try (ResultSet results = txn.execute(p)) {
      while (results.next()) {
        long singerId = results.getLong(0);
        String firstName = results.getString(1);
        String lastName = results.getString(2);
        System.out.println("[" + singerId + "] " + firstName + " " + lastName);
      }
    }
  }
  // [END partition_read]
}
 

void partitionReadUsingIndex() {
  // [START partition_read_using_index]
  final BatchReadOnlyTransaction txn =
      batchClient.batchReadOnlyTransaction(TimestampBound.strong());
  List<Partition> partitions =
      txn.partitionReadUsingIndex(
          PartitionOptions.getDefaultInstance(),
          "Singers",
          "SingerId",
          KeySet.all(),
          Arrays.asList("SingerId", "FirstName", "LastName"));

  for (Partition p : partitions) {
    try (ResultSet results = txn.execute(p)) {
      while (results.next()) {
        long singerId = results.getLong(0);
        String firstName = results.getString(1);
        String lastName = results.getString(2);
        System.out.println("[" + singerId + "] " + firstName + " " + lastName);
      }
    }
  }
  // [END partition_read_using_index]
}
 

private ResultSet executeRead(String tableName, KeySet keys, Iterable<String> columns,
		SpannerReadOptions options) {

	long startTime = LOGGER.isDebugEnabled() ? System.currentTimeMillis() : 0;

	ReadContext readContext = (options != null && options.getTimestampBound() != null)
			? getReadContext(options.getTimestampBound())
			: getReadContext();

	final ResultSet resultSet = options != null && options.getIndex() != null
			? readContext.readUsingIndex(tableName, options.getIndex(), keys, columns, options.getOptions())
			: readContext.read(tableName, keys, columns, options == null ? ArrayUtils.toArray() : options.getOptions());

	if (LOGGER.isDebugEnabled()) {
		StringBuilder logs = logColumns(tableName, keys, columns);
		logReadOptions(options, logs);
		LOGGER.debug(logs.toString());

		LOGGER.debug("Read elapsed milliseconds: " + (System.currentTimeMillis() - startTime));
	}

	return resultSet;
}
 

@Test
public void greenPathTransaction() {
	TestEntity entity1 = new TestEntity();
	TestEntity entity2 = new TestEntity();
	this.transactionalService.doInTransaction(entity1, entity2);
	verify(this.transactionManager, times(1)).begin();
	verify(this.transactionManager, times(1)).commit();
	verify(this.transactionManager, times(0)).rollback();

	verify(this.databaseClient, times(1)).transactionManager(); // only 1 transaction

	verify(this.transactionContext, times(2)).buffer(INSERT_MUTATION);
	verify(this.transactionContext, times(1))
			.read(
					eq("custom_test_table"),
					eq(KeySet.singleKey(Key.of("abc"))),
					Mockito.any(Iterable.class),
					Mockito.any());
	verify(this.transactionContext, times(1)).buffer(Arrays.asList(DELETE_MUTATION));
	verify(this.transactionContext, times(1)).buffer(UPSERT_MUTATION);
	verify(this.transactionContext, times(1)).executeUpdate(eq(DML_STATEMENT));
}
 

@Test
public void findMultipleKeysTest() {
	ResultSet results = mock(ResultSet.class);
	ReadOption readOption = mock(ReadOption.class);
	SpannerReadOptions options = new SpannerReadOptions().addReadOption(readOption)
			.setTimestampBound(TimestampBound.ofMinReadTimestamp(Timestamp.ofTimeMicroseconds(333L)));
	KeySet keySet = KeySet.singleKey(Key.of("key"));
	when(this.readContext.read(any(), any(), any(), any())).thenReturn(results);
	when(this.databaseClient.singleUse(eq(TimestampBound.ofMinReadTimestamp(Timestamp.ofTimeMicroseconds(333L)))))
			.thenReturn(this.readContext);

	verifyAfterEvents(new AfterReadEvent(Collections.emptyList(), keySet, options),
			() -> this.spannerTemplate.read(TestEntity.class, keySet, options), x -> {
				verify(this.objectMapper, times(1)).mapToList(same(results),
						eq(TestEntity.class), isNull(), eq(false));
				verify(this.readContext, times(1)).read(eq("custom_test_table"), same(keySet),
						any(), same(readOption));
			});
	verify(this.databaseClient, times(1))
			.singleUse(TimestampBound.ofMinReadTimestamp(Timestamp.ofTimeMicroseconds(333L)));
}
 

@Test
public void findMultipleKeysWithIndexTest() {
	ResultSet results = mock(ResultSet.class);
	ReadOption readOption = mock(ReadOption.class);
	SpannerReadOptions options = new SpannerReadOptions().addReadOption(readOption)
			.setIndex("index");
	KeySet keySet = KeySet.singleKey(Key.of("key"));
	when(this.readContext.readUsingIndex(any(), any(), any(), any(), any()))
			.thenReturn(results);
	this.spannerTemplate.read(TestEntity.class, keySet, options);
	verify(this.objectMapper, times(1)).mapToList(same(results),
			eq(TestEntity.class), isNull(), eq(false));
	verify(this.readContext, times(1)).readUsingIndex(eq("custom_test_table"),
			eq("index"), same(keySet), any(), same(readOption));
	verify(this.databaseClient, times(1)).singleUse();
}
 

@Test
public void findKeySetTestEager() {
	SpannerTemplate spyTemplate = spy(this.spannerTemplate);
	KeySet keys = KeySet.newBuilder().addKey(Key.of("key1")).addKey(Key.of("key2"))
			.build();
	spyTemplate.read(ParentEntity.class, keys);
	Statement statement = Statement.newBuilder("SELECT other, id, custom_col, id_2, " +
			"ARRAY (SELECT AS STRUCT deleted, id3, id, id_2 FROM child_test_table " +
			"WHERE (child_test_table.id = parent_test_table.id " +
			"AND child_test_table.id_2 = parent_test_table.id_2) AND (deleted = false)) AS childEntities " +
			"FROM parent_test_table WHERE (id = @tag0) OR (id_2 = @tag1)")
			.bind("tag0").to("key1").bind("tag1").to("key2").build();
	verify(spyTemplate, times(1)).query(eq(ParentEntity.class), eq(statement), any());
	verify(this.databaseClient, times(1)).singleUse();
}
 

@Test
public void doWithoutTransaction() {
	TestEntity entity1 = new TestEntity();
	TestEntity entity2 = new TestEntity();
	this.transactionalService.doWithoutTransaction(entity1, entity2);

	verify(this.transactionManager, Mockito.never()).begin();
	verify(this.transactionManager, Mockito.never()).commit();
	verify(this.transactionManager, Mockito.never()).rollback();
	verify(this.databaseClient, Mockito.never()).transactionManager(); // only 1 transaction

	verify(this.transactionContext, Mockito.never()).buffer(Mockito.any(List.class));
	verify(this.transactionContext, Mockito.never())
			.read(Mockito.anyString(), Mockito.any(KeySet.class), Mockito.any(Iterable.class), Mockito.any());
}
 

@Test
public void deleteByKeyTest() {
	Key key = Key.of("key");
	Mutation mutation = Mutation.delete("custom_test_table", key);
	KeySet keys = KeySet.newBuilder().addKey(key).build();
	List<Mutation> mutations = Collections.singletonList(mutation);
	when(this.mutationFactory.delete(eq(TestEntity.class), same(key)))
			.thenReturn(mutation);

	verifyBeforeAndAfterEvents(new BeforeDeleteEvent(mutations, null, keys, TestEntity.class),
			new AfterDeleteEvent(mutations, null, keys, TestEntity.class),
			() -> this.spannerTemplate.delete(TestEntity.class, key), x -> x.verify(this.databaseClient, times(1))
					.write(eq(Collections.singletonList(mutation))));
}
 

@Test
public void rollBackTransaction() {
	TestEntity entity1 = new TestEntity();
	TestEntity entity2 = new TestEntity();
	Exception exception = null;
	try {
		this.transactionalService.doInTransactionWithException(entity1, entity2);
	}
	catch (Exception ex) {
		exception = ex;
	}
	assertThat(exception).isNotNull();

	verify(this.transactionManager, times(1)).begin();
	verify(this.transactionManager, times(0)).commit();
	verify(this.transactionManager, times(1)).rollback();
	verify(this.databaseClient, times(1)).transactionManager(); // only 1 transaction

	verify(this.transactionContext, times(2)).buffer(INSERT_MUTATION);
	verify(this.transactionContext, times(1))
			.read(
					eq("custom_test_table"),
					eq(KeySet.singleKey(Key.of("abc"))),
					Mockito.any(Iterable.class),
					Mockito.any());
	verify(this.transactionContext, times(1)).buffer(Arrays.asList(DELETE_MUTATION));
	verify(this.transactionContext, times(1)).buffer(UPSERT_MUTATION);
}
 

@Test
public void deleteKeysTest() {
	KeySet keys = KeySet.newBuilder().addKey(Key.of("key1")).addKey(Key.of("key2"))
			.build();
	Mutation mutation = Mutation.delete("custom_test_table", keys);
	List<Mutation> mutations = Collections.singletonList(mutation);
	when(this.mutationFactory.delete(eq(TestEntity.class), same(keys)))
			.thenReturn(mutation);

	verifyBeforeAndAfterEvents(new BeforeDeleteEvent(mutations, null, keys, TestEntity.class),
			new AfterDeleteEvent(mutations, null, keys, TestEntity.class),
			() -> this.spannerTemplate.delete(TestEntity.class, keys), x -> x.verify(this.databaseClient, times(1))
					.write(eq(Collections.singletonList(mutation))));
}
 

public void runTemplateExample(Singer singer) {
  // Delete all of the rows in the Singer table.
  this.spannerTemplate.delete(Singer.class, KeySet.all());

  // Insert a singer into the Singers table.
  this.spannerTemplate.insert(singer);

  // Read all of the singers in the Singers table.
  List<Singer> allSingers = this.spannerTemplate
      .query(Singer.class, Statement.of("SELECT * FROM Singers"), null);
}
 

@Test
public void deleteKeysTest() {
	KeySet keySet = KeySet.newBuilder().addKey(Key.of("key1")).addKey(Key.of("key2"))
			.build();
	Mutation mutation = this.spannerMutationFactory.delete(TestEntity.class, keySet);
	assertThat(mutation.getTable()).isEqualTo("custom_test_table");
	assertThat(mutation.getOperation()).isEqualTo(Op.DELETE);
	List<String> keys = new ArrayList<>();
	mutation.getKeySet().getKeys().forEach((key) -> {
		keys.add((String) (key.getParts().iterator().next()));
	});
	assertThat(keys).containsExactlyInAnyOrder("key1", "key2");
}
 

@Test
public void findAllByIdTest() {
	List<Key> unconvertedKey = Arrays.asList(Key.of("key1"), Key.of("key2"));

	when(this.entityProcessor.convertToKey(eq(Key.of("key1")))).thenReturn(Key.of("key1"));
	when(this.entityProcessor.convertToKey(eq(Key.of("key2")))).thenReturn(Key.of("key2"));
	when(this.template.read(eq(Object.class), (KeySet) any())).thenAnswer((invocation) -> {
		KeySet keys = invocation.getArgument(1);
		assertThat(keys.getKeys()).containsExactlyInAnyOrder(Key.of("key2"), Key.of("key1"));
		return null;
	});

	new SimpleSpannerRepository<Object, Key>(this.template, Object.class)
			.findAllById(unconvertedKey);
}
 

private static long sizeOf(KeySet keySet) {
  long result = 0;
  for (Key k : keySet.getKeys()) {
    result += sizeOf(k);
  }
  for (KeyRange kr : keySet.getRanges()) {
    result += sizeOf(kr);
  }
  return result;
}
 

private static void readUsingIndex(PrintWriter pw) {
  ResultSet resultSet =
      SpannerClient.getDatabaseClient()
          .singleUse()
          .readUsingIndex(
              "Albums",
              "AlbumsByAlbumTitle",
              KeySet.all(),
              Arrays.asList("AlbumId", "AlbumTitle"));
  while (resultSet.next()) {
    pw.printf("%d %s\n", resultSet.getLong(0), resultSet.getString(1));
  }
}