Support batch-size for batch mode

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,11 +162,14 @@ 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()),
@@ -176,7 +183,11 @@ public <K, V> Map<K, Set<V>> parallelBatchFilterFromOpa(Map<K, ? extends Collect
                     uri,
                     deserializer)
                 .stream()
-                .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
+                .collect(flatteningToImmutableSetMultimap(Map.Entry::getKey, entry -> entry.getValue().stream()))
+                .asMap()
+                .entrySet()
+                .stream()
+                .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/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(batchUri -> configBuilder.put("opa.policy.batch-size", batchUri.toString()));
         return configBuilder.buildOrThrow();
     }