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

def testEventuallyConsistentGlobalQueryResult(self):
        class TestModel(ndb.Model):
            pass

        user_key = ndb.Key('User', 'ryan')

        # Put two entities
        ndb.put_multi([
            TestModel(parent=user_key),
            TestModel(parent=user_key)
        ])

        # Global query doesn't see the data.
        self.assertEqual(0, TestModel.query().count(3))
        # Ancestor query does see the data.
        self.assertEqual(2, TestModel.query(ancestor=user_key).count(3))
# [END HRD_example_1]

    # [START HRD_example_2] 

def test_set_from_run_result(self):
    request = _gen_request()
    result_summary = task_result.new_result_summary(request)
    to_run = task_to_run.new_task_to_run(request, 0)
    run_result = task_result.new_run_result(
        request, to_run, 'localhost', 'abc', {})
    run_result.started_ts = utils.utcnow()
    self.assertTrue(result_summary.need_update_from_run_result(run_result))
    result_summary.modified_ts = utils.utcnow()
    run_result.modified_ts = utils.utcnow()
    run_result.dead_after_ts = utils.utcnow() + datetime.timedelta(
        seconds=request.bot_ping_tolerance_secs)
    ndb.transaction(lambda: ndb.put_multi((result_summary, run_result)))

    self.assertTrue(result_summary.need_update_from_run_result(run_result))
    ndb.transaction(
        lambda: result_summary.set_from_run_result(run_result, request))
    ndb.transaction(lambda: ndb.put_multi([result_summary]))

    self.assertFalse(result_summary.need_update_from_run_result(run_result)) 

def test_run_result_timeout(self):
    request = _gen_request()
    result_summary = task_result.new_result_summary(request)
    result_summary.modified_ts = utils.utcnow()
    ndb.transaction(result_summary.put)
    to_run = task_to_run.new_task_to_run(request, 0)
    run_result = task_result.new_run_result(
        request, to_run, 'localhost', 'abc', {})
    run_result.state = task_result.State.TIMED_OUT
    run_result.duration = 0.1
    run_result.exit_code = -1
    run_result.started_ts = utils.utcnow()
    run_result.completed_ts = run_result.started_ts
    run_result.modified_ts = run_result.started_ts
    ndb.transaction(
        lambda: result_summary.set_from_run_result(run_result, request))
    ndb.transaction(lambda: ndb.put_multi((run_result, result_summary)))
    run_result = run_result.key.get()
    result_summary = result_summary.key.get()
    self.assertEqual(True, run_result.failure)
    self.assertEqual(True, result_summary.failure) 

def PutResource(url, etag, content):
  """Persist a resource."""
  key = ndb.Key(Resource, url, namespace=settings.PLAYGROUND_NAMESPACE)
  keys = ndb.Query(ancestor=key).fetch(keys_only=True)
  ndb.delete_multi(keys)
  resource = Resource(id=url, etag=etag,
                      namespace=settings.PLAYGROUND_NAMESPACE)
  if len(content) <= _MAX_RAW_PROPERTY_BYTES:
    resource.content = content
    resource.put()
    return
  chunks = [content[i:i + _MAX_RAW_PROPERTY_BYTES]
            for i in range(0, len(content), _MAX_RAW_PROPERTY_BYTES)]
  entities = [ResourceChunk(id=i + 1, parent=resource.key, content=chunks[i])
              for i in range(0, len(chunks))]
  entities.append(resource)
  ndb.put_multi(entities) 

def AdoptProjects(dst_user_id, src_user_id):
  """Transfer project ownership to a new user."""

  @ndb.transactional(xg=True)
  def _AdoptProject(project_key, dst_user_key, src_user_key):
    prj, dst_user, src_user = ndb.get_multi([project_key, dst_user_key,
                                             src_user_key])
    if project_key not in src_user.projects:
      # another concurrent request and transaction beat us to it
      return
    src_user.projects.remove(project_key)
    dst_user.projects.append(project_key)
    prj.owner = dst_user_key.id()
    prj.writers.remove(src_user_key.id())
    prj.writers.append(dst_user_key.id())
    ndb.put_multi([prj, dst_user, src_user])

  src_user = GetUser(src_user_id)
  if not src_user or not src_user.projects:
    return
  dst_user = GetOrCreateUser(dst_user_id)
  for project_key in src_user.projects:
    # slow, but transactionable
    _AdoptProject(project_key, dst_user.key, src_user.key) 

