Search in sources :

Example 6 with Key

use of com.google.datastore.v1.Key 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());
}
Also used : DBException(com.yahoo.ycsb.DBException) Credential(com.google.api.client.auth.oauth2.Credential) GeneralSecurityException(java.security.GeneralSecurityException) DatastoreOptions(com.google.datastore.v1.client.DatastoreOptions) IOException(java.io.IOException)

Example 7 with Key

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

the class DatastoreTest method buildResponsesForQueryPagination.

private List<RunQueryResponse> buildResponsesForQueryPagination() {
    Entity entity4 = Entity.newBuilder(KEY4).set("value", StringValue.of("value")).build();
    Entity entity5 = Entity.newBuilder(KEY5).set("value", "value").build();
    datastore.add(ENTITY3, entity4, entity5);
    List<RunQueryResponse> responses = new ArrayList<>();
    Query<Key> query = Query.newKeyQueryBuilder().build();
    RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder();
    query.populatePb(requestPb);
    QueryResultBatch queryResultBatchPb = RunQueryResponse.newBuilder().mergeFrom(((DatastoreImpl) datastore).runQuery(requestPb.build())).getBatch();
    QueryResultBatch queryResultBatchPb1 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(0, 1)).setEndCursor(queryResultBatchPb.getEntityResultsList().get(0).getCursor()).build();
    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb1).build());
    QueryResultBatch queryResultBatchPb2 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.NOT_FINISHED).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(1, 3)).setEndCursor(queryResultBatchPb.getEntityResultsList().get(2).getCursor()).build();
    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb2).build());
    QueryResultBatch queryResultBatchPb3 = QueryResultBatch.newBuilder().mergeFrom(queryResultBatchPb).setMoreResults(QueryResultBatch.MoreResultsType.NO_MORE_RESULTS).clearEntityResults().addAllEntityResults(queryResultBatchPb.getEntityResultsList().subList(3, 5)).setEndCursor(queryResultBatchPb.getEntityResultsList().get(4).getCursor()).build();
    responses.add(RunQueryResponse.newBuilder().setBatch(queryResultBatchPb3).build());
    return responses;
}
Also used : QueryResultBatch(com.google.datastore.v1.QueryResultBatch) RunQueryResponse(com.google.datastore.v1.RunQueryResponse) RunQueryRequest(com.google.datastore.v1.RunQueryRequest) ArrayList(java.util.ArrayList)

Example 8 with Key

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

the class DatastoreTest method testLookupEventualConsistency.

@Test
public void testLookupEventualConsistency() {
    ReadOptions readOption = ReadOptions.newBuilder().setReadConsistencyValue(ReadConsistency.EVENTUAL_VALUE).build();
    com.google.datastore.v1.Key key = com.google.datastore.v1.Key.newBuilder().setPartitionId(PartitionId.newBuilder().setProjectId(PROJECT_ID).build()).addPath(com.google.datastore.v1.Key.PathElement.newBuilder().setKind("kind1").setName("name").build()).build();
    LookupRequest lookupRequest = LookupRequest.newBuilder().setReadOptions(readOption).addKeys(key).build();
    EasyMock.expect(rpcMock.lookup(lookupRequest)).andReturn(LookupResponse.newBuilder().build()).times(3);
    EasyMock.replay(rpcFactoryMock, rpcMock);
    Datastore datastore = rpcMockOptions.getService();
    datastore.get(KEY1, ReadOption.eventualConsistency());
    datastore.get(ImmutableList.of(KEY1), ReadOption.eventualConsistency());
    datastore.fetch(ImmutableList.of(KEY1), ReadOption.eventualConsistency());
    EasyMock.verify(rpcFactoryMock, rpcMock);
}
Also used : ReadOptions(com.google.datastore.v1.ReadOptions) LookupRequest(com.google.datastore.v1.LookupRequest) Test(org.junit.Test)

Example 9 with Key

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

the class DatastoreTest method testNewBatch.

