Enhance API Key Querying

docs/changelog/103192.yaml

@@ -0,0 +1,5 @@
+pr: 103192
+summary: Query api key api improvements
+area: Security
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -173,6 +173,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion NODE_STATS_REQUEST_SIMPLIFIED = def(8_561_00_0);
     public static final TransportVersion TEXT_EXPANSION_TOKEN_PRUNING_CONFIG_ADDED = def(8_562_00_0);
     public static final TransportVersion ESQL_ASYNC_QUERY = def(8_563_00_0);
+    public static final TransportVersion QUERY_API_KEY_AGGS_ADDED = def(8_564_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/search/aggregations/bucket/filter/FilterAggregationBuilder.java

@@ -41,7 +41,7 @@
 public class FilterAggregationBuilder extends AbstractAggregationBuilder<FilterAggregationBuilder> {
     public static final String NAME = "filter";
 
-    private final QueryBuilder filter;
+    private QueryBuilder filter;
 
     /**
      * @param name
@@ -59,6 +59,14 @@ public FilterAggregationBuilder(String name, QueryBuilder filter) {
         this.filter = filter;
     }
 
+    public FilterAggregationBuilder filter(QueryBuilder filter) {
+        if (filter == null) {
+            throw new IllegalArgumentException("[filter] must not be null: [" + name + "]");
+        }
+        this.filter = filter;
+        return this;
+    }
+
     protected FilterAggregationBuilder(
         FilterAggregationBuilder clone,
         AggregatorFactories.Builder factoriesBuilder,

server/src/main/java/org/elasticsearch/search/aggregations/bucket/filter/FiltersAggregator.java

@@ -61,7 +61,7 @@ public abstract class FiltersAggregator extends BucketsAggregator {
 
     public static class KeyedFilter implements Writeable, ToXContentFragment {
         private final String key;
-        private final QueryBuilder filter;
+        private QueryBuilder filter;
 
         public KeyedFilter(String key, QueryBuilder filter) {
             if (key == null) {
@@ -92,6 +92,10 @@ public String key() {
             return key;
         }
 
+        public void filter(QueryBuilder filter) {
+            this.filter = filter;
+        }
+
         public QueryBuilder filter() {
             return filter;
         }

server/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceAggregationBuilder.java

@@ -192,7 +192,6 @@ public Set<String> metricNames() {
     private String format = null;
     private Object missing = null;
     private ZoneId timeZone = null;
-    protected ValuesSourceConfig config;
 
     protected ValuesSourceAggregationBuilder(String name) {
         super(name);
@@ -209,7 +208,6 @@ protected ValuesSourceAggregationBuilder(
         this.format = clone.format;
         this.missing = clone.missing;
         this.timeZone = clone.timeZone;
-        this.config = clone.config;
         this.script = clone.script;
     }
 

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -728,6 +728,14 @@ public SearchSourceBuilder collapse(CollapseBuilder collapse) {
         return this;
     }
 
+    public SearchSourceBuilder aggregationsBuilder(AggregatorFactories.Builder aggregations) {
+        if (this.aggregations != null) {
+            throw new IllegalStateException("cannot override aggregation build [");
+        }
+        this.aggregations = aggregations;
+        return this;
+    }
+
     /**
      * Add an aggregation to perform as part of the search.
      */

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

@@ -14,6 +14,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.searchafter.SearchAfterBuilder;
 import org.elasticsearch.search.sort.FieldSortBuilder;
 
@@ -27,6 +28,8 @@ public final class QueryApiKeyRequest extends ActionRequest {
     @Nullable
     private final QueryBuilder queryBuilder;
     @Nullable
+    private final AggregatorFactories.Builder aggsBuilder;
+    @Nullable
     private final Integer from;
     @Nullable
     private final Integer size;
@@ -42,18 +45,20 @@ public QueryApiKeyRequest() {
     }
 
     public QueryApiKeyRequest(QueryBuilder queryBuilder) {
-        this(queryBuilder, null, null, null, null, false);
+        this(queryBuilder, null, null, null, null, null, false);
     }
 
     public QueryApiKeyRequest(
         @Nullable QueryBuilder queryBuilder,
+        @Nullable AggregatorFactories.Builder aggsBuilder,
         @Nullable Integer from,
         @Nullable Integer size,
         @Nullable List<FieldSortBuilder> fieldSortBuilders,
         @Nullable SearchAfterBuilder searchAfterBuilder,
         boolean withLimitedBy
     ) {
         this.queryBuilder = queryBuilder;
+        this.aggsBuilder = aggsBuilder;
         this.from = from;
         this.size = size;
         this.fieldSortBuilders = fieldSortBuilders;
@@ -77,12 +82,21 @@ public QueryApiKeyRequest(StreamInput in) throws IOException {
         } else {
             this.withLimitedBy = false;
         }
+        if (in.getTransportVersion().onOrAfter(TransportVersions.QUERY_API_KEY_AGGS_ADDED)) {
+            this.aggsBuilder = in.readOptionalWriteable(AggregatorFactories.Builder::new);
+        } else {
+            this.aggsBuilder = null;
+        }
     }
 
     public QueryBuilder getQueryBuilder() {
         return queryBuilder;
     }
 
+    public AggregatorFactories.Builder getAggsBuilder() {
+        return aggsBuilder;
+    }
+
     public Integer getFrom() {
         return from;
     }
@@ -139,5 +153,8 @@ public void writeTo(StreamOutput out) throws IOException {
         if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_5_0)) {
             out.writeBoolean(withLimitedBy);
         }
+        if (out.getTransportVersion().onOrAfter(TransportVersions.QUERY_API_KEY_AGGS_ADDED)) {
+            out.writeOptionalWriteable(aggsBuilder);
+        }
     }
 }

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

