Python google.appengine.ext.db.put() Examples

def _save_states(self, state, serialized_readers_entity):
    """Run transaction to save state.

    Args:
      state: a model.MapreduceState entity.
      serialized_readers_entity: a model._HugeTaskPayload entity containing
        json serialized input readers.

    Returns:
      False if a fatal error is encountered and this task should be dropped
    immediately. True if transaction is successful. None if a previous
    attempt of this same transaction has already succeeded.
    """
    mr_id = state.key().id_or_name()
    fresh_state = model.MapreduceState.get_by_job_id(mr_id)
    if not self._check_mr_state(fresh_state, mr_id):
      return False
    if fresh_state.active_shards != 0:
      logging.warning(
          "Mapreduce %s already has active shards. Looks like spurious task "
          "execution.", mr_id)
      return None
    config = util.create_datastore_write_config(state.mapreduce_spec)
    db.put([state, serialized_readers_entity], config=config)
    return True 

def testAcquireLeaseSuccess(self):
    # lease acquired a long time ago.
    slice_start_time = datetime.datetime(2000, 1, 1)
    self.shard_state.slice_start_time = slice_start_time
    self.shard_state.slice_request_id = self.PREVIOUS_REQUEST_ID
    self.shard_state.put()
    handler, tstate = self._create_handler()
    with mock.patch("google.appengine.api"
                    ".logservice.fetch") as fetch:
      mock_request_log = mock.Mock()
      mock_request_log.finished = True
      fetch.return_value = [mock_request_log]
      handler._try_acquire_lease(self.shard_state, tstate)

    shard_state = model.ShardState.get_by_shard_id(self.shard_id)
    self.assertTrue(shard_state.active)
    self.assertEqual(self.CURRENT_SLICE_ID, shard_state.slice_id)
    self.assertEqual(self.CURRENT_REQUEST_ID, shard_state.slice_request_id)
    self.assertTrue(shard_state.slice_start_time > slice_start_time)
    self.assertEqual(shard_state, self.shard_state) 

def testSmoke(self):
    """Test main execution path of entity scanning.

    No processing rate limit.
    """
    e1 = TestEntity()
    e1.put()

    e2 = TestEntity()
    e2.put()

    self._handle_request()

    self.assertEquals([str(e1.key()), str(e2.key())],
                      TestHandler.processed_keys)

    # we should have finished
    self.verify_shard_state(
        model.ShardState.get_by_shard_id(self.shard_id),
        active=False, processed=2, input_finished=True,
        result_status=model.ShardState.RESULT_SUCCESS)

    tasks = self.taskqueue.GetTasks("default")
    self.assertEquals(0, len(tasks)) 

def testCompletedState(self):
    self.shard_state.input_finished = True
    self.shard_state.active = False
    self.shard_state.put()

    e1 = TestEntity()
    e1.put()

    self._handle_request(expect_finalize=False)

    # completed state => no data processed
    self.assertEquals([], TestHandler.processed_keys)

    self.verify_shard_state(
        model.ShardState.get_by_shard_id(self.shard_id),
        active=False,
        input_finished=True)

    self.assertEquals(0, len(self.taskqueue.GetTasks("default"))) 

def testAllInputProcessedStopsProcessing(self):
    self.shard_state.input_finished = True
    self.shard_state.put()

    e1 = TestEntity()
    e1.put()

    self.handler.post()

    # completed state => no data processed
    self.assertEquals([], TestHandler.processed_keys)

    # but shard is done now
    self.verify_shard_state(
        model.ShardState.get_by_shard_id(self.shard_id),
        active=False,
        result_status=model.ShardState.RESULT_SUCCESS,
        input_finished=True) 

def testRequestHasTimedOut(self):
    slice_start_time = datetime.datetime(2000, 1, 1)
    self.shard_state.slice_start_time = slice_start_time
    self.shard_state.slice_request_id = self.PREVIOUS_REQUEST_ID
    self.shard_state.put()
    handler, tstate = self._create_handler()
    # Lease has ended.
    self.assertEqual(0,
                     handler._wait_time(self.shard_state,
                                        parameters._LEASE_DURATION_SEC))
    # Logs API doesn't think the request has ended.
    self.assertFalse(handler._has_old_request_ended(self.shard_state))
    # But request has timed out.
    self.assertEqual(0, handler._wait_time(
        self.shard_state, parameters._MAX_LEASE_DURATION_SEC))

    # acquire lease should succeed.
    handler._try_acquire_lease(self.shard_state, tstate)

    shard_state = model.ShardState.get_by_shard_id(self.shard_id)
    self.assertTrue(shard_state.active)
    self.assertEqual(self.CURRENT_SLICE_ID, shard_state.slice_id)
    self.assertEqual(self.CURRENT_REQUEST_ID, shard_state.slice_request_id)
    self.assertTrue(shard_state.slice_start_time > slice_start_time) 