@Test
public void testNewBatch() {
    Batch batch = datastore.newBatch();
    Entity entity1 = Entity.newBuilder(ENTITY1).clear().build();
    Entity entity2 = Entity.newBuilder(ENTITY2).clear().setNull("bla").build();
    Entity entity4 = Entity.newBuilder(KEY4).set("value", StringValue.of("value")).build();
    Entity entity5 = Entity.newBuilder(KEY5).set("value", "value").build();
    List<Entity> entities = batch.add(entity4, PARTIAL_ENTITY2, entity5);
    Entity entity6 = entities.get(1);
    assertSame(entity4, entities.get(0));
    assertEquals(PARTIAL_ENTITY2.getProperties(), entity6.getProperties());
    assertEquals(PARTIAL_ENTITY2.getKey().getProjectId(), entity6.getKey().getProjectId());
    assertEquals(PARTIAL_ENTITY2.getKey().getNamespace(), entity6.getKey().getNamespace());
    assertEquals(PARTIAL_ENTITY2.getKey().getAncestors(), entity6.getKey().getAncestors());
    assertEquals(PARTIAL_ENTITY2.getKey().getKind(), entity6.getKey().getKind());
    assertEquals(PARTIAL_ENTITY2.getKey(), IncompleteKey.newBuilder(entity6.getKey()).build());
    assertNotEquals(PARTIAL_ENTITY2.getKey().getPath(), entity6.getKey().getPath());
    assertNotEquals(PARTIAL_ENTITY2.getKey(), entity6.getKey());
    assertSame(entity5, entities.get(2));
    batch.addWithDeferredIdAllocation(PARTIAL_ENTITY3);
    batch.put(ENTITY3, entity1, entity2);
    Batch.Response response = batch.submit();
    entities = datastore.fetch(KEY1, KEY2, KEY3, entity4.getKey(), entity5.getKey(), entity6.getKey());
    assertEquals(entity1, entities.get(0));
    assertEquals(entity2, entities.get(1));
    assertEquals(ENTITY3, entities.get(2));
    assertEquals(entity4, entities.get(3));
    assertEquals(entity5, entities.get(4));
    assertEquals(entity6, entities.get(5));
    assertEquals(6, entities.size());
    List<Key> generatedKeys = response.getGeneratedKeys();
    assertEquals(1, generatedKeys.size());
    assertEquals(PARTIAL_ENTITY3.getProperties(), datastore.get(generatedKeys.get(0)).getProperties());
    assertEquals(PARTIAL_ENTITY3.getKey(), IncompleteKey.newBuilder(generatedKeys.get(0)).build());
    try {
        batch.submit();
        fail("Expecting a failure");
    } catch (DatastoreException ex) {
    // expected to fail
    }
    verifyNotUsable(batch);
    batch = datastore.newBatch();
    batch.delete(entity4.getKey(), entity5.getKey());
    batch.update(ENTITY1, ENTITY2, ENTITY3);
    batch.submit();
    entities = datastore.fetch(KEY1, KEY2, KEY3, entity4.getKey(), entity5.getKey());
    assertEquals(ENTITY1, entities.get(0));
    assertEquals(ENTITY2, entities.get(1));
    assertEquals(ENTITY3, entities.get(2));
    assertNull(entities.get(3));
    assertNull(entities.get(4));
    assertEquals(5, entities.size());
}
Also used : QueryResultBatch(com.google.datastore.v1.QueryResultBatch) Test(org.junit.Test)

Example 10 with Key

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

the class DatastoreV1 method isValidKey.

/**
   * 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());
}
Also used : PathElement(com.google.datastore.v1.Key.PathElement)

Aggregations

Test (org.junit.Test)10 Key (com.google.datastore.v1.Key)8 DatastoreHelper.makeKey (com.google.datastore.v1.client.DatastoreHelper.makeKey)7 DeleteKey (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteKey)7 DatastoreV1.isValidKey (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.isValidKey)7 Entity (com.google.datastore.v1.Entity)6 Mutation (com.google.datastore.v1.Mutation)5 DeleteEntity (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntity)5 ArrayList (java.util.ArrayList)3 CommitRequest (com.google.datastore.v1.CommitRequest)2 QueryResultBatch (com.google.datastore.v1.QueryResultBatch)2 DatastoreException (com.google.datastore.v1.client.DatastoreException)2 Status (com.yahoo.ycsb.Status)2 StringByteIterator (com.yahoo.ycsb.StringByteIterator)2 DatastoreWriterFn (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DatastoreWriterFn)2 DeleteEntityFn (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteEntityFn)2 DeleteKeyFn (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.DeleteKeyFn)2 UpsertFn (org.apache.beam.sdk.io.gcp.datastore.DatastoreV1.UpsertFn)2 Credential (com.google.api.client.auth.oauth2.Credential)1 PathElement (com.google.datastore.v1.Key.PathElement)1