add a structured type to batchGet in OpenAPI V3 spec

datahub-project · Jul 19, 2024 · df460bc · df460bc
1 parent 44cdb04
commit df460bc
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 17 deletions.
diff --git a/metadata-integration/java/datahub-event/src/main/resources/MetadataChangeProposal.avsc b/metadata-integration/java/datahub-event/src/main/resources/MetadataChangeProposal.avsc
@@ -1,14 +1,14 @@
 {
   "type" : "record",
   "name" : "MetadataChangeProposal",
-  "namespace" : "com.linkedin.pegasus2avro.mxe",
+  "namespace" : "com.linkedin.mxe",
   "doc" : "Kafka event for proposing a metadata change for an entity. A corresponding MetadataChangeLog is emitted when the change is accepted and committed, otherwise a FailedMetadataChangeProposal will be emitted instead.",
   "fields" : [ {
     "name" : "auditHeader",
     "type" : [ "null", {
       "type" : "record",
       "name" : "KafkaAuditHeader",
-      "namespace" : "com.linkedin.events",
+      "namespace" : "com.linkedin.avro2pegasus.events",
       "doc" : "This header records information about the context of an event as it is emitted into kafka and is intended to be used by the kafka audit application.  For more information see go/kafkaauditheader",
       "fields" : [ {
         "name" : "time",
@@ -74,7 +74,7 @@
     "doc" : "Urn of the entity being written",
     "default" : null,
     "java" : {
-      "class" : "com.linkedin.pegasus2avro.common.urn.Urn"
+      "class" : "com.linkedin.common.urn.Urn"
     }
   }, {
     "name" : "entityKeyAspect",
@@ -99,7 +99,7 @@
     "type" : {
       "type" : "enum",
       "name" : "ChangeType",
-      "namespace" : "com.linkedin.pegasus2avro.events.metadata",
+      "namespace" : "com.linkedin.events.metadata",
       "doc" : "Descriptor for a change action",
       "symbols" : [ "UPSERT", "CREATE", "UPDATE", "DELETE", "PATCH", "RESTATE", "CREATE_ENTITY" ],
       "symbolDocs" : {

diff --git a/...ervice/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/OpenAPIV3Generator.java b/...ervice/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/OpenAPIV3Generator.java
@@ -82,6 +82,20 @@ public static OpenAPI generateOpenApiSpec(EntityRegistry entityRegistry) {
         "SystemMetadata", new Schema().type(TYPE_OBJECT).additionalProperties(true));
     components.addSchemas("SortOrder", new Schema()._enum(List.of("ASCENDING", "DESCENDING")));
     components.addSchemas("AspectPatch", buildAspectPatchSchema());
+    components.addSchemas(
+        "BatchGetRequestBody",
+        new Schema<>()
+            .type(TYPE_OBJECT)
+            .description("Request body for batch get aspects.")
+            .properties(
+                Map.of(
+                    "headers",
+                    new Schema<>()
+                        .type(TYPE_OBJECT)
+                        .additionalProperties(new Schema<>().type(TYPE_STRING))
+                        .description("System headers for the operation.")
+                        .nullable(true)))
+            .nullable(true));
     entityRegistry
         .getAspectSpecs()
         .values()
@@ -645,28 +659,19 @@ private static Schema buildEntityScrollSchema(final EntitySpec entity) {
   private static Schema buildEntityBatchGetRequestSchema(
       final EntitySpec entity, Set<String> aspectNames) {
 
-    final Schema stringTypeSchema = new Schema<>();
-    stringTypeSchema.setType(TYPE_STRING);
-    final Map<String, Schema> headers =
-        Map.of(
-            "headers",
-            new Schema<>()
-                .type(TYPE_OBJECT)
-                .additionalProperties(stringTypeSchema)
-                .description("System headers for the operation.")
-                .nullable(true));
-
     final Map<String, Schema> properties =
         entity.getAspectSpecMap().entrySet().stream()
             .filter(a -> aspectNames.contains(a.getValue().getName()))
             .collect(
                 Collectors.toMap(
-                    Map.Entry::getKey, a -> new Schema().type(TYPE_OBJECT).properties(headers)));
+                    Map.Entry::getKey,
+                    a -> new Schema().$ref("#/components/schemas/BatchGetRequestBody")));
     properties.put(
         PROPERTY_URN,
         new Schema<>().type(TYPE_STRING).description("Unique id for " + entity.getName()));
 
-    properties.put(entity.getKeyAspectName(), new Schema().type(TYPE_OBJECT).properties(headers));
+    properties.put(
+        entity.getKeyAspectName(), new Schema().$ref("#/components/schemas/BatchGetRequestBody"));
 
     return new Schema<>()
         .type(TYPE_OBJECT)

diff --git a/...ce/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/OpenAPIV3GeneratorTest.java b/...ce/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/OpenAPIV3GeneratorTest.java
@@ -88,5 +88,18 @@ public void testOpenApiSpecBuilder() throws Exception {
     Schema fabricType = openAPI.getComponents().getSchemas().get("FabricType");
     assertEquals("string", fabricType.getType());
     assertFalse(fabricType.getEnum().isEmpty());
+
+    Map<String, Schema> batchProperties =
+        openAPI
+            .getComponents()
+            .getSchemas()
+            .get("BatchGetContainerEntityRequest_v3")
+            .getProperties();
+    batchProperties.entrySet().stream()
+        .filter(entry -> !entry.getKey().equals("urn"))
+        .forEach(
+            entry ->
+                assertEquals(
+                    "#/components/schemas/BatchGetRequestBody", entry.getValue().get$ref()));
   }
 }