Python google.appengine.ext.ndb.get_context() Examples

def _CreateRuleChangeSet(self, rules_future, new_policy):
    """Creates a RuleChangeSet and trigger an attempted commit."""
    rules = rules_future.get_result()
    # If there are no rules to be created (rare but possible), we can just drop
    # the change set entirely.
    if not rules:
      return

    keys = [rule.key for rule in rules]
    change = rule_models.RuleChangeSet(
        rule_keys=keys, change_type=new_policy, parent=self.blockable.key)
    change.put()

    # Attempt to commit the change set in a deferred task.
    # NOTE: If we're in a transaction, we should only send out the
    # request to Bit9 once the RuleChangeSet has been successfully created.
    # If we're not in a transaction, this executes immediately.
    ndb.get_context().call_on_commit(self._TriggerCommit) 

def setUp(self):
        super(BaseTest, self).setUp()

        root_path = '.'
        application_id = 'graphene-gae-test'

        # First, create an instance of the Testbed class.
        self.testbed = testbed.Testbed()
        self.testbed.activate()
        self.testbed.setup_env(app_id=application_id, overwrite=True)
        policy = datastore_stub_util.PseudoRandomHRConsistencyPolicy(probability=self.datastore_probability)
        self.testbed.init_datastore_v3_stub(root_path=root_path, consistency_policy=policy, require_indexes=True)
        self.testbed.init_app_identity_stub()
        self.testbed.init_blobstore_stub()
        self.testbed.init_memcache_stub()
        self.testbed.init_taskqueue_stub(root_path=root_path)
        self.testbed.init_urlfetch_stub()
        self.storage = cloudstorage_stub.CloudStorageStub(self.testbed.get_stub('blobstore').storage)
        self.testbed.init_mail_stub()
        self.testbed.init_user_stub()
        self.taskqueue_stub = self.testbed.get_stub(testbed.TASKQUEUE_SERVICE_NAME)

        ndb.get_context().clear_cache()
        ndb.get_context().set_cache_policy(lambda x: True) 

def setUp(self):
        # First, create an instance of the Testbed class.
        self.testbed = testbed.Testbed()
        # Then activate the testbed, which prepares the service stubs for use.
        self.testbed.activate()
        # Next, declare which service stubs you want to use.
        self.testbed.init_datastore_v3_stub()
        self.testbed.init_memcache_stub()
        # Clear ndb's in-context cache between tests.
        # This prevents data from leaking between tests.
        # Alternatively, you could disable caching by
        # using ndb.get_context().set_cache_policy(False)
        ndb.get_context().clear_cache()

# [END datastore_example_test]

    # [START datastore_example_teardown] 

def process_export():
    bucket = int(flask.request.values['bucket'])
    filename = '/ilps-search-log.appspot.com/search_log.%d.gz' % bucket
    with gcs.open(filename, 'w' , 'text/plain', {'content-encoding': 'gzip'}) as f:
        bucket_size = int(flask.request.values['bucket_size'])
        offset = bucket * bucket_size
        with gzip.GzipFile('', fileobj=f, mode='wb') as gz:
            ndb.get_context().clear_cache()
            for s in Session.query(Session.shared == True).iter(batch_size=10,
                    offset=offset, limit=bucket_size):
                ndb.get_context().clear_cache()
                gc.collect()
                s.user_id = ''
                print >>gz, json.dumps(s.to_dict(), default=util.default,
                        ensure_ascii=False).encode('utf-8')
    response = 'Written: %s' % str(blobstore.create_gs_key('/gs' + filename))
    app.logger.info(response)
    return response, 200 

def _tidy_stale_BotTaskDimensions(now):
  """Removes all stale BotTaskDimensions entities.

  This also cleans up entities for bots that were deleted or that died, as the
  corresponding will become stale after _EXTEND_VALIDITY.
  """
  qit = BotTaskDimensions.query(BotTaskDimensions.valid_until_ts < now).iter(
      batch_size=64, keys_only=True)
  btd = []
  while (yield qit.has_next_async()):
    key = qit.next()
    # This function takes care of confirming that the entity is indeed
    # expired.
    res = yield _remove_old_entity_async(key, now)
    btd.append(res)
    if res:
      bot_id = res.parent().string_id()
      yield ndb.get_context().memcache_delete(bot_id, namespace='task_queues')
      logging.debug('- BTD: %d for bot %s', res.integer_id(), bot_id)
  raise ndb.Return(btd) 