def test_append_output_max_chunk(self):
    # Ensures that data is dropped.
    # Force tedious chunking.
    self.mock(task_result.TaskOutput, 'CHUNK_SIZE', 2)
    self.mock(task_result.TaskOutput, 'PUT_MAX_CHUNKS', 16)
    self.assertEqual(2*16, task_result.TaskOutput.PUT_MAX_CONTENT())

    run_result = _gen_run_result()

    calls = []
    self.mock(logging, 'warning', lambda *args: calls.append(args))
    max_chunk = 'x' * task_result.TaskOutput.PUT_MAX_CONTENT()
    entities = run_result.append_output(max_chunk, 0)
    self.assertEqual(task_result.TaskOutput.PUT_MAX_CHUNKS, len(entities))
    ndb.put_multi(entities)
    self.assertEqual([], calls)

    # Try with PUT_MAX_CONTENT + 1 bytes, so the last byte is discarded.
    entities = run_result.append_output(max_chunk + 'x', 0)
    self.assertEqual(task_result.TaskOutput.PUT_MAX_CHUNKS, len(entities))
    ndb.put_multi(entities)
    self.assertEqual(1, len(calls))
    self.assertTrue(calls[0][0].startswith('Dropping '), calls[0][0])
    self.assertEqual(1, calls[0][1]) 

def store_specs_from_subscription(subscription_key, week_start, specs):
    """
    Idempotent function to store meeting specs for this week.
    """
    try:
        current_specs = MeetingSpec.query(
            MeetingSpec.meeting_subscription == subscription_key,
            MeetingSpec.datetime > week_start
        ).fetch()
    except NeedIndexError:
        current_specs = []

    if current_specs:
        return

    ndb.put_multi(specs)
    return specs 

def testGet_Active_MultipleOverlapping(self):

    # Create two active Alerts which overlap at the current time.
    alert_entity_1 = _CreateAlert(-10)
    alert_entity_2 = _CreateAlert(-5, end_hours=5)
    ndb.put_multi([alert_entity_1, alert_entity_2])
    self.assertEntityCount(alert.Alert, 2)

    with self.LoggedInUser():
      response = self.testapp.get(self.ROUTE)

    alert_dict = alert_entity_2.to_dict()

    self.assertEntityCount(alert.Alert, 2)
    self.assertMemcacheContains(
        alerts._CreateMemcacheKey('appdetail', 'windows'), alert_dict)
    self.assertEqual(httplib.OK, response.status_int)
    self.assertResponseContains(response, alert_dict) 

def post(self):
        # Force ndb to use v1 of the model by re-loading it.
        reload(models_v1)

        # Save some example data.
        ndb.put_multi([
            models_v1.Picture(author='Alice', name='Sunset'),
            models_v1.Picture(author='Bob', name='Sunrise')
        ])

        self.response.write("""
        Entities created. <a href="/">View entities</a>.
        """)
# [END add_entities]


# [START update_schema] 

def testMultipleEvents_DifferentUserTxns(self):
    self.PatchSetting(
        'EVENT_CREATION', constants.EVENT_CREATION.EXECUTING_USER)
    event1 = self._CreateEvent('the-sha256')
    event2 = event1.copy()
    event2[EVENT_UPLOAD.EXECUTING_USER] = 'anotheruser'
    event2[EVENT_UPLOAD.LOGGED_IN_USERS] = ['anotheruser']

    with mock.patch.object(sync.ndb, 'put_multi_async') as put_multi_mock:
      # It ain't pretty but it works: Wrap put_multi_async such that it behaves
      # like a synchronous version and we have the ability to track its calls.
      def fake_put_multi_async(seq):
        result = ndb.put_multi(seq)
        # Return an ndb.Future that is guaranteed to be done.
        return datastore_utils.GetNoOpFuture(result)
      put_multi_mock.side_effect = fake_put_multi_async

      request_json = {EVENT_UPLOAD.EVENTS: [event1, event2]}
      self.testapp.post_json('/my-uuid', request_json)

      # 1 from creating Certificate entities + 2 from events
      self.assertEqual(3, put_multi_mock.call_count)

    self.assertBigQueryInsertions([TABLE.BINARY] + [TABLE.EXECUTION] * 2) 

