Skip to content

Commit

Permalink
Introduce request attributes in RestClient
Browse files Browse the repository at this point in the history
This commit introduces request attributes in the RestClient and
underlying infrastructure (i.e. HttpRequest).

Closes gh-32027
  • Loading branch information
poutsma committed Jun 11, 2024
1 parent c36e270 commit 60b5bbe
Show file tree
Hide file tree
Showing 21 changed files with 385 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -1007,7 +1007,8 @@ method parameters:
is supported for non-String values.

| `@RequestAttribute`
| Provide an `Object` to add as a request attribute. Only supported by `WebClient`.
| Provide an `Object` to add as a request attribute. Only supported by `RestClient`
and `WebClient`.

| `@RequestBody`
| Provide the body of the request either as an Object to be serialized, or a
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@

import java.io.IOException;
import java.net.URI;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.springframework.http.HttpMethod;
import org.springframework.http.client.ClientHttpRequest;
Expand Down Expand Up @@ -46,6 +48,9 @@ public class MockClientHttpRequest extends MockHttpOutputMessage implements Clie

private boolean executed = false;

@Nullable
Map<String, Object> attributes;


/**
* Create a {@code MockClientHttpRequest} with {@link HttpMethod#GET GET} as
Expand Down Expand Up @@ -115,6 +120,16 @@ public boolean isExecuted() {
return this.executed;
}

@Override
public Map<String, Object> getAttributes() {
Map<String, Object> attributes = this.attributes;
if (attributes == null) {
attributes = new ConcurrentHashMap<>();
this.attributes = attributes;
}
return attributes;
}

/**
* Set the {@link #isExecuted() executed} flag to {@code true} and return the
* configured {@link #setResponse(ClientHttpResponse) response}.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
package org.springframework.http;

import java.net.URI;
import java.util.Map;

/**
* Represents an HTTP request message, consisting of a
Expand All @@ -41,4 +42,10 @@ public interface HttpRequest extends HttpMessage {
*/
URI getURI();

/**
* Return a mutable map of request attributes for this request.
* @since 6.2
*/
Map<String, Object> getAttributes();

}
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@

import java.io.IOException;
import java.io.OutputStream;
import java.util.LinkedHashMap;
import java.util.Map;

