Add sort and pagination support for QueryApiKey API

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/QueryApiKeyRequest.java

@@ -13,32 +13,79 @@
 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;
 
 public final class QueryApiKeyRequest extends ActionRequest {
 
     @Nullable
     private final QueryBuilder queryBuilder;
+    private final int from;
+    private final int 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 = queryBuilder;
+        this(queryBuilder, -1, -1, null, null);
     }
 
     public QueryApiKeyRequest(StreamInput in) throws IOException {
         super(in);
-        queryBuilder = in.readOptionalNamedWriteable(QueryBuilder.class);
+        this.queryBuilder = in.readOptionalNamedWriteable(QueryBuilder.class);
+        this.from = in.readInt();
+        this.size = in.readInt();
+        if (in.readBoolean()) {
+            this.fieldSortBuilders = in.readList(FieldSortBuilder::new);
+        } else {
+            this.fieldSortBuilders = null;
+        }
+        this.searchAfterBuilder = in.readOptionalWriteable(SearchAfterBuilder::new);
+    }
+
+    public QueryApiKeyRequest(
+        @Nullable QueryBuilder queryBuilder,
+        @Nullable Integer from,
+        @Nullable Integer size,
+        @Nullable List<FieldSortBuilder> fieldSortBuilders,
+        @Nullable SearchAfterBuilder searchAfterBuilder
+    ) {
+        this.queryBuilder = queryBuilder;
+        this.from = from == null ? -1 : from;
+        this.size = size == null ? -1 : size;
+        this.fieldSortBuilders = fieldSortBuilders;
+        this.searchAfterBuilder = searchAfterBuilder;
     }
 
     public QueryBuilder getQueryBuilder() {
         return queryBuilder;
     }
 
+    public int getFrom() {
+        return from;
+    }
+
+    public int getSize() {
+        return size;
+    }
+
+    public List<FieldSortBuilder> getFieldSortBuilders() {
+        return fieldSortBuilders;
+    }
+
+    public SearchAfterBuilder getSearchAfterBuilder() {
+        return searchAfterBuilder;
+    }
+
     public boolean isFilterForCurrentUser() {
         return filterForCurrentUser;
     }
@@ -56,5 +103,14 @@ public ActionRequestValidationException validate() {
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         out.writeOptionalNamedWriteable(queryBuilder);
+        out.writeInt(from);
+        out.writeInt(size);
+        if (fieldSortBuilders == null) {
+            out.writeBoolean(false);
+        } else {
+            out.writeBoolean(true);
+            out.writeList(fieldSortBuilders);
+        }
+        out.writeOptionalWriteable(searchAfterBuilder);
     }
 }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/QueryApiKeyResponse.java

