Java Code Examples for com.google.cloud.datastore.Entity#Builder

/** Example of updating multiple entities. */
// [TARGET update(Entity...)]
// [VARIABLE "my_key_name1"]
// [VARIABLE "my_key_name2"]
public void multipleUpdateEntities(String keyName1, String keyName2) {
  Datastore datastore = transaction.getDatastore();
  // [START multipleUpdateEntities]
  Key key1 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName1);
  Entity.Builder entityBuilder1 = Entity.newBuilder(key1);
  entityBuilder1.set("propertyName", "value3");
  Entity entity1 = entityBuilder1.build();

  Key key2 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName2);
  Entity.Builder entityBuilder2 = Entity.newBuilder(key2);
  entityBuilder2.set("propertyName", "value4");
  Entity entity2 = entityBuilder2.build();

  transaction.update(entity1, entity2);
  transaction.commit();
  // [END multipleUpdateEntities]
}
 

@Test
public void testUnsupportedTypeWrite() {
	TestItemUnsupportedFields item = new TestItemUnsupportedFields();
	item.setStringField("string value");
	item.setUnsupportedField(new TestItemUnsupportedFields.NewType(true));

	DatastoreEntityConverter entityConverter = new DefaultDatastoreEntityConverter(
			new DatastoreMappingContext(), new TwoStepsConversions(new DatastoreCustomConversions(
			Arrays.asList(
					getIntegerToNewTypeConverter(),
					getNewTypeToIntegerConverter()
					)), null, datastoreMappingContext));
	Entity.Builder builder = getEntityBuilder();
	entityConverter.write(item, builder);
	Entity entity = builder.build();

	assertThat(entity.getLong("unsupportedField")).as("validate custom conversion")
			.isEqualTo(1L);
	assertThat(entity.getString("stringField")).as("validate string field")
			.isEqualTo("string value");

	TestItemUnsupportedFields readItem =
			entityConverter.read(TestItemUnsupportedFields.class, entity);

	assertThat(item.equals(readItem)).as("read object should be equal to original").isTrue();
}
 

/** Example of updating multiple entities. */
// [TARGET update(Entity...)]
// [VARIABLE "my_key_name_1"]
// [VARIABLE "my_key_name_2"]
public void batchUpdateEntities(String keyName1, String keyName2) {
  // [START batchUpdateEntities]
  Key key1 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName1);
  Entity.Builder entityBuilder1 = Entity.newBuilder(key1);
  entityBuilder1.set("propertyName", "updatedValue1");
  Entity entity1 = entityBuilder1.build();

  Key key2 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName2);
  Entity.Builder entityBuilder2 = Entity.newBuilder(key2);
  entityBuilder2.set("propertyName", "updatedValue2");
  Entity entity2 = entityBuilder2.build();

  datastore.update(entity1, entity2);
  // [END batchUpdateEntities]
}
 

/** Example of putting multiple entities. */
// [TARGET put(FullEntity...)]
// [VARIABLE "my_key_name1"]
// [VARIABLE "my_key_name2"]
public void multiplePutEntities(String keyName1, String keyName2) {
  Datastore datastore = transaction.getDatastore();
  // [START multiplePutEntities]
  Key key1 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName1);
  Entity.Builder entityBuilder1 = Entity.newBuilder(key1);
  entityBuilder1.set("propertyName", "value1");
  Entity entity1 = entityBuilder1.build();

  Key key2 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName2);
  Entity.Builder entityBuilder2 = Entity.newBuilder(key2);
  entityBuilder2.set("propertyName", "value2");
  Entity entity2 = entityBuilder2.build();

  transaction.put(entity1, entity2);
  transaction.commit();
  // [END multiplePutEntities]
}
 

/** Example of adding a single entity. */
// [TARGET add(FullEntity)]
// [VARIABLE "my_key_name"]
public void addSingleEntity(String keyName) {
  // [START addSingleEntity]
  Key key = datastore.newKeyFactory().setKind("MyKind").newKey(keyName);
  Entity.Builder entityBuilder = Entity.newBuilder(key);
  entityBuilder.set("propertyName", "value");
  Entity entity = entityBuilder.build();
  try {
    datastore.add(entity);
  } catch (DatastoreException ex) {
    if ("ALREADY_EXISTS".equals(ex.getReason())) {
      // entity.getKey() already exists
    }
  }
  // [END addSingleEntity]
}
 

