com.mongodb.WriteResult Java Examples

@Test
public void testPersist() {
	DBCollection coll = createMockCollection();
	ToDo td = new ToDo();
	td.setTitle("This is a test");
	td.setId("aaaaaaaaaaaaaaaaaaaaaaa1");
	expect(coll.insert(isA(DBObject.class))).andAnswer(new IAnswer<WriteResult>() {
		@Override
		public WriteResult answer() throws Throwable {
			DBObject obj = (DBObject)getCurrentArguments()[0];
			obj.put("_id", new ObjectId("aaaaaaaaaaaaaaaaaaaaaaa1"));
			return null;
		}
	});
	replay(coll);
	MongoStore store = new MongoStore(coll);
	assertEquals(td, store.persist(td));
	verify(coll);
}
 

@Override
public Object afterMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
    Object ret) throws Throwable {
    AbstractSpan activeSpan = ContextManager.activeSpan();
    CommandResult cresult = null;
    if (ret instanceof WriteResult) {
        WriteResult wresult = (WriteResult) ret;
        cresult = wresult.getCachedLastError();
    } else if (ret instanceof AggregationOutput) {
        AggregationOutput aresult = (AggregationOutput) ret;
        cresult = aresult.getCommandResult();
    }
    if (null != cresult && !cresult.ok()) {
        activeSpan.log(cresult.getException());
    }
    ContextManager.stopSpan();
    return ret;
}
 

@Override
public boolean saveResourcePlan(ResourcePlan resourcePlan) {
    boolean result = false;
    if(resourcePlan.getAddTime() == 0) { //insert
        resourcePlan.setAddTime(new Date().getTime());
        resourcePlan.setModTime(new Date().getTime());

        mongoTemplate.save(resourcePlan, Constant.COL_NAME_RESOURCE_PLAN);
        result = Preconditions.isNotBlank(resourcePlan.getId());

    } else {                            //update
        Query query = new Query().addCriteria(Criteria.where("_id").is(resourcePlan.getId()));
        Update update = new Update();
        update.set("startPageNum", resourcePlan.getStartPageNum());
        update.set("endPageNum", resourcePlan.getEndPageNum());
        update.set("modTime", new Date().getTime());

        WriteResult writeResult = mongoTemplate.updateFirst(query, update, Constant.COL_NAME_RESOURCE_PLAN);
        result = writeResult!=null && writeResult.getN() > 0;
    }

    return result;
}
 

/**
 * 更改密码
 * @param code
 * @param password
 * @param responseMsgVO
 */
public void changePwd(String code, String password, ResponseMsgVO responseMsgVO) {
    String username=  UserContextHandler.getUsername();


    UserDO userDO = userRepository.findByUsername(username);
    log.info("用户:{}更改密码",userDO.getUsername());
    String key = CODE_PREFIX + CodeTypeEnum.CHANGE_PWD.getValue() + userDO.getEmail();
    String value = redisTemplate.opsForValue().get(key);
    if (!StringUtils.equals(code, value)) {
        responseMsgVO.buildWithMsgAndStatus(PosCodeEnum.PARAM_ERROR, "验证码错误");
        return;
    }

    if (!checkPassWordLegal(password, responseMsgVO)){
        return;
    }

    // 更新到mongo数据库
    BCryptPasswordEncoder passwordEncoder = new BCryptPasswordEncoder();
    WriteResult result = userRepository.changePassword(username, passwordEncoder.encode(password));
    if (result.getN() == 0) {
        responseMsgVO.buildWithMsgAndStatus(PosCodeEnum.PARAM_ERROR, "更新失败");
        return;
    }
    responseMsgVO.buildOK();
}
 

/**
 * Save an entity to mongoDB. 
 * If no id in it, then insert the entity.
 * Else, check the id type, to confirm do save or insert.
 * @param t 
 * @return 
 */
