From 7b96273a28c060d82ebc710154fa1c3c84ba0c0f Mon Sep 17 00:00:00 2001
From: AFine-gs <69924417+AFine-gs@users.noreply.github.com>
Date: Fri, 29 Sep 2023 13:10:24 -0400
Subject: [PATCH] add legend response format header to http response (#2286)

add legend response format header to http response
---
 .../main/java/org/finos/legend/engine/server/Server.java    | 3 ++-
 .../engine/plan/execution/api/result/ResultManager.java     | 5 ++++-
 .../engine/plan/execution/result/StreamingResult.java       | 2 ++
 .../shared/runtime/write/ExternalFormatSerializeResult.java | 6 ++++++
 .../engine/external/format/arrow/ArrowRuntimeExtension.java | 2 +-
 .../src/test/java/TestArrowNodeExecutor.java                | 2 +-
 .../external/format/flatdata/FlatDataRuntimeExtension.java  | 2 +-
 .../external/format/json/JsonSchemaRuntimeExtension.java    | 2 +-
 8 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/legend-engine-config/legend-engine-server/src/main/java/org/finos/legend/engine/server/Server.java b/legend-engine-config/legend-engine-server/src/main/java/org/finos/legend/engine/server/Server.java
index b4236e94c01..60fa1a7cfdc 100644
--- a/legend-engine-config/legend-engine-server/src/main/java/org/finos/legend/engine/server/Server.java
+++ b/legend-engine-config/legend-engine-server/src/main/java/org/finos/legend/engine/server/Server.java
@@ -80,6 +80,7 @@
 import org.finos.legend.engine.plan.execution.api.concurrent.ConcurrentExecutionNodeExecutorPoolInfo;
 import org.finos.legend.engine.plan.execution.api.concurrent.ParallelGraphFetchExecutionExecutorPoolInfo;
 import org.finos.legend.engine.plan.execution.api.request.RequestContextHelper;
+import org.finos.legend.engine.plan.execution.api.result.ResultManager;
 import org.finos.legend.engine.plan.execution.concurrent.ParallelGraphFetchExecutionExecutorPool;
 import org.finos.legend.engine.plan.execution.graphFetch.GraphFetchExecutionConfiguration;
 import org.finos.legend.engine.plan.execution.service.api.ServiceModelingApi;
@@ -410,7 +411,7 @@ private void enableCors(Environment environment)
         corsFilter.setInitParameter(CrossOriginFilter.ALLOWED_METHODS_PARAM, "GET,PUT,POST,DELETE,OPTIONS");
         corsFilter.setInitParameter(CrossOriginFilter.ALLOWED_ORIGINS_PARAM, "*");
         corsFilter.setInitParameter(CrossOriginFilter.ALLOWED_TIMING_ORIGINS_PARAM, "*");
-        corsFilter.setInitParameter(CrossOriginFilter.ALLOWED_HEADERS_PARAM, "X-Requested-With,Content-Type,Accept,Origin,Access-Control-Allow-Credentials,x-b3-parentspanid,x-b3-sampled,x-b3-spanid,x-b3-traceid," + RequestContextHelper.LEGEND_REQUEST_ID + RequestContextHelper.LEGEND_USE_PLAN_CACHE);
+        corsFilter.setInitParameter(CrossOriginFilter.ALLOWED_HEADERS_PARAM, "X-Requested-With,Content-Type,Accept,Origin,Access-Control-Allow-Credentials,x-b3-parentspanid,x-b3-sampled,x-b3-spanid,x-b3-traceid," + RequestContextHelper.LEGEND_REQUEST_ID + "," + RequestContextHelper.LEGEND_USE_PLAN_CACHE + "," + ResultManager.LEGEND_RESPONSE_FORMAT);
         corsFilter.setInitParameter(CrossOriginFilter.CHAIN_PREFLIGHT_PARAM, "false");
         corsFilter.addMappingForUrlPatterns(EnumSet.of(DispatcherType.REQUEST), false, "*");
     }
diff --git a/legend-engine-core/legend-engine-core-executionPlan-execution/legend-engine-executionPlan-execution-api/src/main/java/org/finos/legend/engine/plan/execution/api/result/ResultManager.java b/legend-engine-core/legend-engine-core-executionPlan-execution/legend-engine-executionPlan-execution-api/src/main/java/org/finos/legend/engine/plan/execution/api/result/ResultManager.java
index 9fa9ec08097..dea6d0a7794 100644
--- a/legend-engine-core/legend-engine-core-executionPlan-execution/legend-engine-executionPlan-execution-api/src/main/java/org/finos/legend/engine/plan/execution/api/result/ResultManager.java
+++ b/legend-engine-core/legend-engine-core-executionPlan-execution/legend-engine-executionPlan-execution-api/src/main/java/org/finos/legend/engine/plan/execution/api/result/ResultManager.java
@@ -40,6 +40,9 @@ public class ResultManager
     private static final Logger LOGGER = org.slf4j.LoggerFactory.getLogger(ResultManager.class);
     private static final JsonStringEncoder jsonStringEncoder = JsonStringEncoder.getInstance();
 
