@Override
public Map<String, SessionData> getAllSessions() {
  final String originalNamespace = NamespaceManager.get();
  NamespaceManager.set("");
  try {
    DatastoreService ds = DatastoreServiceFactory.getDatastoreService();
    PreparedQuery pq = ds.prepare(new Query(SESSION_ENTITY_TYPE));
    Map<String, SessionData> sessions = new HashMap<>();
    for (Entity entity : pq.asIterable()) {
      sessions.put(entity.getKey().getName(), createSessionFromEntity(entity));
    }
    return sessions;
  } finally {
    NamespaceManager.set(originalNamespace);
  }
}
 

@Test
public void testCommitTx() throws Exception {
    AsyncDatastoreService service = DatastoreServiceFactory.getAsyncDatastoreService();
    Transaction tx = waitOnFuture(service.beginTransaction(TransactionOptions.Builder.withDefaults()));
    Key key;
    try {
        Future<Key> fKey = service.put(tx, new Entity("AsyncTx"));
        key = waitOnFuture(fKey);
        waitOnFuture(tx.commitAsync());
    } catch (Exception e) {
        waitOnFuture(tx.rollbackAsync());
        throw e;
    }

    if (key != null) {
        Assert.assertNotNull(getSingleEntity(service, key));
    }
}
 

/**
 * Returns the Account key associated with the specified authorization key.
 * @param pm               reference to the persistence manager
 * @param authorizationKey authorization key to return the account key for
 * @return the Account key associated with the specified authorization key; or <code>null</code> if the authorization key is invalid
 */
public static Key getAccountKeyByAuthKey( final PersistenceManager pm, final String authorizationKey ) {
	final String memcacheKey = CACHE_KEY_AUTH_KEY_ACCOUNT_KEY_PREFIX + authorizationKey;
	final String accountKeyString = (String) memcacheService.get( memcacheKey );
	if ( accountKeyString != null )
		return KeyFactory.stringToKey( accountKeyString );
	
	final Query q = new Query( Account.class.getSimpleName() );
	q.setFilter( new FilterPredicate( "authorizationKey", FilterOperator.EQUAL, authorizationKey ) );
	q.setKeysOnly();
	final List< Entity > entityList = DatastoreServiceFactory.getDatastoreService().prepare( q ).asList( FetchOptions.Builder.withDefaults() );
	if ( entityList.isEmpty() )
		return null;
	
	final Key accountKey = entityList.get( 0 ).getKey();
	try {
		memcacheService.put( memcacheKey, KeyFactory.keyToString( accountKey ) );
	}
	catch ( final MemcacheServiceException mse ) {
		LOGGER.log( Level.WARNING, "Failed to put key to memcache: " + memcacheKey, mse );
		// Ignore memcache errors, do not prevent serving user request
	}
	
	return accountKey;
}
 

@Test
public void testMiscOps() throws Exception {
    AsyncDatastoreService service = DatastoreServiceFactory.getAsyncDatastoreService();

    DatastoreAttributes attributes = waitOnFuture(service.getDatastoreAttributes());
    Assert.assertNotNull(attributes);
    Assert.assertNotNull(attributes.getDatastoreType());

    Map<Index, Index.IndexState> indexes = waitOnFuture(service.getIndexes());
    Assert.assertNotNull(indexes);

    Transaction tx = waitOnFuture(service.beginTransaction());
    try {
        String txId = tx.getId();
        Assert.assertNotNull(txId);
        Assert.assertEquals(txId, tx.getId());

        String appId = tx.getApp();
        Assert.assertNotNull(appId);
        Assert.assertEquals(appId, tx.getApp());
    } finally {
        tx.rollback();
    }
}
 