public WriteResult save(T t){
    WriteResult wr;
    BuguEntity ent = (BuguEntity)t;
    if(StringUtil.isEmpty(ent.getId())){
        wr = insert(t);
    }
    else{
        Field idField = FieldsCache.getInstance().getIdField(clazz);
        Id idAnnotation = idField.getAnnotation(Id.class);
        if(idAnnotation.type()==IdType.USER_DEFINE){
            if(this.exists(Operator.ID, ent.getId())){
                wr = doSave(ent);
            }else{
                wr = insert(t);
            }
        }
        else{
            wr = doSave(ent);
        }
    }
    return wr;
}
 

@Override
public Boolean updateRetry(final String id, final Integer retry, final String appName) {
    if (StringUtils.isBlank(id) || StringUtils.isBlank(appName) || Objects.isNull(retry)) {
        return Boolean.FALSE;
    }
    final String mongoTableName = RepositoryPathUtils.buildMongoTableName(appName);
    Query query = new Query();
    query.addCriteria(new Criteria("transId").is(id));
    Update update = new Update();
    update.set("lastTime", DateUtils.getCurrentDateTime());
    update.set("retriedCount", retry);
    final WriteResult writeResult = mongoTemplate.updateFirst(query, update,
            MongoAdapter.class, mongoTableName);
    if (writeResult.getN() <= 0) {
        throw new RuntimeException("更新数据异常!");
    }
    return Boolean.TRUE;
}
 

@Override
public Boolean updateRetry(final String id, final Integer retry, final String applicationName) {
    if (StringUtils.isBlank(id)
            || StringUtils.isBlank(applicationName)
            || Objects.isNull(retry)) {
        return Boolean.FALSE;
    }
    final String mongoTableName = RepositoryPathUtils.buildMongoTableName(applicationName);

    Query query = new Query();
    query.addCriteria(new Criteria("transId").is(id));
    Update update = new Update();
    update.set("lastTime", DateUtils.getCurrentDateTime());
    update.set("retriedCount", retry);
    final WriteResult writeResult = mongoTemplate.updateFirst(query, update,
            MongoAdapter.class, mongoTableName);
    if (writeResult.getN() <= 0) {
        throw new TransactionRuntimeException("更新数据异常!");
    }
    return Boolean.TRUE;
}
 

/**
 * set the state of balancer.
 *
 * @param dbConn
 * @param state
 * @return Error description, or null if no errors
 */
private static String setBalancerState(DB dbConn, boolean state) {
    DBObject balancer = new BasicDBObject(ID, "balancer");
    DBObject updateObj = new BasicDBObject();
    String stopped = state ? "false" : "true";
    updateObj.put("$set", new BasicDBObject("stopped", stopped));
    WriteResult wresult = dbConn.getSisterDB("config").getCollection("settings").update(balancer, updateObj, true, false);
    if (wresult != null) {
        CommandResult result = wresult.getLastError();
        if (!result.ok()) {
            LOG.error("Error setting balancer state to {}: {}", state, result.getErrorMessage());
            return result.getErrorMessage();
        }
    }
    return null;
}
 

/**
 * 更新 TM中的消息状态
 *
 * @param transactionMsg
 */
@Override
public Boolean updateTxTransactionMsgStatus(TransactionMsg transactionMsg) {
    try {
        Query query = new Query();
        query.addCriteria(new Criteria("groupId").is(transactionMsg.getGroupId()).and("subTaskId").is(transactionMsg.getSubTaskId()));
        Update update = Update.update("consumed", transactionMsg.getConsumed());
        String message = transactionMsg.getMessage();
        if (StringUtils.isNotBlank(message)) {
            update.set("message", message);
        }
        update.set("updateTime", Timestamp.valueOf(DateUtils.getCurrentDateTime()).getTime());
        final WriteResult writeResult = mongoTemplate.updateFirst(query, update, TransactionMsg.class, CollectionNameEnum.TransactionMsg.name());
        return writeResult.getN() > 0;
    } catch (Exception e) {
        //TODO 处理异常
        LogUtil.error(LOGGER, e::getLocalizedMessage);
        return false;
    }
}
 

