def inner():
   ctx = tasklets.get_context()
   self.assertTrue(ctx is not ctx1)
   key = model.Key('Account', 1)
   ent = yield key.get_async()
   self.assertTrue(tasklets.get_context() is ctx)
   self.assertTrue(ent is None)
   raise tasklets.Return(42)

 def outer():
   ctx1 = tasklets.get_context()
   @tasklets.tasklet
   def inner():
     ctx2 = tasklets.get_context()
     self.assertTrue(ctx1 is not ctx2)
     self.assertTrue(isinstance(ctx2._conn,
                                datastore_rpc.TransactionalConnection))
     return 42
   a = yield tasklets.get_context().transaction(inner)
   ctx1a = tasklets.get_context()
   self.assertTrue(ctx1 is ctx1a)
   raise tasklets.Return(a)

  def count_async(self, limit, **q_options):
    """Count the number of query results, up to a limit.

    This is the asynchronous version of Query.count().
    """
    assert 'offset' not in q_options, q_options
    assert 'limit' not in q_options, q_options
    if (self.__filters is not None and
        isinstance(self.__filters, DisjunctionNode)):
      # _MultiQuery isn't yet smart enough to support the trick below,
      # so just fetch results and count them.
      results = yield self.fetch_async(limit, **q_options)
      raise tasklets.Return(len(results))

    # Issue a special query requesting 0 results at a given offset.
    # The skipped_results count will tell us how many hits there were
    # before that offset without fetching the items.
    q_options['offset'] = limit
    q_options['limit'] = 0
    options = _make_options(q_options)
    conn = tasklets.get_context()._conn
    dsqry = self._get_query(conn)
    rpc = dsqry.run_async(conn, options)
    total = 0
    while rpc is not None:
      batch = yield rpc
      rpc = batch.next_batch_async(options)
      total += batch.skipped_results
    raise tasklets.Return(total)

 def callback():
   ctx = tasklets.get_context()
   self.assertTrue(key not in ctx._cache)  # Whitebox.
   e = yield key.get_async()
   self.assertTrue(key in ctx._cache)  # Whitebox.
   e.bar = 2
   yield e.put_async()

  def get_async(self, **ctx_options):
    """Return a Future whose result is the entity for this Key.

    If no such entity exists, a Future is still returned, and the
    Future's eventual return result be None.
    """
    from ndb import tasklets
    return tasklets.get_context().get(self, **ctx_options)

  def map_async(self, callback, merge_future=None, **q_options):
    """Map a callback function or tasklet over the query results.

    This is the asynchronous version of Query.map().
    """
    return tasklets.get_context().map_query(self, callback,
                                            options=_make_options(q_options),
                                            merge_future=merge_future)

 def allocate_ids_async(cls, size=None, max=None, parent=None):
   from ndb import tasklets
   if parent is None:
     pairs = []
   else:
     pairs = parent.pairs()
   pairs.append((cls.GetKind(), None))
   key = Key(pairs=pairs)
   return tasklets.get_context().allocate_ids(key, size=size, max=max)

    def setup_context_cache(self):
        """Set up the context cache.

        We only need cache active when testing the cache, so the default
        behavior is to disable it to avoid misleading test results. Override
        this when needed.
        """
        ctx = tasklets.get_context()
        ctx.set_cache_policy(False)
        ctx.set_memcache_policy(False)

  def delete_async(self, **ctx_options):
    """Schedule deletion of the entity for this Key.

    This returns a Future, whose result becomes available once the
    deletion is complete.  If no such entity exists, a Future is still
    returned.  In all cases the Future's result is None (i.e. there is
    no way to tell whether the entity existed or not).
    """
    from ndb import tasklets
    return tasklets.get_context().delete(self, **ctx_options)

  def __init__(self, query, options=None):
    """Constructor.  Takes a Query and optionally a QueryOptions.

    This is normally called by Query.iter() or Query.__iter__().
    """
    ctx = tasklets.get_context()
    callback = None
    if options is not None and options.produce_cursors:
      callback = self._extended_callback
    self._iter = ctx.iter_query(query, callback=callback, options=options)
    self._fut = None

  def SetupContextCache(self):
    """Set up the context cache.

    We only need cache active when testing the cache, so the default behavior
    is to disable it to avoid misleading test results. Override this when
    needed.
    """
    from ndb import tasklets
    ctx = tasklets.get_context()
    ctx.set_cache_policy(lambda key: False)
    ctx.set_memcache_policy(lambda key: False)

 def add_context_wrapper(*args):
   __ndb_debug__ = utils.func_info(func)
   tasklets.Future.clear_all_pending()
   # Reset context; a new one will be created on the first call to
   # get_context().
   tasklets.set_context(None)
   ctx = tasklets.get_context()
   try:
     return tasklets.synctasklet(func)(*args)
   finally:
     eventloop.run()  # Ensure writes are flushed, etc.

  def testAddContextDecorator(self):
    class Demo(object):
      @context.toplevel
      def method(self, arg):
        return (tasklets.get_context(), arg)

      @context.toplevel
      def method2(self, **kwds):
        return (tasklets.get_context(), kwds)
    a = Demo()
    old_ctx = tasklets.get_context()
    ctx, arg = a.method(42)
    self.assertTrue(isinstance(ctx, context.Context))
    self.assertEqual(arg, 42)
    self.assertTrue(ctx is not old_ctx)

    old_ctx = tasklets.get_context()
    ctx, kwds = a.method2(foo='bar', baz='ding')
    self.assertTrue(isinstance(ctx, context.Context))
    self.assertEqual(kwds, dict(foo='bar', baz='ding'))
    self.assertTrue(ctx is not old_ctx)

 def testExplicitTransactionClearsDefaultContext(self):
   old_ctx = tasklets.get_context()
   @tasklets.synctasklet
   def outer():
     ctx1 = tasklets.get_context()
     @tasklets.tasklet
     def inner():
       ctx = tasklets.get_context()
       self.assertTrue(ctx is not ctx1)
       key = model.Key('Account', 1)
       ent = yield key.get_async()
       self.assertTrue(tasklets.get_context() is ctx)
       self.assertTrue(ent is None)
       raise tasklets.Return(42)
     fut = ctx1.transaction(inner)
     self.assertEqual(tasklets.get_context(), ctx1)
     val = yield fut
     self.assertEqual(tasklets.get_context(), ctx1)
     raise tasklets.Return(val)
   val = outer()
   self.assertEqual(val, 42)
   self.assertTrue(tasklets.get_context() is old_ctx)

 def count_async(self, limit, options=None):
   conn = tasklets.get_context()._conn
   options = QueryOptions(offset=limit, limit=0, config=options)
   dsqry, post_filters = self._get_query(conn)
   if post_filters:
     raise datastore_errors.BadQueryError(
       'Post-filters are not supported for count().')
   rpc = dsqry.run_async(conn, options)
   total = 0
   while rpc is not None:
     batch = yield rpc
     rpc = batch.next_batch_async(options)
     total += batch.skipped_results
   raise tasklets.Return(total)

  def run(self):
    global cache_policy, memcache_policy, datastore_policy
    ctx = tasklets.get_context()
    ctx.set_cache_policy(cache_policy)
    ctx.set_memcache_policy(memcache_policy)
    ctx.set_datastore_policy(datastore_policy)

    id = threading.current_thread().ident

    try:
      for run in range(1, RUNS + 1):
        workload(id, run).check_success()
    except Exception, e:
      logger.exception('Thread %d run %d raised %s: %s',
                       id, run, e.__class__.__name__, e)

    def search(cls, params):
        """Returns (records, cursor).

        Arguments
            args - Dictionary with Darwin Core concept keys
            keywords - list of keywords to search on
        """        
        ctx = tasklets.get_context()
        ctx.set_memcache_policy(False)

        qry = RecordIndex.query()
        
        # Add darwin core name filters
        args = params['args']
        if len(args) > 0:
            gql = 'SELECT * FROM RecordIndex WHERE'
            for k,v in args.iteritems():
                gql = "%s %s = '%s' AND " % (gql, k, v)
            gql = gql[:-5] # Removes trailing AND
            logging.info(gql)
            qry = query.parse_gql(gql)[0]
            
        # Add full text keyword filters
        keywords = params['keywords']
        for keyword in keywords:
            qry = qry.filter(RecordIndex.corpus == keyword)        

        logging.info('QUERY='+str(qry))

        # Setup query paging
        limit = params['limit']
        cursor = params['cursor']        
        if cursor:
            logging.info('Cursor')
            index_keys, next_cursor, more = qry.fetch_page(limit, start_cursor=cursor, keys_only=True)
            record_keys = [x.parent() for x in index_keys]
        else:
            logging.info('No cursor')
            index_keys, next_cursor, more = qry.fetch_page(limit, keys_only=True)
            record_keys = [x.parent() for x in index_keys]
            
        # Return results
        return (model.get_multi(record_keys), next_cursor, more)

  def count_async(self, limit, **q_options):
    """Count the number of query results, up to a limit.

    This is the asynchronous version of Query.count().
    """
    assert 'offset' not in q_options, q_options
    assert 'limit' not in q_options, q_options
    if (self.__filters is not None and
        (isinstance(self.__filters, DisjunctionNode) or
         self.__filters._post_filters() is not None)):
      results = yield self.fetch_async(limit, **q_options)
      raise tasklets.Return(len(results))
    q_options['offset'] = limit
    q_options['limit'] = 0
    options = _make_options(q_options)
    conn = tasklets.get_context()._conn
    dsqry, post_filters = self._get_query(conn)
    rpc = dsqry.run_async(conn, options)
    total = 0
    while rpc is not None:
      batch = yield rpc
      rpc = batch.next_batch_async(options)
      total += batch.skipped_results
    raise tasklets.Return(total)

  def count_async(self, limit=None, **q_options):
    """Count the number of query results, up to a limit.

    This is the asynchronous version of Query.count().
    """
    # TODO: Support offset by incorporating it to the limit.
    assert 'offset' not in q_options, q_options
    assert 'limit' not in q_options, q_options
    if limit is None:
      limit = _MAX_LIMIT
    if (self.__filters is not None and
        isinstance(self.__filters, DisjunctionNode)):
      # _MultiQuery does not support iterating over result batches,
      # so just fetch results and count them.
      # TODO: Use QueryIterator to avoid materializing the results list.
      q_options.setdefault('prefetch_size', limit)
      q_options.setdefault('batch_size', limit)
      q_options.setdefault('keys_only', True)
      results = yield self.fetch_async(limit, **q_options)
      raise tasklets.Return(len(results))

    # Issue a special query requesting 0 results at a given offset.
    # The skipped_results count will tell us how many hits there were
    # before that offset without fetching the items.
    q_options['offset'] = limit
    q_options['limit'] = 0
    options = _make_options(q_options)
    conn = tasklets.get_context()._conn
    dsquery = self._get_query(conn)
    rpc = dsquery.run_async(conn, options)
    total = 0
    while rpc is not None:
      batch = yield rpc
      rpc = batch.next_batch_async(options)
      total += batch.skipped_results
    raise tasklets.Return(total)

 def setup_context():
   ctx = tasklets.get_context()
   ctx.set_datastore_policy(True)
   ctx.set_memcache_policy(True)
   ctx.set_cache_policy(False)
   return ctx