@Before
public void setUp() throws Exception {

  MockitoAnnotations.initMocks(this);
  helper.setUp();
  ds = DatastoreServiceFactory.getDatastoreService();

  //  Set up some fake HTTP requests
  when(mockRequest.getRequestURI()).thenReturn(FAKE_URL);
  when(mockRequest.getParameter("guestbookName")).thenReturn("default2");
  when(mockRequest.getParameter("content")).thenReturn(testPhrase);

  stringWriter = new StringWriter();
  when(mockResponse.getWriter()).thenReturn(new PrintWriter(stringWriter));

  servletUnderTest = new SignGuestbookServlet();

  ObjectifyService.init();
  ObjectifyService.register(Guestbook.class);
  ObjectifyService.register(Greeting.class);

  closeable = ObjectifyService.begin();

  cleanDatastore(ds, "default");
}
 

private static <T extends TempData> List<T> getAllTempData(Class<T> type, boolean unreadOnly) {
    try {
        DatastoreService ds = DatastoreServiceFactory.getDatastoreService();
        String kind = getKind(type);
        Query query = new Query(kind);
        if (unreadOnly) {
            query.setFilter(new Query.FilterPredicate(TEMP_DATA_READ_PROPERTY, Query.FilterOperator.EQUAL, false));
        } else {
            query.addSort("timestamp", Query.SortDirection.ASCENDING);
        }
        PreparedQuery pq = ds.prepare(query);
        Iterator<Entity> iter = pq.asIterator();
        List<T> result = new ArrayList<>();
        while (iter.hasNext()) {
            Entity entity = iter.next();
            T data = readTempData(type, entity, ds);
            result.add(data);
        }
        return result;
    } catch (Exception e) {
        throw new IllegalStateException(e);
    }
}
 

@Test
public void creatingAnEntityInASpecificEntityGroup() throws Exception {
  String boardName = "my-message-board";

  // [START creating_an_entity_in_a_specific_entity_group]
  DatastoreService datastore = DatastoreServiceFactory.getDatastoreService();

  String messageTitle = "Some Title";
  String messageText = "Some message.";
  Date postDate = new Date();

  Key messageBoardKey = KeyFactory.createKey("MessageBoard", boardName);

  Entity message = new Entity("Message", messageBoardKey);
  message.setProperty("message_title", messageTitle);
  message.setProperty("message_text", messageText);
  message.setProperty("post_date", postDate);

  Transaction txn = datastore.beginTransaction();
  datastore.put(txn, message);

  txn.commit();
  // [END creating_an_entity_in_a_specific_entity_group]
}
 

@Test
public void crossGroupTransactions() throws Exception {
  // [START cross-group_XG_transactions_using_the_Java_low-level_API]
  DatastoreService datastore = DatastoreServiceFactory.getDatastoreService();
  TransactionOptions options = TransactionOptions.Builder.withXG(true);
  Transaction txn = datastore.beginTransaction(options);

  Entity a = new Entity("A");
  a.setProperty("a", 22);
  datastore.put(txn, a);

  Entity b = new Entity("B");
  b.setProperty("b", 11);
  datastore.put(txn, b);

  txn.commit();
  // [END cross-group_XG_transactions_using_the_Java_low-level_API]
}
 

@Override
public void doPost(HttpServletRequest req, HttpServletResponse resp)
    throws IOException {
  final UserService userService = UserServiceFactory.getUserService();
  final User user = userService.getCurrentUser();

  final String guestbookName = req.getParameter("guestbookName");
  final Key guestbookKey = KeyFactory.createKey("Guestbook", guestbookName);
  final String content = req.getParameter("content");
  final Date date = new Date();
  final Entity greeting = new Entity("Greeting", guestbookKey);
  greeting.setProperty("user", user);
  greeting.setProperty("date", date);
  greeting.setProperty("content", content);

  final DatastoreService datastore = DatastoreServiceFactory.getDatastoreService();
  datastore.put(greeting);

  resp.sendRedirect("/guestbook.jsp?guestbookName=" + guestbookName);
}
 

