feat(graphql/ml): Add custom properties to ml entities (#12152)

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

@@ -0,0 +1,12 @@
+package com.linkedin.datahub.graphql.types.mappers;
+
+import com.linkedin.common.urn.Urn;
+import com.linkedin.datahub.graphql.QueryContext;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/** Made for models that are embedded in other models and thus do not encode their own URN. */
+public interface EmbeddedModelMapper<I, O> {
+  O apply(
+      @Nullable final QueryContext context, @Nonnull final I input, @Nonnull final Urn entityUrn);
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureMapper.java

@@ -75,7 +75,8 @@ public class MLFeatureMapper implements ModelMapper<EntityResponse, MLFeature> {
             mlFeature.setOwnership(
                 OwnershipMapper.map(context, new Ownership(dataMap), entityUrn)));
     mappingHelper.mapToResult(
-        context, ML_FEATURE_PROPERTIES_ASPECT_NAME, MLFeatureMapper::mapMLFeatureProperties);
+        ML_FEATURE_PROPERTIES_ASPECT_NAME,
+        (entity, dataMap) -> mapMLFeatureProperties(context, entity, dataMap, entityUrn));
     mappingHelper.mapToResult(
         INSTITUTIONAL_MEMORY_ASPECT_NAME,
         (mlFeature, dataMap) ->
@@ -138,10 +139,13 @@ public class MLFeatureMapper implements ModelMapper<EntityResponse, MLFeature> {
   private static void mapMLFeatureProperties(
       @Nullable final QueryContext context,
       @Nonnull MLFeature mlFeature,
-      @Nonnull DataMap dataMap) {
+      @Nonnull DataMap dataMap,
+      @Nonnull Urn entityUrn) {
     MLFeatureProperties featureProperties = new MLFeatureProperties(dataMap);
-    mlFeature.setFeatureProperties(MLFeaturePropertiesMapper.map(context, featureProperties));
+    com.linkedin.datahub.graphql.generated.MLFeatureProperties graphqlProperties =
-    mlFeature.setProperties(MLFeaturePropertiesMapper.map(context, featureProperties));
+        MLFeaturePropertiesMapper.map(context, featureProperties, entityUrn);
+    mlFeature.setFeatureProperties(graphqlProperties);
+    mlFeature.setProperties(graphqlProperties);
     mlFeature.setDescription(featureProperties.getDescription());
     if (featureProperties.getDataType() != null) {
       mlFeature.setDataType(MLFeatureDataType.valueOf(featureProperties.getDataType().toString()));

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeaturePropertiesMapper.java

@@ -1,29 +1,34 @@
 package com.linkedin.datahub.graphql.types.mlmodel.mappers;
 
+import com.linkedin.common.urn.Urn;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.Dataset;
 import com.linkedin.datahub.graphql.generated.MLFeatureDataType;
 import com.linkedin.datahub.graphql.generated.MLFeatureProperties;
-import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
+import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
+import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
 import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
-import lombok.NonNull;
 
 public class MLFeaturePropertiesMapper
-    implements ModelMapper<com.linkedin.ml.metadata.MLFeatureProperties, MLFeatureProperties> {
+    implements EmbeddedModelMapper<
+        com.linkedin.ml.metadata.MLFeatureProperties, MLFeatureProperties> {
 
   public static final MLFeaturePropertiesMapper INSTANCE = new MLFeaturePropertiesMapper();
 
   public static MLFeatureProperties map(
       @Nullable QueryContext context,
-      @NonNull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties) {
+      @Nonnull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties,
-    return INSTANCE.apply(context, mlFeatureProperties);
+      @Nonnull Urn entityUrn) {
+    return INSTANCE.apply(context, mlFeatureProperties, entityUrn);
   }
 
   @Override
   public MLFeatureProperties apply(
       @Nullable QueryContext context,
-      @NonNull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties) {
+      @Nonnull final com.linkedin.ml.metadata.MLFeatureProperties mlFeatureProperties,
+      @Nonnull Urn entityUrn) {
     final MLFeatureProperties result = new MLFeatureProperties();
 
     result.setDescription(mlFeatureProperties.getDescription());
@@ -45,6 +50,9 @@ public class MLFeaturePropertiesMapper
               .collect(Collectors.toList()));
     }
 
+    result.setCustomProperties(
+        CustomPropertiesMapper.map(mlFeatureProperties.getCustomProperties(), entityUrn));
+
     return result;
   }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureTableMapper.java

@@ -76,7 +76,7 @@ public class MLFeatureTableMapper implements ModelMapper<EntityResponse, MLFeatu
     mappingHelper.mapToResult(ML_FEATURE_TABLE_KEY_ASPECT_NAME, this::mapMLFeatureTableKey);
     mappingHelper.mapToResult(
         ML_FEATURE_TABLE_PROPERTIES_ASPECT_NAME,
-        (entity, dataMap) -> this.mapMLFeatureTableProperties(context, entity, dataMap, entityUrn));
+        (entity, dataMap) -> mapMLFeatureTableProperties(context, entity, dataMap, entityUrn));
     mappingHelper.mapToResult(
         INSTITUTIONAL_MEMORY_ASPECT_NAME,
         (mlFeatureTable, dataMap) ->
@@ -146,10 +146,10 @@ public class MLFeatureTableMapper implements ModelMapper<EntityResponse, MLFeatu
       @Nonnull DataMap dataMap,
       Urn entityUrn) {
     MLFeatureTableProperties featureTableProperties = new MLFeatureTableProperties(dataMap);
-    mlFeatureTable.setFeatureTableProperties(
+    com.linkedin.datahub.graphql.generated.MLFeatureTableProperties graphqlProperties =
-        MLFeatureTablePropertiesMapper.map(context, featureTableProperties, entityUrn));
+        MLFeatureTablePropertiesMapper.map(context, featureTableProperties, entityUrn);
-    mlFeatureTable.setProperties(
+    mlFeatureTable.setFeatureTableProperties(graphqlProperties);
-        MLFeatureTablePropertiesMapper.map(context, featureTableProperties, entityUrn));
+    mlFeatureTable.setProperties(graphqlProperties);
     mlFeatureTable.setDescription(featureTableProperties.getDescription());
   }
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureTablePropertiesMapper.java

@@ -8,26 +8,30 @@ import com.linkedin.datahub.graphql.generated.MLFeature;
 import com.linkedin.datahub.graphql.generated.MLFeatureTableProperties;
 import com.linkedin.datahub.graphql.generated.MLPrimaryKey;
 import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
+import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
 import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
-import lombok.NonNull;
 
-public class MLFeatureTablePropertiesMapper {
+public class MLFeatureTablePropertiesMapper
+    implements EmbeddedModelMapper<
+        com.linkedin.ml.metadata.MLFeatureTableProperties, MLFeatureTableProperties> {
 
   public static final MLFeatureTablePropertiesMapper INSTANCE =
       new MLFeatureTablePropertiesMapper();
 
   public static MLFeatureTableProperties map(
       @Nullable final QueryContext context,
-      @NonNull final com.linkedin.ml.metadata.MLFeatureTableProperties mlFeatureTableProperties,
+      @Nonnull final com.linkedin.ml.metadata.MLFeatureTableProperties mlFeatureTableProperties,
-      Urn entityUrn) {
+      @Nonnull Urn entityUrn) {
     return INSTANCE.apply(context, mlFeatureTableProperties, entityUrn);
   }
 
-  public static MLFeatureTableProperties apply(
+  @Override
+  public MLFeatureTableProperties apply(
       @Nullable final QueryContext context,
-      @NonNull final com.linkedin.ml.metadata.MLFeatureTableProperties mlFeatureTableProperties,
+      @Nonnull final com.linkedin.ml.metadata.MLFeatureTableProperties mlFeatureTableProperties,
-      Urn entityUrn) {
+      @Nonnull Urn entityUrn) {
     final MLFeatureTableProperties result = new MLFeatureTableProperties();
 
     result.setDescription(mlFeatureTableProperties.getDescription());

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelGroupMapper.java

@@ -75,9 +75,8 @@ public class MLModelGroupMapper implements ModelMapper<EntityResponse, MLModelGr
     mappingHelper.mapToResult(
         ML_MODEL_GROUP_KEY_ASPECT_NAME, MLModelGroupMapper::mapToMLModelGroupKey);
     mappingHelper.mapToResult(
-        context,
         ML_MODEL_GROUP_PROPERTIES_ASPECT_NAME,
-        MLModelGroupMapper::mapToMLModelGroupProperties);
+        (entity, dataMap) -> mapToMLModelGroupProperties(context, entity, dataMap, entityUrn));
     mappingHelper.mapToResult(
         STATUS_ASPECT_NAME,
         (mlModelGroup, dataMap) ->
@@ -136,9 +135,13 @@ public class MLModelGroupMapper implements ModelMapper<EntityResponse, MLModelGr
   }
 
   private static void mapToMLModelGroupProperties(
-      @Nullable final QueryContext context, MLModelGroup mlModelGroup, DataMap dataMap) {
+      @Nullable final QueryContext context,
+      MLModelGroup mlModelGroup,
+      DataMap dataMap,
+      @Nonnull Urn entityUrn) {
     MLModelGroupProperties modelGroupProperties = new MLModelGroupProperties(dataMap);
-    mlModelGroup.setProperties(MLModelGroupPropertiesMapper.map(context, modelGroupProperties));
+    mlModelGroup.setProperties(
+        MLModelGroupPropertiesMapper.map(context, modelGroupProperties, entityUrn));
     if (modelGroupProperties.getDescription() != null) {
       mlModelGroup.setDescription(modelGroupProperties.getDescription());
     }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelGroupPropertiesMapper.java

@@ -1,27 +1,31 @@
 package com.linkedin.datahub.graphql.types.mlmodel.mappers;
 
+import com.linkedin.common.urn.Urn;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.MLModelGroupProperties;
-import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
+import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
+import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
-import lombok.NonNull;
 
 public class MLModelGroupPropertiesMapper
-    implements ModelMapper<
+    implements EmbeddedModelMapper<
         com.linkedin.ml.metadata.MLModelGroupProperties, MLModelGroupProperties> {
 
   public static final MLModelGroupPropertiesMapper INSTANCE = new MLModelGroupPropertiesMapper();
 
   public static MLModelGroupProperties map(
       @Nullable QueryContext context,
-      @NonNull final com.linkedin.ml.metadata.MLModelGroupProperties mlModelGroupProperties) {
+      @Nonnull final com.linkedin.ml.metadata.MLModelGroupProperties mlModelGroupProperties,
-    return INSTANCE.apply(context, mlModelGroupProperties);
+      @Nonnull Urn entityUrn) {
+    return INSTANCE.apply(context, mlModelGroupProperties, entityUrn);
   }
 
   @Override
   public MLModelGroupProperties apply(
       @Nullable QueryContext context,
-      @NonNull final com.linkedin.ml.metadata.MLModelGroupProperties mlModelGroupProperties) {
+      @Nonnull final com.linkedin.ml.metadata.MLModelGroupProperties mlModelGroupProperties,
+      @Nonnull Urn entityUrn) {
     final MLModelGroupProperties result = new MLModelGroupProperties();
 
     result.setDescription(mlModelGroupProperties.getDescription());
@@ -30,6 +34,9 @@ public class MLModelGroupPropertiesMapper
     }
     result.setCreatedAt(mlModelGroupProperties.getCreatedAt());
 
+    result.setCustomProperties(
+        CustomPropertiesMapper.map(mlModelGroupProperties.getCustomProperties(), entityUrn));
+
     return result;
   }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelPropertiesMapper.java

@@ -7,25 +7,27 @@ import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.MLModelGroup;
 import com.linkedin.datahub.graphql.generated.MLModelProperties;
 import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
+import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
 import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
-import lombok.NonNull;
 
-public class MLModelPropertiesMapper {
+public class MLModelPropertiesMapper
+    implements EmbeddedModelMapper<com.linkedin.ml.metadata.MLModelProperties, MLModelProperties> {
 
   public static final MLModelPropertiesMapper INSTANCE = new MLModelPropertiesMapper();
 
   public static MLModelProperties map(
       @Nullable final QueryContext context,
-      @NonNull final com.linkedin.ml.metadata.MLModelProperties mlModelProperties,
+      @Nonnull final com.linkedin.ml.metadata.MLModelProperties mlModelProperties,
       Urn entityUrn) {
     return INSTANCE.apply(context, mlModelProperties, entityUrn);
   }
 
   public MLModelProperties apply(
       @Nullable final QueryContext context,
-      @NonNull final com.linkedin.ml.metadata.MLModelProperties mlModelProperties,
+      @Nonnull final com.linkedin.ml.metadata.MLModelProperties mlModelProperties,
-      Urn entityUrn) {
+      @Nonnull Urn entityUrn) {
     final MLModelProperties result = new MLModelProperties();
 
     result.setDate(mlModelProperties.getDate());

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLPrimaryKeyMapper.java

@@ -74,9 +74,8 @@ public class MLPrimaryKeyMapper implements ModelMapper<EntityResponse, MLPrimary
     mappingHelper.mapToResult(
         ML_PRIMARY_KEY_KEY_ASPECT_NAME, MLPrimaryKeyMapper::mapMLPrimaryKeyKey);
     mappingHelper.mapToResult(
-        context,
         ML_PRIMARY_KEY_PROPERTIES_ASPECT_NAME,
-        MLPrimaryKeyMapper::mapMLPrimaryKeyProperties);
+        (entity, dataMap) -> mapMLPrimaryKeyProperties(context, entity, dataMap, entityUrn));
     mappingHelper.mapToResult(
         INSTITUTIONAL_MEMORY_ASPECT_NAME,
         (mlPrimaryKey, dataMap) ->
@@ -132,11 +131,15 @@ public class MLPrimaryKeyMapper implements ModelMapper<EntityResponse, MLPrimary
   }
 
   private static void mapMLPrimaryKeyProperties(
-      @Nullable final QueryContext context, MLPrimaryKey mlPrimaryKey, DataMap dataMap) {
+      @Nullable final QueryContext context,
+      MLPrimaryKey mlPrimaryKey,
+      DataMap dataMap,
+      @Nonnull Urn entityUrn) {
     MLPrimaryKeyProperties primaryKeyProperties = new MLPrimaryKeyProperties(dataMap);
-    mlPrimaryKey.setPrimaryKeyProperties(
+    com.linkedin.datahub.graphql.generated.MLPrimaryKeyProperties graphqlProperties =
-        MLPrimaryKeyPropertiesMapper.map(context, primaryKeyProperties));
+        MLPrimaryKeyPropertiesMapper.map(context, primaryKeyProperties, entityUrn);
-    mlPrimaryKey.setProperties(MLPrimaryKeyPropertiesMapper.map(context, primaryKeyProperties));
+    mlPrimaryKey.setPrimaryKeyProperties(graphqlProperties);
+    mlPrimaryKey.setProperties(graphqlProperties);
     mlPrimaryKey.setDescription(primaryKeyProperties.getDescription());
     if (primaryKeyProperties.getDataType() != null) {
       mlPrimaryKey.setDataType(

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLPrimaryKeyPropertiesMapper.java

@@ -1,30 +1,34 @@
 package com.linkedin.datahub.graphql.types.mlmodel.mappers;
 
+import com.linkedin.common.urn.Urn;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.Dataset;
 import com.linkedin.datahub.graphql.generated.MLFeatureDataType;
 import com.linkedin.datahub.graphql.generated.MLPrimaryKeyProperties;
-import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
+import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
+import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
 import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
-import lombok.NonNull;
 
 public class MLPrimaryKeyPropertiesMapper
-    implements ModelMapper<
+    implements EmbeddedModelMapper<
         com.linkedin.ml.metadata.MLPrimaryKeyProperties, MLPrimaryKeyProperties> {
 
   public static final MLPrimaryKeyPropertiesMapper INSTANCE = new MLPrimaryKeyPropertiesMapper();
 
   public static MLPrimaryKeyProperties map(
       @Nullable QueryContext context,
-      @NonNull final com.linkedin.ml.metadata.MLPrimaryKeyProperties mlPrimaryKeyProperties) {
+      @Nonnull final com.linkedin.ml.metadata.MLPrimaryKeyProperties mlPrimaryKeyProperties,
-    return INSTANCE.apply(context, mlPrimaryKeyProperties);
+      @Nonnull Urn entityUrn) {
+    return INSTANCE.apply(context, mlPrimaryKeyProperties, entityUrn);
   }
 
   @Override
   public MLPrimaryKeyProperties apply(
       @Nullable QueryContext context,
-      @NonNull final com.linkedin.ml.metadata.MLPrimaryKeyProperties mlPrimaryKeyProperties) {
+      @Nonnull final com.linkedin.ml.metadata.MLPrimaryKeyProperties mlPrimaryKeyProperties,
+      @Nonnull Urn entityUrn) {
     final MLPrimaryKeyProperties result = new MLPrimaryKeyProperties();
 
     result.setDescription(mlPrimaryKeyProperties.getDescription());
@@ -45,6 +49,9 @@ public class MLPrimaryKeyPropertiesMapper
                 })
             .collect(Collectors.toList()));
 
+    result.setCustomProperties(
+        CustomPropertiesMapper.map(mlPrimaryKeyProperties.getCustomProperties(), entityUrn));
+
     return result;
   }
 }

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

@@ -9834,6 +9834,8 @@ description: String
     createdAt: Long
 
     version: VersionTag
+
+    customProperties: [CustomPropertiesEntry!]
 }
 
 """
@@ -10028,6 +10030,8 @@ type MLFeatureProperties {
     version: VersionTag
 
     sources: [Dataset]
+
+    customProperties: [CustomPropertiesEntry!]
 }
 
 """
@@ -10171,6 +10175,8 @@ description: String
     version: VersionTag
 
     sources: [Dataset]
+
+    customProperties: [CustomPropertiesEntry!]
 }
 
 """