+    public static final String LEGEND_RESPONSE_FORMAT = "x-legend-response-format";
+
+
     public static Response manageResult(MutableList<CommonProfile> pm, Result result, LoggingEventType loggingEventType)
     {
         return manageResult(pm, result, SerializationFormat.defaultFormat, loggingEventType);
@@ -68,7 +71,7 @@ private static Response manageResultWithCustomErrorCodeImpl(MutableList<CommonPr
         }
         else if (result instanceof StreamingResult)
         {
-            return Response.ok(new StreamingResultHandler((StreamingResult) result, format)).build();
+            return Response.ok(new StreamingResultHandler((StreamingResult) result, format)).header(LEGEND_RESPONSE_FORMAT, ((StreamingResult) result).resultFormat).build();
         }
         else if (result instanceof ConstantResult)
         {
diff --git a/legend-engine-core/legend-engine-core-executionPlan-execution/legend-engine-executionPlan-execution/src/main/java/org/finos/legend/engine/plan/execution/result/StreamingResult.java b/legend-engine-core/legend-engine-core-executionPlan-execution/legend-engine-executionPlan-execution/src/main/java/org/finos/legend/engine/plan/execution/result/StreamingResult.java
index a8efd0ca930..62a55912e5e 100644
--- a/legend-engine-core/legend-engine-core-executionPlan-execution/legend-engine-executionPlan-execution/src/main/java/org/finos/legend/engine/plan/execution/result/StreamingResult.java
+++ b/legend-engine-core/legend-engine-core-executionPlan-execution/legend-engine-executionPlan-execution/src/main/java/org/finos/legend/engine/plan/execution/result/StreamingResult.java
@@ -26,6 +26,8 @@
 
 public abstract class StreamingResult extends Result
 {
+    public String resultFormat = "FormatNotSet";
+
     public abstract Builder getResultBuilder();
 
     public abstract Serializer getSerializer(SerializationFormat format);
diff --git a/legend-engine-core/legend-engine-core-executionPlan-execution/legend-engine-external-shared-format-runtime/src/main/java/org/finos/legend/engine/external/shared/runtime/write/ExternalFormatSerializeResult.java b/legend-engine-core/legend-engine-core-executionPlan-execution/legend-engine-external-shared-format-runtime/src/main/java/org/finos/legend/engine/external/shared/runtime/write/ExternalFormatSerializeResult.java
index 697bb270bbf..1c41eec625f 100644
--- a/legend-engine-core/legend-engine-core-executionPlan-execution/legend-engine-external-shared-format-runtime/src/main/java/org/finos/legend/engine/external/shared/runtime/write/ExternalFormatSerializeResult.java
+++ b/legend-engine-core/legend-engine-core-executionPlan-execution/legend-engine-external-shared-format-runtime/src/main/java/org/finos/legend/engine/external/shared/runtime/write/ExternalFormatSerializeResult.java
@@ -28,10 +28,16 @@ public class ExternalFormatSerializeResult extends StreamingResult
     private final Result childResult;
 
     public ExternalFormatSerializeResult(ExternalFormatWriter externalFormatWriter, Result childResult)
+    {
+        this(externalFormatWriter, childResult, "");
+    }
+
+    public ExternalFormatSerializeResult(ExternalFormatWriter externalFormatWriter, Result childResult, String resultFormat)
     {
         super(Lists.mutable.empty());
         this.externalFormatWriter = externalFormatWriter;
         this.childResult = childResult;
+        this.resultFormat = resultFormat;
     }
 
     @Override
diff --git a/legend-engine-xts-arrow/legend-engine-xt-arrow-runtime/src/main/java/org/finos/legend/engine/external/format/arrow/ArrowRuntimeExtension.java b/legend-engine-xts-arrow/legend-engine-xt-arrow-runtime/src/main/java/org/finos/legend/engine/external/format/arrow/ArrowRuntimeExtension.java
index 3aba69cc8e5..372a03f85ac 100644
--- a/legend-engine-xts-arrow/legend-engine-xt-arrow-runtime/src/main/java/org/finos/legend/engine/external/format/arrow/ArrowRuntimeExtension.java
+++ b/legend-engine-xts-arrow/legend-engine-xt-arrow-runtime/src/main/java/org/finos/legend/engine/external/format/arrow/ArrowRuntimeExtension.java
@@ -65,7 +65,7 @@ public Result executeExternalizeTDSExecutionNode(ExternalFormatExternalizeTDSExe
     private Result streamArrowFromRelational(RelationalResult relationalResult) throws SQLException, IOException
     {
 
-        return new ExternalFormatSerializeResult(new ArrowDataWriter(relationalResult.getResultSet()), relationalResult);
+        return new ExternalFormatSerializeResult(new ArrowDataWriter(relationalResult.getResultSet()), relationalResult, CONTENT_TYPE);
 
     }
 
diff --git a/legend-engine-xts-arrow/legend-engine-xt-arrow-runtime/src/test/java/TestArrowNodeExecutor.java b/legend-engine-xts-arrow/legend-engine-xt-arrow-runtime/src/test/java/TestArrowNodeExecutor.java
index dd39c8d4504..a2cf57af287 100644
--- a/legend-engine-xts-arrow/legend-engine-xt-arrow-runtime/src/test/java/TestArrowNodeExecutor.java
+++ b/legend-engine-xts-arrow/legend-engine-xt-arrow-runtime/src/test/java/TestArrowNodeExecutor.java
@@ -66,7 +66,7 @@ public void testExternalize() throws Exception
             RelationalResult result = new RelationalResult(FastList.newListWith(new RelationalExecutionActivity("SELECT * FROM testtable", null)), mockExecutionNode, FastList.newListWith(new SQLResultColumn("testInt", "INTEGER"), new SQLResultColumn("testString", "VARCHAR"), new SQLResultColumn("testDate", "TIMESTAMP"), new SQLResultColumn("testBool", "TIMESTAMP")), null, "GMT", conn, null, null, null, new RequestContext());
 
             ExternalFormatSerializeResult nodeExecute = (ExternalFormatSerializeResult) extension.executeExternalizeTDSExecutionNode(node, result, null, null);
-
+            Assert.assertEquals(nodeExecute.resultFormat, "application/x.arrow");
             nodeExecute.stream(outputStream, SerializationFormat.DEFAULT);
             assertArrow(outputStream, "TESTINT\tTESTSTRING\tTESTDATE\tTESTBOOL\n" +
                     "1\tA\t1577854800000\ttrue\n" +
diff --git a/legend-engine-xts-flatdata/legend-engine-xt-flatdata-runtime/src/main/java/org/finos/legend/engine/external/format/flatdata/FlatDataRuntimeExtension.java b/legend-engine-xts-flatdata/legend-engine-xt-flatdata-runtime/src/main/java/org/finos/legend/engine/external/format/flatdata/FlatDataRuntimeExtension.java
index 71fba91607a..8f55d6d1d6c 100644
--- a/legend-engine-xts-flatdata/legend-engine-xt-flatdata-runtime/src/main/java/org/finos/legend/engine/external/format/flatdata/FlatDataRuntimeExtension.java
+++ b/legend-engine-xts-flatdata/legend-engine-xt-flatdata-runtime/src/main/java/org/finos/legend/engine/external/format/flatdata/FlatDataRuntimeExtension.java
@@ -85,7 +85,7 @@ public Result executeExternalizeExecutionNode(ExternalFormatExternalizeExecution
             FlatDataContext<?> context = specifics.createContext();
             FlatDataWriter<?> serializer = new FlatDataWriter(context, inputStream);
 
-            return new ExternalFormatSerializeResult(serializer, result);
+            return new ExternalFormatSerializeResult(serializer, result, CONTENT_TYPE);
         }
         catch (Exception e)
         {
diff --git a/legend-engine-xts-json/legend-engine-xt-json-runtime/src/main/java/org/finos/legend/engine/external/format/json/JsonSchemaRuntimeExtension.java b/legend-engine-xts-json/legend-engine-xt-json-runtime/src/main/java/org/finos/legend/engine/external/format/json/JsonSchemaRuntimeExtension.java
index da8d9e28384..2920438da88 100644
--- a/legend-engine-xts-json/legend-engine-xt-json-runtime/src/main/java/org/finos/legend/engine/external/format/json/JsonSchemaRuntimeExtension.java
+++ b/legend-engine-xts-json/legend-engine-xt-json-runtime/src/main/java/org/finos/legend/engine/external/format/json/JsonSchemaRuntimeExtension.java
@@ -104,7 +104,7 @@ public Result executeExternalizeExecutionNode(ExternalFormatExternalizeExecution
             {
                 IJsonExternalizeExecutionNodeSpecifics nodeSpecifics = (IJsonExternalizeExecutionNodeSpecifics) specificsClass.getConstructor().newInstance();
                 JsonDataWriter<?> jsonDataWriter = new JsonDataWriter<>(nodeSpecifics, extractStreamFromResult(result), context);
-                return new ExternalFormatSerializeResult(jsonDataWriter, result);
+                return new ExternalFormatSerializeResult(jsonDataWriter, result, CONTENT_TYPE);
             }
             else
             {