Examples with BatchEntityResponse com.linkedin.restli.internal.client.response.BatchEntityResponse used on opensource projects

Example 1 with BatchEntityResponse

use of com.linkedin.restli.internal.client.response.BatchEntityResponse in project rest.li by linkedin.

the class BatchEntityResponseDecoder method wrapResponse.

@Override
public BatchKVResponse<K, EntityResponse<V>> wrapResponse(DataMap dataMap, Map<String, String> headers, ProtocolVersion version) throws InstantiationException, IllegalAccessException, InvocationTargetException, NoSuchMethodException, IOException {
    if (dataMap == null) {
        return null;
    }
    final DataMap mergedResults = new DataMap();
    final DataMap inputResults = dataMap.containsKey(BatchResponse.RESULTS) ? dataMap.getDataMap(BatchResponse.RESULTS) : new DataMap();
    final DataMap inputStatuses = dataMap.containsKey(BatchResponse.STATUSES) ? dataMap.getDataMap(BatchResponse.STATUSES) : new DataMap();
    final DataMap inputErrors = dataMap.containsKey(BatchResponse.ERRORS) ? dataMap.getDataMap(BatchResponse.ERRORS) : new DataMap();
    final Set<String> mergedKeys = new HashSet<String>(inputResults.keySet());
    mergedKeys.addAll(inputStatuses.keySet());
    mergedKeys.addAll(inputErrors.keySet());
    for (String key : mergedKeys) {
        final DataMap entityResponseData = new DataMap();
        final Object entityData = inputResults.get(key);
        if (entityData != null) {
            CheckedUtil.putWithoutChecking(entityResponseData, EntityResponse.ENTITY, entityData);
        }
        final Object statusData = inputStatuses.get(key);
        if (statusData != null) {
            CheckedUtil.putWithoutChecking(entityResponseData, EntityResponse.STATUS, statusData);
        }
        final Object errorData = inputErrors.get(key);
        if (errorData != null) {
            CheckedUtil.putWithoutChecking(entityResponseData, EntityResponse.ERROR, errorData);
        }
        CheckedUtil.putWithoutChecking(mergedResults, key, entityResponseData);
    }
    final DataMap responseData = new DataMap();
    CheckedUtil.putWithoutChecking(responseData, BatchKVResponse.RESULTS, mergedResults);
    CheckedUtil.putWithoutChecking(responseData, BatchKVResponse.ERRORS, inputErrors);
    return new BatchEntityResponse<K, V>(responseData, _keyType, _entityType, _keyParts, _complexKeyType, version);
}

Example 2 with BatchEntityResponse

use of com.linkedin.restli.internal.client.response.BatchEntityResponse in project rest.li by linkedin.

the class MockBatchEntityResponseFactory method createWithCustomTyperefKey.

/**
 * Creates a {@link BatchKVResponse} where the key is a typeref to a custom Java class.
 *
 * @param keyClass the custom Java class
 * @param typerefClass the typeref class (the generated class that extends {@link RecordTemplate})
 * @param valueClass class for the value
 * @param recordTemplates the data that will be returned for a call to {@link com.linkedin.restli.client.response.BatchKVResponse#getResults()}
 * @param statuses The HTTP status codes that will be returned as part of {@link EntityResponse}s returned in {@link com.linkedin.restli.client.response.BatchKVResponse#getResults()}
 * @param errorResponses the data that will be returned for a call to {@link com.linkedin.restli.client.response.BatchKVResponse#getErrors()}
 * @param <K>
 * @param <TK>
 * @param <V>
 */
@SuppressWarnings({ "unchecked", "rawtypes" })
public static <K, TK, V extends RecordTemplate> BatchKVResponse<K, EntityResponse<V>> createWithCustomTyperefKey(Class<K> keyClass, Class<TK> typerefClass, Class<V> valueClass, Map<K, V> recordTemplates, Map<K, HttpStatus> statuses, Map<K, ErrorResponse> errorResponses) {
    ProtocolVersion version = AllProtocolVersions.BASELINE_PROTOCOL_VERSION;
    DataMap batchResponseDataMap = buildDataMap(recordTemplates, statuses, errorResponses, version);
    return new BatchEntityResponse(batchResponseDataMap, TypeSpec.forClassMaybeNull(typerefClass), TypeSpec.forClassMaybeNull(valueClass), Collections.<String, CompoundKey.TypeInfo>emptyMap(), null, version);
}

