com.google.datastore.v1.Key Java Examples

/**
 * Add a greeting to the specified guestbook.
 */
private void addGreeting(String guestbookName, String user, String message)
    throws DatastoreException {
  Entity.Builder greeting = Entity.newBuilder();
  greeting.setKey(makeKey(GUESTBOOK_KIND, guestbookName, GREETING_KIND));
  greeting.putProperties(USER_PROPERTY, makeValue(user).build());
  greeting.putProperties(MESSAGE_PROPERTY, makeValue(message).build());
  greeting.putProperties(DATE_PROPERTY, makeValue(new Date()).build());
  Key greetingKey = insert(greeting.build());
  System.out.println("greeting key is: " + greetingKey);
}
 

/**
 * Gets a list of split keys given a desired number of splits.
 *
 * <p>This list will contain multiple split keys for each split. Only a single split key
 * will be chosen as the split point, however providing multiple keys allows for more uniform
 * sharding.
 *
 * @param numSplits the number of desired splits.
 * @param query the user query.
 * @param partition the partition to run the query in.
 * @param datastore the datastore containing the data.
 * @throws DatastoreException if there was an error when executing the datastore query.
 */
private List<Key> getScatterKeys(
    int numSplits, Query query, PartitionId partition, Datastore datastore)
    throws DatastoreException {
  Query.Builder scatterPointQuery = createScatterQuery(query, numSplits);

  List<Key> keySplits = new ArrayList<Key>();

  QueryResultBatch batch;
  do {
    RunQueryRequest scatterRequest =
        RunQueryRequest.newBuilder()
            .setPartitionId(partition)
            .setQuery(scatterPointQuery)
            .build();
    batch = datastore.runQuery(scatterRequest).getBatch();
    for (EntityResult result : batch.getEntityResultsList()) {
      keySplits.add(result.getEntity().getKey());
    }
    scatterPointQuery.setStartCursor(batch.getEndCursor());
    scatterPointQuery.getLimitBuilder().setValue(
        scatterPointQuery.getLimit().getValue() - batch.getEntityResultsCount());
  } while (batch.getMoreResults() == MoreResultsType.NOT_FINISHED);
  Collections.sort(keySplits, DatastoreHelper.getKeyComparator());
  return keySplits;
}
 

@Override
public int compare(Key thisKey, Key otherKey) {
  if (!thisKey.getPartitionId().equals(otherKey.getPartitionId())) {
    throw new IllegalArgumentException("Cannot compare keys with different partition ids.");
  }

  Iterator<PathElement> thisPath = thisKey.getPathList().iterator();
  Iterator<PathElement> otherPath = otherKey.getPathList().iterator();
  while (thisPath.hasNext()) {
    if (!otherPath.hasNext()) {
      return 1;
    }
    int result = comparePathElement(thisPath.next(), otherPath.next());
    if (result != 0) {
      return result;
    }
  }

  return otherPath.hasNext() ? -1 : 0;
}
 

/**
 * Build an entity for the given ancestorKey, kind, namespace and value.
 *
 * @param largePropertySize if greater than 0, add an unindexed property of the given size.
 */
static Entity makeEntity(
    Long value, Key ancestorKey, String kind, @Nullable String namespace, int largePropertySize) {
  Entity.Builder entityBuilder = Entity.newBuilder();
  Key.Builder keyBuilder = makeKey(ancestorKey, kind, UUID.randomUUID().toString());
  // NOTE: Namespace is not inherited between keys created with DatastoreHelper.makeKey, so
  // we must set the namespace on keyBuilder. TODO: Once partitionId inheritance is added,
  // we can simplify this code.
  if (namespace != null) {
    keyBuilder.getPartitionIdBuilder().setNamespaceId(namespace);
  }

  entityBuilder.setKey(keyBuilder.build());
  entityBuilder.putProperties("value", makeValue(value).build());
  if (largePropertySize > 0) {
    entityBuilder.putProperties(
        "unindexed_value",
        makeValue(new String(new char[largePropertySize]).replace("\0", "A"))
            .setExcludeFromIndexes(true)
            .build());
  }
  return entityBuilder.build();
}
 

@ProcessElement
public void processElement(ProcessContext c) throws InvalidProtocolBufferException {
  String entityJson = c.element();
  Entity.Builder entityBuilder = Entity.newBuilder();
  entityJsonParser.merge(entityJson, entityBuilder);

  // Build entity who's key has an empty project Id.
  // This allows DatastoreIO to handle what project Entities are loaded into
  Key k = entityBuilder.build().getKey();
  entityBuilder.setKey(Key.newBuilder()
      .addAllPath(k.getPathList())
      .setPartitionId(PartitionId.newBuilder()
          .setProjectId("")
          .setNamespaceId(k.getPartitionId().getNamespaceId())));

  c.output(entityBuilder.build());
}
 