def setUpValidState(self, hooks_class_name=None):
    self.mapper_spec = model.MapperSpec(
        handler_spec=self.HANLDER_SPEC,
        input_reader_spec=self.INPUT_READER_SPEC,
        params={"entity_kind": self.ENTITY_KIND},
        shard_count=self.SHARD_COUNT,
        output_writer_spec=self.OUTPUT_WRITER_SPEC)

    for _ in range(10):
      TestEntity().put()

    # Use StartJobHandler for setup.
    self.mr_id = handlers.StartJobHandler._start_map(
        self.NAME, self.mapper_spec,
        mapreduce_params=self.MAPREDUCE_SPEC_PARAMS,
        queue_name=self.QUEUE,
        hooks_class_name=hooks_class_name) 

def testRequestHasNotEnd(self):
    # Previous request's lease has timed out but the request has not.
    now = datetime.datetime.now()
    old = (now -
           datetime.timedelta(seconds=parameters._LEASE_DURATION_SEC + 1))
    self.shard_state.slice_start_time = old
    self.shard_state.slice_request_id = self.PREVIOUS_REQUEST_ID
    self.shard_state.put()
    handler, _ = self._create_handler()
    # Lease has ended.
    self.assertEqual(0,
                     handler._wait_time(self.shard_state,
                                        parameters._LEASE_DURATION_SEC),
                     lambda: now)
    # Logs API doesn't think the request has ended.
    self.assertFalse(handler._has_old_request_ended(self.shard_state))
    # Request has not timed out.
    self.assertTrue(handler._wait_time(
        self.shard_state,
        parameters._MAX_LEASE_DURATION_SEC,
        lambda: now))
    handler.post()
    self.assertEqual(httplib.SERVICE_UNAVAILABLE, handler.response.status) 

def testParameterValidationSuccess(self):
    """Tests validating user-supplied parameters."""
    TestEntity().put()
    self.handler.request.set("mapper_params.one", ["red", "blue"])
    self.handler.request.set("mapper_params.two", "green")
    self.handler.request.set("mapper_params_validator",
                             __name__ + ".test_param_validator_success")
    self.handler.post()

    state = model.MapreduceState.get_by_job_id(
        self.handler.json_response["mapreduce_id"])
    params = state.mapreduce_spec.mapper.params

    self.assertEquals(["red", "blue"], params["one"])
    self.assertEquals("green", params["two"])

    # From the validator function
    self.assertEquals("good", params["test"]) 

def run(self):
    if model.MissingDataUsersSecondary.all().count() > 0:
      raise Error('Must clear MissingDataUsersSecondary before refreshing')

    # Load the whole data model. As of 2014-05-20, this takes around
    # 30 seconds, out of our allotted 10 minutes.
    logging.info('Load all users')
    users = [u.email for u in model.User.all()
             if not (u.occupation and u.employer)]

    logging.info('Load all Pledges')
    pledges = list(model.Pledge.all())
    logging.info('Load all WpPledges')
    wp_pledges = list(model.WpPledge.all())
    logging.info('Done loading')

    pledge_sum = defaultdict(int)
    for p in pledges + wp_pledges:
      pledge_sum[p.email] += p.amountCents

    users = [u for u in users if pledge_sum[u] >= 20000]
    users = [model.MissingDataUsersSecondary(email=u, amountCents=pledge_sum[u])
             for u in users]
    db.put(users)
    logging.info('Done') 

def testMapreduceParameters(self):
    """Tests propagation of user-supplied mapreduce parameters."""
    TestEntity().put()
    self.handler.request.set("params.one", ["red", "blue"])
    self.handler.request.set("params.two", "green")
    self.handler.request.set("params_validator",
                             __name__ + ".test_param_validator_success")
    self.handler.post()

    state = model.MapreduceState.get_by_job_id(
        self.handler.json_response["mapreduce_id"])
    params = state.mapreduce_spec.params
    self.assertEquals(["red", "blue"], params["one"])
    self.assertEquals("green", params["two"])

    # From the validator function
    self.assertEquals("good", params["test"]) 