Example 3 with BatchEntityResponse

use of com.linkedin.restli.internal.client.response.BatchEntityResponse in project parseq by linkedin.

the class GetRequestGroup method doExecuteBatchGet.

@SuppressWarnings({ "rawtypes", "unchecked" })
private <K, RT extends RecordTemplate> void doExecuteBatchGet(final Client client, final Batch<RestRequestBatchKey, Response<Object>> batch, final Set<Object> ids, final Set<PathSpec> fields, Function<Request<?>, RequestContext> requestContextProvider) {
    final BatchGetEntityRequestBuilder<K, RT> builder = new BatchGetEntityRequestBuilder<>(_baseUriTemplate, _resourceSpec, _requestOptions);
    builder.setHeaders(_headers);
    builder.setCookies(_cookies);
    _queryParams.forEach((key, value) -> builder.setParam(key, value));
    _pathKeys.forEach((key, value) -> builder.pathKey(key, value));
    builder.ids((Set<K>) ids);
    if (fields != null && !fields.isEmpty()) {
        builder.fields(fields.toArray(new PathSpec[fields.size()]));
    }
    final BatchGetEntityRequest<K, RT> batchGet = builder.build();
    client.sendRequest(batchGet, requestContextProvider.apply(batchGet), new Callback<Response<BatchKVResponse<K, EntityResponse<RT>>>>() {

        @Override
        public void onSuccess(Response<BatchKVResponse<K, EntityResponse<RT>>> responseToBatch) {
            final ProtocolVersion version = ProtocolVersionUtil.extractProtocolVersion(responseToBatch.getHeaders());
            batch.entries().stream().forEach(entry -> {
                try {
                    RestRequestBatchKey rrbk = entry.getKey();
                    Request request = rrbk.getRequest();
                    if (request instanceof GetRequest) {
                        successGet((GetRequest) request, responseToBatch, batchGet, entry, version);
                    } else if (request instanceof BatchGetKVRequest) {
                        successBatchGetKV((BatchGetKVRequest) request, responseToBatch, entry, version);
                    } else if (request instanceof BatchGetRequest) {
                        successBatchGet((BatchGetRequest) request, responseToBatch, entry, version);
                    } else if (request instanceof BatchGetEntityRequest) {
                        successBatchGetEntity((BatchGetEntityRequest) request, responseToBatch, entry, version);
                    } else {
                        entry.getValue().getPromise().fail(unsupportedGetRequestType(request));
                    }
                } catch (RemoteInvocationException e) {
                    entry.getValue().getPromise().fail(e);
                }
            });
        }

        @SuppressWarnings({ "deprecation" })
        private void successBatchGetEntity(BatchGetEntityRequest request, Response<BatchKVResponse<K, EntityResponse<RT>>> responseToBatch, Entry<RestRequestBatchKey, BatchEntry<Response<Object>>> entry, final ProtocolVersion version) {
            Set<String> ids = (Set<String>) request.getObjectIds().stream().map(o -> BatchResponse.keyToString(o, version)).collect(Collectors.toSet());
            DataMap dm = filterIdsInBatchResult(responseToBatch.getEntity().data(), ids);
            BatchKVResponse br = new BatchEntityResponse<>(dm, request.getResourceSpec().getKeyType(), request.getResourceSpec().getValueType(), request.getResourceSpec().getKeyParts(), request.getResourceSpec().getComplexKeyType(), version);
            Response rsp = new ResponseImpl(responseToBatch, br);
            entry.getValue().getPromise().done(rsp);
        }

        private void successBatchGet(BatchGetRequest request, Response<BatchKVResponse<K, EntityResponse<RT>>> responseToBatch, Entry<RestRequestBatchKey, BatchEntry<Response<Object>>> entry, final ProtocolVersion version) {
            Set<String> ids = (Set<String>) request.getObjectIds().stream().map(o -> BatchResponse.keyToString(o, version)).collect(Collectors.toSet());
            DataMap dm = filterIdsInBatchResult(responseToBatch.getEntity().data(), ids);
            BatchResponse br = new BatchResponse<>(dm, request.getResponseDecoder().getEntityClass());
            Response rsp = new ResponseImpl(responseToBatch, br);
            entry.getValue().getPromise().done(rsp);
        }

        @SuppressWarnings({ "deprecation" })
        private void successBatchGetKV(BatchGetKVRequest request, Response<BatchKVResponse<K, EntityResponse<RT>>> responseToBatch, Entry<RestRequestBatchKey, BatchEntry<Response<Object>>> entry, final ProtocolVersion version) {
            Set<String> ids = (Set<String>) request.getObjectIds().stream().map(o -> BatchResponse.keyToString(o, version)).collect(Collectors.toSet());
            DataMap dm = filterIdsInBatchResult(responseToBatch.getEntity().data(), ids);
            BatchKVResponse br = new BatchKVResponse(dm, request.getResourceSpec().getKeyType(), request.getResourceSpec().getValueType(), request.getResourceSpec().getKeyParts(), request.getResourceSpec().getComplexKeyType(), version);
            Response rsp = new ResponseImpl(responseToBatch, br);
            entry.getValue().getPromise().done(rsp);
        }

        @SuppressWarnings({ "deprecation" })
        private void successGet(GetRequest request, Response<BatchKVResponse<K, EntityResponse<RT>>> responseToBatch, final BatchGetEntityRequest<K, RT> batchGet, Entry<RestRequestBatchKey, BatchEntry<Response<Object>>> entry, final ProtocolVersion version) throws RemoteInvocationException {
            String idString = BatchResponse.keyToString(request.getObjectId(), version);
            Object id = ResponseUtils.convertKey(idString, request.getResourceSpec().getKeyType(), request.getResourceSpec().getKeyParts(), request.getResourceSpec().getComplexKeyType(), version);
            Response rsp = unbatchResponse(batchGet, responseToBatch, id);
            entry.getValue().getPromise().done(rsp);
        }

        @Override
        public void onError(Throwable e) {
            batch.failAll(e);
        }
    });
}

