com.google.datastore.v1.PartitionId Java Examples

@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 ListenableFuture<QueryResult> executeAsync(final Query statement, final ListenableFuture<TransactionResult> txn) {
  final ListenableFuture<Response> httpResponse = Futures.transformAsync(txn, result -> {
    final String namespace = config.getNamespace();
    final RunQueryRequest.Builder request = RunQueryRequest.newBuilder()
      .setQuery(statement.getPb(namespace != null ? namespace : ""));
    if (namespace != null) {
      request.setPartitionId(PartitionId.newBuilder().setNamespaceId(namespace));
    }
    final ByteString transaction = result.getTransaction();
    if (transaction != null) {
      request.setReadOptions(ReadOptions.newBuilder().setTransaction(transaction));
    }
    final ProtoHttpContent payload = new ProtoHttpContent(request.build());
    return ListenableFutureAdapter.asGuavaFuture(prepareRequest("runQuery", payload).execute());
  }, MoreExecutors.directExecutor());
  return Futures.transformAsync(httpResponse, response -> {
    if (!isSuccessful(response.getStatusCode())) {
      throw new DatastoreException(response.getStatusCode(), response.getResponseBody());
    }
    final RunQueryResponse query = RunQueryResponse.parseFrom(streamResponse(response));
    return Futures.immediateFuture(QueryResult.build(query));
  }, MoreExecutors.directExecutor());
}
 

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

/**
 * 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;
}
 

@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_NonMatchingPartitionId2() {
  PartitionId partitionId1 = PartitionId.newBuilder()
      .setNamespaceId("namespace-id")
      .build();
  PartitionId partitionId2 = PartitionId.newBuilder()
      .setNamespaceId("another-namespace-id")
      .build();
  try {
    makeKey(
        PARENT.toBuilder().setPartitionId(partitionId1).build(),
        CHILD.toBuilder().setPartitionId(partitionId2).build());
    fail("expected IllegalArgumentException");
  } catch (IllegalArgumentException expected) {
  }
}
 

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

private static PartitionId.Builder forNamespace(@Nullable String namespace) {
  PartitionId.Builder partitionBuilder = PartitionId.newBuilder();
  // Namespace either being null or empty represents the default namespace.
  // Datastore Client libraries expect users to not set the namespace proto field in
  // either of these cases.
  if (!Strings.isNullOrEmpty(namespace)) {
    partitionBuilder.setNamespaceId(namespace);
  }
  return partitionBuilder;
}
 

/** Tests {@link SplitQueryFn} when number of query splits is specified. */
@Test
public void testSplitQueryFnWithNumSplits() throws Exception {
  int numSplits = 100;
  when(mockQuerySplitter.getSplits(
          eq(QUERY), any(PartitionId.class), eq(numSplits), any(Datastore.class)))
      .thenReturn(splitQuery(QUERY, numSplits));

  SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, numSplits, mockDatastoreFactory);
  DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
  /**
   * Although Datastore client is marked transient in {@link SplitQueryFn}, when injected through
   * mock factory using a when clause for unit testing purposes, it is not serializable because it
   * doesn't have a no-arg constructor. Thus disabling the cloning to prevent the doFn from being
   * serialized.
   */
  doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
  List<Query> queries = doFnTester.processBundle(QUERY);

  assertEquals(queries.size(), numSplits);

  // Confirms that sub-queries are not equal to original when there is more than one split.
  for (Query subQuery : queries) {
    assertNotEquals(subQuery, QUERY);
  }

  verify(mockQuerySplitter, times(1))
      .getSplits(eq(QUERY), any(PartitionId.class), eq(numSplits), any(Datastore.class));
  verifyZeroInteractions(mockDatastore);
}
 

/** Tests {@link SplitQueryFn} when no query splits is specified. */
@Test
public void testSplitQueryFnWithoutNumSplits() throws Exception {
  // Force SplitQueryFn to compute the number of query splits
  int numSplits = 0;
  int expectedNumSplits = 20;
  long entityBytes = expectedNumSplits * DEFAULT_BUNDLE_SIZE_BYTES;
  // In seconds
  long timestamp = 1234L;

  RunQueryRequest latestTimestampRequest =
      makeRequest(makeLatestTimestampQuery(NAMESPACE), NAMESPACE);
  RunQueryResponse latestTimestampResponse = makeLatestTimestampResponse(timestamp);

  // Per Kind statistics request and response
  RunQueryRequest statRequest = makeRequest(makeStatKindQuery(NAMESPACE, timestamp), NAMESPACE);
  RunQueryResponse statResponse = makeStatKindResponse(entityBytes);

  when(mockDatastore.runQuery(latestTimestampRequest)).thenReturn(latestTimestampResponse);
  when(mockDatastore.runQuery(statRequest)).thenReturn(statResponse);
  when(mockQuerySplitter.getSplits(
          eq(QUERY), any(PartitionId.class), eq(expectedNumSplits), any(Datastore.class)))
      .thenReturn(splitQuery(QUERY, expectedNumSplits));

  SplitQueryFn splitQueryFn = new SplitQueryFn(V_1_OPTIONS, numSplits, mockDatastoreFactory);
  DoFnTester<Query, Query> doFnTester = DoFnTester.of(splitQueryFn);
  doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
  List<Query> queries = doFnTester.processBundle(QUERY);

  assertEquals(expectedNumSplits, queries.size());
  verify(mockQuerySplitter, times(1))
      .getSplits(eq(QUERY), any(PartitionId.class), eq(expectedNumSplits), any(Datastore.class));
  verify(mockDatastore, times(1)).runQuery(latestTimestampRequest);
  verify(mockDatastore, times(1)).runQuery(statRequest);
}
 