def _drop_gracefully(self):
    """Drop worker task gracefully.

    Set current shard_state to failed. Controller logic will take care of
    other shards and the entire MR.
    """
    shard_id = self.request.headers[util._MR_SHARD_ID_TASK_HEADER]
    mr_id = self.request.headers[util._MR_ID_TASK_HEADER]
    shard_state, mr_state = db.get([
        model.ShardState.get_key_by_shard_id(shard_id),
        model.MapreduceState.get_key_by_job_id(mr_id)])

    if shard_state and shard_state.active:
      shard_state.set_for_failure()
      config = util.create_datastore_write_config(mr_state.mapreduce_spec)
      shard_state.put(config=config) 

def testRequiredParams(self):
    """Tests that required parameters are enforced."""
    TestEntity().put()
    self.handler.post()

    self.handler.request.set("name", None)
    self.assertRaises(errors.NotEnoughArgumentsError, self.handler.handle)

    self.handler.request.set("name", "my job")
    self.handler.request.set("mapper_input_reader", None)
    self.assertRaises(errors.NotEnoughArgumentsError, self.handler.handle)

    self.handler.request.set(
        "mapper_input_reader",
        "mapreduce.input_readers.DatastoreInputReader")
    self.handler.request.set("mapper_handler", None)
    self.assertRaises(errors.NotEnoughArgumentsError, self.handler.handle)

    self.handler.request.set("mapper_handler", MAPPER_HANDLER_SPEC)
    self.handler.request.set("mapper_params.entity_kind", None)
    self.assertRaises(input_readers.BadReaderParamsError, self.handler.handle)

    self.handler.request.set("mapper_params.entity_kind",
                             (__name__ + "." + TestEntity.__name__))
    self.handler.post() 

def _save_states(self, state, serialized_readers_entity):
    """Run transaction to save state.

    Args:
      state: a model.MapreduceState entity.
      serialized_readers_entity: a model._HugeTaskPayload entity containing
        json serialized input readers.

    Returns:
      False if a fatal error is encountered and this task should be dropped
    immediately. True if transaction is successful. None if a previous
    attempt of this same transaction has already succeeded.
    """
    mr_id = state.key().id_or_name()
    fresh_state = model.MapreduceState.get_by_job_id(mr_id)
    if not self._check_mr_state(fresh_state, mr_id):
      return False
    if fresh_state.active_shards != 0:
      logging.warning(
          "Mapreduce %s already has active shards. Looks like spurious task "
          "execution.", mr_id)
      return None
    config = util.create_datastore_write_config(state.mapreduce_spec)
    db.put([state, serialized_readers_entity], config=config)
    return True 

def _save_states(self, state, serialized_readers_entity):
    """Run transaction to save state.

    Args:
      state: a model.MapreduceState entity.
      serialized_readers_entity: a model._HugeTaskPayload entity containing
        json serialized input readers.

    Returns:
      False if a fatal error is encountered and this task should be dropped
    immediately. True if transaction is successful. None if a previous
    attempt of this same transaction has already succeeded.
    """
    mr_id = state.key().id_or_name()
    fresh_state = model.MapreduceState.get_by_job_id(mr_id)
    if not self._check_mr_state(fresh_state, mr_id):
      return False
    if fresh_state.active_shards != 0:
      logging.warning(
          "Mapreduce %s already has active shards. Looks like spurious task "
          "execution.", mr_id)
      return None
    config = util.create_datastore_write_config(state.mapreduce_spec)
    db.put([state, serialized_readers_entity], config=config)
    return True 

def _drop_gracefully(self):
    """Drop worker task gracefully.

    Set current shard_state to failed. Controller logic will take care of
    other shards and the entire MR.
    """
    shard_id = self.request.headers[util._MR_SHARD_ID_TASK_HEADER]
    mr_id = self.request.headers[util._MR_ID_TASK_HEADER]
    shard_state, mr_state = db.get([
        model.ShardState.get_key_by_shard_id(shard_id),
        model.MapreduceState.get_key_by_job_id(mr_id)])

    if shard_state and shard_state.active:
      shard_state.set_for_failure()
      config = util.create_datastore_write_config(mr_state.mapreduce_spec)
      shard_state.put(config=config) 

def __setitem__(self, keyname, value):
        """
        Set item in session data.

        Args:
            keyname: They keyname of the mapping.
            value: The value of mapping.
        """

        if self.integrate_flash and (keyname == 'flash'):
            self.flash.msg = value
        else:
            keyname = self._validate_key(keyname)
            self.cache[keyname] = value
            # self._set_memcache() # commented out because this is done in the datestore put
            return self._put(keyname, value) 

