com.google.cloud.spanner.Key Java Examples

/** Example of read only transaction. */
// [TARGET readOnlyTransaction()]
// [VARIABLE my_singer_id]
// [VARIABLE my_album_id]
public String readOnlyTransaction(long singerId, long albumId) {
  // [START readOnlyTransaction]
  String singerColumn = "FirstName";
  String albumColumn = "AlbumTitle";
  String albumTitle = null;
  // ReadOnlyTransaction should be closed to prevent resource leak.
  try (ReadOnlyTransaction txn = dbClient.readOnlyTransaction()) {
    Struct singerRow =
        txn.readRow("Singers", Key.of(singerId), Collections.singleton(singerColumn));
    Struct albumRow =
        txn.readRow("Albums", Key.of(singerId, albumId), Collections.singleton(albumColumn));
    singerRow.getString(singerColumn);
    albumTitle = albumRow.getString(albumColumn);
  }
  // [END readOnlyTransaction]
  return albumTitle;
}
 

@Test
public void testCommitTimestamp() {

	final CommitTimestamps entity = new CommitTimestamps();
	final String id = UUID.randomUUID().toString();
	entity.id = id;

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

	final Timestamp committedAt = databaseClient.write(mutationFactory.insert(entity));

	final CommitTimestamps fetched = spannerOperations.read(CommitTimestamps.class, Key.of(id));
	doWithFields(CommitTimestamps.class,
			f -> assertThat(getField(f, fetched))
					.describedAs("Test of the field %s has tailed", f)
					.isEqualTo(getConverter(f).convert(committedAt)),
			ff -> !ff.isSynthetic() && isNull(ff.getAnnotation(PrimaryKey.class)));
}
 

@Test
public void readOnlyTransactionTest() {

	ReadOnlyTransaction readOnlyTransaction = mock(ReadOnlyTransaction.class);
	when(this.databaseClient.readOnlyTransaction(
			eq(TimestampBound.ofMinReadTimestamp(Timestamp.ofTimeMicroseconds(333)))))
					.thenReturn(readOnlyTransaction);

	String finalResult = this.spannerTemplate
			.performReadOnlyTransaction((spannerOperations) -> {
				List<TestEntity> items = spannerOperations.readAll(TestEntity.class);
				TestEntity item = spannerOperations.read(TestEntity.class,
						Key.of("key"));
				return "all done";
			}, new SpannerReadOptions()
					.setTimestampBound(TimestampBound.ofMinReadTimestamp(Timestamp.ofTimeMicroseconds(333L))));

	assertThat(finalResult).isEqualTo("all done");
	verify(readOnlyTransaction, times(2)).read(eq("custom_test_table"), any(), any());
}
 

Key getId(Object entity) {
	PersistentPropertyAccessor accessor = getOwner().getPropertyAccessor(entity);
	List keyParts = new ArrayList();
	for (SpannerPersistentProperty spannerPersistentProperty : this.primaryKeyColumns) {
		Object value = accessor.getProperty(spannerPersistentProperty);
		if (spannerPersistentProperty.isEmbedded()) {
			Key embeddedKeyParts = this.spannerPersistentEntity
					.getSpannerMappingContext()
					.getPersistentEntity(spannerPersistentProperty.getType())
					.getIdProperty()
					.getId(value);
			for (Object keyPart : embeddedKeyParts.getParts()) {
				keyParts.add(keyPart);
			}
		}
		else if (spannerPersistentProperty.getAnnotatedColumnItemType() == null || value == null) {
			keyParts.add(value);
		}
		else {
			keyParts.add(this.spannerPersistentEntity.getSpannerEntityProcessor().getSpannerWriteConverter()
					.convert(value, SpannerTypeMapper
							.getSimpleJavaClassFor(spannerPersistentProperty.getAnnotatedColumnItemType())));
		}
	}
	return this.spannerPersistentEntity.getSpannerEntityProcessor().convertToKey(keyParts);
}
 