import org.springframework.http.HttpHeaders;
import org.springframework.lang.Nullable;
Expand All @@ -39,6 +41,9 @@ public abstract class AbstractClientHttpRequest implements ClientHttpRequest {
@Nullable
private HttpHeaders readOnlyHeaders;

@Nullable
private Map<String, Object> attributes;


@Override
public final HttpHeaders getHeaders() {
Expand All @@ -60,6 +65,16 @@ public final OutputStream getBody() throws IOException {
return getBodyInternal(this.headers);
}

@Override
public Map<String, Object> getAttributes() {
Map<String, Object> attributes = this.attributes;
if (attributes == null) {
attributes = new LinkedHashMap<>();
this.attributes = attributes;
}
return attributes;
}

@Override
public final ClientHttpResponse execute() throws IOException {
assertNotExecuted();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ public ClientHttpResponse execute(HttpRequest request, byte[] body) throws IOExc
HttpMethod method = request.getMethod();
ClientHttpRequest delegate = requestFactory.createRequest(request.getURI(), method);
request.getHeaders().forEach((key, value) -> delegate.getHeaders().addAll(key, value));
request.getAttributes().forEach((key, value) -> delegate.getAttributes().put(key, value));
if (body.length > 0) {
if (delegate instanceof StreamingHttpOutputMessage streamingOutputMessage) {
streamingOutputMessage.setBody(new StreamingHttpOutputMessage.Body() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
package org.springframework.http.client.support;

import java.net.URI;
import java.util.Map;

import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
Expand Down Expand Up @@ -70,6 +71,14 @@ public URI getURI() {
return this.request.getURI();
}

/**
* Return the attributes of the wrapped request.
*/
@Override
public Map<String, Object> getAttributes() {
return this.request.getAttributes();
}

/**
* Return the headers of the wrapped request.
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,16 @@
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.security.Principal;
import java.util.AbstractCollection;
import java.util.AbstractMap;
import java.util.AbstractSet;
import java.util.Arrays;
import java.util.Collection;
import java.util.Enumeration;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import jakarta.servlet.http.HttpServletRequest;

Expand Down Expand Up @@ -67,6 +72,10 @@ public class ServletServerHttpRequest implements ServerHttpRequest {
@Nullable
private HttpHeaders headers;

@Nullable
private Map<String, Object> attributes;


@Nullable
private ServerHttpAsyncRequestControl asyncRequestControl;

Expand Down Expand Up @@ -207,6 +216,16 @@ public InetSocketAddress getRemoteAddress() {
return new InetSocketAddress(this.servletRequest.getRemoteHost(), this.servletRequest.getRemotePort());
}

@Override
public Map<String, Object> getAttributes() {
Map<String, Object> attributes = this.attributes;
if (attributes == null) {
attributes = new AttributesMap();
this.attributes = attributes;
}
return attributes;
}

@Override
public InputStream getBody() throws IOException {
if (isFormPost(this.servletRequest) && this.servletRequest.getQueryString() == null) {
Expand Down Expand Up @@ -276,4 +295,151 @@ private InputStream getBodyFromServletRequestParameters(HttpServletRequest reque
return new ByteArrayInputStream(bytes);
}


private final class AttributesMap extends AbstractMap<String, Object> {

@Nullable
private transient Set<String> keySet;

@Nullable
private transient Collection<Object> values;

@Nullable
private transient Set<Entry<String, Object>> entrySet;


@Override
public int size() {
int size = 0;
for (Enumeration<?> names = servletRequest.getAttributeNames(); names.hasMoreElements(); names.nextElement()) {
size++;
}
return size;
}

@Override
@Nullable
public Object get(Object key) {
if (key instanceof String name) {
return servletRequest.getAttribute(name);
}
else {
return null;
}
}

@Override
@Nullable
public Object put(String key, Object value) {
Object old = get(key);
servletRequest.setAttribute(key, value);
return old;
}

@Override
@Nullable
public Object remove(Object key) {
if (key instanceof String name) {
Object old = get(key);
servletRequest.removeAttribute(name);
return old;
}
else {
return null;
}
}

@Override
public void clear() {
for (Enumeration<String> names = servletRequest.getAttributeNames(); names.hasMoreElements(); ) {
String name = names.nextElement();
servletRequest.removeAttribute(name);
}
}

@Override
public Set<String> keySet() {
Set<String> keySet = this.keySet;
if (keySet == null) {
keySet = new AbstractSet<>() {
@Override
public Iterator<String> iterator() {
return servletRequest.getAttributeNames().asIterator();
}

@Override
public int size() {
return AttributesMap.this.size();
}
};
this.keySet = keySet;
}
return keySet;
}

@Override
public Collection<Object> values() {
Collection<Object> values = this.values;
if (values == null) {
values = new AbstractCollection<>() {
@Override
public Iterator<Object> iterator() {
Enumeration<String> e = servletRequest.getAttributeNames();
return new Iterator<>() {
@Override
public boolean hasNext() {
return e.hasMoreElements();
}

@Override
public Object next() {
String name = e.nextElement();
return servletRequest.getAttribute(name);
}
};
}

@Override
public int size() {
return AttributesMap.this.size();
}
};
this.values = values;
}
return values;
}

@Override
public Set<Entry<String, Object>> entrySet() {
Set<Entry<String, Object>> entrySet = this.entrySet;
if (entrySet == null) {
entrySet = new AbstractSet<>() {
@Override
public Iterator<Entry<String, Object>> iterator() {
Enumeration<String> e = servletRequest.getAttributeNames();
return new Iterator<>() {
@Override
public boolean hasNext() {
return e.hasMoreElements();
}

@Override
public Entry<String, Object> next() {
String name = e.nextElement();
Object value = servletRequest.getAttribute(name);
return new SimpleImmutableEntry<>(name, value);
}
};
}

@Override
public int size() {
return AttributesMap.this.size();
}
};
this.entrySet = entrySet;
}
return entrySet;
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,9 @@
import java.net.URI;
import java.net.URLDecoder;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.Map;
import java.util.function.Supplier;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

Expand Down Expand Up @@ -68,6 +71,9 @@ public abstract class AbstractServerHttpRequest implements ServerHttpRequest {
@Nullable
private String logPrefix;

@Nullable
private Supplier<Map<String, Object>> attributesSupplier;


/**
* Constructor with the method, URI and headers for the request.
Expand Down Expand Up @@ -122,6 +128,16 @@ public URI getURI() {
return this.uri;
}

@Override
public Map<String, Object> getAttributes() {
if (this.attributesSupplier != null) {
return this.attributesSupplier.get();
}
else {
return Collections.emptyMap();
}
}

@Override
public RequestPath getPath() {
return this.path;
Expand Down Expand Up @@ -230,4 +246,12 @@ protected String initLogPrefix() {
return getId();
}

/**
* Set the attribute supplier.
* <p><strong>Note:</strong> This is exposed mainly for internal framework
* use.
*/
public void setAttributesSupplier(Supplier<Map<String, Object>> attributesSupplier) {
this.attributesSupplier = attributesSupplier;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import java.net.InetSocketAddress;
import java.net.URI;
import java.util.Map;

import reactor.core.publisher.Flux;

Expand Down Expand Up @@ -70,6 +71,11 @@ public URI getURI() {
return getDelegate().getURI();
}

@Override
public Map<String, Object> getAttributes() {
return getDelegate().getAttributes();
}

@Override
public RequestPath getPath() {
return getDelegate().getPath();
Expand Down
Loading

0 comments on commit 60b5bbe

Please sign in to comment.