com.google.cloud.datastore.KeyFactory Java Examples

@Override
public Key getKeyFromId(Object id, String kindName) {
	Assert.notNull(id, "Cannot get key for null ID value.");
	if (id instanceof Key) {
		return (Key) id;
	}
	KeyFactory keyFactory = getKeyFactory();
	keyFactory.setKind(kindName);
	Key key;
	if (id instanceof String) {
		key = keyFactory.newKey((String) id);
	}
	else if (id instanceof Long) {
		key = keyFactory.newKey((long) id);
	}
	else {
		// We will use configurable custom converters to try to convert other types to
		// String or long
		// in the future.
		throw new DatastoreDataException(
				"Keys can only be created using String or long values.");
	}
	return key;
}
 

@Test
public void testSerializationCycle() throws Exception {
  DatastoreSession initialSession = new DatastoreSession(sessionManager);
  initialSession.setValid(true);
  initialSession.setAttribute("count", 5);
  initialSession.setAttribute("map", Collections.singletonMap("key", "value"));

  KeyFactory keyFactory = new KeyFactory("project").setKind("kind");
  List<Entity> attributes = initialSession.saveToEntities(sessionKey, keyFactory);

  DatastoreSession restoredSession = new DatastoreSession(sessionManager);
  restoredSession.restoreFromEntities(sessionKey, attributes);

  assertTrue(restoredSession.getAttribute("count") != null);
  assertTrue(restoredSession.getAttribute("map") != null);

  assertEquals(5, restoredSession.getAttribute("count"));
  assertEquals("value", ((Map)restoredSession.getAttribute("map")).get("key"));
}
 

@Test
public void testAttributesSerializationKey() throws Exception {
  DatastoreSession session = new DatastoreSession(sessionManager);
  session.setValid(true);
  session.setAttribute("count", 2);
  session.setAttribute("map", new HashMap<>());

  KeyFactory factory = new KeyFactory("project").setKind("kind");
  List<Entity> entities = session.saveAttributesToEntity(factory);

  assertTrue(entities.stream()
      .map(BaseEntity::getKey)
      .map(Key::getName)
      .collect(Collectors.toSet())
      .containsAll(Arrays.asList("count", "map")));
}
 

/**
 * Serialize an attribute an embed it into an entity whose key is generated by the provided
 * KeyFactory.
 * @param attributeKeyFactory The KeyFactory to use to create the key for the entity.
 * @param name The name of the attribute to serialize.
 * @return An Entity containing the serialized attribute.
 */
private Entity serializeAttribute(KeyFactory attributeKeyFactory, String name) {
  ByteArrayOutputStream bos = new ByteArrayOutputStream();
  try (ObjectOutputStream oos = new ObjectOutputStream(bos)) {
    oos.writeObject(getAttribute(name));
  } catch (IOException e) {
    throw new UncheckedIOException(e);
  }

  return Entity.newBuilder(attributeKeyFactory.newKey(name))
      .set(SessionMetadata.ATTRIBUTE_VALUE_NAME,
          BlobValue.newBuilder(Blob.copyFrom(bos.toByteArray()))
              .setExcludeFromIndexes(true)
              .build())
      .build();
}
 

@Before
public void setUp() throws Exception {
  MockitoAnnotations.initMocks(this);
  store = new DatastoreStore();
  KeyFactory keyFactory = new KeyFactory("project").setKind("kind");
  key = keyFactory.newKey(keyId);
  attributeKey = keyFactory.newKey("attribute");
  QueryResults<Key> keyQueryResults = new IteratorQueryResults<>(ImmutableList.of(key).iterator());

  when(datastore.newKeyFactory()).thenAnswer((invocation) -> new KeyFactory("project"));
  when(datastore.run(any(KeyQuery.class))).thenReturn(keyQueryResults);

  when(manager.getContext()).thenReturn(new StandardContext());
  when(manager.willAttributeDistribute(anyString(), any())).thenReturn(true);
  when(manager.createEmptySession()).thenReturn(new DatastoreSession(manager));

  store.setDatastore(datastore);
  store.setClock(clock);
  store.setSessionKind("kind");
  store.setManager(manager);
}
 