@Test
public void readPolicy_strong_returnsAllResults() {
  double deadline = 5.0;
  ReadPolicy policy = new ReadPolicy(ReadPolicy.Consistency.STRONG);
  DatastoreServiceConfig datastoreConfig =
      DatastoreServiceConfig.Builder.withReadPolicy(policy).deadline(deadline);
  DatastoreService datastore = DatastoreServiceFactory.getDatastoreService(datastoreConfig);

  Entity parent = new Entity("Person", "a");
  Entity child = new Entity("Person", "b", parent.getKey());
  datastore.put(ImmutableList.<Entity>of(parent, child));

  Query q = new Query("Person").setAncestor(parent.getKey());
  List<Entity> results = datastore.prepare(q).asList(FetchOptions.Builder.withDefaults());
  assertWithMessage("query results").that(results).hasSize(2);
}
 

@Test
public void kindExample_writesEntity() throws Exception {
  // [START kind_example]
  Entity employee = new Entity("Employee", "asalieri");
  employee.setProperty("firstName", "Antonio");
  employee.setProperty("lastName", "Salieri");
  employee.setProperty("hireDate", new Date());
  employee.setProperty("attendedHrTraining", true);

  DatastoreService datastore = DatastoreServiceFactory.getDatastoreService();
  datastore.put(employee);
  // [END kind_example]

  Entity got = datastore.get(employee.getKey());
  assertWithMessage("got.firstName")
      .that((String) got.getProperty("firstName"))
      .isEqualTo("Antonio");
  assertWithMessage("got.lastName")
      .that((String) got.getProperty("lastName"))
      .isEqualTo("Salieri");
  assertWithMessage("got.hireDate").that((Date) got.getProperty("hireDate")).isNotNull();
  assertWithMessage("got.attendedHrTraining")
      .that((boolean) got.getProperty("attendedHrTraining"))
      .isTrue();
}
 

@Before
public void setUp() throws Exception {
  MockitoAnnotations.initMocks(this);
  helper.setUp();
  datastore = DatastoreServiceFactory.getDatastoreService();

  // Add test data.
  ImmutableList.Builder<Entity> people = ImmutableList.builder();
  for (String name : TEST_NAMES) {
    people.add(createPerson(name));
  }
  datastore.put(people.build());

  // Set up a fake HTTP response.
  responseWriter = new StringWriter();
  when(mockResponse.getWriter()).thenReturn(new PrintWriter(responseWriter));

  servletUnderTest = new ListPeopleServlet();
}
 

protected <T> T inTx(Action<T> action) throws Exception {
    AsyncDatastoreService ads = DatastoreServiceFactory.getAsyncDatastoreService();
    Transaction tx = ads.beginTransaction().get();
    boolean ok = false;
    try {
        T result = action.run(ads);
        ok = true;
        return result;
    } finally {
        if (ok)
            tx.commitAsync();
        else
            tx.rollbackAsync();

        sync(); // wait for tx to finish
    }
}
 

@ApiMethod(name = "storeTask")
public void storeTask(TaskBean taskBean) {
    DatastoreService datastoreService = DatastoreServiceFactory.getDatastoreService();
    Transaction txn = datastoreService.beginTransaction();
    try {
        Key taskBeanParentKey = KeyFactory.createKey("TaskBeanParent", "todo.txt");
        Entity taskEntity = new Entity("TaskBean", taskBean.getId(), taskBeanParentKey);
        taskEntity.setProperty("data", taskBean.getData());
        datastoreService.put(taskEntity);
        txn.commit();
    } finally {
        if (txn.isActive()) {
            txn.rollback();
        }
    }
}
 

public static <T extends TempData> T getLastTempData(Class<T> type) {
    try {
        DatastoreService ds = DatastoreServiceFactory.getDatastoreService();
        String kind = getKind(type);
        PreparedQuery pq = ds.prepare(new Query(kind).addSort("timestamp", Query.SortDirection.DESCENDING));
        Iterator<Entity> iter = pq.asIterator();
        if (iter.hasNext()) {
            Entity entity = iter.next();
            return readTempData(type, entity, ds);
        } else {
            return null;
        }
    } catch (Exception e) {
        throw new IllegalStateException(e);
    }
}
 