def testGetVotes_Inactive(self):
    self.assertLen(self.blockable_1.GetVotes(), 0)

    test_utils.CreateVotes(self.blockable_1, 2)

    self.assertLen(self.blockable_1.GetVotes(), 2)

    votes = vote_models.Vote.query().fetch()
    new_votes = []
    for vote in votes:
      new_key = ndb.Key(flat=vote.key.flat()[:-1] + (None,))
      new_votes.append(datastore_utils.CopyEntity(vote, new_key=new_key))
    ndb.delete_multi(vote.key for vote in votes)
    ndb.put_multi(new_votes)

    self.assertLen(self.blockable_1.GetVotes(), 0) 

def init_zoneinfo():
    """
    Add each zone info to the datastore. This will overwrite existing zones.

    This must be called before the AppengineTimezoneLoader will work.
    """
    import os, logging
    from zipfile import ZipFile
    zoneobjs = []

    zoneinfo_path = os.path.abspath(os.path.join(os.path.dirname(__file__),
      'zoneinfo.zip'))

    with ZipFile(zoneinfo_path) as zf:
        for zfi in zf.filelist:
            key = ndb.Key('Zoneinfo', zfi.filename, namespace=NDB_NAMESPACE)
            zobj = Zoneinfo(key=key, data=zf.read(zfi))
            zoneobjs.append(zobj)

    logging.info("Adding %d timezones to the pytz-appengine database" %
        len(zoneobjs)
        )

    ndb.put_multi(zoneobjs) 

def test_append_output_partial_far_split(self):
    # Missing, writing happens on two different TaskOutputChunk entities.
    self.mock(task_result.TaskOutput, 'CHUNK_SIZE', 16)
    run_result = _gen_run_result()
    ndb.put_multi(run_result.append_output(
        'FooBar', 2 * task_result.TaskOutput.CHUNK_SIZE - 3))
    expected_output = (
        '\x00' * (task_result.TaskOutput.CHUNK_SIZE * 2 - 3) + 'FooBar')
    self.assertEqual(expected_output, run_result.get_output(0, 0))
    expected = [
        {
            'chunk': '\x00' * (task_result.TaskOutput.CHUNK_SIZE - 3) + 'Foo',
            'gaps': [0, 13],
        },
        {
            'chunk': 'Bar',
            'gaps': []
        },
    ]
    self.assertTaskOutputChunk(expected) 

def account_rank(organization):
  account_qry = model.Account.query().filter(model.Account.organization == organization)
  account_dbs, account_cursors = util.get_dbs(
    account_qry,
    order='-stars',
    limit=-1,
  )
  updated_dbs = []
  for index, account_db in enumerate(account_dbs, start=1):
    if index < config.MAX_DB_LIMIT:
      account_db.rank = index
    else:
      account_db.rank = 0
    updated_dbs.append(account_db)

  ndb.put_multi(updated_dbs) 

def test_append_output_overwrite(self):
    # Overwrite previously written data.
    # Force tedious chunking.
    self.mock(task_result.TaskOutput, 'CHUNK_SIZE', 2)
    run_result = _gen_run_result()
    ndb.put_multi(run_result.append_output('FooBar', 0))
    ndb.put_multi(run_result.append_output('X', 3))
    self.assertEqual('FooXar', run_result.get_output(0, 0))
    self.assertTaskOutputChunk([
        {
            'chunk': 'Fo',
            'gaps': []
        },
        {
            'chunk': 'oX',
            'gaps': []
        },
        {
            'chunk': 'ar',
            'gaps': []
        },
    ]) 

