Examples with RequestContext com.linkedin.r2.message.RequestContext used on opensource projects

Example 61 with RequestContext

use of com.linkedin.r2.message.RequestContext in project rest.li by linkedin.

the class ScatterGatherBuilder method buildRequestsKV.

@SuppressWarnings({ "unchecked", "deprecation" })
public <K> KVScatterGatherResult<K, T> buildRequestsKV(BatchGetKVRequest<K, T> request, RequestContext requestContext) throws ServiceUnavailableException {
    Set<K> idObjects = (Set<K>) request.getObjectIds();
    MapKeyResult<URI, K> mapKeyResult = mapKeys(request, idObjects);
    Map<URI, Collection<K>> batches = mapKeyResult.getMapResult();
    Collection<KVRequestInfo<K, T>> scatterGatherRequests = new ArrayList<KVRequestInfo<K, T>>(batches.size());
    for (Map.Entry<URI, Collection<K>> batch : batches.entrySet()) {
        BatchGetRequestBuilder<K, T> builder = new BatchGetRequestBuilder<K, T>(request.getBaseUriTemplate(), (Class<T>) request.getResourceProperties().getValueType().getType(), request.getResourceSpec(), request.getRequestOptions());
        builder.ids(batch.getValue());
        for (Map.Entry<String, Object> param : request.getQueryParamsObjects().entrySet()) {
            if (!param.getKey().equals(RestConstants.QUERY_BATCH_IDS_PARAM)) {
                // keep all non-batch query parameters since they could be request specific
                builder.setParam(param.getKey(), param.getValue());
            }
        }
        for (Map.Entry<String, String> header : request.getHeaders().entrySet()) {
            builder.setHeader(header.getKey(), header.getValue());
        }
        RequestContext context = requestContext.clone();
        KeyMapper.TargetHostHints.setRequestContextTargetHost(context, batch.getKey());
        scatterGatherRequests.add(new KVRequestInfo<K, T>(builder.buildKV(), context));
    }
    return new KVScatterGatherResult<K, T>(scatterGatherRequests, mapKeyResult.getUnmappedKeys());
}

Example 62 with RequestContext

use of com.linkedin.r2.message.RequestContext in project rest.li by linkedin.

the class ScatterGatherBuilder method buildRequests.

@SuppressWarnings("deprecation")
public <K> KVScatterGatherResult<K, UpdateStatus> buildRequests(BatchDeleteRequest<K, T> request, RequestContext requestContext) throws ServiceUnavailableException {
    Set<Object> idObjects = request.getObjectIds();
    Collection<K> ids = new HashSet<K>(idObjects.size());
    for (Object o : idObjects) {
        @SuppressWarnings("unchecked") K k = (K) o;
        ids.add(k);
    }
    MapKeyResult<URI, K> mapKeyResult = mapKeys(request, ids);
    Map<URI, Collection<K>> batches = mapKeyResult.getMapResult();
    Collection<KVRequestInfo<K, UpdateStatus>> scatterGatherRequests = new ArrayList<KVRequestInfo<K, UpdateStatus>>(batches.size());
    for (Map.Entry<URI, Collection<K>> batch : batches.entrySet()) {
        TypeSpec<? extends RecordTemplate> value = request.getResourceProperties().getValueType();
        @SuppressWarnings("unchecked") Class<T> valueClass = (Class<T>) ((value == null) ? null : value.getType());
        BatchDeleteRequestBuilder<K, T> builder = new BatchDeleteRequestBuilder<K, T>(request.getBaseUriTemplate(), valueClass, request.getResourceSpec(), request.getRequestOptions());
        builder.ids(batch.getValue());
        for (Map.Entry<String, Object> param : request.getQueryParamsObjects().entrySet()) {
            if (!param.getKey().equals(RestConstants.QUERY_BATCH_IDS_PARAM)) {
                builder.setParam(param.getKey(), param.getValue());
            }
        }
        for (Map.Entry<String, String> header : request.getHeaders().entrySet()) {
            builder.setHeader(header.getKey(), header.getValue());
        }
        RequestContext context = requestContext.clone();
        KeyMapper.TargetHostHints.setRequestContextTargetHost(context, batch.getKey());
        BatchRequest<BatchKVResponse<K, UpdateStatus>> build = builder.build();
        scatterGatherRequests.add(new KVRequestInfo<K, UpdateStatus>(build, context));
    }
    return new KVScatterGatherResult<K, UpdateStatus>(scatterGatherRequests, mapKeyResult.getUnmappedKeys());
}

Example 63 with RequestContext

