feat(graphql): adds container aspect for dataflow and datajob entities

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataflow/mappers/DataFlowMapper.java

@@ -16,6 +16,7 @@
 import com.linkedin.data.DataMap;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
+import com.linkedin.datahub.graphql.generated.Container;
 import com.linkedin.datahub.graphql.generated.DataFlow;
 import com.linkedin.datahub.graphql.generated.DataFlowEditableProperties;
 import com.linkedin.datahub.graphql.generated.DataFlowInfo;
@@ -106,6 +107,7 @@ public DataFlow apply(
         (dataset, dataMap) ->
             dataset.setDataPlatformInstance(
                 DataPlatformInstanceAspectMapper.map(context, new DataPlatformInstance(dataMap))));
+    mappingHelper.mapToResult(context, CONTAINER_ASPECT_NAME, DataFlowMapper::mapContainers);
     mappingHelper.mapToResult(
         BROWSE_PATHS_V2_ASPECT_NAME,
         (dataFlow, dataMap) ->
@@ -206,6 +208,17 @@ private static void mapGlobalTags(
     dataFlow.setTags(globalTags);
   }
 
+  private static void mapContainers(
+      @Nullable final QueryContext context, @Nonnull DataFlow dataFlow, @Nonnull DataMap dataMap) {
+    final com.linkedin.container.Container gmsContainer =
+        new com.linkedin.container.Container(dataMap);
+    dataFlow.setContainer(
+        Container.builder()
+            .setType(EntityType.CONTAINER)
+            .setUrn(gmsContainer.getContainer().toString())
+            .build());
+  }
+
   private static void mapDomains(
       @Nullable final QueryContext context, @Nonnull DataFlow dataFlow, @Nonnull DataMap dataMap) {
     final Domains domains = new Domains(dataMap);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/datajob/mappers/DataJobMapper.java

@@ -18,6 +18,7 @@
 import com.linkedin.data.DataMap;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
+import com.linkedin.datahub.graphql.generated.Container;
 import com.linkedin.datahub.graphql.generated.DataFlow;
 import com.linkedin.datahub.graphql.generated.DataJob;
 import com.linkedin.datahub.graphql.generated.DataJobEditableProperties;
@@ -129,6 +130,14 @@ public DataJob apply(
               } else if (DATA_PLATFORM_INSTANCE_ASPECT_NAME.equals(name)) {
                 result.setDataPlatformInstance(
                     DataPlatformInstanceAspectMapper.map(context, new DataPlatformInstance(data)));
+              } else if (CONTAINER_ASPECT_NAME.equals(name)) {
+                final com.linkedin.container.Container gmsContainer =
+                    new com.linkedin.container.Container(data);
+                result.setContainer(
+                    Container.builder()
+                        .setType(EntityType.CONTAINER)
+                        .setUrn(gmsContainer.getContainer().toString())
+                        .build());
               } else if (BROWSE_PATHS_V2_ASPECT_NAME.equals(name)) {
                 result.setBrowsePathV2(BrowsePathsV2Mapper.map(context, new BrowsePathsV2(data)));
               } else if (SUB_TYPES_ASPECT_NAME.equals(name)) {

datahub-graphql-core/src/main/resources/entity.graphql

@@ -6275,6 +6275,16 @@ type DataFlow implements EntityWithRelationships & Entity & BrowsableEntity {
     """
     dataPlatformInstance: DataPlatformInstance
 
+    """
+    The parent container in which the entity resides
+    """
+    container: Container
+
+    """
+    Recursively get the lineage of containers for this entity
+    """
+    parentContainers: ParentContainersResult
+
     """
     Granular API for querying edges extending from this entity
     """
@@ -6457,6 +6467,16 @@ type DataJob implements EntityWithRelationships & Entity & BrowsableEntity {
     """
     dataPlatformInstance: DataPlatformInstance
 
+    """
+    The parent container in which the entity resides
+    """
+    container: Container
+
+    """
+    Recursively get the lineage of containers for this entity
+    """
+    parentContainers: ParentContainersResult
+
     """
     Additional read write properties associated with the Data Job
     """

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/dataflow/mappers/DataFlowMapperTest.java

@@ -0,0 +1,42 @@
+package com.linkedin.datahub.graphql.types.dataflow.mappers;
+
+import com.linkedin.common.urn.Urn;
+import com.linkedin.datahub.graphql.generated.DataFlow;
+import com.linkedin.entity.Aspect;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.EnvelopedAspect;
+import com.linkedin.entity.EnvelopedAspectMap;
+import com.linkedin.metadata.Constants;
+import java.net.URISyntaxException;
+import java.util.HashMap;
+import java.util.Map;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class DataFlowMapperTest {
+  private static final Urn TEST_DATA_FLOW_URN =
+      Urn.createFromTuple(Constants.DATA_FLOW_ENTITY_NAME, "dataflow1");
+  private static final Urn TEST_CONTAINER_URN =
+      Urn.createFromTuple(Constants.CONTAINER_ENTITY_NAME, "container1");
+
+  @Test
+  public void testMapDataFlowContainer() throws URISyntaxException {
+    com.linkedin.container.Container input = new com.linkedin.container.Container();
+    input.setContainer(TEST_CONTAINER_URN);
+
+    final Map<String, EnvelopedAspect> containerAspect = new HashMap<>();
+    containerAspect.put(
+        Constants.CONTAINER_ASPECT_NAME,
+        new com.linkedin.entity.EnvelopedAspect().setValue(new Aspect(input.data())));
+    final EntityResponse response =
+        new EntityResponse()
+            .setEntityName(Constants.DATA_FLOW_ENTITY_NAME)
+            .setUrn(TEST_DATA_FLOW_URN)
+            .setAspects(new EnvelopedAspectMap(containerAspect));
+
+    final DataFlow actual = DataFlowMapper.map(null, response);
+
+    Assert.assertEquals(actual.getUrn(), TEST_DATA_FLOW_URN.toString());
+    Assert.assertEquals(actual.getContainer().getUrn(), TEST_CONTAINER_URN.toString());
+  }
+}

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/datajob/mappers/DataJobMapperTest.java

@@ -0,0 +1,42 @@
+package com.linkedin.datahub.graphql.types.datajob.mappers;
+
+import com.linkedin.common.urn.Urn;
+import com.linkedin.datahub.graphql.generated.DataJob;
+import com.linkedin.entity.Aspect;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.EnvelopedAspect;
+import com.linkedin.entity.EnvelopedAspectMap;
+import com.linkedin.metadata.Constants;
+import java.net.URISyntaxException;
+import java.util.HashMap;
+import java.util.Map;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class DataJobMapperTest {
+  private static final Urn TEST_DATA_JOB_URN =
+      Urn.createFromTuple(Constants.DATA_JOB_ENTITY_NAME, "datajob1");
+  private static final Urn TEST_CONTAINER_URN =
+      Urn.createFromTuple(Constants.CONTAINER_ENTITY_NAME, "container1");
+
+  @Test
+  public void testMapDataJobContainer() throws URISyntaxException {
+    com.linkedin.container.Container input = new com.linkedin.container.Container();
+    input.setContainer(TEST_CONTAINER_URN);
+
+    final Map<String, EnvelopedAspect> containerAspect = new HashMap<>();
+    containerAspect.put(
+        Constants.CONTAINER_ASPECT_NAME,
+        new com.linkedin.entity.EnvelopedAspect().setValue(new Aspect(input.data())));
+    final EntityResponse response =
+        new EntityResponse()
+            .setEntityName(Constants.DATA_JOB_ENTITY_NAME)
+            .setUrn(TEST_DATA_JOB_URN)
+            .setAspects(new EnvelopedAspectMap(containerAspect));
+
+    final DataJob actual = DataJobMapper.map(null, response);
+
+    Assert.assertEquals(actual.getUrn(), TEST_DATA_JOB_URN.toString());
+    Assert.assertEquals(actual.getContainer().getUrn(), TEST_CONTAINER_URN.toString());
+  }
+}

docs/how/updating-datahub.md

@@ -43,6 +43,7 @@ This file documents any backwards-incompatible changes in DataHub and assists pe
 - OpenAPI Update: PIT Keep Alive parameter added to scroll. NOTE: This parameter requires the `pointInTimeCreationEnabled` feature flag to be enabled and the `elasticSearch.implementation` configuration to be `elasticsearch`. This feature is not supported for OpenSearch at this time and the parameter will not be respected without both of these set.
 - OpenAPI Update 2: Previously there was an incorrectly marked parameter named `sort` on the generic list entities endpoint for v3. This parameter is deprecated and only supports a single string value while the documentation indicates it supports a list of strings. This documentation error has been fixed and the correct field, `sortCriteria`, is now documented which supports a list of strings. 
 - #12223: For dbt Cloud ingestion, the "View in dbt" link will point at the "Explore" page in the dbt Cloud UI. You can revert to the old behavior of linking to the dbt Cloud IDE by setting `external_url_mode: ide".
+- #XXXXX: Data flow and data job entities may additionally produce container aspect that will require a corresponding upgrade of server. Otherwise server can reject the aspect.
 
 ### Breaking Changes
 

metadata-models/src/main/resources/entity-registry.yml

@@ -70,6 +70,7 @@ entities:
       - glossaryTerms
       - institutionalMemory
       - dataPlatformInstance
+      - container
       - browsePathsV2
       - structuredProperties
       - forms
@@ -92,6 +93,7 @@ entities:
       - glossaryTerms
       - institutionalMemory
       - dataPlatformInstance
+      - container
       - browsePathsV2
       - structuredProperties
       - incidentsSummary