@Override
public List<Query> getSplits(
    Query query, PartitionId partition, int numSplits, Datastore datastore)
    throws DatastoreException, IllegalArgumentException {

  List<Query> splits = new ArrayList<Query>(numSplits);
  if (numSplits == 1) {
    splits.add(query);
    return splits;
  }
  validateQuery(query);
  validateSplitSize(numSplits);

  List<Key> scatterKeys = getScatterKeys(numSplits, query, partition, datastore);
  Key lastKey = null;
  for (Key nextKey : getSplitKey(scatterKeys, numSplits)) {
    splits.add(createSplit(lastKey, nextKey, query));
    lastKey = nextKey;
  }
  splits.add(createSplit(lastKey, null, query));
  return splits;
}
 

/**
 * Create a new {@link Query} given the query and range.
 *
 * @param lastKey the previous key. If null then assumed to be the beginning.
 * @param nextKey the next key. If null then assumed to be the end.
 * @param query the desired query.
 */
private Query createSplit(Key lastKey, Key nextKey, Query query) {
  if (lastKey == null && nextKey == null) {
    return query;
  }
  List<Filter> keyFilters = new ArrayList<Filter>();
  if (query.hasFilter()) {
    keyFilters.add(query.getFilter());
  }
  if (lastKey != null) {
    Filter lowerBound = DatastoreHelper.makeFilter(DatastoreHelper.KEY_PROPERTY_NAME,
        PropertyFilter.Operator.GREATER_THAN_OR_EQUAL,
        DatastoreHelper.makeValue(lastKey)).build();
    keyFilters.add(lowerBound);
  }
  if (nextKey != null) {
    Filter upperBound = DatastoreHelper.makeFilter(DatastoreHelper.KEY_PROPERTY_NAME,
        PropertyFilter.Operator.LESS_THAN,
        DatastoreHelper.makeValue(nextKey)).build();
    keyFilters.add(upperBound);
  }
  return Query.newBuilder(query).setFilter(makeAndFilter(keyFilters)).build();
}
 

@Test
public void testMakeKey_PartitionId() {
  PartitionId partitionId = PartitionId.newBuilder()
      .setNamespaceId("namespace-id")
      .build();
  Key parent = PARENT.toBuilder()
      .setPartitionId(partitionId)
      .build();
  assertEquals(
      Key.newBuilder()
          .setPartitionId(partitionId)
          .addPath(PARENT.getPath(0))
          .addPath(Key.PathElement.newBuilder().setKind("Child"))
          .build(),
      makeKey(parent, "Child").build());
}
 

@Test
public void testMakeKey_IdInt() {
  assertEquals(
      Key.newBuilder()
          .addPath(Key.PathElement.newBuilder().setKind("Foo").setId(1))
          .build(),
      makeKey("Foo", 1).build());
}
 

@Test
public void testMakeKey_Incomplete() {
  assertEquals(
      Key.newBuilder()
          .addPath(Key.PathElement.newBuilder().setKind("Foo"))
          .build(),
      makeKey("Foo").build());
}
 

@Test
public void testMakeKey_KindIdKey() {
  // 1 key at the end of the series
  assertEquals(
      Key.newBuilder()
          .addPath(Key.PathElement.newBuilder().setKind("Grandparent").setId(24L))
          .addPath(PARENT.getPath(0))
          .build(),
      makeKey("Grandparent", 24L, PARENT).build());
}
 

/**
 * Given a list of keys and a number of splits find the keys to split on.
 *
 * @param keys the list of keys.
 * @param numSplits the number of splits.
 */
private Iterable<Key> getSplitKey(List<Key> keys, int numSplits) {
  // If the number of keys is less than the number of splits, we are limited in the number of
  // splits we can make.
  if (keys.size() < numSplits - 1) {
    return keys;
  }

  // Calculate the number of keys per split. This should be KEYS_PER_SPLIT, but may
  // be less if there are not KEYS_PER_SPLIT * (numSplits - 1) scatter entities.
  //
  // Consider the following dataset, where - represents an entity and * represents an entity
  // that is returned as a scatter entity:
  // ||---*-----*----*-----*-----*------*----*----||
  // If we want 4 splits in this data, the optimal split would look like:
  // ||---*-----*----*-----*-----*------*----*----||
  //            |          |            |
  // The scatter keys in the last region are not useful to us, so we never request them:
  // ||---*-----*----*-----*-----*------*---------||
  //            |          |            |
  // With 6 scatter keys we want to set scatter points at indexes: 1, 3, 5.
  //
  // We keep this as a double so that any "fractional" keys per split get distributed throughout
  // the splits and don't make the last split significantly larger than the rest.
  double numKeysPerSplit = Math.max(1.0, ((double) keys.size()) / (numSplits - 1));

  List<Key> keysList = new ArrayList<Key>(numSplits - 1);
  // Grab the last sample for each split, otherwise the first split will be too small.
  for (int i = 1; i < numSplits; i++) {
    int splitIndex = (int) Math.round(i * numKeysPerSplit) - 1;
    keysList.add(keys.get(splitIndex));
  }

  return keysList;
}
 