Example 4 with BatchEntityResponse

use of com.linkedin.restli.internal.client.response.BatchEntityResponse in project parseq by linkedin.

the class TestRequest404WithBatching method remove404.

private Object remove404(Object o) {
    if (o instanceof Response) {
        Response r = (Response) o;
        Object entity = r.getEntity();
        if (entity instanceof BatchEntityResponse) {
            BatchEntityResponse ber = (BatchEntityResponse) entity;
            DataMap data = ber.data();
            DataMap errors = (DataMap) data.getDataMap("errors");
            Set<String> keys = new HashSet<>(errors.keySet());
            keys.forEach(key -> {
                DataMap error = errors.getDataMap(key);
                if (error.getInteger("status").equals(404)) {
                    errors.remove(key);
                }
            });
        }
    }
    return o;
}

Example 5 with BatchEntityResponse

use of com.linkedin.restli.internal.client.response.BatchEntityResponse in project rest.li by linkedin.

the class TestDefaultScatterGatherStrategy method createBatchEntityResponse.

private static Response<BatchKVResponse<Long, EntityResponse<TestRecord>>> createBatchEntityResponse(ProtocolVersion version, Set<Long> resultKeys, Set<Long> errorKeys) {
    DataMap resultMap = new DataMap();
    for (Long id : resultKeys) {
        resultMap.put(id.toString(), new EntityResponse<>(TestRecord.class).setEntity(new TestRecord().setId(id)).setStatus(HttpStatus.S_200_OK).data());
    }
    DataMap errorMap = new DataMap();
    for (Long id : errorKeys) {
        errorMap.put(id.toString(), new ErrorResponse().setStatus(HttpStatus.S_404_NOT_FOUND.getCode()).data());
    }
    DataMap responseMap = new DataMap();
    responseMap.put(BatchResponse.RESULTS, resultMap);
    responseMap.put(BatchResponse.ERRORS, errorMap);
    BatchEntityResponse<Long, TestRecord> response = new BatchEntityResponse<>(responseMap, new TypeSpec<>(Long.class), new TypeSpec<>(TestRecord.class), Collections.emptyMap(), null, version);
    return new ResponseImpl<>(HttpStatus.S_200_OK.getCode(), Collections.emptyMap(), Collections.emptyList(), response, null);
}