def setUp(self):
    self.testbed = testbed.Testbed()
    self.testbed.activate()
    self.testbed.init_datastore_v3_stub()
    self.testbed.init_memcache_stub()
    ndb.get_context().clear_cache() 

def _urlfetch(**kwargs):
  """To be mocked in tests."""
  return ndb.get_context().urlfetch(**kwargs) 

def _memcache_set(*args, **kwargs):
  """To be mocked in tests."""
  return ndb.get_context().memcache_set(*args, **kwargs) 

def _urlfetch_async(**kwargs):
  """To be mocked in tests."""
  return ndb.get_context().urlfetch(**kwargs) 

def get_config_by_hash_async(self, content_hash):
    """Returns a config blob by its hash. Memcaches results."""
    assert content_hash
    cache_key = '%sconfig_by_hash/%s' % (MEMCACHE_PREFIX, content_hash)
    ctx = ndb.get_context()
    content = yield ctx.memcache_get(cache_key)
    if content is not None:
      raise ndb.Return(zlib.decompress(content))

    res = yield self._api_call_async('config/%s' % content_hash)
    content = base64.b64decode(res.get('content')) if res else None
    if content is not None:
      yield ctx.memcache_set(cache_key, zlib.compress(content))
    raise ndb.Return(content) 

def _set_ndb_cache_policy():
  """Tell NDB to never cache anything in memcache or in-process.

  This ensures that entities fetched from Datastore input_readers via NDB
  will not bloat up the request memory size and Datastore Puts will avoid
  doing calls to memcache. Without this you get soft memory limit exits,
  which hurts overall throughput.
  """
  ndb_ctx = ndb.get_context()
  ndb_ctx.set_cache_policy(lambda key: False)
  ndb_ctx.set_memcache_policy(lambda key: False) 

def setUp(self):
        self.prepare_environ() 
        from gae.main import app
        self.test_app = webtest.TestApp(app)
        self.testbed = testbed.Testbed()
        self.testbed.activate()
        self.testbed.init_datastore_v3_stub()
        self.testbed.init_memcache_stub()
        ndb.get_context().clear_cache() 

def urlfetch_async(**kwargs):
  """To be mocked in tests."""
  return ndb.get_context().urlfetch(**kwargs) 

def setUp(self):
    self.testbed = testbed.Testbed()
    self.testbed.activate()
    self.testbed.init_datastore_v3_stub()
    self.testbed.init_memcache_stub()
    ndb.get_context().clear_cache() 

def setUp(self):
    self.testbed = testbed.Testbed()
    self.testbed.activate()
    self.testbed.init_datastore_v3_stub()
    self.testbed.init_memcache_stub()

    ndb.get_context().clear_cache() 

def _FutureFactory(future_cls, *args, **kwargs):
  """Create a Future object."""

  # NOTE: Futures do not execute callbacks from within the transaction
  # scope that they were created in. Since we rely on this behavior in some
  # places, we wrap all manually-created Futures to achieve this behavior.
  ctx = ndb.get_context()

  class _TxnPreservingFuture(future_cls):
    """Executes the Future's callbacks inside the current txn."""

    def add_callback(self, callback, *args, **kwds):
      @functools.wraps(callback)
      def _Wrapped(*inner_args, **inner_kwargs):
        with NdbContext(ctx):
          callback(*inner_args, **inner_kwargs)
      return super(
          _TxnPreservingFuture, self).add_callback(_Wrapped, *args, **kwds)

    def add_immediate_callback(self, callback, *args, **kwds):
      @functools.wraps(callback)
      def _Wrapped(*inner_args, **inner_kwargs):
        with NdbContext(ctx):
          callback(*inner_args, **inner_kwargs)
      return super(
          _TxnPreservingFuture, self).add_immediate_callback(
              _Wrapped, *args, **kwds)

  return _TxnPreservingFuture(*args, **kwargs) 