@Test
public void testMakeKey_KindIdKeyKind() {
  // 1 key in the middle of the series
  assertEquals(
      Key.newBuilder()
          .addPath(Key.PathElement.newBuilder().setKind("Grandparent").setId(24L))
          .addPath(PARENT.getPath(0))
          .addPath(Key.PathElement.newBuilder().setKind("Child"))
          .build(),
      makeKey("Grandparent", 24L, PARENT, "Child").build());
}
 

@Test
public void testMakeKey_IdLong() {
  assertEquals(
      Key.newBuilder()
          .addPath(Key.PathElement.newBuilder().setKind("Foo").setId(1))
          .build(),
      makeKey("Foo", 1L).build());
}
 

@Test
public void testMakeKey_IdShort() {
  assertEquals(
      Key.newBuilder()
          .addPath(Key.PathElement.newBuilder().setKind("Foo").setId(1))
          .build(),
      makeKey("Foo", (short) 1).build());
}
 

@Test
public void testMakeKey_KeyKindIdKey() {
  // 1 key at the beginning and 1 key at the end of the series
  assertEquals(
      Key.newBuilder()
          .addPath(GRANDPARENT.getPath(0))
          .addPath(Key.PathElement.newBuilder().setKind("Parent").setId(23L))
          .addPath(CHILD.getPath(0))
          .build(),
      makeKey(GRANDPARENT, "Parent", 23, CHILD).build());
}
 

/**
 * @return the key contained in value
 * @throws IllegalArgumentException if the value does not contain a key.
 */
public static Key getKey(Value value) {
  if (value.getValueTypeCase() != ValueTypeCase.KEY_VALUE) {
    throw new IllegalArgumentException("Value does not contain a key.");
  }
  return value.getKeyValue();
}
 

@Test
public void testMakeKey_Key() {
  // Just 1 key
  assertEquals(
      Key.newBuilder()
          .addPath(CHILD.getPath(0))
          .build(),
      makeKey(CHILD).build());
}
 

@Test
public void testMakeKey_Name() {
  assertEquals(
      Key.newBuilder()
          .addPath(Key.PathElement.newBuilder().setKind("Foo").setName("hi"))
          .build(),
      makeKey("Foo", "hi").build());
}
 

@Test
public void testMakeKey_KindNameKind() {
  assertEquals(
      Key.newBuilder()
          .addPath(Key.PathElement.newBuilder().setKind("Foo").setName("hi"))
          .addPath(Key.PathElement.newBuilder().setKind("Bar"))
          .build(),
      makeKey("Foo", "hi", "Bar").build());
}
 

@Test
public void testMakeKey_KeyKey() {
  // Just 2 keys
  assertEquals(
      Key.newBuilder()
          .addPath(PARENT.getPath(0))
          .addPath(CHILD.getPath(0))
          .build(),
      makeKey(PARENT, CHILD).build());
}
 

@Test
public void testMakeKey_KeyKind() {
  // 1 key at the beginning of the series
  assertEquals(
      Key.newBuilder()
          .addPath(PARENT.getPath(0))
          .addPath(Key.PathElement.newBuilder().setKind("Child"))
          .build(),
      makeKey(PARENT, "Child").build());
}
 

@Test
public void testMakeKey_KeyKeyKey() {
  // Just 3 keys
  assertEquals(
      Key.newBuilder()
          .addPath(GRANDPARENT.getPath(0))
          .addPath(PARENT.getPath(0))
          .addPath(CHILD.getPath(0))
          .build(),
      makeKey(GRANDPARENT, PARENT, CHILD).build());
}
 

/**
  * Insert an entity into the datastore.
  *
  * The entity must have no ids.
  *
  * @return The key for the inserted entity.
  * @throws DatastoreException on error
  */
 private Key insert(Entity entity) throws DatastoreException {
   CommitRequest request = CommitRequest.newBuilder()
.addMutations(Mutation.newBuilder()
    .setInsert(entity))
       .setMode(CommitRequest.Mode.NON_TRANSACTIONAL)
.build();
   return datastore.commit(request).getMutationResults(0).getKey();
 }
 

