Examples with Batch com.linkedin.parseq.batching.Batch used on opensource projects

Example 1 with Batch

use of com.linkedin.parseq.batching.Batch 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);
        }
    });
}