@Before
public void testData() {
  entities = new ArrayList<>();
  entities.add(Entity.newBuilder()
      .setKey(Key.newBuilder()
          .setPartitionId(PartitionId.newBuilder()
              .setProjectId("my-project")
              .setNamespaceId("some-namespace"))
          .addPath(PathElement.newBuilder()
              .setId(1234)
              .setKind("monkey")))
      .putProperties("someString", Value.newBuilder()
          .setStringValue("someValue")
          .build())
      .build());
  entities.add(Entity.newBuilder()
      .setKey(Key.newBuilder()
          .setPartitionId(PartitionId.newBuilder()
              .setProjectId("my-project")
              .setNamespaceId("some-namespace"))
          .addPath(PathElement.newBuilder()
              .setName("SomeName")
              .setKind("monkey")))
      .putProperties("someString", Value.newBuilder()
          .setIntegerValue(100L)
          .build())
      .putProperties("someSubRecord", Value.newBuilder()
          .setEntityValue(Entity.newBuilder()
              .putProperties("someSubFloat", Value.newBuilder()
                  .setDoubleValue(0.234)
                  .build()))
          .build())
      .putProperties("someArray", Value.newBuilder()
          .setArrayValue(ArrayValue.newBuilder()
              .addValues(Value.newBuilder()
                  .setIntegerValue(1234L))
              .addValues(Value.newBuilder()
                  .setIntegerValue(1234L))
              .addValues(Value.newBuilder()
                  .setArrayValue(ArrayValue.newBuilder()
                      .addValues(Value.newBuilder()
                          .setStringValue("Some nested string in a nested array"))))
              .build())
          .build())
      .build());

  entitiesJson = new LinkedList<>();
  entitiesJson.add(
      "{\"key\":{\"partitionId\":{\"projectId\":\"my-project\","
          + "\"namespaceId\":\"some-namespace\"},\"path\":"
          + "[{\"kind\":\"monkey\",\"id\":\"1234\"}]},"
          + "\"properties\":{\"someString\":{\"stringValue\":\"someValue\"}}}");

}
 

/**
 * Make a key from the specified path of kind/id-or-name pairs
 * and/or Keys.
 *
 * <p>The id-or-name values must be either String, Long, Integer or Short.
 *
 * <p>The last id-or-name value may be omitted, in which case an entity without
 * an id is created (for use with automatic id allocation).
 *
 * <p>The PartitionIds of all Keys in the path must be equal. The returned
 * Key.Builder will use this PartitionId.
 */
public static Key.Builder makeKey(Object... elements) {
  Key.Builder key = Key.newBuilder();
  PartitionId partitionId = null;
  for (int pathIndex = 0; pathIndex < elements.length; pathIndex += 2) {
    PathElement.Builder pathElement = PathElement.newBuilder();
    Object element =  elements[pathIndex];
    if (element instanceof Key) {
      Key subKey = (Key) element;
      if (partitionId == null) {
        partitionId = subKey.getPartitionId();
      } else if (!partitionId.equals(subKey.getPartitionId())) {
        throw new IllegalArgumentException("Partition IDs did not match, found: "
            + partitionId + " and " + subKey.getPartitionId());
      }
      key.addAllPath(((Key) element).getPathList());
      // We increment by 2, but since we got a Key argument we're only consuming 1 element in this
      // iteration of the loop. Decrement the index so that when we jump by 2 we end up in the
      // right spot.
      pathIndex--;
    } else {
      String kind;
      try {
        kind = (String) element;
      } catch (ClassCastException e) {
        throw new IllegalArgumentException("Expected string or Key, got: " + element.getClass());
      }
      pathElement.setKind(kind);
      if (pathIndex + 1 < elements.length) {
        Object value = elements[pathIndex + 1];
        if (value instanceof String) {
          pathElement.setName((String) value);
        } else if (value instanceof Long) {
          pathElement.setId((Long) value);
        } else if (value instanceof Integer) {
          pathElement.setId((Integer) value);
        } else if (value instanceof Short) {
          pathElement.setId((Short) value);
        } else {
          throw new IllegalArgumentException(
              "Expected string or integer, got: " + value.getClass());
        }
      }
      key.addPath(pathElement);
    }
  }
  if (partitionId != null && !partitionId.equals(PartitionId.getDefaultInstance())) {
    key.setPartitionId(partitionId);
  }
  return key;
}
 

/**
 * Returns a list of sharded {@link Query}s for the given query.
 *
 * <p>This will create up to the desired number of splits, however it may return less splits if
 * the desired number of splits is unavailable. This will happen if the number of split points
 * provided by the underlying Datastore is less than the desired number, which will occur if the
 * number of results for the query is too small.
 *
 * @param query the query to split.
 * @param partition the partition the query is running in.
 * @param numSplits the desired number of splits.
 * @param datastore the datastore to run on.
 * @throws DatastoreException if there was a datastore error while generating query splits.
 * @throws IllegalArgumentException if the given query or numSplits was invalid.
 */
List<Query> getSplits(Query query, PartitionId partition, int numSplits, Datastore datastore)
    throws DatastoreException;