Search in sources :

Example 1 with Datastore

use of com.google.datastore.v1.client.Datastore in project YCSB by brianfrankcooper.

the class GoogleDatastoreClient method read.

@Override
public Status read(String table, String key, Set<String> fields, HashMap<String, ByteIterator> result) {
    LookupRequest.Builder lookupRequest = LookupRequest.newBuilder();
    lookupRequest.addKeys(buildPrimaryKey(table, key));
    lookupRequest.getReadOptionsBuilder().setReadConsistency(this.readConsistency);
    // Note above, datastore lookupRequest always reads the entire entity, it
    // does not support reading a subset of "fields" (properties) of an entity.
    logger.debug("Built lookup request as: " + lookupRequest.toString());
    LookupResponse response = null;
    try {
        response = datastore.lookup(lookupRequest.build());
    } catch (DatastoreException exception) {
        logger.error(String.format("Datastore Exception when reading (%s): %s %s", exception.getMessage(), exception.getMethodName(), exception.getCode()));
        // will bubble up to the user as part of the YCSB Status "name".
        return new Status("ERROR-" + exception.getCode(), exception.getMessage());
    }
    if (response.getFoundCount() == 0) {
        return new Status("ERROR-404", "Not Found, key is: " + key);
    } else if (response.getFoundCount() > 1) {
        // entity back. Unexpected State.
        return Status.UNEXPECTED_STATE;
    }
    Entity entity = response.getFound(0).getEntity();
    logger.debug("Read entity: " + entity.toString());
    Map<String, Value> properties = entity.getProperties();
    Set<String> propertiesToReturn = (fields == null ? properties.keySet() : fields);
    for (String name : propertiesToReturn) {
        if (properties.containsKey(name)) {
            result.put(name, new StringByteIterator(properties.get(name).getStringValue()));
        }
    }
    return Status.OK;
}
Also used : Status(com.yahoo.ycsb.Status) StringByteIterator(com.yahoo.ycsb.StringByteIterator) DatastoreException(com.google.datastore.v1.client.DatastoreException)

Example 2 with Datastore

use of com.google.datastore.v1.client.Datastore in project YCSB by brianfrankcooper.

the class GoogleDatastoreClient method doSingleItemMutation.

private Status doSingleItemMutation(String table, String key, @Nullable HashMap<String, ByteIterator> values, MutationType mutationType) {
    // First build the key.
    Key.Builder datastoreKey = buildPrimaryKey(table, key);
    // Build a commit request in non-transactional mode.
    // Single item mutation to google datastore
    // is always atomic and strongly consistent. Transaction is only necessary
    // for multi-item mutation, or Read-modify-write operation.
    CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
    commitRequest.setMode(Mode.NON_TRANSACTIONAL);
    if (mutationType == MutationType.DELETE) {
        commitRequest.addMutationsBuilder().setDelete(datastoreKey);
    } else {
        // If this is not for delete, build the entity.
        Entity.Builder entityBuilder = Entity.newBuilder();
        entityBuilder.setKey(datastoreKey);
        for (Entry<String, ByteIterator> val : values.entrySet()) {
            entityBuilder.getMutableProperties().put(val.getKey(), Value.newBuilder().setStringValue(val.getValue().toString()).build());
        }
        Entity entity = entityBuilder.build();
        logger.debug("entity built as: " + entity.toString());
        if (mutationType == MutationType.UPSERT) {
            commitRequest.addMutationsBuilder().setUpsert(entity);
        } else if (mutationType == MutationType.UPDATE) {
            commitRequest.addMutationsBuilder().setUpdate(entity);
        } else {
            throw new RuntimeException("Impossible MutationType, code bug.");
        }
    }
    try {
        datastore.commit(commitRequest.build());
        logger.debug("successfully committed.");
    } catch (DatastoreException exception) {
        // Catch all Datastore rpc errors.
        // Log the exception, the name of the method called and the error code.
        logger.error(String.format("Datastore Exception when committing (%s): %s %s", exception.getMessage(), exception.getMethodName(), exception.getCode()));
        // will bubble up to the user as part of the YCSB Status "name".
        return new Status("ERROR-" + exception.getCode(), exception.getMessage());
    }
    return Status.OK;
}
Also used : Status(com.yahoo.ycsb.Status) StringByteIterator(com.yahoo.ycsb.StringByteIterator) ByteIterator(com.yahoo.ycsb.ByteIterator) DatastoreException(com.google.datastore.v1.client.DatastoreException)

Example 3 with Datastore

use of com.google.datastore.v1.client.Datastore in project google-cloud-java by GoogleCloudPlatform.

the class DatastoreTest method testRetryableException.

