Skip to content

Commit

Permalink
wip fixing tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@srinjoyray
srinjoyray committed Nov 6, 2024
1 parent 2c718f9 commit 6b973ca
Show file tree
Hide file tree
Showing 18 changed files with 99 additions and 99 deletions.
9 changes: 5 additions & 4 deletions api/src/main/java/com/google/appengine/api/datastore/AdminDatastoreService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -309,14 +309,15 @@ private static List<DatastoreV3Pb.Query.Builder> convertQueryToPbs(
List<MultiQueryBuilder> queriesToRun = QuerySplitHelper.splitQuery(query);
// All Filters should be in queriesToRun
query.setFilter(null);
query.getFilterPredicates().clear();
// query.getFilterPredicates().clear();
List<DatastoreV3Pb.Query.Builder> resultQueries = new ArrayList<DatastoreV3Pb.Query.Builder>();
for (MultiQueryBuilder multiQuery : queriesToRun) {
for (List<List<FilterPredicate>> parallelQueries : multiQuery) {
for (List<FilterPredicate> singleQuery : parallelQueries) {
Query newQuery = new Query(query);
newQuery.getFilterPredicates().addAll(singleQuery);
DatastoreV3Pb.Query.Builder queryProto = QueryTranslator.convertToPb(newQuery, fetchOptions);
// Query newQuery = new Query(query);
query.getFilterPredicates().clear();
query.getFilterPredicates().addAll(singleQuery);
DatastoreV3Pb.Query.Builder queryProto = QueryTranslator.convertToPb(query, fetchOptions);
resultQueries.add(queryProto);
}
}
Expand Down
1 change: 1 addition & 0 deletions api/src/main/java/com/google/appengine/api/datastore/AsyncCloudDatastoreV1ServiceImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import com.google.appengine.api.datastore.FutureHelper.MultiFuture;
import com.google.appengine.api.datastore.Index.IndexState;
import com.google.appengine.api.utils.FutureWrapper;
import com.google.apphosting.datastore.proto2api.DatastoreV3Pb.GetResponse;
import com.google.common.base.Preconditions;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableList;
Expand Down
85 changes: 42 additions & 43 deletions api/src/main/java/com/google/appengine/api/datastore/AsyncDatastoreServiceImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,8 @@
import com.google.apphosting.datastore.proto2api.DatastoreV3Pb.PutResponse;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
// import com.google.io.protocol.ProtocolMessage;
import com.google.protobuf.Message;
import com.google.protobuf.MessageLiteOrBuilder;
import com.google.storage.onestore.v3.proto2api.OnestoreEntity.CompositeIndex;
import com.google.storage.onestore.v3.proto2api.OnestoreEntity.EntityProto;
import com.google.storage.onestore.v3.proto2api.OnestoreEntity.Reference;
Expand Down Expand Up @@ -79,13 +79,13 @@ class AsyncDatastoreServiceImpl extends BaseAsyncDatastoreServiceImpl {
*/
private abstract class V3Batcher<
S extends Message,
R extends Message,
R extends MessageLiteOrBuilder,
F,
T extends Message>
extends BaseRpcBatcher<S, R, F, T> {
@Override
final R newBatch(R baseBatch) {
return (R) baseBatch.toBuilder().clone().build();
return (R) baseBatch.getDefaultInstanceForType().toBuilder().clone();
}
}

Expand All @@ -95,7 +95,7 @@ final R newBatch(R baseBatch) {
* @param <S> the response message type
* @param <R> the request message type
*/
private abstract class V3KeyBatcher<S extends Message, R extends Message>
private abstract class V3KeyBatcher<S extends Message, R extends Message.Builder>
extends V3Batcher<S, R, Key, Reference> {
@Override
final Object getGroup(Key value) {
Expand All @@ -108,11 +108,11 @@ final Reference toPb(Key value) {
}
}

private final V3KeyBatcher<DeleteResponse, DeleteRequest> deleteBatcher =
new V3KeyBatcher<DeleteResponse, DeleteRequest>() {
private final V3KeyBatcher<DeleteResponse, DeleteRequest.Builder> deleteBatcher =
new V3KeyBatcher<DeleteResponse, DeleteRequest.Builder>() {
@Override
void addToBatch(Reference value, DeleteRequest batch) {
batch.toBuilder().addKey(value).build();
void addToBatch(Reference value, DeleteRequest.Builder batch) {
batch.addKey(value);
}

@Override
Expand All @@ -121,17 +121,17 @@ int getMaxCount() {
}

@Override
protected Future<DeleteResponse> makeCall(DeleteRequest batch) {
protected Future<DeleteResponse> makeCall(DeleteRequest.Builder batch) {
return makeAsyncCall(
apiConfig, DatastoreService_3.Method.Delete, batch, DeleteResponse.getDefaultInstance());
apiConfig, DatastoreService_3.Method.Delete, batch, DeleteResponse.newBuilder());
}
};

private final V3KeyBatcher<GetResponse, GetRequest> getByKeyBatcher =
new V3KeyBatcher<GetResponse, GetRequest>() {
private final V3KeyBatcher<GetResponse, GetRequest.Builder> getByKeyBatcher =
new V3KeyBatcher<GetResponse, GetRequest.Builder>() {
@Override
void addToBatch(Reference value, GetRequest batch) {
batch.toBuilder().addKey(value).build();
void addToBatch(Reference value, GetRequest.Builder batch) {
batch.addKey(value);
}

@Override
Expand All @@ -140,13 +140,13 @@ int getMaxCount() {
}

@Override
protected Future<GetResponse> makeCall(GetRequest batch) {
return makeAsyncCall(apiConfig, DatastoreService_3.Method.Get, batch, GetResponse.getDefaultInstance());
protected Future<GetResponse> makeCall(GetRequest.Builder batch) {
return makeAsyncCall(apiConfig, DatastoreService_3.Method.Get, batch, GetResponse.newBuilder());
}
};

private final V3Batcher<GetResponse, GetRequest, Reference, Reference> getByReferenceBatcher =
new V3Batcher<GetResponse, GetRequest, Reference, Reference>() {
private final V3Batcher<GetResponse, GetRequest.Builder, Reference, Reference> getByReferenceBatcher =
new V3Batcher<GetResponse, GetRequest.Builder, Reference, Reference>() {
@Override
final Object getGroup(Reference value) {
return value.getPath().getElement(0);
Expand All @@ -158,8 +158,8 @@ final Reference toPb(Reference value) {
}

@Override
void addToBatch(Reference value, GetRequest batch) {
batch.toBuilder().addKey(value).build();
void addToBatch(Reference value, GetRequest.Builder batch) {
batch.addKey(value);
}

@Override
Expand All @@ -168,21 +168,21 @@ int getMaxCount() {
}

@Override
protected Future<GetResponse> makeCall(GetRequest batch) {
return makeAsyncCall(apiConfig, DatastoreService_3.Method.Get, batch, GetResponse.getDefaultInstance());
protected Future<GetResponse> makeCall(GetRequest.Builder batch) {
return makeAsyncCall(apiConfig, DatastoreService_3.Method.Get, batch, GetResponse.newBuilder());
}
};

private final V3Batcher<PutResponse, PutRequest, Entity, EntityProto> putBatcher =
new V3Batcher<PutResponse, PutRequest, Entity, EntityProto>() {
private final V3Batcher<PutResponse, PutRequest.Builder, Entity, EntityProto> putBatcher =
new V3Batcher<PutResponse, PutRequest.Builder, Entity, EntityProto>() {
@Override
Object getGroup(Entity value) {
return value.getKey().getRootKey();
}

@Override
void addToBatch(EntityProto value, PutRequest batch) {
batch.toBuilder().addEntity(value);
void addToBatch(EntityProto value, PutRequest.Builder batch) {
batch.addEntity(value);
}

@Override
Expand All @@ -191,8 +191,8 @@ int getMaxCount() {
}

@Override
protected Future<PutResponse> makeCall(PutRequest batch) {
return makeAsyncCall(apiConfig, DatastoreService_3.Method.Put, batch, PutResponse.getDefaultInstance());
protected Future<PutResponse> makeCall(PutRequest.Builder batch) {
return makeAsyncCall(apiConfig, DatastoreService_3.Method.Put, batch, PutResponse.newBuilder());
}

@Override
Expand Down Expand Up @@ -247,7 +247,7 @@ protected TransactionImpl.InternalTransaction doBeginTransaction(TransactionOpti

Future<DatastoreV3Pb.Transaction> future =
DatastoreApiHelper.makeAsyncCall(
apiConfig, DatastoreService_3.Method.BeginTransaction, request.build(), remoteTxn.buildPartial());
apiConfig, DatastoreService_3.Method.BeginTransaction, request, remoteTxn);

return new InternalTransactionV3(apiConfig, request.getApp(), future);
}
Expand All @@ -271,9 +271,9 @@ protected final Future<Map<Key, Entity>> doBatchGet(
txn == null && datastoreServiceConfig.getReadPolicy().getConsistency() != EVENTUAL;

// Batch and issue the request(s).
Iterator<GetRequest> batches =
Iterator<GetRequest.Builder> batches =
getByKeyBatcher.getBatches(
keysToGet, baseReq.build(), baseReq.getKeyCount(), shouldUseMultipleBatches);
keysToGet, baseReq, baseReq.getKeyCount(), shouldUseMultipleBatches);
List<Future<GetResponse>> futures = getByKeyBatcher.makeCalls(batches);

return registerInTransaction(
Expand Down Expand Up @@ -339,9 +339,9 @@ private void aggregate(
}

// Some keys were deferred. Issue followup requests, and loop again.
Iterator<GetRequest> followupBatches =
Iterator<GetRequest.Builder> followupBatches =
getByReferenceBatcher.getBatches(
deferredRefs, baseReq.build(), baseReq.getKeyCount(), shouldUseMultipleBatches);
deferredRefs, baseReq, baseReq.getKeyCount(), shouldUseMultipleBatches);
currentFutures = getByReferenceBatcher.makeCalls(followupBatches);
}
}
Expand Down Expand Up @@ -438,8 +438,8 @@ protected Future<List<Key>> doBatchPut(@Nullable Transaction txn, final List<Ent
}
boolean group = !baseReq.hasTransaction(); // Do not group when inside a transaction.
final List<Integer> order = Lists.newArrayListWithCapacity(entities.size());
Iterator<PutRequest> batches =
putBatcher.getBatches(entities, baseReq.build(), baseReq.getEntityCount(), group, order);
Iterator<PutRequest.Builder> batches =
putBatcher.getBatches(entities, baseReq, baseReq.getEntityCount(), group, order);
List<Future<PutResponse>> futures = putBatcher.makeCalls(batches);

return registerInTransaction(
Expand Down Expand Up @@ -474,8 +474,8 @@ protected Future<Void> doBatchDelete(@Nullable Transaction txn, Collection<Key>
baseReq.setTransaction(InternalTransactionV3.toProto(txn));
}
boolean group = !baseReq.hasTransaction(); // Do not group inside a transaction.
Iterator<DeleteRequest> batches =
deleteBatcher.getBatches(keys, baseReq.build(), baseReq.getKeyCount(), group);
Iterator<DeleteRequest.Builder> batches =
deleteBatcher.getBatches(keys, baseReq, baseReq.getKeyCount(), group);
List<Future<DeleteResponse>> futures = deleteBatcher.makeCalls(batches);
return registerInTransaction(
txn,
Expand Down Expand Up @@ -527,7 +527,7 @@ public Future<KeyRange> allocateIds(final Key parent, final String kind, long nu
req.setModelKey(allocateIdsRef);
AllocateIdsResponse.Builder resp = AllocateIdsResponse.newBuilder();
Future<AllocateIdsResponse> future =
makeAsyncCall(apiConfig, DatastoreService_3.Method.AllocateIds, req.build(), resp.buildPartial());
makeAsyncCall(apiConfig, DatastoreService_3.Method.AllocateIds, req, resp);
return new FutureWrapper<AllocateIdsResponse, KeyRange>(future) {
@Override
protected KeyRange wrap(AllocateIdsResponse resp) throws Exception {
Expand All @@ -553,7 +553,7 @@ public Future<KeyRangeState> allocateIdRange(final KeyRange range) {
req.setMax(end);
AllocateIdsResponse.Builder resp = AllocateIdsResponse.newBuilder();
Future<AllocateIdsResponse> future =
makeAsyncCall(apiConfig, DatastoreService_3.Method.AllocateIds, req.build(), resp.build());
makeAsyncCall(apiConfig, DatastoreService_3.Method.AllocateIds, req, resp);
return new FutureWrapper<AllocateIdsResponse, KeyRangeState>(future) {
@SuppressWarnings("deprecation")
@Override
Expand Down Expand Up @@ -590,13 +590,12 @@ protected Throwable convertException(Throwable cause) {

@Override
public Future<Map<Index, IndexState>> getIndexes() {
StringProto req =
StringProto.Builder req =
StringProto.newBuilder()
.setValue(datastoreServiceConfig.getAppIdNamespace().getAppId())
.build();
.setValue(datastoreServiceConfig.getAppIdNamespace().getAppId());
return new FutureWrapper<CompositeIndices, Map<Index, IndexState>>(
makeAsyncCall(
apiConfig, DatastoreService_3.Method.GetIndices, req, CompositeIndices.getDefaultInstance())) {
apiConfig, DatastoreService_3.Method.GetIndices, req, CompositeIndices.newBuilder())) {
@Override
protected Map<Index, IndexState> wrap(CompositeIndices indices) throws Exception {
Map<Index, IndexState> answer = new LinkedHashMap<Index, IndexState>();
Expand Down
1 change: 1 addition & 0 deletions api/src/main/java/com/google/appengine/api/datastore/BaseAsyncDatastoreServiceImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import com.google.appengine.api.datastore.DatastoreAttributes.DatastoreType;
import com.google.appengine.api.datastore.TransactionOptions.Mode;
import com.google.appengine.api.utils.FutureWrapper;
import com.google.apphosting.datastore.proto2api.DatastoreV3Pb.GetResponse;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
Expand Down
1 change: 1 addition & 0 deletions api/src/main/java/com/google/appengine/api/datastore/BaseQueryResultsSource.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

package com.google.appengine.api.datastore;

import com.google.apphosting.datastore.proto2api.DatastoreV3Pb;
import com.google.common.collect.Sets;
import java.util.Collection;
import java.util.List;
Expand Down
16 changes: 8 additions & 8 deletions api/src/main/java/com/google/appengine/api/datastore/DatastoreApiHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,36 +119,36 @@ static RuntimeException createV1Exception(Code code, String message, Throwable c
}
}

static <T extends Message> Future<T> makeAsyncCall(
static <T extends Message, S extends Message.Builder> Future<T> makeAsyncCall(
ApiConfig apiConfig,
final DatastoreService_3.Method method,
MessageLite request,
final T responseProto) {
MessageLite.Builder request,
final S responseProto) {
Future<byte[]> response =
ApiProxy.makeAsyncCall(
DATASTORE_V3_PACKAGE, method.name(), request.toByteArray(), apiConfig);
DATASTORE_V3_PACKAGE, method.name(), request.buildPartial().toByteArray(), apiConfig);
return new FutureWrapper<byte[], T>(response) {
@Override
protected T wrap(byte[] responseBytes) throws InvalidProtocolBufferException {
T updatedResponseProto = null;
// This null check is mainly for the benefit of unit tests
// (specifically ones using EasyMock, where the default behavior
// is to return null).
if (responseBytes != null && responseProto != null) {

try {
updatedResponseProto = (T) responseProto.getParserForType().parseFrom(responseBytes);
responseProto.clear();
responseProto.mergeFrom(responseBytes);
}
catch(InvalidProtocolBufferException e) {
throw new InvalidProtocolBufferException(
String.format("Invalid %s.%s response", DATASTORE_V3_PACKAGE, method.name()));
}
List<String> initializationErrors = updatedResponseProto.findInitializationErrors();
List<String> initializationErrors = responseProto.findInitializationErrors();
if (initializationErrors != null && !initializationErrors.isEmpty()) {
throw new InvalidProtocolBufferException(initializationErrors.toString());
}
}
return updatedResponseProto;
return (T) responseProto.build();
}

@Override
Expand Down
2 changes: 1 addition & 1 deletion api/src/main/java/com/google/appengine/api/datastore/IndexComponentsOnlyQuery.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,7 @@ public boolean equals(@Nullable Object o) {

IndexComponentsOnlyQuery that = (IndexComponentsOnlyQuery) o;

if (!query.equals(that.query)) {
if (!query.build().equals(that.query.build())) {
return false;
}

Expand Down
10 changes: 5 additions & 5 deletions api/src/main/java/com/google/appengine/api/datastore/InternalTransactionV3.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,8 +58,8 @@ private long getHandle() {
}

// extracted method to facilitate testing
<T extends Message> Future<Void> makeAsyncCall(
DatastoreService_3.Method method, MessageLite request, T response) {
<T extends Message, S extends Message.Builder> Future<Void> makeAsyncCall(
DatastoreService_3.Method method, Message.Builder request, S response) {
Future<T> resultProto = DatastoreApiHelper.makeAsyncCall(apiConfig, method, request, response);
return new FutureWrapper<T, Void>(resultProto) {
@Override
Expand All @@ -74,17 +74,17 @@ protected Throwable convertException(Throwable cause) {
};
}

private <T extends Message> Future<Void> makeAsyncTxnCall(
private <T extends Message.Builder> Future<Void> makeAsyncTxnCall(
DatastoreService_3.Method method, T response) {
DatastoreV3Pb.Transaction.Builder txn = DatastoreV3Pb.Transaction.newBuilder();
txn.setApp(app);
txn.setHandle(getHandle());
return makeAsyncCall(method, txn.build(), response);
return makeAsyncCall(method, txn, response);
}

@Override
public Future<Void> doCommitAsync() {
return makeAsyncTxnCall(DatastoreService_3.Method.Commit, CommitResponse.getDefaultInstance());
return makeAsyncTxnCall(DatastoreService_3.Method.Commit, CommitResponse.newBuilder());
}

@Override
Expand Down
1 change: 1 addition & 0 deletions api/src/main/java/com/google/appengine/api/datastore/QueryResultsSourceCloudDatastoreV1.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

package com.google.appengine.api.datastore;

import com.google.apphosting.datastore.proto2api.DatastoreV3Pb;
import com.google.datastore.v1.QueryResultBatch;
import com.google.datastore.v1.RunQueryRequest;
import com.google.datastore.v1.RunQueryResponse;
Expand Down
2 changes: 1 addition & 1 deletion api/src/main/java/com/google/appengine/api/datastore/QueryResultsSourceV3.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ public Future<QueryResult> makeNextCall(
if (offsetOrNull != null) {
req.setOffset(offsetOrNull);
}
return makeAsyncCall(apiConfig, Method.Next, req.build(), DatastoreV3Pb.QueryResult.getDefaultInstance());
return makeAsyncCall(apiConfig, Method.Next, req, DatastoreV3Pb.QueryResult.newBuilder());
}

@Override
Expand Down
3 changes: 1 addition & 2 deletions api/src/main/java/com/google/appengine/api/datastore/QueryRunnerV3.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ public QueryResultsSource runQuery(FetchOptions fetchOptions, Query query, Trans

Future<DatastoreV3Pb.QueryResult> result =
DatastoreApiHelper.makeAsyncCall(
apiConfig, Method.RunQuery, queryProto.build(), DatastoreV3Pb.QueryResult.getDefaultInstance());
apiConfig, Method.RunQuery, queryProto, DatastoreV3Pb.QueryResult.newBuilder());

// Adding more info to DatastoreNeedIndexException if thrown
result =
Expand All @@ -69,7 +69,6 @@ protected DatastoreV3Pb.QueryResult wrap(DatastoreV3Pb.QueryResult result)
return result;
}
};

return new QueryResultsSourceV3(
datastoreServiceConfig.getDatastoreCallbacks(),
fetchOptions,
Expand Down
Loading

0 comments on commit 6b973ca

Please sign in to comment.