@Test
public void testMakeKey_MultiLevelKey() {
  // 1 key with 3 elements
  assertEquals(
      Key.newBuilder()
          .addPath(GRANDPARENT.getPath(0))
          .addPath(PARENT.getPath(0))
          .addPath(CHILD.getPath(0))
          .build(),
      makeKey(makeKey(GRANDPARENT, PARENT, CHILD).build()).build());
}
 

@Test
public void testDataStoreV1SqlWriteRead() {
  BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(new DataStoreV1TableProvider());
  String projectId = options.getProject();

  String createTableStatement =
      "CREATE EXTERNAL TABLE TEST( \n"
          + "   `__key__` VARBINARY, \n"
          + "   `content` VARCHAR \n"
          + ") \n"
          + "TYPE 'datastoreV1' \n"
          + "LOCATION '"
          + projectId
          + "/"
          + KIND
          + "'";
  sqlEnv.executeDdl(createTableStatement);

  Key ancestor = makeKey(KIND, UUID.randomUUID().toString()).build();
  Key itemKey = makeKey(ancestor, KIND, UUID.randomUUID().toString()).build();
  String insertStatement =
      "INSERT INTO TEST VALUES ( \n" + keyToSqlByteString(itemKey) + ", \n" + "'2000' \n" + ")";

  BeamSqlRelUtils.toPCollection(writePipeline, sqlEnv.parseQuery(insertStatement));
  writePipeline.run().waitUntilFinish();

  String selectTableStatement = "SELECT * FROM TEST";
  PCollection<Row> output =
      BeamSqlRelUtils.toPCollection(readPipeline, sqlEnv.parseQuery(selectTableStatement));

  assertThat(output.getSchema(), equalTo(SOURCE_SCHEMA));

  PipelineResult.State state = readPipeline.run().waitUntilFinish(Duration.standardMinutes(5));
  assertThat(state, equalTo(State.DONE));
}
 

@Test
public void testWriteRead_viaCoreBeamIO() {
  String projectId = options.getProject();
  Key ancestor = makeKey(KIND, UUID.randomUUID().toString()).build();
  Key itemKey =
      makeKey(ancestor, KIND, UUID.randomUUID().toString())
          .setPartitionId(PartitionId.newBuilder().setProjectId(projectId).build())
          .build();
  Row testWriteRow =
      Row.withSchema(SOURCE_SCHEMA).addValues(itemKey.toByteArray(), "4000").build();

  writePipeline
      .apply(Create.of(testWriteRow).withRowSchema(SOURCE_SCHEMA))
      .apply(RowToEntity.create("__key__", KIND))
      .apply(DatastoreIO.v1().write().withProjectId(projectId));
  writePipeline.run().waitUntilFinish();

  Query.Builder query = Query.newBuilder();
  query.addKindBuilder().setName(KIND);
  query.setFilter(makeFilter("__key__", Operator.EQUAL, makeValue(itemKey)));

  DatastoreV1.Read read =
      DatastoreIO.v1().read().withProjectId(projectId).withQuery(query.build());
  PCollection<Row> rowsRead =
      readPipeline.apply(read).apply(EntityToRow.create(SOURCE_SCHEMA, "__key__"));

  PAssert.that(rowsRead).containsInAnyOrder(testWriteRow);
  readPipeline.run().waitUntilFinish();
}
 

@Test
public void testMakeKey_MultiLevelKeyKey() {
  // 1 key with 2 elements
  assertEquals(
      Key.newBuilder()
          .addPath(GRANDPARENT.getPath(0))
          .addPath(PARENT.getPath(0))
          .addPath(CHILD.getPath(0))
          .build(),
      makeKey(makeKey(GRANDPARENT, PARENT).build(), CHILD).build());
}
 

/**
 * Returns true if a Cloud Datastore key is complete. A key is complete if its last element has
 * either an id or a name.
 */
static boolean isValidKey(Key key) {
  List<PathElement> elementList = key.getPathList();
  if (elementList.isEmpty()) {
    return false;
  }
  PathElement lastElement = elementList.get(elementList.size() - 1);
  return (lastElement.getId() != 0 || !lastElement.getName().isEmpty());
}
 

@Override
public Mutation apply(Key key) {
  // Verify that the entity to delete has a complete key.
  checkArgument(
      isValidKey(key),
      "Keys to be deleted from the Cloud Datastore must be complete:\n%s",
      key);

  return makeDelete(key).build();
}