Gh-2989: Add configurable GetAllElements Limit to Tinkerpop

library/tinkerpop/src/main/java/uk/gov/gchq/gaffer/tinkerpop/GafferPopGraph.java

@@ -16,6 +16,7 @@
 
 package uk.gov.gchq.gaffer.tinkerpop;
 
+import org.apache.commons.collections4.IterableUtils;
 import org.apache.commons.configuration2.Configuration;
 import org.apache.tinkerpop.gremlin.process.computer.GraphComputer;
 import org.apache.tinkerpop.gremlin.process.traversal.TraversalStrategies;
@@ -44,6 +45,7 @@
 import uk.gov.gchq.gaffer.operation.data.ElementSeed;
 import uk.gov.gchq.gaffer.operation.data.EntitySeed;
 import uk.gov.gchq.gaffer.operation.graph.SeededGraphFilters.IncludeIncomingOutgoingType;
+import uk.gov.gchq.gaffer.operation.impl.Limit;
 import uk.gov.gchq.gaffer.operation.impl.add.AddElements;
 import uk.gov.gchq.gaffer.operation.impl.get.GetAdjacentIds;
 import uk.gov.gchq.gaffer.operation.impl.get.GetAllElements;
@@ -170,6 +172,16 @@ public class GafferPopGraph implements org.apache.tinkerpop.gremlin.structure.Gr
      */
     public static final String OP_OPTIONS = "gaffer.operation.options";
 
+    /**
+     * Configuration key for the max number of elements returned by getAllElements
+     */
+    public static final String GET_ALL_ELEMENTS_LIMIT = "gaffer.elements.getalllimit";
+
+    /**
+     * Default value for the max number of elements returned by getAllElements
+     */
+    public static final int DEFAULT_GET_ALL_ELEMENTS_LIMIT = 5000;
+
     public static final String USER_ID = "gaffer.userId";
 
     public static final String DATA_AUTHS = "gaffer.dataAuths";