@Test
public void getIdTest() {
	TestEntity t = new TestEntity();
	t.id = "aaa";
	t.embeddedColumns = new EmbeddedColumns();
	t.embeddedColumns.id2 = "2";
	t.id3 = 3L;

	Key expectedKey = Key.newBuilder()
			.append(t.id)
			.appendObject(t.embeddedColumns.id2)
			.append(t.id3)
			.build();

	assertThat(this.spannerSchemaUtils.getKey(t)).isEqualTo(expectedKey);
}
 

private static long sizeOf(Key k) {
  long result = 0;
  for (Object part : k.getParts()) {
    if (part == null) {
      continue;
    }
    if (part instanceof Boolean) {
      result += 1;
    } else if (part instanceof Long) {
      result += 8;
    } else if (part instanceof Double) {
      result += 8;
    } else if (part instanceof String) {
      result += ((String) part).length();
    } else if (part instanceof ByteArray) {
      result += ((ByteArray) part).length();
    } else if (part instanceof Timestamp) {
      result += 12;
    } else if (part instanceof Date) {
      result += 12;
    }
  }
  return result;
}
 

@Test
public void findAllPageableTest() {
	Pageable pageable = mock(Pageable.class);
	Sort sort = mock(Sort.class);
	when(pageable.getSort()).thenReturn(sort);
	when(pageable.getOffset()).thenReturn(3L);
	when(pageable.getPageSize()).thenReturn(5);
	when(this.template.queryAll(eq(Object.class), any())).thenAnswer((invocation) -> {
		SpannerPageableQueryOptions spannerQueryOptions = invocation.getArgument(1);
		assertThat(spannerQueryOptions.getSort()).isSameAs(sort);
		assertThat(spannerQueryOptions.getOffset()).isEqualTo(3);
		assertThat(spannerQueryOptions.getLimit()).isEqualTo(5);
		return new ArrayList<>();
	});
	new SimpleSpannerRepository<Object, Key>(this.template, Object.class)
			.findAll(pageable);
	verify(this.template, times(1)).queryAll(eq(Object.class), any());
}
 

@Override
public Key convertToKey(Object key) {
	Assert.notNull(key, "Key of an entity to be written cannot be null!");

	Key k;
	boolean isIterable = Iterable.class.isAssignableFrom(key.getClass());
	boolean isArray = Object[].class.isAssignableFrom(key.getClass());
	if ((isIterable || isArray) && !ByteArray.class.isAssignableFrom(key.getClass())) {
		Key.Builder kb = Key.newBuilder();
		for (Object keyPart : (isArray ? (Arrays.asList((Object[]) key))
				: ((Iterable) key))) {
			kb.appendObject(convertKeyPart(keyPart));
		}
		k = kb.build();
		if (k.size() == 0) {
			throw new SpannerDataException(
					"A key must have at least one component, but 0 were given.");
		}
	}
	else {
		k = Key.class.isAssignableFrom(key.getClass()) ? (Key) key
				: Key.of(convertKeyPart(key));
	}
	return k;
}
 

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

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

private static void cleanupPrepared(TransactionContext transaction, String xid) {
  boolean foundRecords = false;
  KeySet.Builder builder = KeySet.newBuilder();
  try (ResultSet rs = transaction.executeQuery(getPreparedMutationsStatement(xid))) {
    while (rs.next()) {
      foundRecords = true;
      long number = rs.getLong(0);
      builder.addKey(Key.of(xid, number));
    }
  }
  if (foundRecords) {
    Mutation delete =
        Mutation.delete(CloudSpannerXAConnection.XA_PREPARED_MUTATIONS_TABLE, builder.build());
    transaction.buffer(delete);
  }
}
 

@Test
public void getEntityInformationTest() {
	EntityInformation<TestEntity, Key> entityInformation = this.spannerRepositoryFactory
			.getEntityInformation(TestEntity.class);
	assertThat(entityInformation.getJavaType()).isEqualTo(TestEntity.class);
	assertThat(entityInformation.getIdType()).isEqualTo(Key.class);

	TestEntity t = new TestEntity();
	t.id = "a";
	t.id2 = 3L;
	assertThat(entityInformation.getId(t))
			.isEqualTo(Key.newBuilder().append(t.id).append(t.id2).build());
}
 