private WriteResult execute(DBObject condition){
    List ids = null;
    if(dao.hasCustomListener){
        ids = dao.getCollection().distinct(Operator.ID, condition);
    }
    if(isolated){
        condition.put(Operator.ISOLATED, 1);
    }
    WriteResult wr = dao.getCollection().update(condition, modifier, upsert, multi);
    if(dao.hasCustomListener && ids != null){
        DBObject in = new BasicDBObject(Operator.IN, ids);
        DBCursor cursor = dao.getCollection().find(new BasicDBObject(Operator.ID, in));
        List<T> list = MapperUtil.toList(dao.getEntityClass(), cursor);
        for(T t : list){
            dao.notifyUpdated((BuguEntity)t);
        }
    }
    return wr;
}
 

@Override
public boolean remove(String taskTrackerNodeGroup, String jobId) {
    String tableName = JobQueueUtils.getExecutableQueueName(taskTrackerNodeGroup);
    Query<JobPo> query = template.createQuery(tableName, JobPo.class);
    query.field("jobId").equal(jobId);
    WriteResult wr = template.delete(query);
    return wr.getN() == 1;
}
 

@Override
public WriteResult changeUserInfo(UserDO userDO) {
    Query query = new Query();
    query.addCriteria(new Criteria("username").is(userDO.getUsername()));
    Update update = new Update();
    if (userDO.getBlogLink()!=null){
        update.set("blogLink", userDO.getBlogLink());
    }
    if (userDO.getAvatar()!=null){
        update.set("avatar", userDO.getAvatar());
    }
    return this.mongoTemplate.upsert(query, update, UserDO.class);
}
 

@Test public void doBatchWithRetry() throws Exception {
  setupReturns();
  setupRowMeta();
  dbOutput.m_batch = new ArrayList<DBObject>();
  dbOutput.m_batch.add( new BasicDBObject( ImmutableMap.of( "foo", "fooval", "bar", "barval", "baz", "bazval" ) ) );
  List<Object[]> batchRows = new ArrayList<Object[]>();
  batchRows.add( rowData );

  List<DBObject> batchCopy = new ArrayList( dbOutput.m_batch );

  dbOutput.m_batchRows = batchRows;
  when( stepMetaInterface.getWriteRetries() ).thenReturn( "1" );
  when( stepMetaInterface.getWriteRetryDelay() ).thenReturn( "0" );
  WriteResult result = mock( WriteResult.class );
  CommandResult commandResult = mock( CommandResult.class );
  when( commandResult.ok() ).thenReturn( true );
  when( mongoCollectionWrapper.save( dbOutput.m_batch.get( 0 ) ) ).thenReturn( result );

  doThrow( mock( MongoException.class ) ).when( mongoCollectionWrapper ).insert( anyList() );
  dbOutput.init( stepMetaInterface, stepDataInterace );
  dbOutput.doBatch();

  // should attempt insert once, falling back to save on retry
  verify( mongoCollectionWrapper, times( 1 ) ).insert( anyList() );
  verify( mongoCollectionWrapper, times( 1 ) ).save( batchCopy.get( 0 ) );

  // batch should be cleared.
  assertThat( dbOutput.m_batch.size(), equalTo( 0 ) );
  assertThat( dbOutput.m_batchRows.size(), equalTo( 0 ) );
}
 