@@ -213,6 +225,7 @@ public enum DefaultIdManager {
     private final Map<String, String> opOptions;
     private final User defaultUser;
     private final ServiceRegistry serviceRegistry;
+    private final int getAllElementsLimit;
 
     private static final Logger LOGGER = LoggerFactory.getLogger(GafferPopGraph.class);
 
@@ -237,6 +250,8 @@ public GafferPopGraph(final Configuration configuration, final Graph graph) {
                 .dataAuths(configuration().getStringArray(DATA_AUTHS))
                 .build();
 
+        getAllElementsLimit = configuration.getInteger(GET_ALL_ELEMENTS_LIMIT, DEFAULT_GET_ALL_ELEMENTS_LIMIT);
+
         // Set the graph variables to current config
         variables = new GafferPopGraphVariables();
         setDefaultVariables(variables);
@@ -367,10 +382,13 @@ public Iterator<Vertex> vertices(final Object... vertexIds) {
 
         final Output<Iterable<? extends Element>> getOperation;
         if (getAll) {
-            getOperation = new GetAllElements.Builder()
-                    .view(new View.Builder()
-                            .entities(graph.getSchema().getEntityGroups())
+            getOperation = new Builder()
+                    .first(new GetAllElements.Builder()
+                            .view(new View.Builder()
+                                    .entities(graph.getSchema().getEntityGroups())
+                                    .build())
                             .build())
+                    .then(new Limit<>(getAllElementsLimit, true))
                     .build();
         } else {
             getOperation = new GetElements.Builder()
@@ -393,6 +411,12 @@ public Iterator<Vertex> vertices(final Object... vertexIds) {
                 .map(e -> (Vertex) e)
                 .iterator();
 
+        final int resultSize = IterableUtils.size(translatedResults);
+        if (getAll && resultSize >= getAllElementsLimit) {
+            LOGGER.warn("Result size is equal to configured limit ({}). Results may have been truncated", getAllElementsLimit);
+        }
+
+
         return translatedResults.iterator();
     }
 
@@ -518,11 +542,14 @@ public Iterator<Edge> edges(final Object... elementIds) {
 
         final Output<Iterable<? extends Element>> getOperation;
         if (getAll) {
-            getOperation = new GetAllElements.Builder()
-                    .view(new View.Builder()
+            getOperation = new Builder()
+                .first(new GetAllElements.Builder()
+                        .view(new View.Builder()
                             .edges(graph.getSchema().getEdgeGroups())
                             .build())
-                    .build();
+                        .build())
+                .then(new Limit<>(getAllElementsLimit, true))
+                .build();
         } else {
             getOperation = new GetElements.Builder()
                     .input(getElementSeeds(Arrays.asList(elementIds)))
@@ -545,6 +572,11 @@ public Iterator<Edge> edges(final Object... elementIds) {
                 .map(e -> (Edge) e)
                 .iterator();
 
+        final int resultSize = IterableUtils.size(translatedResults);
+        if (getAll && resultSize >= getAllElementsLimit) {
+            LOGGER.warn("Result size is equal to configured limit ({}). Results may have been truncated", getAllElementsLimit);
+        }
+
         return translatedResults.iterator();
     }
 
@@ -913,6 +945,7 @@ private void setDefaultVariables(final GafferPopGraphVariables variables) {
         variables.set(GafferPopGraphVariables.OP_OPTIONS, Collections.unmodifiableMap(opOptions));
         variables.set(GafferPopGraphVariables.USER_ID, defaultUser.getUserId());
         variables.set(GafferPopGraphVariables.DATA_AUTHS, configuration().getStringArray(DATA_AUTHS));
+        variables.set(GafferPopGraphVariables.GET_ALL_ELEMENTS_LIMIT, getAllElementsLimit);
     }
 
     /**

library/tinkerpop/src/main/java/uk/gov/gchq/gaffer/tinkerpop/GafferPopGraphVariables.java

@@ -45,6 +45,11 @@ public final class GafferPopGraphVariables implements Graph.Variables {
      */
     public static final String USER_ID = "userId";
 
+    /**
+     * The max number of elements that can be returned by GetAllElements
+     */
+    public static final String GET_ALL_ELEMENTS_LIMIT = "getAllElementsLimit";
+
     private final Map<String, Object> variables;
 
     public GafferPopGraphVariables(final Map<String, Object> variables) {

library/tinkerpop/src/test/java/uk/gov/gchq/gaffer/tinkerpop/GafferPopGraphIT.java

@@ -74,10 +74,11 @@ public void shouldConstructGafferPopGraphWithOnlyConfig() {
         // Then
         final Map<String, Object> variables = graph.variables().asMap();
         assertThat(variables.get(GafferPopGraphVariables.USER_ID)).isEqualTo(expectedUser.getUserId());
+        assertThat(variables.get(GafferPopGraphVariables.GET_ALL_ELEMENTS_LIMIT)).isEqualTo(1);
 
         final Map<String, String> opOptions = (Map<String, String>) variables.get(GafferPopGraphVariables.OP_OPTIONS);
         assertThat(opOptions).containsEntry("key1", "value1").containsEntry("key2", "value2").hasSize(2);
-        assertThat(variables.size()).isEqualTo(3);
+        assertThat(variables.size()).isEqualTo(4);
     }
 
     @Test
@@ -91,10 +92,11 @@ public void shouldConstructGafferPopGraphWithConfigFile() {
         // Then
         final Map<String, Object> variables = graph.variables().asMap();
         assertThat(variables.get(GafferPopGraphVariables.USER_ID)).isEqualTo(expectedUser.getUserId());
+        assertThat(variables.get(GafferPopGraphVariables.GET_ALL_ELEMENTS_LIMIT)).isEqualTo(GafferPopGraph.DEFAULT_GET_ALL_ELEMENTS_LIMIT);
 
         final Map<String, String> opOptions = (Map<String, String>) variables.get(GafferPopGraphVariables.OP_OPTIONS);
         assertThat(opOptions).containsEntry("key1", "value1").hasSize(1);
-        assertThat(variables.size()).isEqualTo(3);
+        assertThat(variables.size()).isEqualTo(4);
     }
 
     @Test
@@ -110,15 +112,16 @@ public void shouldConstructGafferPopGraph() {
         final Map<String, Object> variables = graph.variables().asMap();
         assertThat(variables)
             .containsEntry(GafferPopGraphVariables.DATA_AUTHS, expectedUser.getDataAuths().toArray())
-            .containsEntry(GafferPopGraphVariables.USER_ID, expectedUser.getUserId());
+            .containsEntry(GafferPopGraphVariables.USER_ID, expectedUser.getUserId())
+            .containsEntry(GafferPopGraphVariables.GET_ALL_ELEMENTS_LIMIT, GafferPopGraph.DEFAULT_GET_ALL_ELEMENTS_LIMIT);
 
         final Map<String, String> opOptions = (Map<String, String>) variables.get(GafferPopGraphVariables.OP_OPTIONS);
         assertThat(opOptions).containsEntry("key1", "value1").containsEntry("key2", "value2").hasSize(2);
-        assertThat(variables.size()).isEqualTo(3);
+        assertThat(variables.size()).isEqualTo(4);
     }
 
     @Test
-    public void shouldThrowUnsupportedExceptionForNoGraphId() {
+    public void shouldThrowIllegalArgumentExceptionForNoGraphId() {
 
         // Given/Then
         assertThatExceptionOfType(IllegalArgumentException.class)
@@ -249,14 +252,26 @@ public void shouldGetAllVertices() {
         final Iterator<Vertex> vertices = graph.vertices();
 
         // Then
-        final List<Vertex> verticesList = new ArrayList<>();
-        while (vertices.hasNext()) {
-            verticesList.add(vertices.next());
-        }
-        assertThat(verticesList).contains(
+        assertThat(vertices)
+            .toIterable()
+            .contains(
                 new GafferPopVertex(SOFTWARE_NAME_GROUP, VERTEX_1, graph),
-                new GafferPopVertex(SOFTWARE_NAME_GROUP, VERTEX_2, graph)
-        );
+                new GafferPopVertex(SOFTWARE_NAME_GROUP, VERTEX_2, graph));
+    }
+
+    @Test
+    public void shouldTruncateGetAllVertices() {
+        // Given
+        final Graph gafferGraph = getGafferGraph();
+        final GafferPopGraph graph = GafferPopGraph.open(TEST_CONFIGURATION_2, gafferGraph);
+
+        // When
+        addSoftwareVertex(graph);
+        graph.addVertex(T.label, PERSON_GROUP, T.id, VERTEX_2, NAME_PROPERTY, "Gaffer");
+        final Iterator<Vertex> vertices = graph.vertices();
+
+        // Then
+        assertThat(vertices).toIterable().hasSize(1);
     }
 
     @Test
@@ -439,6 +454,23 @@ public void shouldGetAllEdges() {
         assertThat(edges).toIterable().contains(edgeToAdd1, edgeToAdd2);
     }
 
+    @Test
+    public void shouldTruncateGetAllEdges() {
+        // Given
+        final Graph gafferGraph = getGafferGraph();
+        final GafferPopGraph graph = GafferPopGraph.open(TEST_CONFIGURATION_2, gafferGraph);
+        final GafferPopEdge edgeToAdd1 = new GafferPopEdge(CREATED_EDGE_GROUP, VERTEX_1, VERTEX_2, graph);
+        final GafferPopEdge edgeToAdd2 = new GafferPopEdge(DEPENDS_ON_EDGE_GROUP, VERTEX_2, VERTEX_1, graph);
+        graph.addEdge(edgeToAdd1);
+        graph.addEdge(edgeToAdd2);
+
+        // When
+        final Iterator<Edge> edges = graph.edges();
+
+        // Then
+        assertThat(edges).toIterable().hasSize(1);
+    }
+
     @Test
     public void shouldGetAllEdgesInGroup() {
         // Given

library/tinkerpop/src/test/java/uk/gov/gchq/gaffer/tinkerpop/util/GafferPopTestUtil.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 Crown Copyright
+ * Copyright 2023-2024 Crown Copyright
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -51,6 +51,7 @@ private GafferPopTestUtil() {
             this.setProperty(GafferPopGraph.DATA_AUTHS, new String[]{AUTH_1, AUTH_2});
             this.setProperty(GafferPopGraph.GRAPH_ID, "Graph1");
             this.setProperty(GafferPopGraph.STORE_PROPERTIES, GafferPopGraphIT.class.getClassLoader().getResource("gaffer/store.properties").getPath());
+            this.setProperty(GafferPopGraph.GET_ALL_ELEMENTS_LIMIT, 1);
         }
     };