private <RT extends RecordTemplate> Set<String> extractIds(BatchRequest<RT> request) { return request.getObjectIds().stream().map(Object::toString).collect(Collectors.toSet()); } }
@Override public boolean equals(Object o) { if (this == o) return true; if (o == null || !(o instanceof BatchingKey)) return false; @SuppressWarnings("unchecked") BatchingKey<T,R> that = (BatchingKey<T,R>) o; EqualsBuilder builder = new EqualsBuilder(); builder.append(_request.getBaseUriTemplate(), that._request.getBaseUriTemplate()); builder.append(_request.getPathKeys(), that._request.getPathKeys()); builder.append(_request.getResourceProperties(), that._request.getResourceProperties()); builder.append(_request.getRequestOptions(), that._request.getRequestOptions()); builder.append(_queryParams, that._queryParams); builder.append(_batchFields, that._batchFields ); if (_batchFields) builder.append(_request.getFields(), that._request.getFields()); return builder.isEquals(); }
return new GetRequestBuilder(batchRequest.getBaseUriTemplate(), batchRequest.getResourceSpec().getValueClass(), batchRequest.getResourceSpec(), batchRequest.getRequestOptions()); return new DeleteRequestBuilder(batchRequest.getBaseUriTemplate(), batchRequest.getResourceSpec().getValueClass(), batchRequest.getResourceSpec(), batchRequest.getRequestOptions()); return new UpdateRequestBuilder(batchRequest.getBaseUriTemplate(), batchRequest.getResourceSpec().getValueClass(), batchRequest.getResourceSpec(), batchRequest.getRequestOptions()); return new PartialUpdateRequestBuilder(batchRequest.getBaseUriTemplate(), batchRequest.getResourceSpec().getValueClass(), batchRequest.getResourceSpec(), batchRequest.getRequestOptions());
Set<PathSpec> fields = batchFields ? new HashSet<PathSpec>() : firstRequest.getFields(); final Set<Object> requestIds = request.getObjectIds(); final Set<PathSpec> requestFields = request.getFields();
@SuppressWarnings({"rawtypes", "unchecked"}) private <K> Request buildScatterBatchRequestByKeys(BatchRequest<?> batchRequest, Set<K> keys, Map<K, ?> body) { final BatchKVRequestBuilder builder = getBatchBuilder(batchRequest, keys, body); // keep all non-batch query parameters batchRequest.getQueryParamsObjects().entrySet().stream() .filter(queryParam -> !queryParam.getKey().equals(RestConstants.QUERY_BATCH_IDS_PARAM)) .forEach(queryParam -> builder.setParam(queryParam.getKey(), queryParam.getValue())); // keep all headers batchRequest.getHeaders().forEach(builder::setHeader); return builder.build(); }
@SuppressWarnings({"unchecked", "rawtypes", "deprecation"}) private <T> T constructResponseFromDataMap(BatchRequest<T> request, ProtocolVersion protocolVersion, DataMap data) { if (request instanceof BatchGetRequest) { return (T) new BatchResponse(data, request.getResponseDecoder().getEntityClass()); } else { return (T) new BatchKVResponse(data, request.getResourceSpec().getKeyType(), request.getResourceSpec().getValueType(), request.getResourceSpec().getKeyParts(), request.getResourceSpec().getComplexKeyType(), protocolVersion); } }
@Override public Set<PathSpec> getFields() { return super.getFields(); } }
/** * Creates a map of those queryParams which are required to be the same for all batchable requests. * * @param request the BatchGetRequest to pull query params from * @return Map which contains all query params save for Batch_Ids and Fields */ public static Map<String, Object> getQueryParamsForBatchingKey(BatchRequest<?> request) { final Map<String, Object> params = new HashMap<String, Object>(request.getQueryParamsObjects()); params.remove(RestConstants.QUERY_BATCH_IDS_PARAM); params.remove(RestConstants.FIELDS_PARAM); return params; }
private void checkBatchRequest(BatchRequest<?> request) { if (request.getMethod() == ResourceMethod.BATCH_CREATE ) { throw new UnsupportedOperationException("BATCH_CREATE is not supported for scatter-gather!"); } }
@Override public int hashCode() { HashCodeBuilder builder = new HashCodeBuilder(); builder.append( _request.getBaseUriTemplate() ) .append( _request.getPathKeys() ) .append( _request.getResourceProperties() ) .append( _request.getRequestOptions() ) .append( _queryParams ); if(_batchFields && null != _request.getFields()) builder.append( _request.getFields() ); return builder.toHashCode(); } }
return new BatchGetRequestBuilder(batchRequest.getBaseUriTemplate(), batchRequest.getResourceSpec().getValueClass(), batchRequest.getResourceSpec(), batchRequest.getRequestOptions()).ids(keys); return new BatchGetEntityRequestBuilder(batchRequest.getBaseUriTemplate(), batchRequest.getResourceSpec(), batchRequest.getRequestOptions()).ids(keys); return new BatchDeleteRequestBuilder(batchRequest.getBaseUriTemplate(), batchRequest.getResourceSpec().getValueClass(), batchRequest.getResourceSpec(), batchRequest.getRequestOptions()).ids(keys); return new BatchUpdateRequestBuilder(batchRequest.getBaseUriTemplate(), batchRequest.getResourceSpec().getValueClass(), batchRequest.getResourceSpec(), batchRequest.getRequestOptions()).inputs(body); return new BatchPartialUpdateRequestBuilder(batchRequest.getBaseUriTemplate(), batchRequest.getResourceSpec().getValueClass(), batchRequest.getResourceSpec(), batchRequest.getRequestOptions()).inputs(body);
@SuppressWarnings({"rawtypes", "unchecked"}) private <K> Request unbatchRequestByKey(BatchRequest<?> batchRequest, K key) { final SingleEntityRequestBuilder builder = getBuilder(batchRequest); // For BATCH_UPDATE and BATCH_PARTIAL_UPDATE, the generated Request is missing body, // but that is sufficient for us to implement getUris where we only care about individual request URI. builder.id(key); // keep all non-batch query parameters batchRequest.getQueryParamsObjects().entrySet().stream() .filter(queryParam -> !queryParam.getKey().equals(RestConstants.QUERY_BATCH_IDS_PARAM)) .forEach(queryParam -> builder.setParam(queryParam.getKey(), queryParam.getValue())); // keep all headers batchRequest.getHeaders().forEach(builder::setHeader); return builder.build(); }
@Override public Set<PathSpec> getFields() { return super.getFields(); } }
/** * @return the IDs of the objects in this request. The IDs are the keys with their original types (non-coerced) */ public Set<Object> getObjectIds() { @SuppressWarnings({"unchecked"}) Collection<Object> ids = (Collection<Object>) getQueryParamsObjects().get(RestConstants.QUERY_BATCH_IDS_PARAM); if (ids == null || ids.isEmpty()) { return Collections.emptySet(); } return new HashSet<Object>(ids); } }
/** * This method provides granular exception messages for developers to know why a request * intended to be added to a batch, might fail to be batched. * * @param request the request to validate is the same or not * @throws IllegalArgumentException if there are any problems */ public <T> void validate(BatchRequest<T> request) { if( !request.getBaseUriTemplate().equals(_request.getBaseUriTemplate()) || !request.getPathKeys().equals(_request.getPathKeys()) ) { throw new IllegalArgumentException("Requests must have same base URI template and path keys to batch"); } if( !request.getResourceProperties().equals(_request.getResourceProperties()) ) { throw new IllegalArgumentException("Requests must be for the same resource to batch"); } if( !request.getRequestOptions().equals(_request.getRequestOptions()) ) { throw new IllegalArgumentException("Requests must have the same RestliRequestOptions to batch!"); } // Enforce uniformity of query params excluding ids and fields final Map<String, Object> queryParams = BatchGetRequestUtil.getQueryParamsForBatchingKey(request); if ( !queryParams.equals(_queryParams) ) { throw new IllegalArgumentException("Requests must have same parameters to batch"); } if (!_batchFields) { if( !request.getFields().equals(_request.getFields()) ) { throw new IllegalArgumentException("Requests must have same fields to batch"); } } }
private static Tuple3<Set<Object>, Set<PathSpec>, Boolean> reduceIds(Tuple3<Set<Object>, Set<PathSpec>, Boolean> state, Request<?> request) { if (request instanceof GetRequest) { GetRequest<?> getRequest = (GetRequest<?>)request; state._1().add(getRequest.getObjectId()); return state; } else if (request instanceof BatchRequest) { BatchRequest<?> batchRequest = (BatchRequest<?>)request; state._1().addAll(batchRequest.getObjectIds()); return state; } else { throw unsupportedGetRequestType(request); } }
@Override public Set<PathSpec> getFields() { return super.getFields(); } }
/** * {@inheritDoc} * * Note that if the custom ScatterGatherStrategy overrides this method to associate each URI with a set of partition * Ids to bypass the partitioning by D2 later, it should also override {@link ScatterGatherStrategy#onAllResponsesReceived(Request, * ProtocolVersion, Map, Map, Map, Callback)} to handle custom response gathering. */ @Override @SuppressWarnings("rawtypes") public <K, T> List<URIKeyPair<K>> getUris(Request<T> request, ProtocolVersion version) { BatchRequest<T> batchRequest = safeCastRequest(request); @SuppressWarnings("unchecked") Set<K> keys = (Set<K>) batchRequest.getObjectIds(); return keys.stream() .map(key -> { Request unbatchRequestByKey = unbatchRequestByKey(batchRequest, key); URI requestUri = RestliUriBuilderUtil.createUriBuilder(unbatchRequestByKey, RestConstants.D2_URI_PREFIX, version).build(); return new URIKeyPair<>(key, requestUri); }) .collect(Collectors.toList()); }
Set<K> failedKeys = (Set<K>)((BatchRequest<T>)req.getRequest()).getObjectIds(); gatherException(gatheredResponseDataMap, failedKeys, e, protocolVersion); });