Support splitting OPA filter requests in batch mode by batch size

Author: thinaih

plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaConfig.java

@@ -16,6 +16,7 @@
 import io.airlift.configuration.Config;
 import io.airlift.configuration.ConfigDescription;
 import jakarta.validation.constraints.NotNull;
+import jakarta.validation.constraints.Positive;
 
 import java.net.URI;
 import java.util.Optional;
@@ -31,6 +32,7 @@ public class OpaConfig
     private Optional<URI> opaRowFiltersUri = Optional.empty();
     private Optional<URI> opaColumnMaskingUri = Optional.empty();
     private Optional<URI> opaBatchColumnMaskingUri = Optional.empty();
+    private Optional<Integer> opaBatchSize = Optional.empty();
 
     @NotNull
     public URI getOpaUri()
@@ -46,6 +48,19 @@ public OpaConfig setOpaUri(@NotNull URI opaUri)
         return this;
     }
 
+    @Config("opa.policy.batch-size")
+    @ConfigDescription("Size of a single batch for OPA requests")
+    public OpaConfig setOpaBatchSize(Integer batchSize)
+    {
+        this.opaBatchSize = Optional.ofNullable(batchSize);
+        return this;
+    }
+
+    public Optional<@Positive Integer> getOpaBatchSize()
+    {
+        return this.opaBatchSize;
+    }
+
     @NotNull
     public Optional<URI> getOpaBatchUri()
     {

plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaHttpClient.java

@@ -49,6 +49,8 @@
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static com.google.common.collect.ImmutableSetMultimap.flatteningToImmutableSetMultimap;
+import static com.google.common.collect.Lists.partition;
 import static com.google.common.net.HttpHeaders.CONTENT_TYPE;
 import static com.google.common.net.MediaType.JSON_UTF_8;
 import static io.airlift.http.client.FullJsonResponseHandler.createFullJsonResponseHandler;
@@ -66,6 +68,7 @@ public class OpaHttpClient
     private final boolean logRequests;
     private final boolean logResponses;
     private static final Logger log = Logger.get(OpaHttpClient.class);
+    private final Optional<Integer> opaBatchSize;
 
     @Inject
     public OpaHttpClient(
@@ -79,6 +82,7 @@ public OpaHttpClient(
         this.executor = requireNonNull(executor, "executor is null");
         this.logRequests = config.getLogRequests();
         this.logResponses = config.getLogResponses();
+        this.opaBatchSize = config.getOpaBatchSize();
     }
 
     public <T> FluentFuture<T> submitOpaRequest(OpaQueryInput input, URI uri, JsonCodec<T> deserializer)
@@ -158,25 +162,30 @@ public <T> Set<T> batchFilterFromOpa(Collection<T> items, Function<List<T>, OpaQ
 
     public <K, V> Map<K, Set<V>> parallelBatchFilterFromOpa(Map<K, ? extends Collection<V>> items, BiFunction<K, List<V>, OpaQueryInput> requestBuilder, URI uri, JsonCodec<? extends OpaBatchQueryResult> deserializer)
     {
-        List<Map.Entry<K, ImmutableList<V>>> parallelRequestItems = items.entrySet()
+        List<Map.Entry<K, List<V>>> parallelRequestItems = items.entrySet()
                 .stream()
                 .filter(entry -> !entry.getValue().isEmpty())
-                .map(entry -> Map.entry(entry.getKey(), ImmutableList.copyOf(entry.getValue())))
+                .flatMap(entry -> partition(ImmutableList.copyOf(entry.getValue()), this.opaBatchSize.orElse(entry.getValue().size()))
+                        .stream()
+                        .map(partition -> Map.entry(entry.getKey(), partition)))
                 .collect(toImmutableList());
+
         return parallelRequest(
-                    parallelRequestItems,
-                    entry -> requestBuilder.apply(entry.getKey(), entry.getValue()),
-                    (entry, result) ->
-                            Optional.of(requireNonNullElse(result.result(), ImmutableList.<Integer>of()))
-                                    .flatMap(indices -> indices.isEmpty() ? Optional.empty() : Optional.of(indices))
-                                    .map(indices -> indices.stream()
-                                        .map(index -> entry.getValue().get(index))
-                                        .collect(toImmutableSet()))
-                                    .map(values -> Map.entry(entry.getKey(), values)),
-                    uri,
-                    deserializer)
+                parallelRequestItems,
+                entry -> requestBuilder.apply(entry.getKey(), entry.getValue()),
+                (entry, result) ->
+                        Optional.of(requireNonNullElse(result.result(), ImmutableList.<Integer>of()))
+                                .flatMap(indices -> indices.isEmpty() ? Optional.empty() : Optional.of(indices))
+                                .map(indices -> indices.stream().map(index -> entry.getValue().get(index)).collect(toImmutableSet()))
+                                .map(values -> Map.entry(entry.getKey(), values)),
+                uri,
+                deserializer)
+                .stream()
+                .collect(flatteningToImmutableSetMultimap(Map.Entry::getKey, entry -> entry.getValue().stream()))
+                .asMap()
+                .entrySet()
                 .stream()
-                .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+                .collect(toImmutableMap(Map.Entry::getKey, entry -> ImmutableSet.copyOf(entry.getValue())));
     }
 
     private <T> T parseOpaResponse(FullJsonResponseHandler.JsonResponse<T> response, URI uri)

plugin/trino-opa/src/test/java/io/trino/plugin/opa/RequestTestUtilities.java

@@ -34,17 +34,20 @@ private RequestTestUtilities() {}
 
     private static final JsonMapper jsonMapper = new JsonMapper();
 
+    public static JsonNode toJsonNode(String item)
+    {
+        try {
+            return jsonMapper.readTree(item);
+        }
+        catch (IOException e) {
+            throw new IllegalStateException("Cannot parse to a JsonNode", e);
+        }
+    }
+
     public static void assertStringRequestsEqual(Set<String> expectedRequests, Collection<JsonNode> actualRequests, String extractPath)
     {
         Set<JsonNode> parsedExpectedRequests = expectedRequests.stream()
-                .map(expectedRequest -> {
-                    try {
-                        return jsonMapper.readTree(expectedRequest);
-                    }
-                    catch (IOException e) {
-                        throw new AssertionError("Cannot parse expected request", e);
-                    }
-                })
+                .map(RequestTestUtilities::toJsonNode)
                 .collect(toImmutableSet());
         Set<JsonNode> extractedActualRequests = actualRequests.stream().map(node -> node.at(extractPath)).collect(toImmutableSet());
         assertThat(extractedActualRequests).containsExactlyInAnyOrderElementsOf(parsedExpectedRequests);

plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestHelpers.java

@@ -92,6 +92,7 @@ public static Map<String, String> opaConfigToDict(OpaConfig config)
         config.getOpaRowFiltersUri().ifPresent(rowFiltersUri -> configBuilder.put("opa.policy.row-filters-uri", rowFiltersUri.toString()));
         config.getOpaColumnMaskingUri().ifPresent(columnMaskingUri -> configBuilder.put("opa.policy.column-masking-uri", columnMaskingUri.toString()));
         config.getOpaBatchColumnMaskingUri().ifPresent(batchColumnMaskingUri -> configBuilder.put("opa.policy.batch-column-masking-uri", batchColumnMaskingUri.toString()));
+        config.getOpaBatchSize().ifPresent(batchSize -> configBuilder.put("opa.policy.batch-size", batchSize.toString()));
         return configBuilder.buildOrThrow();
     }