def test_append_output_reverse_order_second_chunk(self):
    # Write the data in reverse order in multiple calls.
    self.mock(task_result.TaskOutput, 'CHUNK_SIZE', 16)
    run_result = _gen_run_result()
    ndb.put_multi(run_result.append_output(
        'Wow', task_result.TaskOutput.CHUNK_SIZE + 11))
    ndb.put_multi(run_result.append_output(
        'Foo', task_result.TaskOutput.CHUNK_SIZE + 8))
    ndb.put_multi(run_result.append_output(
        'Baz', task_result.TaskOutput.CHUNK_SIZE + 0))
    ndb.put_multi(
        run_result.append_output('Bar', task_result.TaskOutput.CHUNK_SIZE + 4))
    expected_output = (
        task_result.TaskOutput.CHUNK_SIZE * '\x00' + 'Baz\x00Bar\x00FooWow')
    self.assertEqual(expected_output, run_result.get_output(0, 0))
    self.assertTaskOutputChunk([{
        'chunk': 'Baz\x00Bar\x00FooWow',
        'gaps': [3, 4, 7, 8]
    }]) 

def _GetRepoCollections():
  repo_collections = []
  for uri, description in REPO_COLLECTIONS:
    repo_collection = model.GetOrInsertRepoCollection(uri, description)
    task = taskqueue.add(queue_name='repo',
                         url='/_playground_tasks/populate_repo_collection',
                         params={
                             'repo_collection_url': repo_collection.key.id(),
                         }, transactional=True)
    shared.w('adding task {} to populate repo collection {!r}'.format(task.name,
                                                                      uri))
    repo_collections.append(repo_collection)
  ndb.put_multi(repo_collections)
  return repo_collections 

def commit(self):
    assert not self.committed
    puts = [
      c.entity.make_historical_copy(c.deletion, c.comment)
      for c in self.changes
    ]
    ndb.put_multi(puts + [self.replication_state])
    for cb in _commit_callbacks:
      cb(self.replication_state.auth_db_rev)
    self.committed = True 

def test_revision_revision_exists(self):
    self.mock(gitiles, 'get_archive', mock.Mock())
    with open(TEST_ARCHIVE_PATH, 'r') as test_archive_file:
      gitiles.get_archive.return_value = test_archive_file.read()

    loc = gitiles.Location(
        hostname='localhost',
        project='project',
        treeish='master',
        path='/')
    cs = storage.ConfigSet(
        id='config_set',
        latest_revision=None,
        location=str(loc),
    )
    rev = storage.Revision(
        parent=cs.key,
        id='deadbeef',
    )
    ndb.put_multi([cs, rev])

    gitiles_import._import_revision(
        'config_set',
        loc,
        self.test_commit,
        False,
        self.test_project_id)

    cs_fresh = cs.key.get()
    self.assertEqual(cs_fresh.latest_revision, self.test_commit.sha) 

def commit(self):
    assert not self.committed
    puts = [
      c.entity.make_historical_copy(c.deletion, c.comment)
      for c in self.changes
    ]
    ndb.put_multi(puts + [self.replication_state])
    for cb in _commit_callbacks:
      cb(self.replication_state.auth_db_rev)
    self.committed = True 

def commit(self):
    assert not self.committed
    puts = [
      c.entity.make_historical_copy(c.deletion, c.comment)
      for c in self.changes
    ]
    ndb.put_multi(puts + [self.replication_state])
    for cb in _commit_callbacks:
      cb(self.replication_state.auth_db_rev)
    self.committed = True 

def _populate_shortlinks(self):
    test_shortlinks = [

      ShortLink(id=321,
                organization='1.com',
                owner='jay@1.com',
                shortpath='wiki',
                shortpath_prefix='wiki',
                destination_url='http://wiki.com'),

      ShortLink(organization='1.com',
                owner='jay@1.com',
                shortpath='sfdc/%s',
                shortpath_prefix='sfdc',
                destination_url='http://sfdc.com/search/%s'),

      ShortLink(organization='1.com',
                owner='bay@1.com',
                shortpath='slack',
                shortpath_prefix='slack',
                destination_url='http://slack.com'),


      ShortLink(organization='2.com',
                owner='jay@2.com',
                shortpath='drive',
                shortpath_prefix='drive',
                destination_url='http://drive3.com')
    ]

    ndb.put_multi(test_shortlinks) 

