Examples with Datastore com.google.datastore.v1.client.Datastore used on opensource projects

Example 6 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);
}

Example 7 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);
}

Example 8 with Datastore

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

the class DatastoreTest method testNonRetryableException.

@Test
public void testNonRetryableException() throws Exception {
    LookupRequest requestPb = LookupRequest.newBuilder().addKeys(KEY1.toPb()).build();
    EasyMock.expect(rpcMock.lookup(requestPb)).andThrow(new DatastoreException(DatastoreException.UNKNOWN_CODE, "denied", "PERMISSION_DENIED")).times(1);
    EasyMock.replay(rpcFactoryMock, rpcMock);
    Datastore datastore = rpcMockOptions.getService();
    thrown.expect(DatastoreException.class);
    thrown.expectMessage("denied");
    datastore.get(KEY1);
    EasyMock.verify(rpcFactoryMock, rpcMock);
}

Example 9 with Datastore

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

the class GoogleDatastoreClient method init.

/**
   * Initialize any state for this DB. Called once per DB instance; there is
   * one DB instance per client thread.
   */
@Override
public void init() throws DBException {
    String debug = getProperties().getProperty("googledatastore.debug", null);
    if (null != debug && "true".equalsIgnoreCase(debug)) {
        logger.setLevel(Level.DEBUG);
    }
    // We need the following 3 essential properties to initialize datastore:
    //
    // - DatasetId,
    // - Path to private key file,
    // - Service account email address.
    String datasetId = getProperties().getProperty("googledatastore.datasetId", null);
    if (datasetId == null) {
        throw new DBException("Required property \"datasetId\" missing.");
    }
    String privateKeyFile = getProperties().getProperty("googledatastore.privateKeyFile", null);
    if (privateKeyFile == null) {
        throw new DBException("Required property \"privateKeyFile\" missing.");
    }
    String serviceAccountEmail = getProperties().getProperty("googledatastore.serviceAccountEmail", null);
    if (serviceAccountEmail == null) {
        throw new DBException("Required property \"serviceAccountEmail\" missing.");
    }
    // Below are properties related to benchmarking.
    String readConsistencyConfig = getProperties().getProperty("googledatastore.readConsistency", null);
    if (readConsistencyConfig != null) {
        try {
            this.readConsistency = ReadConsistency.valueOf(readConsistencyConfig.trim().toUpperCase());
        } catch (IllegalArgumentException e) {
            throw new DBException("Invalid read consistency specified: " + readConsistencyConfig + ". Expecting STRONG or EVENTUAL.");
        }
    }
    //
    // Entity Grouping Mode (googledatastore.entitygroupingmode), see
    // documentation in conf/googledatastore.properties.
    //
    String entityGroupingConfig = getProperties().getProperty("googledatastore.entityGroupingMode", null);
    if (entityGroupingConfig != null) {
        try {
            this.entityGroupingMode = EntityGroupingMode.valueOf(entityGroupingConfig.trim().toUpperCase());
        } catch (IllegalArgumentException e) {
            throw new DBException("Invalid entity grouping mode specified: " + entityGroupingConfig + ". Expecting ONE_ENTITY_PER_GROUP or " + "MULTI_ENTITY_PER_GROUP.");
        }
    }
    this.rootEntityName = getProperties().getProperty("googledatastore.rootEntityName", "YCSB_ROOT_ENTITY");
    try {
        // Setup the connection to Google Cloud Datastore with the credentials
        // obtained from the configure.
        DatastoreOptions.Builder options = new DatastoreOptions.Builder();
        Credential credential = DatastoreHelper.getServiceAccountCredential(serviceAccountEmail, privateKeyFile);
        logger.info("Using JWT Service Account credential.");
        logger.info("DatasetID: " + datasetId + ", Service Account Email: " + serviceAccountEmail + ", Private Key File Path: " + privateKeyFile);
        datastore = DatastoreFactory.get().create(options.credential(credential).projectId(datasetId).build());
    } catch (GeneralSecurityException exception) {
        throw new DBException("Security error connecting to the datastore: " + exception.getMessage(), exception);
    } catch (IOException exception) {
        throw new DBException("I/O error connecting to the datastore: " + exception.getMessage(), exception);
    }
    logger.info("Datastore client instance created: " + datastore.toString());
}

Example 10 with Datastore

use of com.google.datastore.v1.client.Datastore in project beam by apache.

the class V1ReadIT method writeEntitiesToDatastore.

// Creates entities and write them to datastore
private static void writeEntitiesToDatastore(V1TestOptions options, String project, String ancestor, long numEntities) throws Exception {
    Datastore datastore = getDatastore(options, project);
    // Write test entities to datastore
    V1TestWriter writer = new V1TestWriter(datastore, new UpsertMutationBuilder());
    Key ancestorKey = makeAncestorKey(options.getNamespace(), options.getKind(), ancestor);
    for (long i = 0; i < numEntities; i++) {
        Entity entity = makeEntity(i, ancestorKey, options.getKind(), options.getNamespace(), 0);
        writer.write(entity);
    }
    writer.close();
}