com.google.cloud.datastore.Datastore Java Examples

@Override
public List<String> retrieveClaims(int limit) {
  // Get claim saved in Datastore
  Datastore datastore = getDatastoreInstance();
  Query<Entity> query = Query.newEntityQueryBuilder().setKind(claimsKind).setLimit(limit).build();
  QueryResults<Entity> results = datastore.run(query);

  List<String> claims = new ArrayList<>();
  while (results.hasNext()) {
    Entity entity = results.next();
    String claim = entity.getString("claim");
    if (claim != null) {
      claims.add(claim);
    }
  }
  return claims;
}
 

/** Example of getting entities for several keys. */
// [TARGET get(Key...)]
// [VARIABLE "my_first_key_name"]
// [VARIABLE "my_second_key_name"]
public List<Entity> getMultiple(String firstKeyName, String secondKeyName) {
  Datastore datastore = transaction.getDatastore();
  // TODO change so that it's not necessary to hold the entities in a list for integration testing
  // [START getMultiple]
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key firstKey = keyFactory.newKey(firstKeyName);
  Key secondKey = keyFactory.newKey(secondKeyName);
  Iterator<Entity> entitiesIterator = transaction.get(firstKey, secondKey);
  List<Entity> entities = Lists.newArrayList();
  while (entitiesIterator.hasNext()) {
    Entity entity = entitiesIterator.next();
    // do something with the entity
    entities.add(entity);
  }
  transaction.commit();
  // [END getMultiple]
  return entities;
}
 

/** Example of verifying if a transaction is active. */
// [TARGET active()]
public Key active() {
  Datastore datastore = transaction.getDatastore();
  // [START active]
  // create an entity
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key key = datastore.allocateId(keyFactory.newKey());
  Entity entity = Entity.newBuilder(key).set("description", "active()").build();
  // calling transaction.active() now would return true
  try {
    // add the entity and commit
    transaction.put(entity);
    transaction.commit();
  } finally {
    // if committing succeeded
    // then transaction.isActive() will be false
    if (transaction.isActive()) {
      // otherwise it's true and we need to rollback
      transaction.rollback();
    }
  }
  // [END active]
  return key;
}
 

/** Example of verifying if a transaction is active. */
// [TARGET isActive()]
public Key isActive() {
  Datastore datastore = transaction.getDatastore();
  // [START isActive]
  // create an entity
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key key = datastore.allocateId(keyFactory.newKey());
  Entity entity = Entity.newBuilder(key).set("description", "active()").build();
  // calling transaction.active() now would return true
  try {
    // add the entity and commit
    transaction.put(entity);
    transaction.commit();
  } finally {
    // if committing succeeded
    // then transaction.active() will be false
    if (transaction.isActive()) {
      // otherwise it's true and we need to rollback
      transaction.rollback();
    }
  }
  // [END isActive]
  return key;
}
 

/** Example of putting multiple entities with deferred id allocation. */
// [TARGET putWithDeferredIdAllocation(FullEntity...)]
public List<Key> multiplePutEntitiesDeferredId() {
  Datastore datastore = transaction.getDatastore();
  // [START multiplePutEntitiesDeferredId]
  IncompleteKey key1 = datastore.newKeyFactory().setKind("MyKind").newKey();
  FullEntity.Builder entityBuilder1 = FullEntity.newBuilder(key1);
  entityBuilder1.set("propertyName", "value1");
  FullEntity entity1 = entityBuilder1.build();

  IncompleteKey key2 = datastore.newKeyFactory().setKind("MyKind").newKey();
  FullEntity.Builder entityBuilder2 = FullEntity.newBuilder(key2);
  entityBuilder2.set("propertyName", "value2");
  FullEntity entity2 = entityBuilder2.build();

  transaction.putWithDeferredIdAllocation(entity1, entity2);
  Response response = transaction.commit();
  // [END multiplePutEntitiesDeferredId]
  return response.getGeneratedKeys();
}
 

public static void wipeDatastore() {
  Datastore datastore = getDatastore();
  QueryResults<Key> guestbooks =
      datastore.run(Query.newKeyQueryBuilder().setKind("Greeting").build());
  ArrayList<Key> keys = Lists.newArrayList(guestbooks);

  if (!keys.isEmpty()) {
    datastore.delete(keys.toArray(new Key[keys.size()]));
  }
}
 

