Skip to content

Commit

Permalink
Add sort and pagination support for QueryApiKey API (#76144) (#76518)
Browse files Browse the repository at this point in the history
This PR adds support for sort and pagination similar to those used with
regular search API. Similar to the query field, the sort field also
supports only a subset of what is available for regular search.
  • Loading branch information
ywangd committed Aug 14, 2021
1 parent 211a958 commit 377517a
Show file tree
Hide file tree
Showing 15 changed files with 834 additions and 164 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -108,20 +108,28 @@ public Map<String, Object> getMetadata() {

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject()
.field("id", id)
.field("name", name)
.field("creation", creation.toEpochMilli());
builder.startObject();
innerToXContent(builder, params);
return builder.endObject();
}

public XContentBuilder innerToXContent(XContentBuilder builder, Params params) throws IOException {
builder
.field("id", id)
.field("name", name)
.field("creation", creation.toEpochMilli());
if (expiration != null) {
builder.field("expiration", expiration.toEpochMilli());
}
builder.field("invalidated", invalidated)
.field("username", username)
.field("realm", realm)
.field("metadata", (metadata == null ? org.elasticsearch.core.Map.of() : metadata));
return builder.endObject();
builder
.field("invalidated", invalidated)
.field("username", username)
.field("realm", realm)
.field("metadata", (metadata == null ? org.elasticsearch.core.Map.of() : metadata));
return builder;
}


@Override
public void writeTo(StreamOutput out) throws IOException {
if (out.getVersion().onOrAfter(Version.V_7_5_0)) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,32 +13,83 @@
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.core.Nullable;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.search.searchafter.SearchAfterBuilder;
import org.elasticsearch.search.sort.FieldSortBuilder;

import java.io.IOException;
import java.util.List;

import static org.elasticsearch.action.ValidateActions.addValidationError;

public final class QueryApiKeyRequest extends ActionRequest {

@Nullable
private final QueryBuilder queryBuilder;
@Nullable
private final Integer from;
@Nullable
private final Integer size;
@Nullable
private final List<FieldSortBuilder> fieldSortBuilders;
@Nullable
private final SearchAfterBuilder searchAfterBuilder;
private boolean filterForCurrentUser;

public QueryApiKeyRequest() {
this((QueryBuilder) null);
}

public QueryApiKeyRequest(QueryBuilder queryBuilder) {
this(queryBuilder, null, null, null, null);
}

public QueryApiKeyRequest(
@Nullable QueryBuilder queryBuilder,
@Nullable Integer from,
@Nullable Integer size,
@Nullable List<FieldSortBuilder> fieldSortBuilders,
@Nullable SearchAfterBuilder searchAfterBuilder
) {
this.queryBuilder = queryBuilder;
this.from = from;
this.size = size;
this.fieldSortBuilders = fieldSortBuilders;
this.searchAfterBuilder = searchAfterBuilder;
}

public QueryApiKeyRequest(StreamInput in) throws IOException {
super(in);
queryBuilder = in.readOptionalNamedWriteable(QueryBuilder.class);
this.queryBuilder = in.readOptionalNamedWriteable(QueryBuilder.class);
this.from = in.readOptionalVInt();
this.size = in.readOptionalVInt();
if (in.readBoolean()) {
this.fieldSortBuilders = in.readList(FieldSortBuilder::new);
} else {
this.fieldSortBuilders = null;
}
this.searchAfterBuilder = in.readOptionalWriteable(SearchAfterBuilder::new);
}

public QueryBuilder getQueryBuilder() {
return queryBuilder;
}

public Integer getFrom() {
return from;
}

public Integer getSize() {
return size;
}

public List<FieldSortBuilder> getFieldSortBuilders() {
return fieldSortBuilders;
}

public SearchAfterBuilder getSearchAfterBuilder() {
return searchAfterBuilder;
}

public boolean isFilterForCurrentUser() {
return filterForCurrentUser;
}
Expand All @@ -49,12 +100,28 @@ public void setFilterForCurrentUser() {

@Override
public ActionRequestValidationException validate() {
return null;
ActionRequestValidationException validationException = null;
if (from != null && from < 0) {
validationException = addValidationError("[from] parameter cannot be negative but was [" + from + "]", validationException);
}
if (size != null && size < 0) {
validationException = addValidationError("[size] parameter cannot be negative but was [" + size + "]", validationException);
}
return validationException;
}

@Override
public void writeTo(StreamOutput out) throws IOException {
super.writeTo(out);
out.writeOptionalNamedWriteable(queryBuilder);
out.writeOptionalVInt(from);
out.writeOptionalVInt(size);
if (fieldSortBuilders == null) {
out.writeBoolean(false);
} else {
out.writeBoolean(true);
out.writeList(fieldSortBuilders);
}
out.writeOptionalWriteable(searchAfterBuilder);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,10 @@
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.io.stream.Writeable;
import org.elasticsearch.common.lucene.Lucene;
import org.elasticsearch.common.xcontent.ToXContentObject;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.core.Nullable;
import org.elasticsearch.xpack.core.security.action.ApiKey;

import java.io.IOException;
Expand All @@ -27,36 +29,46 @@
*/
public final class QueryApiKeyResponse extends ActionResponse implements ToXContentObject, Writeable {

private final ApiKey[] foundApiKeysInfo;
private final long total;
private final Item[] items;

public QueryApiKeyResponse(StreamInput in) throws IOException {
super(in);
this.foundApiKeysInfo = in.readArray(ApiKey::new, ApiKey[]::new);
this.total = in.readLong();
this.items = in.readArray(Item::new, Item[]::new);
}

public QueryApiKeyResponse(Collection<ApiKey> foundApiKeysInfo) {
Objects.requireNonNull(foundApiKeysInfo, "found_api_keys_info must be provided");
this.foundApiKeysInfo = foundApiKeysInfo.toArray(new ApiKey[0]);
public QueryApiKeyResponse(long total, Collection<Item> items) {
this.total = total;
Objects.requireNonNull(items, "items must be provided");
this.items = items.toArray(new Item[0]);
}

public static QueryApiKeyResponse emptyResponse() {
return new QueryApiKeyResponse(Collections.emptyList());
return new QueryApiKeyResponse(0, Collections.emptyList());
}

public ApiKey[] getApiKeyInfos() {
return foundApiKeysInfo;
public long getTotal() {
return total;
}

public Item[] getItems() {
return items;
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject()
.array("api_keys", (Object[]) foundApiKeysInfo);
.field("total", total)
.field("count", items.length)
.array("api_keys", (Object[]) items);
return builder.endObject();
}

@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeArray(foundApiKeysInfo);
out.writeLong(total);
out.writeArray(items);
}

@Override
Expand All @@ -66,17 +78,81 @@ public boolean equals(Object o) {
if (o == null || getClass() != o.getClass())
return false;
QueryApiKeyResponse that = (QueryApiKeyResponse) o;
return Arrays.equals(foundApiKeysInfo, that.foundApiKeysInfo);
return total == that.total && Arrays.equals(items, that.items);
}

@Override
public int hashCode() {
return Arrays.hashCode(foundApiKeysInfo);
int result = Objects.hash(total);
result = 31 * result + Arrays.hashCode(items);
return result;
}

@Override
public String toString() {
return "QueryApiKeyResponse [foundApiKeysInfo=" + foundApiKeysInfo + "]";
return "QueryApiKeyResponse{" + "total=" + total + ", items=" + Arrays.toString(items) + '}';
}

public static class Item implements ToXContentObject, Writeable {
private final ApiKey apiKey;
@Nullable
private final Object[] sortValues;

public Item(ApiKey apiKey, @Nullable Object[] sortValues) {
this.apiKey = apiKey;
this.sortValues = sortValues;
}

public Item(StreamInput in) throws IOException {
this.apiKey = new ApiKey(in);
this.sortValues = in.readOptionalArray(Lucene::readSortValue, Object[]::new);
}

public ApiKey getApiKey() {
return apiKey;
}

public Object[] getSortValues() {
return sortValues;
}

@Override
public void writeTo(StreamOutput out) throws IOException {
apiKey.writeTo(out);
out.writeOptionalArray(Lucene::writeSortValue, sortValues);
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject();
apiKey.innerToXContent(builder, params);
if (sortValues != null && sortValues.length > 0) {
builder.array("_sort", sortValues);
}
builder.endObject();
return builder;
}

@Override
public boolean equals(Object o) {
if (this == o)
return true;
if (o == null || getClass() != o.getClass())
return false;
Item item = (Item) o;
return Objects.equals(apiKey, item.apiKey) && Arrays.equals(sortValues, item.sortValues);
}

@Override
public int hashCode() {
int result = Objects.hash(apiKey);
result = 31 * result + Arrays.hashCode(sortValues);
return result;
}

@Override
public String toString() {
return "Item{" + "apiKey=" + apiKey + ", sortValues=" + Arrays.toString(sortValues) + '}';
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,18 @@
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.MatchAllQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.SearchModule;
import org.elasticsearch.search.searchafter.SearchAfterBuilder;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.elasticsearch.test.ESTestCase;

import java.io.ByteArrayInputStream;
import java.io.IOException;

import static org.hamcrest.Matchers.containsString;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.nullValue;
Expand Down Expand Up @@ -56,5 +61,39 @@ public void testReadWrite() throws IOException {
assertThat((BoolQueryBuilder) deserialized.getQueryBuilder(), equalTo(boolQueryBuilder2));
}
}

final QueryApiKeyRequest request3 = new QueryApiKeyRequest(
QueryBuilders.matchAllQuery(),
42,
20,
org.elasticsearch.core.List.of(new FieldSortBuilder("name"),
new FieldSortBuilder("creation_time").setFormat("strict_date_time").order(SortOrder.DESC),
new FieldSortBuilder("username")),
new SearchAfterBuilder().setSortValues(new String[] { "key-2048", "2021-07-01T00:00:59.000Z" }));
try (BytesStreamOutput out = new BytesStreamOutput()) {
request3.writeTo(out);
try (StreamInput in = new NamedWriteableAwareStreamInput(out.bytes().streamInput(), writableRegistry())) {
final QueryApiKeyRequest deserialized = new QueryApiKeyRequest(in);
assertThat(deserialized.getQueryBuilder().getClass(), is(MatchAllQueryBuilder.class));
assertThat(deserialized.getFrom(), equalTo(request3.getFrom()));
assertThat(deserialized.getSize(), equalTo(request3.getSize()));
assertThat(deserialized.getFieldSortBuilders(), equalTo(request3.getFieldSortBuilders()));
assertThat(deserialized.getSearchAfterBuilder(), equalTo(request3.getSearchAfterBuilder()));
}
}
}

public void testValidate() {
final QueryApiKeyRequest request1 =
new QueryApiKeyRequest(null, randomIntBetween(0, Integer.MAX_VALUE), randomIntBetween(0, Integer.MAX_VALUE), null, null);
assertThat(request1.validate(), nullValue());

final QueryApiKeyRequest request2 =
new QueryApiKeyRequest(null, randomIntBetween(Integer.MIN_VALUE, -1), randomIntBetween(0, Integer.MAX_VALUE), null, null);
assertThat(request2.validate().getMessage(), containsString("[from] parameter cannot be negative"));

final QueryApiKeyRequest request3 =
new QueryApiKeyRequest(null, randomIntBetween(0, Integer.MAX_VALUE), randomIntBetween(Integer.MIN_VALUE, -1), null, null);
assertThat(request3.validate().getMessage(), containsString("[size] parameter cannot be negative"));
}
}

0 comments on commit 377517a

Please sign in to comment.