@Test
public void testSetIdPropertyLongerKey() {
	this.thrown.expect(SpannerDataException.class);
	this.thrown.expectMessage(
			"The number of key parts is not equal to the number of primary key properties");

	SpannerPersistentEntity entity = new SpannerMappingContext()
			.getPersistentEntity(MultiIdsEntity.class);

	PersistentProperty idProperty = entity.getIdProperty();

	MultiIdsEntity t = new MultiIdsEntity();
	entity.getPropertyAccessor(t).setProperty(idProperty, Key.of("blah", 123L, 123.45D, "abc"));
}
 

@Test
public void keyWritingTestCase() {
	try {
		Key key = this.spannerEntityWriter.convertToKey(this.objectToTest);
		assertThat(key).isEqualTo(this.expectedResult);
	}
	catch (Exception ex) {
		assertThat(ex.getClass())
				.as("Unexpected exception: " + ex + "\nexpected: " + this.expectedResult)
				.isEqualTo(this.expectedResult);
	}
}
 

@Test
public void testSetIdProperty() {
	SpannerPersistentEntity entity = new SpannerMappingContext()
			.getPersistentEntity(MultiIdsEntity.class);

	PersistentProperty idProperty = entity.getIdProperty();

	MultiIdsEntity t = new MultiIdsEntity();
	entity.getPropertyAccessor(t).setProperty(idProperty, Key.of("blah", 123L, 123.45D));

	assertThat(t.id).isEqualTo("blah");
	assertThat(t.id2).isEqualTo(123L);
	assertThat(t.id3).isEqualTo(123.45D);
}
 

private void encodeKey(OrderedCode orderedCode, String tableName, Key key) {
  List<SpannerSchema.KeyPart> parts = schema.getKeyParts(tableName);
  Iterator<Object> it = key.getParts().iterator();
  for (SpannerSchema.KeyPart part : parts) {
    Object value = it.next();
    if (value == null) {
      if (part.isDesc()) {
        orderedCode.writeInfinityDecreasing();
      } else {
        orderedCode.writeInfinity();
      }
    } else {
      if (value instanceof Boolean) {
        writeNumber(orderedCode, part, (long) ((Boolean) value ? 0 : 1));
      } else if (value instanceof Long) {
        writeNumber(orderedCode, part, (long) value);
      } else if (value instanceof Double) {
        writeNumber(orderedCode, part, Double.doubleToLongBits((double) value));
      } else if (value instanceof String) {
        writeString(orderedCode, part, (String) value);
      } else if (value instanceof ByteArray) {
        writeBytes(orderedCode, part, (ByteArray) value);
      } else if (value instanceof Timestamp) {
        writeTimestamp(orderedCode, part, (Timestamp) value);
      } else if (value instanceof Date) {
        writeNumber(orderedCode, part, encodeDate((Date) value));
      } else {
        throw new IllegalArgumentException("Unknown key part " + value);
      }
    }
  }
}
 

@Test
public void testEmbeddedParentKeys() {
	GrandParentEmbedded grandParentEmbedded = new GrandParentEmbedded();
	grandParentEmbedded.id = "1";

	ParentEmbedded parentEmbedded = new ParentEmbedded();
	parentEmbedded.grandParentEmbedded = grandParentEmbedded;
	parentEmbedded.id2 = 2;
	parentEmbedded.id3 = 3L;

	ChildEmbedded childEmbedded = new ChildEmbedded();
	childEmbedded.parentEmbedded = parentEmbedded;
	childEmbedded.id4 = "4";

	// intentionally null, which is a supported key component type.
	childEmbedded.id5 = null;

	Key key = (Key) this.spannerMappingContext
			.getPersistentEntity(ChildEmbedded.class)
			.getIdentifierAccessor(childEmbedded).getIdentifier();

	assertThat(key).isEqualTo(
			Key.newBuilder()
					.append("1")
					.append("2")
					.append("3")
					.append("4")
					.appendObject(null)
					.build());
}
 