/** Example of committing a transaction. */
// [TARGET commit()]
public Key commit() {
  Datastore datastore = transaction.getDatastore();
  // [START commit]
  // create an entity
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key key = datastore.allocateId(keyFactory.newKey());
  Entity entity = Entity.newBuilder(key).set("description", "commit()").build();

  // add the entity and commit
  try {
    transaction.put(entity);
    transaction.commit();
  } catch (DatastoreException ex) {
    // handle exception
  }
  // [END commit]

  return key;
}
 

/** Example of adding multiple entities with deferred id allocation. */
// [TARGET addWithDeferredIdAllocation(FullEntity...)]
public List<Key> multipleAddEntitiesDeferredId() {
  Datastore datastore = transaction.getDatastore();
  // [START multipleAddEntitiesDeferredId]
  IncompleteKey key1 = datastore.newKeyFactory().setKind("MyKind").newKey();
  FullEntity.Builder entityBuilder1 = FullEntity.newBuilder(key1);
  entityBuilder1.set("propertyName", "value1");
  FullEntity entity1 = entityBuilder1.build();

  IncompleteKey key2 = datastore.newKeyFactory().setKind("MyKind").newKey();
  FullEntity.Builder entityBuilder2 = FullEntity.newBuilder(key2);
  entityBuilder2.set("propertyName", "value2");
  FullEntity entity2 = entityBuilder2.build();

  transaction.addWithDeferredIdAllocation(entity1, entity2);
  Response response = transaction.commit();
  // [END multipleAddEntitiesDeferredId]
  return response.getGeneratedKeys();
}
 

/** Example of fetching a list of entities for several keys. */
// [TARGET fetch(Key...)]
// [VARIABLE "my_first_key_name"]
// [VARIABLE "my_second_key_name"]
public List<Entity> fetchEntitiesWithKeys(String firstKeyName, String secondKeyName) {
  Datastore datastore = transaction.getDatastore();
  // [START fetchEntitiesWithKeys]
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key firstKey = keyFactory.newKey(firstKeyName);
  Key secondKey = keyFactory.newKey(secondKeyName);
  List<Entity> entities = transaction.fetch(firstKey, secondKey);
  for (Entity entity : entities) {
    // do something with the entity
  }
  transaction.commit();
  // [END fetchEntitiesWithKeys]
  return entities;
}
 

@Override
public List<String> retrieveTokens(int limit) {
  // Get token saved in Datastore
  Datastore datastore = getDatastoreInstance();
  Query<Entity> query = Query.newEntityQueryBuilder().setKind(tokensKind).setLimit(limit).build();
  QueryResults<Entity> results = datastore.run(query);

  List<String> tokens = new ArrayList<>();
  while (results.hasNext()) {
    Entity entity = results.next();
    String token = entity.getString("token");
    if (token != null) {
      tokens.add(token);
    }
  }
  return tokens;
}
 

/** Example of putting multiple entities. */
// [TARGET put(FullEntity...)]
// [VARIABLE "my_key_name1"]
// [VARIABLE "my_key_name2"]
public void multiplePutEntities(String keyName1, String keyName2) {
  Datastore datastore = transaction.getDatastore();
  // [START multiplePutEntities]
  Key key1 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName1);
  Entity.Builder entityBuilder1 = Entity.newBuilder(key1);
  entityBuilder1.set("propertyName", "value1");
  Entity entity1 = entityBuilder1.build();

  Key key2 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName2);
  Entity.Builder entityBuilder2 = Entity.newBuilder(key2);
  entityBuilder2.set("propertyName", "value2");
  Entity entity2 = entityBuilder2.build();

  transaction.put(entity1, entity2);
  transaction.commit();
  // [END multiplePutEntities]
}
 

/** Example of running a query to find all entities with an ancestor. */
// [TARGET run(Query)]
// [VARIABLE "my_parent_key_name"]
public List<Entity> run(String parentKeyName) {
  Datastore datastore = transaction.getDatastore();
  // [START run]
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("ParentKind");
  Key parentKey = keyFactory.newKey(parentKeyName);
  // Build a query
  Query<Entity> query =
      Query.newEntityQueryBuilder()
          .setKind("MyKind")
          .setFilter(PropertyFilter.hasAncestor(parentKey))
          .build();
  QueryResults<Entity> results = transaction.run(query);
  List<Entity> entities = Lists.newArrayList();
  while (results.hasNext()) {
    Entity result = results.next();
    // do something with result
    entities.add(result);
  }
  transaction.commit();
  // [END run]
  return entities;
}
 