/**
 * Save the specified Session into this Store. Any previously saved information for
 * the associated session identifier is replaced.
 *
 * <p>Attempt to serialize the session and send it to the datastore.</p>
 *
 * @throws IOException If an error occurs during the serialization of the session.
 *
 * @param session Session to be saved
 */
@Override
public void save(Session session) throws IOException {
  log.debug("Persisting session: " + session.getId());

  if (!(session instanceof DatastoreSession)) {
    throw new IOException(
        "The session must be an instance of DatastoreSession to be serialized");
  }
  DatastoreSession datastoreSession = (DatastoreSession) session;
  Key sessionKey = newKey(session.getId());
  KeyFactory attributeKeyFactory = datastore.newKeyFactory()
      .setKind(sessionKind)
      .addAncestor(PathElement.of(sessionKind, sessionKey.getName()));

  List<Entity> entities = serializeSession(datastoreSession, sessionKey, attributeKeyFactory);

  TraceContext datastoreSaveContext = startSpan("Storing the session in the Datastore");
  datastore.put(entities.toArray(new FullEntity[0]));
  datastore.delete(datastoreSession.getSuppressedAttributes().stream()
      .map(attributeKeyFactory::newKey)
      .toArray(Key[]::new));
  endSpan(datastoreSaveContext);
}
 

@Test
public void testNamespaceRunQuery() {
  setUpQueryTests();
  // [START datastore_namespace_run_query]
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("__namespace__");
  Key startNamespace = keyFactory.newKey("g");
  Key endNamespace = keyFactory.newKey("h");
  Query<Key> query = Query.newKeyQueryBuilder()
      .setKind("__namespace__")
      .setFilter(CompositeFilter.and(
          PropertyFilter.gt("__key__", startNamespace),
          PropertyFilter.lt("__key__", endNamespace)))
      .build();
  List<String> namespaces = new ArrayList<>();
  QueryResults<Key> results = datastore.run(query);
  while (results.hasNext()) {
    namespaces.add(results.next().getName());
  }
  // [END datastore_namespace_run_query]
  assertEquals(ImmutableList.of("ghijklmnop"), namespaces);
}
 

private Entity getShardCounter(final String location) {
  KeyFactory keyFactory = datastore.newKeyFactory().addAncestors(
      NXRM_ROOT,
      PathElement.of(METRICS_STORE, 1L)
  );
  Key key = keyFactory.setNamespace(namespace)
      .setKind(SHARD).newKey(location);
  Entity exists = datastore.get(key);
  if (exists != null) {
    log.trace("counter for {} already present", location);
    return exists;
  }

  log.debug("creating metrics store counter shard for {}", location);
  // otherwise make it
  Entity entity = Entity.newBuilder(key)
      .set(SIZE, LongValue.newBuilder(0L).build())
      .set(COUNT, LongValue.newBuilder(0L).build())
      .build();

  return datastore.put(entity);
}
 

/** Example of getting multiple entity objects. */
// [TARGET get(Iterable, ReadOption...)]
// [VARIABLE "my_first_key_name"]
// [VARIABLE "my_second_key_name"]
public List<Entity> getEntitiesWithKeys(String firstKeyName, String secondKeyName) {
  // TODO change so that it's not necessary to hold the entities in a list for integration testing
  // [START getEntitiesWithKeys]
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key firstKey = keyFactory.newKey(firstKeyName);
  Key secondKey = keyFactory.newKey(secondKeyName);
  Iterator<Entity> entitiesIterator = datastore.get(Lists.newArrayList(firstKey, secondKey));
  List<Entity> entities = Lists.newArrayList();
  while (entitiesIterator.hasNext()) {
    Entity entity = entitiesIterator.next();
    // do something with the entity
    entities.add(entity);
  }
  // [END getEntitiesWithKeys]
  return entities;
}
 