def _import_backup(gs_handle):
  """Import backup from `gs_handle` to the current project."""
  bucket_name, path = parse_gs_handle(gs_handle)
  file_content = get_gs_object(bucket_name, path)
  entities = parse_backup_info_file(file_content)
  original_backup_info = entities.next()
  entity = datastore.Entity(BackupInformation.kind())
  entity.update(original_backup_info)
  backup_info = BackupInformation.from_entity(entity)
  if original_backup_info.key().app() != os.getenv('APPLICATION_ID'):
    backup_info.original_app = original_backup_info.key().app()

  def tx():
    backup_info.put(force_writes=True)
    kind_files_models = []
    for entity in entities:
      kind_files = backup_info.create_kind_backup_files(
          entity.key().name(), entity['files'])
      kind_files_models.append(kind_files)
    db.put(kind_files_models, force_writes=True)
  db.run_in_transaction(tx)
  return str(backup_info.key()) 

def _perform_backup_complete(
    operation, job_id, kind, backup_info_pk, gcs_path_prefix, filenames, queue):
  backup_info = BackupInformation.get(backup_info_pk)
  if backup_info:
    if job_id in backup_info.active_jobs:
      backup_info.active_jobs.remove(job_id)
      backup_info.completed_jobs = list(
          set(backup_info.completed_jobs + [job_id]))


    filenames = [GCSUtil.add_gs_prefix_if_missing(name) for name in filenames]
    kind_backup_files = backup_info.get_kind_backup_files([kind])[0]
    if kind_backup_files:
      kind_backup_files.files = list(set(kind_backup_files.files + filenames))
    else:
      kind_backup_files = backup_info.create_kind_backup_files(kind, filenames)
    db.put((backup_info, kind_backup_files), force_writes=True)
    if operation.status == utils.DatastoreAdminOperation.STATUS_COMPLETED:
      deferred.defer(finalize_backup_info, backup_info.key(),
                     gcs_path_prefix,
                     _url=config.DEFERRED_PATH,
                     _queue=queue,
                     _transactional=True)
  else:
    logging.warn('BackupInfo was not found for %s', backup_info_pk) 

def create_key(self):
        """
        Creates a unique key for the session.
        """
        self.session_key = time.time()
        valid = False
        while valid == False:
            # verify session_key is unique
            if memcache.get("_AppEngineUtilities_Session_" + str(self.session_key)):
                self.session_key = self.session_key + 0.001
            else:
                query = _AppEngineUtilities_Session.all()
                query.filter("session_key = ", self.session_key)
                results = query.fetch(1)
                if len(results) > 0:
                    self.session_key = self.session_key + 0.001
                else:
                    try:
                        self.put()
                        memcache.set("_AppEngineUtilities_Session_" + str(self.session_key), self)
                    except:
                        self.dirty = True
                        memcache.set("_AppEngineUtilities_Session_" + str(self.session_key), self)
                    valid = True 

def put(self):
        """
        Extend put so that it writes vaules to memcache as well as the datastore,
        and keeps them in sync, even when datastore writes fails.
        """
        if self.session_key:
            memcache.set("_AppEngineUtilities_Session_" + str(self.session_key), self)
        else:
            # new session, generate a new key, which will handle the put and set the memcache
            self.create_key()

        self.last_activity = datetime.datetime.now()

        try:
            self.dirty = False
            logging.info("doing a put")
            db.put(self)
            memcache.set("_AppEngineUtilities_Session_" + str(self.session_key), self)
        except:
            self.dirty = True
            memcache.set("_AppEngineUtilities_Session_" + str(self.session_key), self)

        return self 

def _create_and_save_state(cls, mapreduce_spec, _app):
    """Save mapreduce state to datastore.

    Save state to datastore so that UI can see it immediately.

    Args:
      mapreduce_spec: model.MapreduceSpec,
      _app: app id if specified. None otherwise.

    Returns:
      The saved Mapreduce state.
    """
    state = model.MapreduceState.create_new(mapreduce_spec.mapreduce_id)
    state.mapreduce_spec = mapreduce_spec
    state.active = True
    state.active_shards = 0
    if _app:
      state.app_id = _app
    config = util.create_datastore_write_config(mapreduce_spec)
    state.put(config=config)
    return state 