/** Example of adding multiple entities. */
// [TARGET add(FullEntity...)]
// [VARIABLE "my_key_name1"]
// [VARIABLE "my_key_name2"]
public void multipleAddEntities(String keyName1, String keyName2) {
  Datastore datastore = transaction.getDatastore();
  // [START multipleAddEntities]
  Key key1 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName1);
  Entity.Builder entityBuilder1 = Entity.newBuilder(key1);
  entityBuilder1.set("propertyName", "value1");
  Entity entity1 = entityBuilder1.build();

  Key key2 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName2);
  Entity.Builder entityBuilder2 = Entity.newBuilder(key2);
  entityBuilder2.set("propertyName", "value2");
  Entity entity2 = entityBuilder2.build();

  transaction.add(entity1, entity2);
  transaction.commit();
  // [END multipleAddEntities]
}
 

/**
 * Creates and returns an {@link EntityManager} using the specified connection parameters.
 * 
 * @param parameters
 *          the connection parameters
 * @return a new {@link EntityManager} created using the specified connection parameters.
 * @throws EntityManagerException
 *           if any error occurs while creating the EntityManager.
 */
public EntityManager createEntityManager(ConnectionParameters parameters) {
  try {
    DatastoreOptions.Builder datastoreOptionsBuilder = DatastoreOptions.newBuilder();
    datastoreOptionsBuilder.setHost(parameters.getServiceURL());
    datastoreOptionsBuilder.setTransportOptions(getHttpTransportOptions(parameters));
    String projectId = parameters.getProjectId();
    if (!Utility.isNullOrEmpty(projectId)) {
      datastoreOptionsBuilder.setProjectId(projectId);
    }
    String namespace = parameters.getNamespace();
    if (namespace != null) {
      datastoreOptionsBuilder.setNamespace(namespace);
    }
    datastoreOptionsBuilder.setCredentials(getCredentials(parameters));
    Datastore datastore = datastoreOptionsBuilder.build().getService();
    return new DefaultEntityManager(datastore);
  } catch (Exception exp) {
    throw new EntityManagerFactoryException(exp);
  }
}
 

@Override
public void save(Message message) {
  // Save message to "messages"
  Datastore datastore = getDatastoreInstance();
  Key key = datastore.allocateId(keyFactory.newKey());

  Entity.Builder messageEntityBuilder = Entity.newBuilder(key)
      .set("messageId", message.getMessageId());

  if (message.getData() != null) {
    messageEntityBuilder = messageEntityBuilder.set("data", message.getData());
  }

  if (message.getPublishTime() != null) {
    messageEntityBuilder = messageEntityBuilder.set("publishTime", message.getPublishTime());
  }
  datastore.put(messageEntityBuilder.build());
}
 

/** Example of updating multiple entities. */
// [TARGET update(Entity...)]
// [VARIABLE "my_key_name1"]
// [VARIABLE "my_key_name2"]
public void multipleUpdateEntities(String keyName1, String keyName2) {
  Datastore datastore = transaction.getDatastore();
  // [START multipleUpdateEntities]
  Key key1 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName1);
  Entity.Builder entityBuilder1 = Entity.newBuilder(key1);
  entityBuilder1.set("propertyName", "value3");
  Entity entity1 = entityBuilder1.build();

  Key key2 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName2);
  Entity.Builder entityBuilder2 = Entity.newBuilder(key2);
  entityBuilder2.set("propertyName", "value4");
  Entity entity2 = entityBuilder2.build();

  transaction.update(entity1, entity2);
  transaction.commit();
  // [END multipleUpdateEntities]
}
 

@Test
public void deleteExpiredDatastoreTestNamespaces() {
  final Datastore datastore = DatastoreOptions.newBuilder()
      .setProjectId("styx-oss-test")
      .build()
      .getService();

  var expiredNamespaces = new ArrayList<String>();
  datastore.run(KeyQuery.newKeyQueryBuilder().setKind("__namespace__").build())
      .forEachRemaining(k -> {
        if (k.hasName() && isExpiredTestNamespace(k.getName(), NOW)) {
          expiredNamespaces.add(k.getName());
        }
      });

  for (var namespace : expiredNamespaces) {
    log.info("Deleting expired datastore test namespace: {}", namespace);
    try {
      deleteDatastoreNamespace(datastore, namespace);
    } catch (Exception e) {
      log.error("Failed to delete expired datastore test namespace: {}", namespace);
    }
  }
}
 