def test_get_shortlinks_for_user(self):
    ndb.put_multi([models.ShortLink(id=1,
                                    created=datetime.datetime(2018, 10, 1),
                                    organization='googs.com',
                                    owner='kay@googs.com',
                                    shortpath='there',
                                    destination_url='http://example.com'),
                   models.ShortLink(id=2,
                                    created=datetime.datetime(2018, 11, 1),
                                    organization='googs.com',
                                    owner='jay@googs.com',
                                    shortpath='here',
                                    destination_url='http://gmail.com'),
                   models.ShortLink(id=3,
                                    organization='widgets.com',
                                    owner='el@widgets.com',
                                    shortpath='elsewhere',
                                    destination_url='http://drive.com')])

    response = self.testapp.get('/_/api/links',
                                headers={'TROTTO_USER_UNDER_TEST': 'kay@googs.com'})

    self.assertEqual([{'oid': 1,
                       'created': '2018-10-01 00:00:00',
                       'mine': True,
                       'owner': 'kay@googs.com',
                       'shortpath': 'there',
                       'destination_url': 'http://example.com',
                       'visits_count': 0},
                      {'oid': 2,
                       'created': '2018-11-01 00:00:00',
                       'mine': False,
                       'owner': 'jay@googs.com',
                       'shortpath': 'here',
                       'destination_url': 'http://gmail.com',
                       'visits_count': 0}],
                     json.loads(response.text)) 

def _flush_ndb_puts(self, items, options):
    """Flush all NDB puts to datastore."""
    assert ndb is not None
    ndb.put_multi(items, config=self._create_config(options)) 

def get(self):
    policies_future = policy_models.Bit9Policy.query().fetch_async()

    active_policies = (
        api.Policy.query().filter(api.Policy.total_computers > 0)
        .execute(bit9_utils.CONTEXT))
    local_policies = {
        policy.key.id(): policy for policy in policies_future.get_result()}
    policies_to_update = []
    for policy in active_policies:
      try:
        level = constants.BIT9_ENFORCEMENT_LEVEL.MAP_FROM_INTEGRAL_LEVEL[
            policy.enforcement_level]
      except KeyError:
        logging.warning(
            'Unknown enforcement level "%s". Skipping...',
            policy.enforcement_level)
        continue
      local_policy = local_policies.get(str(policy.id))

      if local_policy is None:
        new_policy = policy_models.Bit9Policy(
            id=str(policy.id), name=policy.name, enforcement_level=level)
        policies_to_update.append(new_policy)
      else:
        dirty = False
        if local_policy.name != policy.name:
          local_policy.name = policy.name
          dirty = True
        if local_policy.enforcement_level != level:
          local_policy.enforcement_level = level
          dirty = True
        if dirty:
          policies_to_update.append(local_policy)

    if policies_to_update:
      logging.info('Updating %s policies', len(policies_to_update))
      ndb.put_multi(policies_to_update) 

def _flush_ndb_puts(self, items, options):
    """Flush all NDB puts to datastore."""
    assert ndb is not None
    ndb.put_multi(items, config=self._create_config(options)) 

def testMultipleEvents_RetriedTxn(self):
    user = test_utils.CreateUser()
    blockable = test_utils.CreateSantaBlockable()
    host = test_utils.CreateSantaHost()

    event_key = ndb.Key(user_models.User, user.key.id(),
                        host_models.SantaHost, host.key.id(),
                        binary_models.SantaBlockable, blockable.key.id(),
                        event_models.SantaEvent, '1')

    # This Event will already exist in the datastore but calling
    # _DedupeExistingAndPut on it will simulate an identical Event being synced.
    event = test_utils.CreateSantaEvent(
        key=event_key, blockable=blockable, host_id=host.key.id(),
        executing_user=user.nickname, count=10)

    # Simulate a retried transaction.
    # We use an exception to exit before committing the put_multi because,
    # according to the docs: "There is no mechanism to force a retry."
    # See https://cloud.google.com/appengine/docs/python/ndb/transactions
    with mock.patch.object(sync.ndb, 'put_multi_async', side_effect=Exception):
      with self.assertRaises(Exception):
        sync.EventUploadHandler()._DedupeExistingAndPut([event]).get_result()

    # And now retry...
    sync.EventUploadHandler()._DedupeExistingAndPut([event]).get_result()

    # Test that no additional event counts were added as a result of the retry.
    put_event = event_key.get()
    self.assertEqual(20, put_event.count)

    self.assertNoBigQueryInsertions() 