/** Example of fetching a list of Entity objects. */
// [TARGET fetch(Iterable, ReadOption...)]
// [VARIABLE "my_first_key_name"]
// [VARIABLE "my_second_key_name"]
public List<Entity> fetchEntitiesWithKeys(String firstKeyName, String secondKeyName) {
  // [START fetchEntitiesWithKeys]
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key firstKey = keyFactory.newKey(firstKeyName);
  Key secondKey = keyFactory.newKey(secondKeyName);
  List<Entity> entities = datastore.fetch(Lists.newArrayList(firstKey, secondKey));
  for (Entity entity : entities) {
    // do something with the entity
  }
  // [END fetchEntitiesWithKeys]
  return entities;
}
 

/** Example of getting entities for several keys. */
// [TARGET get(Key...)]
// [VARIABLE "my_first_key_name"]
// [VARIABLE "my_second_key_name"]
public List<Entity> getMultiple(String firstKeyName, String secondKeyName) {
  Datastore datastore = transaction.getDatastore();
  // TODO change so that it's not necessary to hold the entities in a list for integration testing
  // [START getMultiple]
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key firstKey = keyFactory.newKey(firstKeyName);
  Key secondKey = keyFactory.newKey(secondKeyName);
  Iterator<Entity> entitiesIterator = transaction.get(firstKey, secondKey);
  List<Entity> entities = Lists.newArrayList();
  while (entitiesIterator.hasNext()) {
    Entity entity = entitiesIterator.next();
    // do something with the entity
    entities.add(entity);
  }
  transaction.commit();
  // [END getMultiple]
  return entities;
}
 

/** Example of verifying if a transaction is active. */
// [TARGET isActive()]
public Key isActive() {
  Datastore datastore = transaction.getDatastore();
  // [START isActive]
  // create an entity
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key key = datastore.allocateId(keyFactory.newKey());
  Entity entity = Entity.newBuilder(key).set("description", "active()").build();
  // calling transaction.active() now would return true
  try {
    // add the entity and commit
    transaction.put(entity);
    transaction.commit();
  } finally {
    // if committing succeeded
    // then transaction.active() will be false
    if (transaction.isActive()) {
      // otherwise it's true and we need to rollback
      transaction.rollback();
    }
  }
  // [END isActive]
  return key;
}
 

/** Example of verifying if a transaction is active. */
// [TARGET active()]
public Key active() {
  Datastore datastore = transaction.getDatastore();
  // [START active]
  // create an entity
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key key = datastore.allocateId(keyFactory.newKey());
  Entity entity = Entity.newBuilder(key).set("description", "active()").build();
  // calling transaction.active() now would return true
  try {
    // add the entity and commit
    transaction.put(entity);
    transaction.commit();
  } finally {
    // if committing succeeded
    // then transaction.isActive() will be false
    if (transaction.isActive()) {
      // otherwise it's true and we need to rollback
      transaction.rollback();
    }
  }
  // [END active]
  return key;
}
 

/** Example of rolling back a transaction. */
// [TARGET rollback()]
public Key rollback() {
  Datastore datastore = transaction.getDatastore();
  // [START rollback]
  // create an entity
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key key = datastore.allocateId(keyFactory.newKey());
  Entity entity = Entity.newBuilder(key).set("description", "rollback()").build();

  // add the entity and rollback
  transaction.put(entity);
  transaction.rollback();
  // calling transaction.commit() now would fail
  // [END rollback]
  return key;
}
 

/** Example of committing a transaction. */
// [TARGET commit()]
public Key commit() {
  Datastore datastore = transaction.getDatastore();
  // [START commit]
  // create an entity
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key key = datastore.allocateId(keyFactory.newKey());
  Entity entity = Entity.newBuilder(key).set("description", "commit()").build();

  // add the entity and commit
  try {
    transaction.put(entity);
    transaction.commit();
  } catch (DatastoreException ex) {
    // handle exception
  }
  // [END commit]

  return key;
}
 