def setup_func():
    """set up test fixtures"""
    ndb.get_context().set_cache_policy(False) 

def _set_ndb_cache_policy():
  """Tell NDB to never cache anything in memcache or in-process.

  This ensures that entities fetched from Datastore input_readers via NDB
  will not bloat up the request memory size and Datastore Puts will avoid
  doing calls to memcache. Without this you get soft memory limit exits,
  which hurts overall throughput.
  """
  ndb_ctx = ndb.get_context()
  ndb_ctx.set_cache_policy(lambda key: False)
  ndb_ctx.set_memcache_policy(lambda key: False) 

def _get_pending_auth_db_transaction():
  """Used internally to keep track of changes done in the transaction.

  Returns:
    Instance of _AuthDBTransaction (stored in the transaction context).
  """
  # Use transaction context to store the object. Note that each transaction
  # retry gets its own new transaction context which is what we need,
  # see ndb/context.py, 'transaction' tasklet, around line 982 (for SDK 1.9.6).
  assert ndb.in_transaction()
  ctx = ndb.get_context()
  txn = getattr(ctx, '_auth_db_transaction', None)
  if txn:
    return txn

  # Prepare next AuthReplicationState (auth_db_rev +1).
  state = replication_state_key().get()
  if not state:
    primary_id = app_identity.get_application_id() if is_primary() else None
    state = AuthReplicationState(
        key=replication_state_key(),
        primary_id=primary_id,
        auth_db_rev=0)
  # Assert Primary or Standalone. Replicas can't increment auth db revision.
  if not is_primary() and state.primary_id:
    raise ValueError('Can\'t modify Auth DB on Replica')
  state.auth_db_rev += 1
  state.modified_ts = utils.utcnow()

  # Store the state in the transaction context. Used in replicate_auth_db(...)
  # later.
  txn = _AuthDBTransaction(state)
  ctx._auth_db_transaction = txn
  return txn 

def setUp(self):
    super(ConfigTest, self).setUp()
    # Disable in-memory NDB cache, it messes with cache related test cases.
    ndb.get_context().set_cache_policy(lambda _: False) 

def _get_pending_auth_db_transaction():
  """Used internally to keep track of changes done in the transaction.

  Returns:
    Instance of _AuthDBTransaction (stored in the transaction context).
  """
  # Use transaction context to store the object. Note that each transaction
  # retry gets its own new transaction context which is what we need,
  # see ndb/context.py, 'transaction' tasklet, around line 982 (for SDK 1.9.6).
  assert ndb.in_transaction()
  ctx = ndb.get_context()
  txn = getattr(ctx, '_auth_db_transaction', None)
  if txn:
    return txn

  # Prepare next AuthReplicationState (auth_db_rev +1).
  state = replication_state_key().get()
  if not state:
    primary_id = app_identity.get_application_id() if is_primary() else None
    state = AuthReplicationState(
        key=replication_state_key(),
        primary_id=primary_id,
        auth_db_rev=0)
  # Assert Primary or Standalone. Replicas can't increment auth db revision.
  if not is_primary() and state.primary_id:
    raise ValueError('Can\'t modify Auth DB on Replica')
  state.auth_db_rev += 1
  state.modified_ts = utils.utcnow()

  # Store the state in the transaction context. Used in replicate_auth_db(...)
  # later.
  txn = _AuthDBTransaction(state)
  ctx._auth_db_transaction = txn
  return txn 

def get_self_config_async(path, message_factory):
  """Parses a config file in the app's config set into a protobuf message."""
  cache_key = 'get_self_config_async(%r)' % path
  ctx = ndb.get_context()
  cached = yield ctx.memcache_get(cache_key)
  if cached:
    msg = message_factory()
    msg.ParseFromString(cached)
    raise ndb.Return(msg)

  cs = get_self_config_set()
  messages = yield get_latest_messages_async([cs], path, message_factory)
  msg = messages[cs]
  yield ctx.memcache_set(cache_key, msg.SerializeToString(), time=60)
  raise ndb.Return(msg) 