@@ -7,19 +7,22 @@
 
 package org.elasticsearch.xpack.core.security.action.apikey;
 
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.action.ActionResponse;
 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.core.Nullable;
+import org.elasticsearch.search.aggregations.Aggregations;
+import org.elasticsearch.search.aggregations.InternalAggregations;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.Collections;
+import java.util.List;
 import java.util.Objects;
 
 /**
@@ -30,21 +33,28 @@ public final class QueryApiKeyResponse extends ActionResponse implements ToXCont
 
     private final long total;
     private final Item[] items;
+    private final @Nullable Aggregations aggregations;
 
     public QueryApiKeyResponse(StreamInput in) throws IOException {
         super(in);
         this.total = in.readLong();
         this.items = in.readArray(Item::new, Item[]::new);
+        if (in.getTransportVersion().onOrAfter(TransportVersions.QUERY_API_KEY_AGGS_ADDED)) {
+            this.aggregations = in.readOptionalWriteable(InternalAggregations::readFrom);
+        } else {
+            this.aggregations = null;
+        }
     }
 
-    public QueryApiKeyResponse(long total, Collection<Item> items) {
+    public QueryApiKeyResponse(long total, Collection<Item> items, @Nullable Aggregations aggregations) {
         this.total = total;
         Objects.requireNonNull(items, "items must be provided");
         this.items = items.toArray(new Item[0]);
+        this.aggregations = aggregations;
     }
 
     public static QueryApiKeyResponse emptyResponse() {
-        return new QueryApiKeyResponse(0, Collections.emptyList());
+        return new QueryApiKeyResponse(0, List.of(), null);
     }
 
     public long getTotal() {
@@ -59,36 +69,47 @@ public int getCount() {
         return items.length;
     }
 
+    public Aggregations getAggregations() {
+        return aggregations;
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject().field("total", total).field("count", items.length).array("api_keys", (Object[]) items);
+        if (aggregations != null) {
+            aggregations.toXContent(builder, params);
+        }
         return builder.endObject();
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeLong(total);
         out.writeArray(items);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.QUERY_API_KEY_AGGS_ADDED)) {
+            out.writeOptionalWriteable((InternalAggregations) aggregations);
+        }
     }
 
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         QueryApiKeyResponse that = (QueryApiKeyResponse) o;
-        return total == that.total && Arrays.equals(items, that.items);
+        return total == that.total && Arrays.equals(items, that.items) && Objects.equals(aggregations, that.aggregations);
     }
 
     @Override
     public int hashCode() {
         int result = Objects.hash(total);
         result = 31 * result + Arrays.hashCode(items);
+        result = 31 * result + Objects.hash(aggregations);
         return result;
     }
 
     @Override
     public String toString() {
-        return "QueryApiKeyResponse{" + "total=" + total + ", items=" + Arrays.toString(items) + '}';
+        return "QueryApiKeyResponse{" + "total=" + total + ", items=" + Arrays.toString(items) + ", aggs=" + aggregations + '}';
     }
 
     public static class Item implements ToXContentObject, Writeable {

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

@@ -65,6 +65,7 @@ public void testReadWrite() throws IOException {
 
         final QueryApiKeyRequest request3 = new QueryApiKeyRequest(
             QueryBuilders.matchAllQuery(),
+            null,
             42,
             20,
             List.of(
@@ -91,6 +92,7 @@ public void testReadWrite() throws IOException {
 
     public void testValidate() {
         final QueryApiKeyRequest request1 = new QueryApiKeyRequest(
+            null,
             null,
             randomIntBetween(0, Integer.MAX_VALUE),
             randomIntBetween(0, Integer.MAX_VALUE),
@@ -101,6 +103,7 @@ public void testValidate() {
         assertThat(request1.validate(), nullValue());
 
         final QueryApiKeyRequest request2 = new QueryApiKeyRequest(
+            null,
             null,
             randomIntBetween(Integer.MIN_VALUE, -1),
             randomIntBetween(0, Integer.MAX_VALUE),
@@ -111,6 +114,7 @@ public void testValidate() {
         assertThat(request2.validate().getMessage(), containsString("[from] parameter cannot be negative"));
 
         final QueryApiKeyRequest request3 = new QueryApiKeyRequest(
+            null,
             null,
             randomIntBetween(0, Integer.MAX_VALUE),
             randomIntBetween(Integer.MIN_VALUE, -1),

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

@@ -34,7 +34,7 @@ protected Writeable.Reader<QueryApiKeyResponse> instanceReader() {
     @Override
     protected QueryApiKeyResponse createTestInstance() {
         final List<QueryApiKeyResponse.Item> items = randomList(0, 3, this::randomItem);
-        return new QueryApiKeyResponse(randomIntBetween(items.size(), 100), items);
+        return new QueryApiKeyResponse(randomIntBetween(items.size(), 100), items, null);
     }
 
     @Override
@@ -43,13 +43,13 @@ protected QueryApiKeyResponse mutateInstance(QueryApiKeyResponse instance) {
         switch (randomIntBetween(0, 3)) {
             case 0:
                 items.add(randomItem());
-                return new QueryApiKeyResponse(instance.getTotal(), items);
+                return new QueryApiKeyResponse(instance.getTotal(), items, null);
             case 1:
                 if (false == items.isEmpty()) {
-                    return new QueryApiKeyResponse(instance.getTotal(), items.subList(1, items.size()));
+                    return new QueryApiKeyResponse(instance.getTotal(), items.subList(1, items.size()), null);
                 } else {
                     items.add(randomItem());
-                    return new QueryApiKeyResponse(instance.getTotal(), items);
+                    return new QueryApiKeyResponse(instance.getTotal(), items, null);
                 }
             case 2:
                 if (false == items.isEmpty()) {
@@ -58,9 +58,9 @@ protected QueryApiKeyResponse mutateInstance(QueryApiKeyResponse instance) {
                 } else {
                     items.add(randomItem());
                 }
-                return new QueryApiKeyResponse(instance.getTotal(), items);
+                return new QueryApiKeyResponse(instance.getTotal(), items, null);
             default:
-                return new QueryApiKeyResponse(instance.getTotal() + 1, items);
+                return new QueryApiKeyResponse(instance.getTotal() + 1, items, null);
         }
     }
 

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/security/authz/privilege/ManageOwnApiKeyClusterPrivilegeTests.java

@@ -259,7 +259,7 @@ public void testCheckQueryApiKeyRequest() {
         final ClusterPermission clusterPermission = ManageOwnApiKeyClusterPrivilege.INSTANCE.buildPermission(ClusterPermission.builder())
             .build();
 
-        final QueryApiKeyRequest queryApiKeyRequest = new QueryApiKeyRequest(null, null, null, null, null, randomBoolean());
+        final QueryApiKeyRequest queryApiKeyRequest = new QueryApiKeyRequest(null, null, null, null, null, null, randomBoolean());
         if (randomBoolean()) {
             queryApiKeyRequest.setFilterForCurrentUser();
         }
@@ -278,7 +278,7 @@ public void testAuthenticationWithApiKeyAllowsDeniesQueryApiKeyWithLimitedBy() {
             .build();
 
         final boolean withLimitedBy = randomBoolean();
-        final QueryApiKeyRequest queryApiKeyRequest = new QueryApiKeyRequest(null, null, null, null, null, withLimitedBy);
+        final QueryApiKeyRequest queryApiKeyRequest = new QueryApiKeyRequest(null, null, null, null, null, null, withLimitedBy);
         queryApiKeyRequest.setFilterForCurrentUser();
         assertThat(
             clusterPermission.check(QueryApiKeyAction.NAME, queryApiKeyRequest, AuthenticationTestHelper.builder().apiKey().build(false)),

x-pack/plugin/security/build.gradle

@@ -27,6 +27,9 @@ dependencies {
   testImplementation project(path: xpackModule('wildcard'))
   testImplementation project(path: ':modules:legacy-geo')
   testImplementation project(path: ':modules:percolator')
+//  Conflict found for the following module:
+//  - org.ow2.asm:asm between versions 8.0.1 and 7.2
+//  testImplementation project(path: ':modules:lang-painless')
   testImplementation project(path: xpackModule('sql:sql-action'))
   testImplementation project(path: ':modules:analysis-common')
   testImplementation project(path: ':modules:reindex')

x-pack/plugin/security/src/internalClusterTest/java/org/elasticsearch/xpack/security/authc/ApiKeyIntegTests.java

@@ -2852,7 +2852,7 @@ private ApiKey getApiKeyInfo(Client client, String apiKeyId, boolean withLimited
             final PlainActionFuture<QueryApiKeyResponse> future = new PlainActionFuture<>();
             client.execute(
                 QueryApiKeyAction.INSTANCE,
-                new QueryApiKeyRequest(QueryBuilders.idsQuery().addIds(apiKeyId), null, null, null, null, withLimitedBy),
+                new QueryApiKeyRequest(QueryBuilders.idsQuery().addIds(apiKeyId), null, null, null, null, null, withLimitedBy),
                 future
             );
             final QueryApiKeyResponse queryApiKeyResponse = future.actionGet();
@@ -2871,7 +2871,7 @@ private ApiKey[] getAllApiKeyInfo(Client client, boolean withLimitedBy) {
             final PlainActionFuture<QueryApiKeyResponse> future = new PlainActionFuture<>();
             client.execute(
                 QueryApiKeyAction.INSTANCE,
-                new QueryApiKeyRequest(QueryBuilders.matchAllQuery(), null, 1000, null, null, withLimitedBy),
+                new QueryApiKeyRequest(QueryBuilders.matchAllQuery(), null, null, 1000, null, null, withLimitedBy),
                 future
             );
             final QueryApiKeyResponse queryApiKeyResponse = future.actionGet();

x-pack/plugin/security/src/internalClusterTest/java/org/elasticsearch/xpack/security/authc/apikey/ApiKeySingleNodeTests.java

@@ -625,6 +625,7 @@ public void testCreateCrossClusterApiKey() throws IOException {
             null,
             null,
             null,
+            null,
             randomBoolean()
         );
         final QueryApiKeyResponse queryApiKeyResponse = client().execute(QueryApiKeyAction.INSTANCE, queryApiKeyRequest).actionGet();
@@ -722,6 +723,7 @@ public void testUpdateCrossClusterApiKey() throws IOException {
             null,
             null,
             null,
+            null,
             randomBoolean()
         );
         final QueryApiKeyResponse queryApiKeyResponse = client().execute(QueryApiKeyAction.INSTANCE, queryApiKeyRequest).actionGet();