@Test public void testUpdate() throws Exception {
  setupReturns();
  WriteResult result = mock( WriteResult.class );
  CommandResult commandResult = mock( CommandResult.class );
  when( commandResult.ok() ).thenReturn( true );
  when( mongoCollectionWrapper.update( any( DBObject.class ), any( DBObject.class ), anyBoolean(), anyBoolean() ) )
    .thenReturn( result );
  when( stepMetaInterface.getUpdate() ).thenReturn( true );

  // flag a field for update = "foo"
  MongoDbOutputMeta.MongoField mongoField = mongoFields.get( 0 );
  mongoField.m_updateMatchField = true;

  setupRowMeta();
  dbOutput.init( stepMetaInterface, stepDataInterace );
  assertTrue( dbOutput.processRow( stepMetaInterface, stepDataInterace ) );
  ArgumentCaptor<BasicDBObject> updateQueryCaptor = ArgumentCaptor.forClass( BasicDBObject.class );
  ArgumentCaptor<BasicDBObject> insertCaptor = ArgumentCaptor.forClass( BasicDBObject.class );

  // update is executed
  verify( mongoCollectionWrapper )
    .update( updateQueryCaptor.capture(), insertCaptor.capture(), anyBoolean(), anyBoolean() );
  // updated field is expected
  assertThat( updateQueryCaptor.getValue(), equalTo( new BasicDBObject( "foo", "foo" ) ) );
  // insert document is expected
  assertThat( insertCaptor.getValue(),
    equalTo( new BasicDBObject( ( ImmutableMap.of( "foo", "foo", "bar", "bar", "baz", "baz" ) ) ) ) );
}
 

public void sendJSON(Object obj, HttpServletRequest request, HttpServletResponse response) throws IOException {
	boolean gzip = false;
	if (request != null) {
		String acceptedEncodings = request.getHeader("accept-encoding");
		gzip = acceptedEncodings != null && acceptedEncodings.indexOf("gzip") != -1;
	}
	response.setCharacterEncoding("UTF-8");
	response.setContentType("application/json");
	if (obj instanceof WriteResult) {
		String error = ((WriteResult) obj).getError();
		if (error != null) {
			obj = error;
		} else {
			obj = "OK";
		}
	}
	OutputStream os = null;
	try {
		byte data[] = JSON.serialize(obj).getBytes("UTF-8");
		os = response.getOutputStream();
		if (gzip && data.length >= 860) {
			response.setHeader("Content-Encoding", "gzip");
			GZIPOutputStream gzos = new GZIPOutputStream(os);
			gzos.write(data);
			gzos.finish();
			gzos.close();
		} else {
			os.write(data);
		}
	} catch (Exception e) {
		e.printStackTrace();
		response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR, e.getMessage());
	} finally {
		if (os != null) {
			os.close();
		}
	}
}
 

@Override
public WriteResult changePassword( final String username, final String password) {
    Query query = new Query();
    query.addCriteria(new Criteria("username").is(username));
    Update update = new Update();
    update.set("password", password);
    update.set("lastPasswordResetDate", new Date());
    return this.mongoTemplate.updateFirst(query, update, UserDO.class);
}
 

@Override
public void endWindow()
{
  logger.debug("mongo datalist size: " + dataList.size());
  if (dataList.size() > 0) {
    WriteResult result = dbCollection.insert(dataList, writeConcern);
    logger.debug("Result for MongoDB insert: " + result);
    dataList.clear();
  }
}
 

@Override
public boolean remove(String jobId) {
    Query<JobPo> query = template.createQuery(JobPo.class);
    query.field("jobId").equal(jobId);
    WriteResult wr = template.delete(query);
    return wr.getN() == 1;
}
 

@Override
public boolean remove(String jobClientNodeGroup, String id) {
    Query<JobFeedbackPo> query = createQuery(jobClientNodeGroup);
    query.field("id").equal(id);
    WriteResult wr = template.delete(query);
    return wr.getN() == 1;
}
 

@Override
public boolean remove(String jobId) {
    Query<JobPo> query = template.createQuery(JobPo.class);
    query.field("jobId").equal(jobId);
    WriteResult wr = template.delete(query);
    return wr.getN() == 1;
}
 

@Override
public boolean remove(String jobId) {
    Query<JobPo> query = template.createQuery(JobPo.class);
    query.field("jobId").equal(jobId);
    WriteResult wr = template.delete(query);
    return wr.getN() == 1;
}
 

/**
 * execute the update operation on a single entity.
 * @param id
 * @return 
 */