@Test
public void testRetryableException() throws Exception {
    LookupRequest requestPb = LookupRequest.newBuilder().addKeys(KEY1.toPb()).build();
    LookupResponse responsePb = LookupResponse.newBuilder().addFound(EntityResult.newBuilder().setEntity(ENTITY1.toPb())).build();
    EasyMock.expect(rpcMock.lookup(requestPb)).andThrow(new DatastoreException(14, "UNAVAILABLE", "UNAVAILABLE", null)).andReturn(responsePb);
    EasyMock.replay(rpcFactoryMock, rpcMock);
    Datastore datastore = rpcMockOptions.getService();
    Entity entity = datastore.get(KEY1);
    assertEquals(ENTITY1, entity);
    EasyMock.verify(rpcFactoryMock, rpcMock);
}
Also used : LookupResponse(com.google.datastore.v1.LookupResponse) LookupRequest(com.google.datastore.v1.LookupRequest) Test(org.junit.Test)

Example 4 with Datastore

use of com.google.datastore.v1.client.Datastore in project google-cloud-java by GoogleCloudPlatform.

the class DatastoreTest method testQueryPaginationWithLimit.

@Test
public void testQueryPaginationWithLimit() throws DatastoreException {
    List<RunQueryResponse> responses = buildResponsesForQueryPaginationWithLimit();
    List<ByteString> endCursors = Lists.newArrayListWithCapacity(responses.size());
    for (RunQueryResponse response : responses) {
        EasyMock.expect(rpcMock.runQuery(EasyMock.anyObject(RunQueryRequest.class))).andReturn(response);
        if (response.getBatch().getMoreResults() != QueryResultBatch.MoreResultsType.NOT_FINISHED) {
            endCursors.add(response.getBatch().getEndCursor());
        }
    }
    EasyMock.replay(rpcFactoryMock, rpcMock);
    Datastore datastore = rpcMockOptions.getService();
    int limit = 2;
    int totalCount = 0;
    Iterator<ByteString> cursorIter = endCursors.iterator();
    StructuredQuery<Entity> query = Query.newEntityQueryBuilder().setLimit(limit).build();
    while (true) {
        QueryResults<Entity> results = datastore.run(query);
        int resultCount = 0;
        while (results.hasNext()) {
            results.next();
            resultCount++;
            totalCount++;
        }
        assertTrue(cursorIter.hasNext());
        Cursor expectedEndCursor = Cursor.copyFrom(cursorIter.next().toByteArray());
        assertEquals(expectedEndCursor, results.getCursorAfter());
        if (resultCount < limit) {
            break;
        }
        query = query.toBuilder().setStartCursor(results.getCursorAfter()).build();
    }
    assertEquals(5, totalCount);
    EasyMock.verify(rpcFactoryMock, rpcMock);
}
Also used : RunQueryResponse(com.google.datastore.v1.RunQueryResponse) ByteString(com.google.protobuf.ByteString) RunQueryRequest(com.google.datastore.v1.RunQueryRequest) Test(org.junit.Test)

Example 5 with Datastore

use of com.google.datastore.v1.client.Datastore in project google-cloud-java by GoogleCloudPlatform.

the class DatastoreTest method testRuntimeException.

@Test
public void testRuntimeException() throws Exception {
    LookupRequest requestPb = LookupRequest.newBuilder().addKeys(KEY1.toPb()).build();
    String exceptionMessage = "Artificial runtime exception";
    EasyMock.expect(rpcMock.lookup(requestPb)).andThrow(new RuntimeException(exceptionMessage));
    EasyMock.replay(rpcFactoryMock, rpcMock);
    Datastore datastore = rpcMockOptions.getService();
    thrown.expect(DatastoreException.class);
    thrown.expectMessage(exceptionMessage);
    datastore.get(KEY1);
    EasyMock.verify(rpcFactoryMock, rpcMock);
}
Also used : ByteString(com.google.protobuf.ByteString) LookupRequest(com.google.datastore.v1.LookupRequest) Test(org.junit.Test)

Aggregations

Test (org.junit.Test)11 Entity (com.google.datastore.v1.Entity)5 Query (com.google.datastore.v1.Query)5 Key (com.google.datastore.v1.Key)4 LookupRequest (com.google.datastore.v1.LookupRequest)4 RunQueryRequest (com.google.datastore.v1.RunQueryRequest)4 Datastore (com.google.datastore.v1.client.Datastore)4 DatastoreException (com.google.datastore.v1.client.DatastoreException)4 DatastoreHelper.makeKey (com.google.datastore.v1.client.DatastoreHelper.makeKey)3 DatastoreOptions (com.google.datastore.v1.client.DatastoreOptions)3 DeleteEntity (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntity)3 DeleteKey (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteKey)3 DatastoreV1.isValidKey (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.isValidKey)3 Credential (com.google.api.client.auth.oauth2.Credential)2 GqlQuery (com.google.datastore.v1.GqlQuery)2 ReadOptions (com.google.datastore.v1.ReadOptions)2 RunQueryResponse (com.google.datastore.v1.RunQueryResponse)2 ByteString (com.google.protobuf.ByteString)2 Status (com.yahoo.ycsb.Status)2 StringByteIterator (com.yahoo.ycsb.StringByteIterator)2