def CreateSantaBundle(bundle_binaries=None, **kwargs):
  """Create a SantaBundle entity."""
  defaults = {
      'name': 'bundle_name_%s' % RandomLetters(3),
      'bundle_id': '.'.join(RandomLetters(3) for _ in range(3)),
      'version': '.'.join(RandomDigits(1) for _ in range(3)),
      'short_version': '.'.join(RandomDigits(1) for _ in range(3)),
      'binary_count': len(bundle_binaries) if bundle_binaries else 1,
      'uploaded_dt': Now()
  }
  defaults.update(kwargs.copy())

  santa_bundle = CreateBlockableEntity(package_models.SantaBundle, **defaults)
  santa_bundle.put()

  # Create the SantaBundleBinary entities, if any bundle binaries were
  # specified.
  if bundle_binaries:
    entities = []
    for binary in bundle_binaries:
      entity = package_models.SantaBundleBinary.Generate(
          santa_bundle.key, binary.key, cert_key=binary.cert_key,
          file_name=binary.file_name, rel_path='Content/MacOS')
      entities.append(entity)

    ndb.put_multi(entities)

  return santa_bundle 

def _CommitChangeSet(change_key):
  """Attempts to commit and delete a given RuleChangeSet."""

  change = change_key.get()
  if change is None:
    logging.info('Change no longer exists. (already committed?)')
    return

  logging.info(
      'Committing a %s change set of %s rules for blockable %s',
      change.change_type, len(change.rule_keys), change.blockable_key.id())

  blockable = change.blockable_key.get()
  rules = ndb.get_multi(change.rule_keys)

  if change.change_type == constants.RULE_POLICY.WHITELIST:
    change_func = _Whitelist
  elif change.change_type == constants.RULE_POLICY.BLACKLIST:
    change_func = _Blacklist
  elif change.change_type == constants.RULE_POLICY.REMOVE:
    change_func = _Remove
  elif change.change_type in constants.RULE_POLICY.SET_INSTALLER:
    change_func = _ChangeInstallerState
  else:
    raise NotImplementedError

  # Attempt to perform the change. If something fails, just let the Exception
  # escape and kill the task. A retry will be attempted soon enough via cron.
  change_func(blockable, rules)

  # Clean up if the change went through.
  for rule in rules:
    rule.is_committed = True
  ndb.put_multi(rules)
  change.key.delete() 

def _ChangeModeForHosts(mode, user_keys, honor_lock=True):
  """Performs a client mode change for the specified users' hosts.

  Args:
    mode: The new client_mode to set.
    user_keys: The users whose host modes are to be changed.
    honor_lock: bool, whether the client_mode_lock property will be honored.
  """
  predicates = [
      host_models.SantaHost.primary_user == user_utils.EmailToUsername(key.id())
      for key in user_keys]
  query = host_models.SantaHost.query(ndb.OR(*predicates))
  hosts = query.fetch()
  updated_hosts = []

  for host in hosts:

    # If lock is honored, skip locked users.
    if honor_lock and host.client_mode_lock:
      continue

    # Ignore non-changes also.
    if host.client_mode == mode:
      continue

    # Proceed with the mode change.
    host.client_mode = mode
    host.client_mode_lock = False
    updated_hosts.append(host)

  ndb.put_multi(updated_hosts)
  logging.info(
      'Client mode changed to %s for %d host(s)', mode, len(updated_hosts))