public WriteResult execute(String id){
    Class<T> clazz = dao.getEntityClass();
    DBObject condition = new BasicDBObject(Operator.ID, IdUtil.toDbId(clazz, id));
    WriteResult wr = dao.getCollection().update(condition, modifier, upsert, false); //update one
    if(dao.hasCustomListener){
        BuguEntity entity = (BuguEntity)dao.findOne(id);
        dao.notifyUpdated(entity);
    }
    return wr;
}
 

@Override
public boolean remove(String jobId) {
    Query<JobPo> query = template.createQuery(JobPo.class);
    query.field("jobId").equal(jobId);
    WriteResult wr = template.delete(query);
    return wr.getN() == 1;
}
 

@Override
public <T> int update(Class<T> entityClass, String[] fields, Object[] fieldValues, Serializable id) {
	WriteResult result = operations.updateFirst(new Query(where("_id").is(id)), setFieldValue(fields, fieldValues), entityClass);
	CommandResult commandResult = result.getLastError();
	if(commandResult.ok()){
		return 1;
	}
	throw new DaoException(commandResult.getErrorMessage());
}
 

@Before
public void setup() {
    mockCollection = Mockito.mock(DBCollection.class);
    writeResult = Mockito.mock(WriteResult.class);
    commandResult = Mockito.mock(CommandResult.class);
    testAccessor = new ContainerDocumentAccessor(generatorStrategy, naturalKeyExtractor, mongoTemplate, schemaRepo);
    MockitoAnnotations.initMocks(this);
    when(mockHolder.getContainerDocument(ATTENDANCE)).thenReturn(createContainerDocAttendance());
    entity = createAttendanceEntity();
    when(writeResult.getLastError()).thenReturn(commandResult);
    when(commandResult.ok()).thenReturn(true);
}
 

private WriteResult doSave(BuguEntity ent){
    WriteResult wr = getCollection().save(MapperUtil.toDBObject(ent));
    if(hasCustomListener){
        notifyUpdated(ent);
    }
    return wr;
}
 

/**
 * This method deletes the entity from the collection for a given key.
 *
 * @param tableName Name of the table
 * @param entity    Entity
 * @throws ODataServiceFault
 */
@Override
public boolean deleteEntityInTable(String tableName, ODataEntry entity) throws ODataServiceFault {
    String documentId = entity.getValue(DOCUMENT_ID);
    WriteResult delete = jongo.getCollection(tableName).remove(new ObjectId(documentId));
    int wasDeleted = delete.getN();
    if (wasDeleted == 1) {
        return delete.wasAcknowledged();
    } else {
        throw new ODataServiceFault("Document ID: " + documentId + " does not exist in "
                + "collection: " + tableName + ".");
    }
}
 

/**
 * Batch remove by id.
 * @param idList
 * @return 
 */
public WriteResult remove(List<String> idList){
    int len = idList.size();
    Object[] arr = new Object[len];
    for(int i=0; i<len; i++){
        arr[i] = IdUtil.toDbId(clazz, idList.get(i));
    }
    DBObject in = new BasicDBObject(Operator.IN, arr);
    return removeMulti(new BasicDBObject(Operator.ID, in));
}
 

private WriteResult removeMulti(DBObject condition){
    List<T> list = null;
    if(!listenerList.isEmpty()){
        DBCursor cursor = getCollection().find(condition);
        list = MapperUtil.toList(clazz, cursor);
    }
    WriteResult wr = getCollection().remove(condition);
    if(!listenerList.isEmpty() && list!=null){
        for(T t : list){
            notifyDeleted((BuguEntity)t);
        }
    }
    return wr;
}
 

private WriteResult doInsertWithoutCascade(T t, boolean withoutCascade){
    DBObject dbo = MapperUtil.toDBObject(t, withoutCascade);
    WriteResult wr = getCollection().insert(dbo);
    String id = dbo.get(Operator.ID).toString();
    BuguEntity ent = (BuguEntity)t;
    ent.setId(id);
    if(hasCustomListener){
        notifyInserted(ent);
    }
    return wr;
}