fix(openapiv3): v3 scroll response fix (#10654)

Co-authored-by: Kevin Chun <[email protected]>

metadata-service/openapi-analytics-servlet/build.gradle

@@ -9,6 +9,7 @@ dependencies {
     implementation project(':metadata-service:auth-impl')
     implementation project(':metadata-service:factories')
     implementation project(':metadata-service:openapi-servlet')
+    implementation project(':metadata-service:openapi-servlet:models')
     implementation project(':metadata-models')
 
     implementation externalDependency.springBoot

metadata-service/openapi-entity-servlet/build.gradle

@@ -9,6 +9,7 @@ dependencies {
   implementation project(':metadata-service:auth-impl')
   implementation project(':metadata-service:factories')
   implementation project(':metadata-service:openapi-servlet')
+  implementation project(':metadata-service:openapi-servlet:models')
   implementation project(':metadata-models')
 
   implementation externalDependency.servletApi

metadata-service/openapi-servlet/models/build.gradle

@@ -6,6 +6,14 @@ dependencies {
   implementation project(':entity-registry')
   implementation project(':metadata-operation-context')
   implementation project(':metadata-auth:auth-api')
+  implementation project(':metadata-service:auth-impl')
+  implementation project(':metadata-io')
+
+  implementation externalDependency.springWeb
+  implementation(externalDependency.springDocUI) {
+    exclude group: 'org.springframework.boot'
+  }
+  implementation externalDependency.swaggerAnnotations
 
   implementation externalDependency.jacksonDataBind
   implementation externalDependency.httpClient

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericEntity.java

@@ -0,0 +1,7 @@
+package io.datahubproject.openapi.models;
+
+import java.util.Map;
+
+public interface GenericEntity {
+  Map<String, Object> getAspects();
+}

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericEntityScrollResult.java

@@ -0,0 +1,3 @@
+package io.datahubproject.openapi.models;
+
+public interface GenericEntityScrollResult<T extends GenericEntity> {}

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericScrollResult.java → metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/models/GenericScrollResult.java

@@ -1,4 +1,4 @@
-package io.datahubproject.openapi.v2.models;
+package io.datahubproject.openapi.models;
 
 import java.util.List;
 import lombok.Builder;

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/BatchGetUrnResponse.java

@@ -16,5 +16,5 @@
 public class BatchGetUrnResponse implements Serializable {
   @JsonProperty("entities")
   @Schema(description = "List of entity responses")
-  List<GenericEntity> entities;
+  List<GenericEntityV2> entities;
 }

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntityScrollResultV2.java

@@ -0,0 +1,15 @@
+package io.datahubproject.openapi.v2.models;
+
+import io.datahubproject.openapi.models.GenericEntity;
+import io.datahubproject.openapi.models.GenericEntityScrollResult;
+import java.util.List;
+import lombok.Builder;
+import lombok.Data;
+
+@Data
+@Builder
+public class GenericEntityScrollResultV2<T extends GenericEntity>
+    implements GenericEntityScrollResult<T> {
+  private String scrollId;
+  private List<T> results;
+}

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntity.java → metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v2/models/GenericEntityV2.java

@@ -7,6 +7,7 @@
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.mxe.SystemMetadata;
 import com.linkedin.util.Pair;
+import io.datahubproject.openapi.models.GenericEntity;
 import io.swagger.v3.oas.annotations.media.Schema;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
@@ -23,7 +24,7 @@
 @JsonInclude(JsonInclude.Include.NON_NULL)
 @NoArgsConstructor(force = true, access = AccessLevel.PRIVATE)
 @AllArgsConstructor
-public class GenericEntity {
+public class GenericEntityV2 implements GenericEntity {
   @JsonProperty("urn")
   @Schema(description = "Urn of the entity")
   private String urn;
@@ -32,9 +33,9 @@ public class GenericEntity {
   @Schema(description = "Map of aspect name to aspect")
   private Map<String, Object> aspects;
 
-  public static class GenericEntityBuilder {
+  public static class GenericEntityV2Builder {
 
-    public GenericEntity build(
+    public GenericEntityV2 build(
         ObjectMapper objectMapper, Map<String, Pair<RecordTemplate, SystemMetadata>> aspects) {
       Map<String, Object> jsonObjectMap =
           aspects.entrySet().stream()
@@ -63,7 +64,7 @@ public GenericEntity build(
                   })
               .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
 
-      return new GenericEntity(urn, jsonObjectMap);
+      return new GenericEntityV2(urn, jsonObjectMap);
     }
   }
 }

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v3/models/GenericEntityScrollResultV3.java

@@ -0,0 +1,15 @@
+package io.datahubproject.openapi.v3.models;
+
+import io.datahubproject.openapi.models.GenericEntity;
+import io.datahubproject.openapi.models.GenericEntityScrollResult;
+import java.util.List;
+import lombok.Builder;
+import lombok.Data;
+
+@Data
+@Builder
+public class GenericEntityScrollResultV3<T extends GenericEntity>
+    implements GenericEntityScrollResult<T> {
+  private String scrollId;
+  private List<T> entities;
+}

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v3/models/GenericEntityV3.java

@@ -0,0 +1,77 @@
+package io.datahubproject.openapi.v3.models;
+
+import com.datahub.util.RecordUtils;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.data.template.RecordTemplate;
+import com.linkedin.mxe.SystemMetadata;
+import com.linkedin.util.Pair;
+import io.datahubproject.openapi.models.GenericEntity;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+import lombok.EqualsAndHashCode;
+
+@EqualsAndHashCode(callSuper = true)
+@Data
+@Builder
+@JsonInclude(JsonInclude.Include.NON_NULL)
+@AllArgsConstructor
+public class GenericEntityV3 extends LinkedHashMap<String, Object> implements GenericEntity {
+
+  public GenericEntityV3(Map<? extends String, ?> m) {
+    super(m);
+  }
+
+  @Override
+  public Map<String, Object> getAspects() {
+    return this;
+  }
+
+  public static class GenericEntityV3Builder {
+
+    public GenericEntityV3 build(
+        ObjectMapper objectMapper,
+        @Nonnull Urn urn,
+        Map<String, Pair<RecordTemplate, SystemMetadata>> aspects) {
+      Map<String, Object> jsonObjectMap =
+          aspects.entrySet().stream()
+              .map(
+                  e -> {
+                    try {
+                      Map<String, Object> valueMap =
+                          Map.of(
+                              "value",
+                              objectMapper.readTree(
+                                  RecordUtils.toJsonString(e.getValue().getFirst())
+                                      .getBytes(StandardCharsets.UTF_8)));
+
+                      if (e.getValue().getSecond() != null) {
+                        return Map.entry(
+                            e.getKey(),
+                            Map.of(
+                                "systemMetadata", e.getValue().getSecond(),
+                                "value", valueMap.get("value")));
+                      } else {
+                        return Map.entry(e.getKey(), Map.of("value", valueMap.get("value")));
+                      }
+                    } catch (IOException ex) {
+                      throw new RuntimeException(ex);
+                    }
+                  })
+              .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+      GenericEntityV3 genericEntityV3 = new GenericEntityV3();
+      genericEntityV3.put("urn", urn.toString());
+      genericEntityV3.putAll(jsonObjectMap);
+      return genericEntityV3;
+    }
+  }
+}