/** Example of running a query to find all entities with an ancestor. */
// [TARGET run(Query)]
// [VARIABLE "my_parent_key_name"]
public List<Entity> run(String parentKeyName) {
  Datastore datastore = transaction.getDatastore();
  // [START run]
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("ParentKind");
  Key parentKey = keyFactory.newKey(parentKeyName);
  // Build a query
  Query<Entity> query =
      Query.newEntityQueryBuilder()
          .setKind("MyKind")
          .setFilter(PropertyFilter.hasAncestor(parentKey))
          .build();
  QueryResults<Entity> results = transaction.run(query);
  List<Entity> entities = Lists.newArrayList();
  while (results.hasNext()) {
    Entity result = results.next();
    // do something with result
    entities.add(result);
  }
  transaction.commit();
  // [END run]
  return entities;
}
 

/** Example of fetching a list of entities for several keys. */
// [TARGET fetch(Key...)]
// [VARIABLE "my_first_key_name"]
// [VARIABLE "my_second_key_name"]
public List<Entity> fetchEntitiesWithKeys(String firstKeyName, String secondKeyName) {
  Datastore datastore = transaction.getDatastore();
  // [START fetchEntitiesWithKeys]
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key firstKey = keyFactory.newKey(firstKeyName);
  Key secondKey = keyFactory.newKey(secondKeyName);
  List<Entity> entities = transaction.fetch(firstKey, secondKey);
  for (Entity entity : entities) {
    // do something with the entity
  }
  transaction.commit();
  // [END fetchEntitiesWithKeys]
  return entities;
}
 

/**
 * Serialize a session to a list of Entities that can be stored to the Datastore.
 * @param session The session to serialize.
 * @return A list of one or more entities containing the session and its attributes.
 * @throws IOException If the session cannot be serialized.
 */
@VisibleForTesting
List<Entity> serializeSession(DatastoreSession session, Key sessionKey,
    KeyFactory attributeKeyFactory) throws IOException {
  TraceContext serializationContext = startSpan("Serialization of the session");
  List<Entity> entities = session.saveToEntities(sessionKey, attributeKeyFactory);
  endSpan(serializationContext);
  return entities;
}
 

/**
 * Serialize the session attributes into entities.
 * @param attributeKeyFactory The key builder for the entities.
 * @return A list of entities where the key correspond to the name of the attribute
           and the property `value` to the serialized attribute.
 * @throws IOException If an error occur during the serialization.
 */
@VisibleForTesting
List<Entity> saveAttributesToEntity(KeyFactory attributeKeyFactory) throws
    IOException {
  Stream<Entity> entities = Collections.list(getAttributeNames()).stream()
      .filter(name -> accessedAttributes.contains(name))
      .filter(name -> isAttributeDistributable(name, getAttribute(name)))
      .map(name -> serializeAttribute(attributeKeyFactory, name));

  try {
    return entities.collect(Collectors.toList());
  } catch (UncheckedIOException e) {
    throw e.getCause();
  }
}
 

private List<Key> setUpTransferTests() {
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("People");
  Key from = keyFactory.newKey("from");
  Key to = keyFactory.newKey("to");
  datastore.put(Entity.newBuilder(from).set("balance", 100).build());
  datastore.put(Entity.newBuilder(to).set("balance", 0).build());
  return ImmutableList.of(from, to);
}
 

@Test
public void testKeyWithMultilevelParent() {
  // [START datastore_key_with_multilevel_parent]
  KeyFactory keyFactory = datastore.newKeyFactory()
      .addAncestors(PathElement.of("User", "Alice"), PathElement.of("TaskList", "default"))
      .setKind("Task");
  Key taskKey = keyFactory.newKey("sampleTask");
  // [END datastore_key_with_multilevel_parent]
  assertValidKey(taskKey);
}
 

