Search in sources :

Example 31 with BatchResponse

use of com.linkedin.restli.common.BatchResponse in project rest.li by linkedin.

the class TestScatterGather method testBuildSGGetRequests.

private static void testBuildSGGetRequests(int numEndpoints, ScatterGatherBuilder<Greeting> sg, Long[] ids) throws ServiceUnavailableException {
    Collection<ScatterGatherBuilder.RequestInfo<Greeting>> requests = buildScatterGatherGetRequests(sg, ids);
    Assert.assertEquals(requests.size(), numEndpoints);
    Set<Set<String>> requestIdSets = new HashSet<>();
    Set<Long> requestIds = new HashSet<>();
    for (ScatterGatherBuilder.RequestInfo<Greeting> requestInfo : requests) {
        // URI will be something like "greetings/?ids=21&ids=4&ids=53&ids=60&ids=66&ids=88&ids=93&foo=bar"
        BatchRequest<BatchResponse<Greeting>> request = requestInfo.getBatchRequest();
        Set<String> expectedParams = new HashSet<>();
        expectedParams.add(RestConstants.QUERY_BATCH_IDS_PARAM);
        expectedParams.add("foo");
        expectedParams.add(RestConstants.FIELDS_PARAM);
        Set<PathSpec> expectedFields = Collections.singleton(new PathSpec("message"));
        testRequest(request, expectedParams, expectedFields, null, requestIdSets, requestIds);
    }
    Assert.assertTrue(requestIds.containsAll(Arrays.asList(ids)));
    Assert.assertEquals(requestIds.size(), ids.length);
}
Also used : Greeting(com.linkedin.restli.examples.greetings.api.Greeting) Set(java.util.Set) HashSet(java.util.HashSet) BatchResponse(com.linkedin.restli.common.BatchResponse) PathSpec(com.linkedin.data.schema.PathSpec) HashSet(java.util.HashSet)

Example 32 with BatchResponse

use of com.linkedin.restli.common.BatchResponse 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);
        }
    });
}
Also used : ResponseImpl(com.linkedin.restli.internal.client.ResponseImpl) ResourceMethod(com.linkedin.restli.common.ResourceMethod) LoggerFactory(org.slf4j.LoggerFactory) HashMap(java.util.HashMap) Callback(com.linkedin.common.callback.Callback) Function(java.util.function.Function) BatchEntityResponse(com.linkedin.restli.internal.client.response.BatchEntityResponse) HashSet(java.util.HashSet) HttpCookie(java.net.HttpCookie) DataMap(com.linkedin.data.DataMap) BatchKVResponse(com.linkedin.restli.client.response.BatchKVResponse) RestConstants(com.linkedin.restli.common.RestConstants) Map(java.util.Map) EntityResponse(com.linkedin.restli.common.EntityResponse) PathSpec(com.linkedin.data.schema.PathSpec) RestResponseBuilder(com.linkedin.r2.message.rest.RestResponseBuilder) ProtocolVersionUtil(com.linkedin.restli.internal.common.ProtocolVersionUtil) RecordTemplate(com.linkedin.data.template.RecordTemplate) ProtocolVersion(com.linkedin.restli.common.ProtocolVersion) Logger(org.slf4j.Logger) ResourceSpec(com.linkedin.restli.common.ResourceSpec) BatchEntry(com.linkedin.parseq.batching.BatchImpl.BatchEntry) RemoteInvocationException(com.linkedin.r2.RemoteInvocationException) Set(java.util.Set) Batch(com.linkedin.parseq.batching.Batch) Tuple3(com.linkedin.parseq.function.Tuple3) Tuples(com.linkedin.parseq.function.Tuples) Collectors(java.util.stream.Collectors) Objects(java.util.Objects) List(java.util.List) RequestContext(com.linkedin.r2.message.RequestContext) ErrorResponse(com.linkedin.restli.common.ErrorResponse) HttpStatus(com.linkedin.restli.common.HttpStatus) Entry(java.util.Map.Entry) ResponseUtils(com.linkedin.restli.internal.common.ResponseUtils) BatchResponse(com.linkedin.restli.common.BatchResponse) HashSet(java.util.HashSet) Set(java.util.Set) BatchResponse(com.linkedin.restli.common.BatchResponse) ProtocolVersion(com.linkedin.restli.common.ProtocolVersion) PathSpec(com.linkedin.data.schema.PathSpec) BatchKVResponse(com.linkedin.restli.client.response.BatchKVResponse) DataMap(com.linkedin.data.DataMap) RemoteInvocationException(com.linkedin.r2.RemoteInvocationException) BatchEntry(com.linkedin.parseq.batching.BatchImpl.BatchEntry) ResponseImpl(com.linkedin.restli.internal.client.ResponseImpl) BatchEntityResponse(com.linkedin.restli.internal.client.response.BatchEntityResponse) BatchKVResponse(com.linkedin.restli.client.response.BatchKVResponse) EntityResponse(com.linkedin.restli.common.EntityResponse) ErrorResponse(com.linkedin.restli.common.ErrorResponse) BatchResponse(com.linkedin.restli.common.BatchResponse) BatchEntityResponse(com.linkedin.restli.internal.client.response.BatchEntityResponse) EntityResponse(com.linkedin.restli.common.EntityResponse)

Aggregations

BatchResponse (com.linkedin.restli.common.BatchResponse)32 Test (org.testng.annotations.Test)23 HashMap (java.util.HashMap)11 DataMap (com.linkedin.data.DataMap)10 Greeting (com.linkedin.restli.examples.greetings.api.Greeting)9 RoutingResult (com.linkedin.restli.internal.server.RoutingResult)9 ServerResourceContext (com.linkedin.restli.internal.server.ServerResourceContext)9 ResourceMethodDescriptor (com.linkedin.restli.internal.server.model.ResourceMethodDescriptor)9 Map (java.util.Map)8 BatchKVResponse (com.linkedin.restli.client.response.BatchKVResponse)6 ErrorResponse (com.linkedin.restli.common.ErrorResponse)6 UpdateStatus (com.linkedin.restli.common.UpdateStatus)6 ConcurrentHashMap (java.util.concurrent.ConcurrentHashMap)6 GreetingsBuilders (com.linkedin.restli.examples.greetings.client.GreetingsBuilders)5 EntityResponse (com.linkedin.restli.common.EntityResponse)4 HashSet (java.util.HashSet)4 Callback (com.linkedin.common.callback.Callback)3 Foo (com.linkedin.pegasus.generator.examples.Foo)3 Set (java.util.Set)3 PathSpec (com.linkedin.data.schema.PathSpec)2