def testContentionWhenAcquireLease(self):
    # Shard has moved on AFTER we got shard state.
    self.shard_state.slice_id += 1
    self.shard_state.put()

    # Revert in memory shard state.
    self.shard_state.slice_id -= 1
    handler, tstate = self._create_handler()
    self.assertEqual(
        handler._TASK_DIRECTIVE.RETRY_TASK,
        # Use old shard state.
        handler._try_acquire_lease(self.shard_state, tstate)) 

def _init_shard(self):
    """Init shard state."""
    self.shard_state = model.ShardState.create_new(
        self.mr_spec.mapreduce_id,
        shard_number=self.SHARD_NUMBER)
    self.shard_state.slice_id = self.CURRENT_SLICE_ID
    self.shard_state.put()
    self.shard_id = self.shard_state.shard_id 

def _put(self, keyname, value):
        """
        Insert a keyname/value pair into the datastore for the session.

        Args:
            keyname: The keyname of the mapping.
            value: The value of the mapping.
        """
        if self.writer == "datastore":
            writer = _DatastoreWriter()
        else:
            writer = _CookieWriter()

        writer.put(keyname, value, self) 

def put(self, keyname, value, session):
        """
        Insert a keyname/value pair into the datastore for the session.

        Args:
            keyname: The keyname of the mapping.
            value: The value of the mapping.
        """
        keyname = session._validate_key(keyname)
        if value is None:
            raise ValueError('You must pass a value to put.')

        # datestore write trumps cookie. If there is a cookie value
        # with this keyname, delete it so we don't have conflicting
        # entries.
        if session.cookie_vals.has_key(keyname):
            del(session.cookie_vals[keyname])
            session.output_cookie[session.cookie_name + '_data'] = \
                simplejson.dumps(session.cookie_vals)
            print session.output_cookie.output()

        sessdata = session._get(keyname=keyname)
        if sessdata is None:
            sessdata = _AppEngineUtilities_SessionData()
            sessdata.session_key = session.session.session_key
            sessdata.keyname = keyname
        sessdata.content = pickle.dumps(value)
        # UNPICKLING CACHE session.cache[keyname] = pickle.dumps(value)
        session.cache[keyname] = value
        sessdata.put()
        # todo _set_memcache() should be going away when this is done
        # session._set_memcache() 

def put(self):
        """
        Adds a keyname/value for session to the datastore and memcache
        """
        # update or insert in datastore
        try:
            db.put(self)
            self.dirty = False
        except:
            self.dirty = True

        # update or insert in memcache
        mc_items = memcache.get("_AppEngineUtilities_SessionData_" + str(self.session_key))
        if mc_items:
            value_updated = False
            for item in mc_items:
                if value_updated == True:
                    break
                if item.keyname == self.keyname:
                    logging.info("updating " + self.keyname)
                    item.content = self.content
                    memcache.set("_AppEngineUtilities_SessionData_" + str(self.session_key), mc_items)
                    value_updated = True
                    break
            if value_updated == False:
                #logging.info("adding " + self.keyname)
                mc_items.append(self)
                memcache.set("_AppEngineUtilities_SessionData_" + str(self.session_key), mc_items) 

def testShardStateCollision(self):
    handlers._TEST_INJECTED_FAULTS.add("worker_active_state_collision")

    e1 = TestEntity()
    e1.put()

    self._handle_request(expect_finalize=False)

    # Data will still be processed
    self.assertEquals([str(e1.key())], TestHandler.processed_keys)
    # Shard state should not be overriden, i.e. left active.
    self.verify_shard_state(
        model.ShardState.get_by_shard_id(self.shard_id), active=True)
    self.assertEquals(0, len(self.taskqueue.GetTasks("default"))) 

def testInvalidMRState(self):
    self.createDummyHandler()
    # No mr_state exists.
    self.handler.request.set("mapreduce_id", self.mapreduce_id)
    self.handler.post()
    self.assertEqual(0, len(self.taskqueue.GetTasks(self.QUEUE)))

    # mr_state is not active.
    state = model.MapreduceState.create_new(self.mapreduce_id)
    state.active = False
    state.put()
    self.handler.post()
    self.assertEqual(0, len(self.taskqueue.GetTasks(self.QUEUE))) 

def testNoShardState(self):
    """Correct handling of missing shard state."""
    self.shard_state.delete()
    e1 = TestEntity()
    e1.put()

    self._handle_request(expect_finalize=False)

    # no state => no data processed
    self.assertEquals([], TestHandler.processed_keys)
    self.assertEquals(0, len(self.taskqueue.GetTasks("default")))