@Test
public void testIncompleteKey() {
  // [START datastore_incomplete_key]
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("Task");
  Key taskKey = datastore.allocateId(keyFactory.newKey());
  // [END datastore_incomplete_key]
  assertValidKey(taskKey);
}
 

public static void main(String... args) {
  Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("keyKind");
  Key key = keyFactory.newKey("keyName");
  Entity entity = datastore.get(key);
  if (entity != null) {
    System.out.println("Updating access_time for " + entity.getString("name"));
    entity = Entity.newBuilder(entity).set("access_time", Timestamp.now()).build();
    datastore.update(entity);
  }
}
 

public static void main(String... args) {
  Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("keyKind");
  Key key = keyFactory.newKey("keyName");
  Entity entity =
      Entity.newBuilder(key)
          .set("name", "John Doe")
          .set("age", 30)
          .set("access_time", Timestamp.now())
          .build();
  datastore.put(entity);
}
 

/** Example of creating a {@code KeyFactory}. */
// [TARGET newKeyFactory()]
public KeyFactory createKeyFactory() {
  // [START createKeyFactory]
  KeyFactory keyFactory = datastore.newKeyFactory();
  // [END createKeyFactory]
  return keyFactory;
}
 

/** Example of allocating multiple ids in a single batch. */
// [TARGET allocateId(IncompleteKey...)]
public List<Key> batchAllocateId() {
  // [START batchAllocateId]
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  IncompleteKey incompleteKey1 = keyFactory.newKey();
  IncompleteKey incompleteKey2 = keyFactory.newKey();

  // let cloud datastore automatically assign the ids
  List<Key> keys = datastore.allocateId(incompleteKey1, incompleteKey2);
  // [END batchAllocateId]
  return keys;
}
 

/** Example of allocating an id. */
// [TARGET allocateId(IncompleteKey)]
public Key allocateIdSingle() {
  // [START allocateIdSingle]
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  IncompleteKey incompleteKey = keyFactory.newKey();

  // let cloud datastore automatically assign an id
  Key key = datastore.allocateId(incompleteKey);
  // [END allocateIdSingle]
  return key;
}
 

@Test
public void testDecomposedSessionLoad() throws Exception {
  DatastoreSession session = new DatastoreSession(manager);
  session.setValid(true);
  session.setId(keyId);
  session.setAttribute("count", 2);
  session.setAttribute("map", Collections.singletonMap("key", "value"));

  KeyFactory attributeKeyFactory = datastore.newKeyFactory()
      .setKind("kind")
      .addAncestor(PathElement.of("kind", key.getName()));
  List<Entity> entities = session.saveToEntities(key, attributeKeyFactory);

  QueryResults<Entity> queryResults = new IteratorQueryResults<>(entities.iterator());
  when(datastore.<Entity>run(any())).thenReturn(queryResults);

  Session restored = store.load(keyId);

  assertEquals(keyId, restored.getId());
  assertEquals(2, restored.getSession().getAttribute("count"));
  assertEquals("value",
      ((Map<String, String>)session.getSession().getAttribute("map")).get("key"));
}
 

@Test(expected = IOException.class)
public void testSerializationWithoutMetadata() throws Exception {
  DatastoreSession session = new DatastoreSession(sessionManager);
  Entity attribute = Entity.newBuilder(new KeyFactory("project")
      .setKind("kind")
      .newKey("456"))
      .build();
  session.restoreFromEntities(sessionKey, Collections.singleton(attribute));
}
 

@Test(expected = NotSerializableException.class)
public void testSerializationError() throws Exception {
  DatastoreSession session = spy(new DatastoreSession(sessionManager));
  session.setValid(true);
  session.setAttribute("count", 5);
  when(session.isAttributeDistributable(any(), any())).thenReturn(true);
  when(session.getAttribute("count")).thenReturn(sessionManager);

  session.saveAttributesToEntity(new KeyFactory("project").setKind("kind"));
}