From 7e21d9dcee361002f8cd9d4ebf8ad89d9435eebd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Patrick=20Sodr=C3=A9?= <psodre@gmail.com>
Date: Tue, 9 Mar 2021 09:14:35 -0500
Subject: [PATCH] Add support to annotationTypes using multivalued query
 parameters.

---
 src/main/java/co/zeroae/gate/App.java     | 20 ++++++++++++++------
 src/test/java/co/zeroae/gate/AppTest.java | 22 +++++++++++++++++++---
 2 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/src/main/java/co/zeroae/gate/App.java b/src/main/java/co/zeroae/gate/App.java
index 4194f18..d33540c 100644
--- a/src/main/java/co/zeroae/gate/App.java
+++ b/src/main/java/co/zeroae/gate/App.java
@@ -63,6 +63,7 @@ public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent in
         final APIGatewayProxyResponseEvent response = new APIGatewayProxyResponseEvent()
                 .withHeaders(new HashMap<>());
         final Map<String, String> queryStringParams = Optional.ofNullable(input.getQueryStringParameters()).orElse(new HashMap<>());
+        final Map<String, List<String>> mQueryStringParams = Optional.ofNullable(input.getMultiValueQueryStringParameters()).orElse(new HashMap<>());
         try {
             final String acceptHeader = input.getHeaders().getOrDefault("Accept", "application/json");
             final String responseType = ((Supplier<String>) () -> {
@@ -101,8 +102,10 @@ public APIGatewayProxyResponseEvent handleRequest(APIGatewayProxyRequestEvent in
 
             AWSXRay.beginSubsegment("Gate Export");
             AWSXRay.getCurrentSubsegment().putMetadata("Content-Type", response.getHeaders().get("Content-Type"));
+            final List<String> annotationSelector = mQueryStringParams.get("annotations");
+
             try {
-                return export(exporter, doc, response).withStatusCode(200);
+                return export(exporter, doc, annotationSelector, response).withStatusCode(200);
             } finally {
                 Factory.deleteResource(doc);
                 AWSXRay.endSubsegment();
@@ -166,23 +169,28 @@ private Document execute(FeatureMap docFeatureMap) throws GateException {
     /**
      * @param exporter The document exporter
      * @param doc an instance of gate.Document
+     * @param annotationSelector the List of AnnotationTypes to return
      * @param response The response where we put the exported Document as body
      * @return the modified response
      */
     private APIGatewayProxyResponseEvent export(
             DocumentExporter exporter,
             Document doc,
+            List<String> annotationSelector,
             APIGatewayProxyResponseEvent response
     ) throws IOException {
         final FeatureMap exportOptions = Factory.newFeatureMap();
 
         // Take *all* annotation types.
         final AnnotationSet defaultAnnots = doc.getAnnotations();
-        final HashSet<String> annotationTypes = new HashSet<>();
-        for (Annotation annotation : defaultAnnots.inDocumentOrder()) {
-            annotationTypes.add(annotation.getType());
-        }
-        exportOptions.put("annotationTypes", annotationTypes);
+        final Set<String> excludedTypes = new HashSet<>(defaultAnnots.getAllTypes());
+        Optional.ofNullable(annotationSelector).ifPresent((selector) ->
+                doc.getFeatures().put("gate.cloud.annotationSelectors", annotationSelector));
+        excludedTypes.removeIf((type) ->
+                !Optional.ofNullable(annotationSelector).isPresent() || annotationSelector.contains(":"+ type)
+        );
+
+        defaultAnnots.removeIf((annotation) -> excludedTypes.contains(annotation.getType()));
 
         final ByteArrayOutputStream baos = new ByteArrayOutputStream();
         try {
diff --git a/src/test/java/co/zeroae/gate/AppTest.java b/src/test/java/co/zeroae/gate/AppTest.java
index 8307cf7..eb737f4 100644
--- a/src/test/java/co/zeroae/gate/AppTest.java
+++ b/src/test/java/co/zeroae/gate/AppTest.java
@@ -23,9 +23,7 @@
 import java.io.IOException;
 import java.io.StringReader;
 import java.io.UnsupportedEncodingException;
-import java.util.HashMap;
-import java.util.Random;
-import java.util.UUID;
+import java.util.*;
 
 import static com.github.stefanbirkner.systemlambda.SystemLambda.withEnvironmentVariable;
 import static org.junit.Assert.*;
@@ -83,6 +81,24 @@ public void testGateXMLToDocument() throws Exception {
         assertEquals(input.getBody(), doc.getContent().toString());
     }
 
+    @Test
+    public void testAnnotationSelector() throws Exception {
+        APIGatewayProxyResponseEvent result = app.handleRequest(input, context);
+        assertEquals("application/gate+xml", result.getHeaders().get("Content-Type"));
+
+        Document doc = Utils.xmlToDocument(new StringReader(result.getBody()));
+        assertEquals(7, doc.getAnnotations().getAllTypes().size());
+
+        // Now we downselect to only one field.
+        input.withMultiValueQueryStringParameters(new HashMap<>())
+                .getMultiValueQueryStringParameters()
+               .put("annotations", Collections.singletonList(":Token"));
+        result = app.handleRequest(input, context);
+
+        doc = Utils.xmlToDocument(new StringReader(result.getBody()));
+        assertEquals(1, doc.getAnnotations().getAllTypes().size());
+    }
+
     @Test
     public void testNextAnnotationId() {
         final int nextAnnotationId = 1000 + new Random().nextInt(1000);