@Test
public void testCreateEntityManager_GoodFilePath() {
  EntityManagerFactory emf = EntityManagerFactory.getInstance();
  String projectId = System.getenv(TestUtils.ENV_PROJECT_ID);
  String jsonFile = System.getenv(TestUtils.ENV_CREDENTIALS);
  if (jsonFile == null) {
    System.out.printf("Enviornment variable %s is not set, skipping the test case%n",
        TestUtils.ENV_CREDENTIALS);
    return;
  }
  EntityManager em = emf.createEntityManager(projectId, jsonFile);
  DefaultEntityManager dem = (DefaultEntityManager) em;
  Datastore ds = dem.getDatastore();
  assertTrue(
      ds.getOptions().getProjectId() != null && ds.getOptions().getProjectId().length() != 0
          && ds.getOptions().getNamespace().equals(""));
}
 

@Test
public void testCreateEntityManager_Namespace() {
  EntityManagerFactory emf = EntityManagerFactory.getInstance();
  try {
    String projectId = System.getenv(TestUtils.ENV_PROJECT_ID);
    String jsonFile = System.getenv(TestUtils.ENV_CREDENTIALS);
    if (jsonFile == null) {
      System.out.printf("Enviornment variable %s is not set, skipping the test case%n",
          TestUtils.ENV_CREDENTIALS);
      return;
    }
    EntityManager em = emf.createEntityManager(projectId, new File(jsonFile), "junit");
    DefaultEntityManager dem = (DefaultEntityManager) em;
    Datastore ds = dem.getDatastore();
    assertTrue(
        ds.getOptions().getProjectId() != null && ds.getOptions().getProjectId().length() != 0
            && ds.getOptions().getNamespace().equals("junit"));
  } catch (Exception exp) {
    System.out.println(exp);
    throw exp;
  }
}
 

/** Example of rolling back a transaction. */
// [TARGET rollback()]
public Key rollback() {
  Datastore datastore = transaction.getDatastore();
  // [START rollback]
  // create an entity
  KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
  Key key = datastore.allocateId(keyFactory.newKey());
  Entity entity = Entity.newBuilder(key).set("description", "rollback()").build();

  // add the entity and rollback
  transaction.put(entity);
  transaction.rollback();
  // calling transaction.commit() now would fail
  // [END rollback]
  return key;
}
 

@SuppressWarnings("JavadocMethod")
public static Datastore getDatastore() {
  if (datastore.get() == null) {
    datastore.set(
        DatastoreOptions.newBuilder().setProjectId("your-project-id-here").build().getService());
  }

  return datastore.get();
}
 

static void deleteDatastoreNamespace(Datastore datastore, String namespace) {
  var keys = new ArrayList<Key>();
  var query = datastore.run(KeyQuery.newKeyQueryBuilder().setNamespace(namespace).build());
  query.forEachRemaining(key -> {
    if (!RESERVED_KIND_PATTERN.matcher(key.getKind()).matches()) {
      keys.add(key);
    }
  });
  if (keys.isEmpty()) {
    return;
  }
  log.info("Deleting datastore entities in namespace {}: {}", namespace, keys.size());
  Lists.partition(keys, 500).forEach(keyBatch ->
      datastore.delete(keyBatch.toArray(Key[]::new)));
}
 

/** Example of getting an entity for a given key. */
// [TARGET get(Key)]
// [VARIABLE "my_key_name"]
public Entity get(String keyName) {
  Datastore datastore = transaction.getDatastore();
  // [START get]
  Key key = datastore.newKeyFactory().setKind("MyKind").newKey(keyName);
  Entity entity = transaction.get(key);
  transaction.commit();
  // Do something with the entity
  // [END get]
  return entity;
}
 

@Test
@Parameters(method = "clients")
public void testEmulator(Datastore client) {
  var key = client.newKeyFactory().setKind("foo").newKey("bar");
  var entity = Entity.newBuilder(key).set("baz", "quux").build();
  client.put(entity);
  var readEntity = client.get(key);
  assertThat(readEntity, is(entity));
  datastoreEmulator.reset();
  var readEntityAfterReset = client.get(key);
  assertThat(readEntityAfterReset, is(nullValue()));
}
 