def get_config_hash_async(
      self, config_set, path, revision=None, use_memcache=True):
    """Returns tuple (revision, content_hash). Optionally memcaches results.

    If |revision| is not specified, memcaches for only 1 min.
    """
    assert config_set
    assert path

    get_latest = not revision

    content_hash = None
    if use_memcache:
      cache_key = (
          '%sconfig_hash/%s:%s@%s' %
          (MEMCACHE_PREFIX, config_set, path, revision or '!latest'))
      ctx = ndb.get_context()
      revision, content_hash = (
          (yield ctx.memcache_get(cache_key)) or (revision, None))

    if not content_hash:
      url_path = format_url('config_sets/%s/config/%s', config_set, path)
      params = {'hash_only': True}
      if revision:
        params['revision'] = revision
      res = yield self._api_call_async(url_path, params=params)
      if res:
        revision = res['revision']
        content_hash = res['content_hash']
        if content_hash and use_memcache:
          yield ctx.memcache_set(
              cache_key, (revision, content_hash), time=60 if get_latest else 0)
    raise ndb.Return(revision, content_hash) 

def get_config_by_hash_async(self, content_hash):
    """Returns a config blob by its hash. Memcaches results."""
    assert content_hash
    cache_key = '%sconfig_by_hash/%s' % (MEMCACHE_PREFIX, content_hash)
    ctx = ndb.get_context()
    content = yield ctx.memcache_get(cache_key)
    if content is not None:
      raise ndb.Return(zlib.decompress(content))

    res = yield self._api_call_async('config/%s' % content_hash)
    content = base64.b64decode(res.get('content')) if res else None
    if content is not None:
      yield ctx.memcache_set(cache_key, zlib.compress(content))
    raise ndb.Return(content) 

def _urlfetch_async(**kwargs):
  """To be mocked in tests."""
  return ndb.get_context().urlfetch(**kwargs) 

def _memcache_set(*args, **kwargs):
  """To be mocked in tests."""
  return ndb.get_context().memcache_set(*args, **kwargs) 

def _memcache_get(*args, **kwargs):
  """To be mocked in tests."""
  return ndb.get_context().memcache_get(*args, **kwargs) 

def _get_pending_auth_db_transaction():
  """Used internally to keep track of changes done in the transaction.

  Returns:
    Instance of _AuthDBTransaction (stored in the transaction context).
  """
  # Use transaction context to store the object. Note that each transaction
  # retry gets its own new transaction context which is what we need,
  # see ndb/context.py, 'transaction' tasklet, around line 982 (for SDK 1.9.6).
  assert ndb.in_transaction()
  ctx = ndb.get_context()
  txn = getattr(ctx, '_auth_db_transaction', None)
  if txn:
    return txn

  # Prepare next AuthReplicationState (auth_db_rev +1).
  state = replication_state_key().get()
  if not state:
    primary_id = app_identity.get_application_id() if is_primary() else None
    state = AuthReplicationState(
        key=replication_state_key(),
        primary_id=primary_id,
        auth_db_rev=0)
  # Assert Primary or Standalone. Replicas can't increment auth db revision.
  if not is_primary() and state.primary_id:
    raise ValueError('Can\'t modify Auth DB on Replica')
  state.auth_db_rev += 1
  state.modified_ts = utils.utcnow()

  # Store the state in the transaction context. Used in replicate_auth_db(...)
  # later.
  txn = _AuthDBTransaction(state)
  ctx._auth_db_transaction = txn
  return txn 

def setUp(self):
    super(ConfigTest, self).setUp()
    # Disable in-memory NDB cache, it messes with cache related test cases.
    ndb.get_context().set_cache_policy(lambda _: False) 

def urlfetch_async(**kwargs):
  """To be mocked in tests."""
  return ndb.get_context().urlfetch(**kwargs)