@Override
public String toRequestId(Serializable source, Class<?> entityType) {
	Key id = (Key) source;
	StringJoiner stringJoiner = new StringJoiner(getUrlIdSeparator());
	id.getParts().forEach((p) -> stringJoiner.add(p.toString()));
	return stringJoiner.toString();
}
 

@Test
public void existsByIdTestFound() {
	when(this.entityProcessor.convertToKey(eq(A_KEY))).thenReturn(A_KEY);
	when(this.template.existsById(eq(Object.class), eq(A_KEY))).thenReturn(true);
	assertThat(new SimpleSpannerRepository<Object, Key>(this.template, Object.class)
			.existsById(A_KEY)).isTrue();
}
 

@Test
public void findByIdTest() {
	Object ret = new Object();
	when(this.entityProcessor.convertToKey(eq(A_KEY))).thenReturn(A_KEY);
	when(this.template.read(eq(Object.class), eq(A_KEY))).thenReturn(ret);
	assertThat(new SimpleSpannerRepository<Object, Key>(this.template, Object.class).findById(A_KEY).get())
			.isEqualTo(ret);
	verify(this.template, times(1)).read(eq(Object.class), eq(A_KEY));
}
 

@Test
public void saveAllTest() {
	Object ob = new Object();
	Object ob2 = new Object();
	Iterable<Object> ret = new SimpleSpannerRepository<Object, Key>(this.template,
			Object.class)
			.saveAll(Arrays.asList(ob, ob2));
	assertThat(ret).containsExactlyInAnyOrder(ob, ob2);
	verify(this.template, times(1)).upsertAll(eq(Arrays.asList(ob, ob2)));
}
 

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

@Transactional
public void doInTransaction(TestEntity entity1, TestEntity entity2) {
	this.spannerTemplate.read(TestEntity.class, Key.of("abc"));
	this.spannerTemplate.executeDmlStatement(DML_STATEMENT);
	this.spannerTemplate.insert(entity1);
	this.spannerTemplate.insert(entity2);
	this.spannerTemplate.delete(entity1);
	this.spannerTemplate.upsert(entity2);
}
 

@Transactional
public void doInTransactionWithException(TestEntity entity1, TestEntity entity2) {
	this.spannerTemplate.read(TestEntity.class, Key.of("abc"));
	this.spannerTemplate.insert(entity1);
	this.spannerTemplate.insert(entity2);
	this.spannerTemplate.delete(entity1);
	this.spannerTemplate.upsert(entity2);
	throw new RuntimeException("oops");
}
 

public void doWithoutTransaction(TestEntity entity1, TestEntity entity2) {
	this.spannerTemplate.read(TestEntity.class, Key.of("abc"));
	this.spannerTemplate.insert(entity1);
	this.spannerTemplate.insert(entity2);
	this.spannerTemplate.delete(entity1);
	this.spannerTemplate.upsert(entity2);
}
 

@Test
public void deleteNullEntityTest() {
	this.expectedEx.expect(IllegalArgumentException.class);
	this.expectedEx.expectMessage("A non-null entity is required.");
	new SimpleSpannerRepository<Object, Key>(this.template, Object.class)
			.delete(null);
}
 

@Test
public void deleteNullIdTest() {
	this.expectedEx.expect(IllegalArgumentException.class);
	this.expectedEx.expectMessage("A non-null ID is required.");
	new SimpleSpannerRepository<Object, Key>(this.template, Object.class)
			.deleteById(null);
}
 

@Test
public void existsNullIdTest() {
	this.expectedEx.expect(IllegalArgumentException.class);
	this.expectedEx.expectMessage("A non-null ID is required.");
	new SimpleSpannerRepository<Object, Key>(this.template, Object.class)
			.existsById(null);
}
 

@Test
public void deleteAllNullEntityTest() {
	this.expectedEx.expect(IllegalArgumentException.class);
	this.expectedEx.expectMessage("A non-null list of entities is required.");
	new SimpleSpannerRepository<Object, Key>(this.template, Object.class)
			.deleteAll(null);
}
 

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