@Before
public void setUp() {

  final Datastore datastore = datastoreEmulator.client();
  storage = new AggregateStorage(bigtable, datastore);

  StorageFactory storageFactory = (env, stats) -> storage;
  StatsFactory statsFactory = (env) -> Stats.NOOP;
  StyxScheduler.ExecutorFactory executorFactory = (ts, tf) -> executor;
  StyxScheduler.PublisherFactory publisherFactory = (env) -> Publisher.NOOP;
  StyxScheduler.DockerRunnerFactory dockerRunnerFactory =
      (id, env, states, stats, debug, secretWhitelist, time) -> fakeDockerRunner();
  WorkflowResourceDecorator resourceDecorator = (rs, cfg, res) ->
      Sets.union(res, resourceIdsToDecorateWith);
  StyxScheduler.EventConsumerFactory eventConsumerFactory =
      (env, stats) -> (event, state) ->  transitionedEvents.add(Tuple.of(event, state.state()));
  AuthenticatorFactory authenticatorFactory = (cfg) -> mock(Authenticator.class);

  final ServiceAccountUsageAuthorizer.Factory serviceAccountUsageAuthorizerFactory =
      (cfg, name) -> ServiceAccountUsageAuthorizer.nop();
  styxScheduler = StyxScheduler.newBuilder()
      .setTime(time)
      .setStorageFactory(storageFactory)
      .setDockerRunnerFactory(dockerRunnerFactory)
      .setStatsFactory(statsFactory)
      .setExecutorFactory(executorFactory)
      .setPublisherFactory(publisherFactory)
      .setResourceDecorator(resourceDecorator)
      .setEventConsumerFactory(eventConsumerFactory)
      .setAuthenticatorFactory(authenticatorFactory)
      .setServiceAccountUsageAuthorizerFactory(serviceAccountUsageAuthorizerFactory)
      .build();

  serviceHelper = ServiceHelper.create(styxScheduler, StyxScheduler.SERVICE_NAME)
      .startTimeoutSeconds(30);
}
 

private static AggregateStorage storage(Environment environment, Stats stats  ) {
  final Config config = environment.config();
  final Closer closer = environment.closer();

  final Connection bigTable = closer.register(createBigTableConnection(config));
  final Datastore datastore = createDatastore(config, stats);
  return new AggregateStorage(bigTable, datastore);
}
 

@Test
public void createDatastore() {
  when(config.getString(Connections.DATASTORE_PROJECT)).thenReturn(TEST_PROJECT);
  when(config.getString(Connections.DATASTORE_NAMESPACE)).thenReturn(TEST_NAMESPACE);
  final InstrumentedDatastore instrumentedDatastore = Connections.createDatastore(config, stats);
  final Datastore datastore = instrumentedDatastore.delegate();
  final DatastoreOptions options = datastore.getOptions();
  assertThat(options.getProjectId(), is(TEST_PROJECT));
  assertThat(options.getNamespace(), is(TEST_NAMESPACE));
}
 

/** Example of deleting multiple entities. */
// [TARGET delete(Key...)]
// [VARIABLE "my_key_name1"]
// [VARIABLE "my_key_name2"]
public void multipleDeleteEntities(String keyName1, String keyName2) {
  Datastore datastore = transaction.getDatastore();
  // [START multipleDeleteEntities]
  Key key1 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName1);
  Key key2 = datastore.newKeyFactory().setKind("MyKind").newKey(keyName2);
  transaction.delete(key1, key2);
  transaction.commit();
  // [END multipleDeleteEntities]
}
 

Datastore create(final BlobStoreConfiguration configuration) throws Exception {
  DatastoreOptions.Builder builder = DatastoreOptions.newBuilder().setTransportOptions(transportOptions());

  String credentialFile = configuration.attributes(CONFIG_KEY).get(CREDENTIAL_FILE_KEY, String.class);
  if (StringUtils.hasText(credentialFile)) {
    ServiceAccountCredentials credentials = ServiceAccountCredentials.fromStream(new FileInputStream(credentialFile));
    logger.debug("loaded {} from {} for Google Datastore client", credentials, credentialFile);
    builder.setCredentials(credentials);
    builder.setProjectId(getProjectId(credentialFile));
  }

  return builder.build().getService();
}
 

public static void main(String... args) throws Exception {
  // Instantiates a client
  Datastore datastore = DatastoreOptions.getDefaultInstance().getService();

  // The kind for the new entity
  String kind = "Task";
  // The name/ID for the new entity
  String name = "sampletask1";
  // The Cloud Datastore key for the new entity
  Key taskKey = datastore.newKeyFactory().setKind(kind).newKey(name);

  // Prepares the new entity
  Entity task = Entity.newBuilder(taskKey)
      .set("description", "Buy milk")
      .build();

  // Saves the entity
  datastore.put(task);

  System.out.printf("Saved %s: %s%n", task.getKey().getName(), task.getString("description"));

  //Retrieve entity
  Entity retrieved = datastore.get(taskKey);

  System.out.printf("Retrieved %s: %s%n", taskKey.getName(), retrieved.getString("description"));

}