/** Example of adding multiple entities. */
// [TARGET add(FullEntity...)]
// [VARIABLE "my_key_name1"]
// [VARIABLE "my_key_name2"]
public void multipleAddEntities(String keyName1, String keyName2) {
  Datastore datastore = transaction.getDatastore();
  // [START multipleAddEntities]
  Key key1 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName1);
  Entity.Builder entityBuilder1 = Entity.newBuilder(key1);
  entityBuilder1.set("propertyName", "value1");
  Entity entity1 = entityBuilder1.build();

  Key key2 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName2);
  Entity.Builder entityBuilder2 = Entity.newBuilder(key2);
  entityBuilder2.set("propertyName", "value2");
  Entity entity2 = entityBuilder2.build();

  transaction.add(entity1, entity2);
  transaction.commit();
  // [END multipleAddEntities]
}
 

/**
 * Delete a value stored in the project's datastore.
 *
 * @param sessionId Request from which the session is extracted.
 */
protected void deleteSessionVariables(String sessionId, String... varNames) {
  if (sessionId.equals("")) {
    return;
  }
  Key key = keyFactory.newKey(sessionId);
  Transaction transaction = datastore.newTransaction();
  try {
    Entity stateEntity = transaction.get(key);
    if (stateEntity != null) {
      Entity.Builder builder = Entity.newBuilder(stateEntity);
      StringBuilder delNames = new StringBuilder();
      for (String varName : varNames) {
        delNames.append(varName + " ");
        builder = builder.remove(varName);
      }
      datastore.update(builder.build());
    }
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

@Test
public void shouldReturnGlobalCounterLimit() throws IOException {
  final Key key = globalConfigKey(datastore.newKeyFactory());
  Entity.Builder builder = Entity.newBuilder(key)
      .set(PROPERTY_CONCURRENCY, 4000L);
  datastore.put(builder.build());

  assertEquals(4000L, storage.getLimitForCounter("GLOBAL_STYX_CLUSTER"));
}
 

/**
 * Stores the state value in each key-value pair in the project's datastore.
 *
 * @param sessionId Request from which to extract session.
 * @param varName   the name of the desired session variable
 * @param varValue  the value of the desired session variable
 */
protected void setSessionVariables(String sessionId, Map<String, String> setMap) {
  if (sessionId.equals("")) {
    return;
  }
  Key key = keyFactory.newKey(sessionId);
  Transaction transaction = datastore.newTransaction();
  DateTime dt = DateTime.now(DateTimeZone.UTC);
  dt.toString(dtf);
  try {
    Entity stateEntity = transaction.get(key);
    Entity.Builder seBuilder;
    if (stateEntity == null) {
      seBuilder = Entity.newBuilder(key);
    } else {
      seBuilder = Entity.newBuilder(stateEntity);
    }
    for (String varName : setMap.keySet()) {
      seBuilder.set(varName, setMap.get(varName));
    }
    transaction.put(seBuilder.set("lastModified", dt.toString(dtf)).build());
    transaction.commit();
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

@Test
public void shouldReadBackfillWithMissingReverseField() throws Exception {
  final Backfill backfill = Backfill.newBuilder()
      .id("backfill-2")
      .start(Instant.parse("2017-01-01T00:00:00Z"))
      .end(Instant.parse("2017-01-02T00:00:00Z"))
      .workflowId(WorkflowId.create("component", "workflow2"))
      .concurrency(2)
      .nextTrigger(Instant.parse("2017-01-01T00:00:00Z"))
      .schedule(DAYS)
      .created(currentTime)
      .lastModified(currentTime)
      .build();

  final Key key = DatastoreStorage.backfillKey(datastore.newKeyFactory(), backfill.id());
  Entity.Builder builder = Entity.newBuilder(key)
      .set(PROPERTY_CONCURRENCY, backfill.concurrency())
      .set(PROPERTY_START, instantToTimestamp(backfill.start()))
      .set(PROPERTY_END, instantToTimestamp(backfill.end()))
      .set(PROPERTY_COMPONENT, backfill.workflowId().componentId())
      .set(PROPERTY_WORKFLOW, backfill.workflowId().id())
      .set(PROPERTY_SCHEDULE, backfill.schedule().toString())
      .set(PROPERTY_NEXT_TRIGGER, instantToTimestamp(backfill.nextTrigger()))
      .set(PROPERTY_ALL_TRIGGERED, backfill.allTriggered())
      .set(PROPERTY_HALTED, backfill.halted())
      .set(PROPERTY_CREATED, instantToTimestamp(backfill.created().orElseThrow()))
      .set(PROPERTY_LAST_MODIFIED, instantToTimestamp(backfill.lastModified().orElseThrow()));

  datastore.put(builder.build());

  assertThat(storage.getBackfill(backfill.id()), equalTo(Optional.of(backfill)));
}
 

private void addEntity(String keyName, String keyClass, String property, String value) {
  Key key = datastore.newKeyFactory().setKind(keyClass).newKey(keyName);
  Entity.Builder entityBuilder = Entity.newBuilder(key);
  entityBuilder.set(property, value);
  Entity entity = entityBuilder.build();
  datastore.put(entity);
}
 

static Entity runStateToEntity(KeyFactory keyFactory, WorkflowInstance wfi, RunState state)
    throws JsonProcessingException {
  final Key key = activeWorkflowInstanceKey(keyFactory, wfi);
  final Entity.Builder entity = Entity.newBuilder(key)
      .set(PROPERTY_COMPONENT, wfi.workflowId().componentId())
      .set(PROPERTY_WORKFLOW, wfi.workflowId().id())
      .set(PROPERTY_PARAMETER, wfi.parameter())
      .set(PROPERTY_COUNTER, state.counter());

  entity
      .set(PROPERTY_STATE, state.state().toString())
      .set(PROPERTY_STATE_TIMESTAMP, state.timestamp())
      .set(PROPERTY_STATE_TRIES, state.data().tries())
      .set(PROPERTY_STATE_CONSECUTIVE_FAILURES, state.data().consecutiveFailures())
      .set(PROPERTY_STATE_RETRY_COST, state.data().retryCost())
      // TODO: consider making this list bounded or not storing it here to avoid exceeding entity size limit
      .set(PROPERTY_STATE_MESSAGES, jsonValue(state.data().messages()));

  state.data().retryDelayMillis().ifPresent(v -> entity.set(PROPERTY_STATE_RETRY_DELAY_MILLIS, v));
  state.data().lastExit().ifPresent(v -> entity.set(PROPERTY_STATE_LAST_EXIT, v));
  state.data().trigger().ifPresent(trigger -> {
    entity.set(PROPERTY_STATE_TRIGGER_TYPE, TriggerUtil.triggerType(trigger));
    entity.set(PROPERTY_STATE_TRIGGER_ID, TriggerUtil.triggerId(trigger));
  });
  state.data().executionId().ifPresent(v -> entity.set(PROPERTY_STATE_EXECUTION_ID, v));
  state.data().runnerId().ifPresent(v -> entity.set(PROPERTY_STATE_RUNNER_ID, v));
  if (state.data().triggerParameters().isPresent()) {
    entity.set(PROPERTY_STATE_TRIGGER_PARAMETERS, jsonValue(state.data().triggerParameters().get()));
  }

  if (state.data().executionDescription().isPresent()) {
    entity.set(PROPERTY_STATE_EXECUTION_DESCRIPTION, jsonValue(state.data().executionDescription().get()));
  }
  if (state.data().resourceIds().isPresent()) {
    entity.set(PROPERTY_STATE_RESOURCE_IDS, jsonValue(state.data().resourceIds().get()));
  }

  return entity.build();
}
 

/** Example of putting a single entity. */
// [TARGET put(FullEntity)]
// [VARIABLE "my_key_name"]
public void putSingleEntity(String keyName) {
  Datastore datastore = transaction.getDatastore();
  // [START putSingleEntity]
  Key key = datastore.newKeyFactory().setKind("MyKind").newKey(keyName);
  Entity.Builder entityBuilder = Entity.newBuilder(key);
  entityBuilder.set("propertyName", "value");
  Entity entity = entityBuilder.build();
  transaction.put(entity);
  transaction.commit();
  // [END putSingleEntity]
}
 

@Test
public void testCollectionFieldsUnsupported() {
	this.thrown.expect(DatastoreDataException.class);
	this.thrown.expectMessage("Unable to write collectionOfUnsupportedTypes.unsupportedElts");
	this.thrown.expectMessage("Unable to convert " +
					"class org.springframework.cloud.gcp.data.datastore.core.convert." +
					"TestItemUnsupportedFields$NewType to Datastore supported type.");

	TestItemUnsupportedFields.CollectionOfUnsupportedTypes item = getCollectionOfUnsupportedTypesItem();

	Entity.Builder builder = getEntityBuilder();
	ENTITY_CONVERTER.write(item, builder);
}
 

@Test
public void testCollectionFieldsNulls() {
	TestDatastoreItemCollections item =
			new TestDatastoreItemCollections(
					Arrays.asList(1, 2),
					null,
					null, new boolean[] {true, false}, null, null);

	Entity.Builder builder = getEntityBuilder();
	ENTITY_CONVERTER.write(item, builder);
	Entity entity = builder.build();

	List<Value<?>> intList = entity.getList("intList");
	assertThat(intList.stream().map(Value::get).collect(Collectors.toList()))
			.as("validate int list values").isEqualTo(Arrays.asList(1L, 2L));

	List<Value<?>> stringArray = entity.getList("stringArray");
	assertThat(stringArray)
			.as("validate string array is null").isNull();

	List<Value<?>> beanContext = entity.getList("beanContext");
	assertThat(beanContext)
			.as("validate bean context is null")
			.isNull();

	TestDatastoreItemCollections readItem = ENTITY_CONVERTER.read(TestDatastoreItemCollections.class, entity);
	assertThat(item.equals(readItem)).as("read object should be equal to original").isTrue();

}
 

/**
 * Stores the state value in each key-value pair in the project's datastore.
 *
 * @param sessionId Request from which to extract session.
 * @param varName   the name of the desired session variable
 * @param varValue  the value of the desired session variable
 */
protected void setSessionVariables(String sessionId, Map<String, String> setMap) {
  if (sessionId.equals("")) {
    return;
  }
  Key key = keyFactory.newKey(sessionId);
  Transaction transaction = datastore.newTransaction();
  DateTime dt = DateTime.now(DateTimeZone.UTC);
  dt.toString(dtf);
  try {
    Entity stateEntity = transaction.get(key);
    Entity.Builder seBuilder;
    if (stateEntity == null) {
      seBuilder = Entity.newBuilder(key);
    } else {
      seBuilder = Entity.newBuilder(stateEntity);
    }
    for (String varName : setMap.keySet()) {
      seBuilder.set(varName, setMap.get(varName));
    }
    transaction.put(seBuilder.set("lastModified", dt.toString(dtf)).build());
    transaction.commit();
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

/**
 * Stores the state value in each key-value pair in the project's datastore.
 * @param sessionId Request from which to extract session.
 * @param varName the name of the desired session variable
 * @param varValue the value of the desired session variable
 */
protected void setSessionVariables(String sessionId, Map<String, String> setMap) {
  if (sessionId.equals("")) {
    return;
  }
  Key key = keyFactory.newKey(sessionId);
  Transaction transaction = datastore.newTransaction();
  DateTime dt = DateTime.now(DateTimeZone.UTC);
  dt.toString(dtf);
  try {
    Entity stateEntity = transaction.get(key);
    Entity.Builder seBuilder;
    if (stateEntity == null) {
      seBuilder = Entity.newBuilder(key);
    } else {
      seBuilder = Entity.newBuilder(stateEntity);
    }
    for (String varName : setMap.keySet()) {
      seBuilder.set(varName, setMap.get(varName));
    }
    transaction.put(seBuilder.set("lastModified", dt.toString(dtf)).build());
    transaction.commit();
  } finally {
    if (transaction.isActive()) {
      transaction.rollback();
    }
  }
}
 

@Override
public void saveToken(String token) {
  // Save message to "messages"
  Datastore datastore = getDatastoreInstance();
  Key key = datastore.allocateId(tokensKindKeyFactory.newKey());

  Entity.Builder tokenEntityBuilder = Entity.newBuilder(key).set("token", token);

  datastore.put(tokenEntityBuilder.build());
}
 

@Override
public void save(Message message) {
  // Save message to "messages"
  Datastore datastore = getDatastoreInstance();
  Key key = datastore.allocateId(keyFactory.newKey());

  Entity.Builder messageEntityBuilder =
      Entity.newBuilder(key).set("messageId", message.getMessageId());

  String translated = message.getTranslated();
  if (translated != null) {
    messageEntityBuilder = messageEntityBuilder.set("data", translated);
  }

  if (message.getPublishTime() != null) {
    messageEntityBuilder = messageEntityBuilder.set("publishTime", message.getPublishTime());
  }

  if (message.getSourceLang() != null) {
    messageEntityBuilder = messageEntityBuilder.set("sourceLang", message.getSourceLang());
  }

  if (message.getTargetLang() != null) {
    messageEntityBuilder = messageEntityBuilder.set("targetLang", message.getTargetLang());
  }
  datastore.put(messageEntityBuilder.build());
}
 

private Entity.Builder getEntityBuilder() {
	return Entity.newBuilder(this.datastore.newKeyFactory().setKind("aKind").newKey("1"));
}