@@ -27,20 +27,27 @@
  */
 public final class QueryApiKeyResponse extends ActionResponse implements ToXContentObject, Writeable {
 
+    private final long total;
     private final ApiKey[] foundApiKeysInfo;
 
     public QueryApiKeyResponse(StreamInput in) throws IOException {
         super(in);
+        this.total = in.readLong();
         this.foundApiKeysInfo = in.readArray(ApiKey::new, ApiKey[]::new);
     }
 
-    public QueryApiKeyResponse(Collection<ApiKey> foundApiKeysInfo) {
+    public QueryApiKeyResponse(long total, Collection<ApiKey> foundApiKeysInfo) {
+        this.total = total;
         Objects.requireNonNull(foundApiKeysInfo, "found_api_keys_info must be provided");
         this.foundApiKeysInfo = foundApiKeysInfo.toArray(new ApiKey[0]);
     }
 
     public static QueryApiKeyResponse emptyResponse() {
-        return new QueryApiKeyResponse(Collections.emptyList());
+        return new QueryApiKeyResponse(0, Collections.emptyList());
+    }
+
+    public long getTotal() {
+        return total;
     }
 
     public ApiKey[] getApiKeyInfos() {
@@ -50,12 +57,15 @@ public ApiKey[] getApiKeyInfos() {
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject()
+            .field("total", total)
+            .field("count", foundApiKeysInfo.length)
             .array("api_keys", (Object[]) foundApiKeysInfo);
         return builder.endObject();
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
+        out.writeLong(total);
         out.writeArray(foundApiKeysInfo);
     }
 
@@ -66,17 +76,18 @@ 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(foundApiKeysInfo, that.foundApiKeysInfo);
     }
 
     @Override
     public int hashCode() {
-        return Arrays.hashCode(foundApiKeysInfo);
+        int result = Objects.hash(total);
+        result = 31 * result + Arrays.hashCode(foundApiKeysInfo);
+        return result;
     }
 
     @Override
     public String toString() {
-        return "QueryApiKeyResponse [foundApiKeysInfo=" + foundApiKeysInfo + "]";
+        return "QueryApiKeyResponse{" + "total=" + total + ", foundApiKeysInfo=" + Arrays.toString(foundApiKeysInfo) + '}';
     }
-
 }

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/security/action/apikey/QueryApiKeyRequestTests.java

@@ -14,8 +14,12 @@
 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;
@@ -57,5 +61,25 @@ public void testReadWrite() throws IOException {
                 assertThat((BoolQueryBuilder) deserialized.getQueryBuilder(), equalTo(boolQueryBuilder2));
             }
         }
+
+        final QueryApiKeyRequest request3 = new QueryApiKeyRequest(
+            QueryBuilders.matchAllQuery(),
+            42,
+            20,
+            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()));
+            }
+        }
     }
 }

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/security/action/apikey/QueryApiKeyResponseTests.java

@@ -30,32 +30,34 @@ protected Writeable.Reader<QueryApiKeyResponse> instanceReader() {
     @Override
     protected QueryApiKeyResponse createTestInstance() {
         final List<ApiKey> apiKeys = randomList(0, 3, this::randomApiKeyInfo);
-        return new QueryApiKeyResponse(apiKeys);
+        return new QueryApiKeyResponse(randomIntBetween(apiKeys.size(), 100), apiKeys);
     }
 
     @Override
     protected QueryApiKeyResponse mutateInstance(QueryApiKeyResponse instance) throws IOException {
         final ArrayList<ApiKey> apiKeyInfos =
             Arrays.stream(instance.getApiKeyInfos()).collect(Collectors.toCollection(ArrayList::new));
-        switch (randomIntBetween(0, 2)) {
+        switch (randomIntBetween(0, 3)) {
             case 0:
                 apiKeyInfos.add(randomApiKeyInfo());
-                return new QueryApiKeyResponse(apiKeyInfos);
+                return new QueryApiKeyResponse(instance.getTotal(), apiKeyInfos);
             case 1:
                 if (false == apiKeyInfos.isEmpty()) {
-                    return new QueryApiKeyResponse(apiKeyInfos.subList(1, apiKeyInfos.size()));
+                    return new QueryApiKeyResponse(instance.getTotal(), apiKeyInfos.subList(1, apiKeyInfos.size()));
                 } else {
                     apiKeyInfos.add(randomApiKeyInfo());
-                    return new QueryApiKeyResponse(apiKeyInfos);
+                    return new QueryApiKeyResponse(instance.getTotal(), apiKeyInfos);
                 }
-            default:
+            case 2:
                 if (false == apiKeyInfos.isEmpty()) {
                     final int index = randomIntBetween(0, apiKeyInfos.size() - 1);
                     apiKeyInfos.set(index, randomApiKeyInfo());
                 } else {
                     apiKeyInfos.add(randomApiKeyInfo());
                 }
-                return new QueryApiKeyResponse(apiKeyInfos);
+                return new QueryApiKeyResponse(instance.getTotal(), apiKeyInfos);
+            default:
+                return new QueryApiKeyResponse(instance.getTotal() + 1, apiKeyInfos);
         }
     }