use of com.linkedin.r2.message.RequestContext in project rest.li by linkedin.

the class TestComplexKeysResource method testBatchGetEntityEmpty.

@SuppressWarnings("deprecation")
public void testBatchGetEntityEmpty(BatchGetEntityRequestBuilder<ComplexResourceKey<TwoPartKey, TwoPartKey>, Message> builder) throws Exception {
    final Request<BatchKVResponse<ComplexResourceKey<TwoPartKey, TwoPartKey>, EntityResponse<Message>>> request = builder.build();
    final FutureCallback<RestResponse> callback = new FutureCallback<RestResponse>();
    getClient().sendRestRequest(request, new RequestContext(), callback);
    final RestResponse result = callback.get();
    final DataMap responseMap = DataMapUtils.readMap(result);
    final DataMap resultsMap = responseMap.getDataMap(BatchKVResponse.RESULTS);
    Assert.assertNotNull(resultsMap, "Response does not contain results map");
    Assert.assertTrue(resultsMap.isEmpty());
}

Example 64 with RequestContext

use of com.linkedin.r2.message.RequestContext in project rest.li by linkedin.

the class TestParseqTraceDebugRequestHandler method testResponseStreamingAttachmentsForbidden.

@Test
public void testResponseStreamingAttachmentsForbidden() {
    //This test verifies that the ParseqTraceDebugRequestHandler aborts any potential outgoing response attachments
    //set by a resource method.
    final URI uri = URI.create("http://host/abc/12/__debug/parseqtrace/raw");
    ParseqTraceDebugRequestHandler requestHandler = new ParseqTraceDebugRequestHandler();
    RestRequestBuilder requestBuilder = new RestRequestBuilder(uri);
    RestRequest request = requestBuilder.build();
    RequestContext requestContext = new RequestContext();
    final RequestExecutionCallback<RestResponse> callback = new RequestExecutionCallback<RestResponse>() {

        @Override
        public void onError(Throwable e, RequestExecutionReport executionReport, RestLiAttachmentReader requestAttachmentReader, RestLiResponseAttachments responseAttachments) {
            Assert.fail("Request execution failed unexpectedly.");
        }

        @Override
        public void onSuccess(RestResponse result, RequestExecutionReport executionReport, RestLiResponseAttachments responseAttachments) {
        }
    };
    final RestLiTestAttachmentDataSource testAttachmentDataSource = RestLiTestAttachmentDataSource.createWithRandomPayload("1");
    final RestLiResponseAttachments responseAttachments = new RestLiResponseAttachments.Builder().appendSingleAttachment(testAttachmentDataSource).build();
    requestHandler.handleRequest(request, requestContext, new RestLiDebugRequestHandler.ResourceDebugRequestHandler() {

        @Override
        public void handleRequest(RestRequest request, RequestContext requestContext, RequestExecutionCallback<RestResponse> callback) {
            RestResponse response = EasyMock.createMock(RestResponse.class);
            //Provide some attachments that should be aborted.
            callback.onSuccess(response, new RequestExecutionReportBuilder().build(), responseAttachments);
        }
    }, null, callback);
    Assert.assertTrue(testAttachmentDataSource.dataSourceAborted());
}

Example 65 with RequestContext

use of com.linkedin.r2.message.RequestContext in project rest.li by linkedin.

the class TestParseqTraceDebugRequestHandler method executeRequestThroughParseqDebugHandler.

private void executeRequestThroughParseqDebugHandler(URI uri, RequestExecutionCallback<RestResponse> callback) {
    ParseqTraceDebugRequestHandler requestHandler = new ParseqTraceDebugRequestHandler();
    RestRequestBuilder requestBuilder = new RestRequestBuilder(uri);
    RestRequest request = requestBuilder.build();
    RequestContext requestContext = new RequestContext();
    requestHandler.handleRequest(request, requestContext, new RestLiDebugRequestHandler.ResourceDebugRequestHandler() {

        @Override
        public void handleRequest(RestRequest request, RequestContext requestContext, RequestExecutionCallback<RestResponse> callback) {
            RestResponse response = EasyMock.createMock(RestResponse.class);
            RequestExecutionReportBuilder executionReportBuilder = new RequestExecutionReportBuilder();
            JsonTraceCodec jsonTraceCodec = new JsonTraceCodec();
            Trace t = null;
            try {
                t = jsonTraceCodec.decode(TEST_TRACE);
                executionReportBuilder.setParseqTrace(t);
            } catch (IOException exc) {
            //test will fail later
            }
            callback.onSuccess(response, executionReportBuilder.build(), null);
        }
    }, null, callback);
}