@ApiMethod(name = "clearTasks")
public void clearTasks() {
    DatastoreService datastoreService = DatastoreServiceFactory.getDatastoreService();
    Transaction txn = datastoreService.beginTransaction();
    try {
        Key taskBeanParentKey = KeyFactory.createKey("TaskBeanParent", "todo.txt");
        Query query = new Query(taskBeanParentKey);
        List<Entity> results = datastoreService.prepare(query).asList(FetchOptions.Builder.withDefaults());
        for (Entity result : results) {
            datastoreService.delete(result.getKey());
        }
        txn.commit();
    } finally {
        if (txn.isActive()) {
            txn.rollback();
        }
    }
}
 

@Override
protected void doGet(HttpServletRequest req, HttpServletResponse resp)
    throws ServletException, IOException {
  resp.setContentType("text/plain");

  DatastoreService ds = DatastoreServiceFactory.getDatastoreService();
  Entity e = new Entity("foo");
  e.setProperty("foo", 23);
  Key key = ds.put(e);

  try {
    e = ds.get(key);
  } catch (EntityNotFoundException e1) {
    throw new ServletException(e1);
  }

  e.setProperty("bar", 44);
  ds.put(e);

  Query q = new Query("foo");
  q.addFilter("foo", Query.FilterOperator.GREATER_THAN_OR_EQUAL, 22);
  Iterator<Entity> iter = ds.prepare(q).asIterator();
  iter.next();
}
 

/** Allocates an id. */
public static long allocateId() {
  return RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get())
    ? nextTestId.getAndIncrement()
    : DatastoreServiceFactory.getDatastoreService()
        .allocateIds(APP_WIDE_ALLOCATION_KIND, 1)
        .iterator()
        .next()
        .getId();
}
 

public static void deleteTempData(Class<? extends TempData> type) {
    // check if in-container
    if (isInContainer() == false) {
        return;
    }

    String kind = getKind(type);
    DatastoreService ds = DatastoreServiceFactory.getDatastoreService();

    final List<Entity> list = ds.prepare(new Query(kind)).asList(FetchOptions.Builder.withDefaults());
    for (Entity e : list) {
        deleteTempDataInTx(ds, e, type);
    }
}
 

@Before
public void init() {
  action.gcsService = gcsService;
  action.dryRun = false;
  action.datastoreService = DatastoreServiceFactory.getDatastoreService();
  action.fromTime = now.minusMillis(1);
  action.retrier = new Retrier(new FakeSleeper(new FakeClock()), 1);
  action.diffLister = new GcsDiffFileLister();
  action.diffLister.gcsService = gcsService;
  action.diffLister.gcsBucket = GCS_BUCKET;
  action.diffLister.executor = newDirectExecutorService();
}
 

@Test(expected = IllegalArgumentException.class)
public void testTransactionalTasksMustBeNamelessIterable() {
    Transaction tx = DatastoreServiceFactory.getDatastoreService().beginTransaction();
    try {
        getDefaultQueue().add(tx, Collections.singleton(TaskOptions.Builder.withTaskName("foo")));
    } finally {
        tx.rollback();
    }
}
 

@Before
public void setUp() {
    xmppService = XMPPServiceFactory.getXMPPService();
    datastoreService = DatastoreServiceFactory.getDatastoreService();
    clearTestData();
    initConfig();
}
 

@Override
public void init(FilterConfig config) throws ServletException {
  // initialize local copy of datastore session variables

  datastore = DatastoreServiceFactory.getDatastoreService();
  // Delete all sessions unmodified for over two days
  DateTime dt = DateTime.now(DateTimeZone.UTC);
  Query query = new Query(SESSION_KIND).setFilter(new FilterPredicate(
          "lastModified", FilterOperator.LESS_THAN_OR_EQUAL, dt.minusDays(2).toString(DTF)));
  Iterator<Entity> results = datastore.prepare(query).asIterator();
  while (results.hasNext()) {
    Entity stateEntity = results.next();
    datastore.delete(stateEntity.getKey());
  }
}
 

@Test(expected = IllegalArgumentException.class)
public void testTransactionalTasksMustBeNamelessSingle() {
    Transaction tx = DatastoreServiceFactory.getDatastoreService().beginTransaction();
    try {
        TaskOptions options = TaskOptions.Builder.withTaskName("foo");
        getDefaultQueue().add(tx, new TaskOptions(options));
    } finally {
        tx.rollback();
    }
}
 

@Override
public void init(FilterConfig config) throws ServletException {
  // initialize local copy of datastore session variables

  datastore = DatastoreServiceFactory.getDatastoreService();
  // Delete all sessions unmodified for over two days
  DateTime dt = DateTime.now(DateTimeZone.UTC);
  Query query = new Query(SESSION_KIND).setFilter(new FilterPredicate(
          "lastModified", FilterOperator.LESS_THAN_OR_EQUAL, dt.minusDays(2).toString(DTF)));
  Iterator<Entity> results = datastore.prepare(query).asIterator();
  while (results.hasNext()) {
    Entity stateEntity = results.next();
    datastore.delete(stateEntity.getKey());
  }
}
 

private void doTest() {
  DatastoreService ds = DatastoreServiceFactory.getDatastoreService();
  assertEquals(0, ds.prepare(new Query("yam")).countEntities(withLimit(10)));
  ds.put(new Entity("yam"));
  ds.put(new Entity("yam"));
  assertEquals(2, ds.prepare(new Query("yam")).countEntities(withLimit(10)));
}
 

@Test
public void testEventuallyConsistentGlobalQueryResult() {
  DatastoreService ds = DatastoreServiceFactory.getDatastoreService();
  Key ancestor = KeyFactory.createKey("foo", 3);
  ds.put(new Entity("yam", ancestor));
  ds.put(new Entity("yam", ancestor));
  // Global query doesn't see the data.
  assertEquals(0, ds.prepare(new Query("yam")).countEntities(withLimit(10)));
  // Ancestor query does see the data.
  assertEquals(2, ds.prepare(new Query("yam", ancestor)).countEntities(withLimit(10)));
}
 

public void purgeTestRunRecords() {
    DatastoreService datastoreService = DatastoreServiceFactory. getDatastoreService();
    FilterPredicate testRunFilter = new FilterPredicate(TEST_RUN_ID, FilterOperator.EQUAL, testRunId);
    Query query = new Query(entityName).setFilter(testRunFilter).setKeysOnly();
    for (Entity readRec : datastoreService.prepare(query).asIterable()) {
        datastoreService.delete(readRec.getKey());
    }
}
 

@Test(expected = ApiDeadlineExceededException.class)
public void testDeadlineConfig() {
    DatastoreServiceConfig config = DatastoreServiceConfig.Builder.withDeadline(0.00001);
    DatastoreService ds = DatastoreServiceFactory.getDatastoreService(config);
    assertNotNull(ds);
    Entity g1 = new Entity("test");
    g1.setProperty("deadline", "0.00001");
    ds.put(g1);
}
 

@Test
public void testFutureCancel() {
    clearData();
    AsyncDatastoreService asyncService = DatastoreServiceFactory.getAsyncDatastoreService();
    Entity newRec = new Entity(ASYNC_ENTITY);
    newRec.setProperty("count", -1);
    newRec.setProperty("timestamp", new Date());
    Future<Key> future = asyncService.put(newRec);
    future.cancel(true);

    // The actual call may already succeeded, so just verify that cancel has been called.
    assertTrue(future.isCancelled());
}