feat(graphql/versioning): Add versioning support to graphql; mutations return version set

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/Constants.java

@@ -28,6 +28,7 @@ private Constants() {}
   public static final String INCIDENTS_SCHEMA_FILE = "incident.graphql";
   public static final String CONTRACTS_SCHEMA_FILE = "contract.graphql";
   public static final String CONNECTIONS_SCHEMA_FILE = "connection.graphql";
+  public static final String VERSION_SCHEMA_FILE = "versioning.graphql";
   public static final String BROWSE_PATH_DELIMITER = "/";
   public static final String BROWSE_PATH_V2_DELIMITER = "␟";
   public static final String VERSION_STAMP_FIELD_NAME = "versionStamp";

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java

@@ -124,6 +124,8 @@
 import com.linkedin.datahub.graphql.generated.TestResult;
 import com.linkedin.datahub.graphql.generated.TypeQualifier;
 import com.linkedin.datahub.graphql.generated.UserUsageCounts;
+import com.linkedin.datahub.graphql.generated.VersionProperties;
+import com.linkedin.datahub.graphql.generated.VersionSet;
 import com.linkedin.datahub.graphql.resolvers.MeResolver;
 import com.linkedin.datahub.graphql.resolvers.assertion.AssertionRunEventResolver;
 import com.linkedin.datahub.graphql.resolvers.assertion.DeleteAssertionResolver;
@@ -324,6 +326,7 @@
 import com.linkedin.datahub.graphql.resolvers.user.ListUsersResolver;
 import com.linkedin.datahub.graphql.resolvers.user.RemoveUserResolver;
 import com.linkedin.datahub.graphql.resolvers.user.UpdateUserStatusResolver;
+import com.linkedin.datahub.graphql.resolvers.versioning.VersionsSearchResolver;
 import com.linkedin.datahub.graphql.resolvers.view.CreateViewResolver;
 import com.linkedin.datahub.graphql.resolvers.view.DeleteViewResolver;
 import com.linkedin.datahub.graphql.resolvers.view.ListGlobalViewsResolver;
@@ -381,6 +384,7 @@
 import com.linkedin.datahub.graphql.types.structuredproperty.StructuredPropertyType;
 import com.linkedin.datahub.graphql.types.tag.TagType;
 import com.linkedin.datahub.graphql.types.test.TestType;
+import com.linkedin.datahub.graphql.types.versioning.VersionSetType;
 import com.linkedin.datahub.graphql.types.view.DataHubViewType;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.entity.client.SystemEntityClient;
@@ -537,6 +541,7 @@ public class GmsGraphQLEngine {
   private final IncidentType incidentType;
   private final RestrictedType restrictedType;
   private final DataProcessInstanceType dataProcessInstanceType;
+  private final VersionSetType versionSetType;
 
   private final int graphQLQueryComplexityLimit;
   private final int graphQLQueryDepthLimit;
@@ -658,6 +663,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
     this.incidentType = new IncidentType(entityClient);
     this.restrictedType = new RestrictedType(entityClient, restrictedService);
     this.dataProcessInstanceType = new DataProcessInstanceType(entityClient, featureFlags);
+    this.versionSetType = new VersionSetType(entityClient);
 
     this.graphQLQueryComplexityLimit = args.graphQLQueryComplexityLimit;
     this.graphQLQueryDepthLimit = args.graphQLQueryDepthLimit;
@@ -707,6 +713,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
                 entityTypeType,
                 formType,
                 incidentType,
+                versionSetType,
                 restrictedType,
                 businessAttributeType,
                 dataProcessInstanceType));
@@ -809,6 +816,8 @@ public void configureRuntimeWiring(final RuntimeWiring.Builder builder) {
     configureConnectionResolvers(builder);
     configureDeprecationResolvers(builder);
     configureMetadataAttributionResolver(builder);
+    configureVersionPropertiesResolvers(builder);
+    configureVersionSetResolvers(builder);
   }
 
   private void configureOrganisationRoleResolvers(RuntimeWiring.Builder builder) {
@@ -863,7 +872,8 @@ public GraphQLEngine.Builder builder() {
         .addSchema(fileBasedSchema(ASSERTIONS_SCHEMA_FILE))
         .addSchema(fileBasedSchema(INCIDENTS_SCHEMA_FILE))
         .addSchema(fileBasedSchema(CONTRACTS_SCHEMA_FILE))
-        .addSchema(fileBasedSchema(COMMON_SCHEMA_FILE));
+        .addSchema(fileBasedSchema(COMMON_SCHEMA_FILE))
+        .addSchema(fileBasedSchema(VERSION_SCHEMA_FILE));
 
     for (GmsGraphQLPlugin plugin : this.graphQLPlugins) {
       List<String> pluginSchemaFiles = plugin.getSchemaFiles();
@@ -1050,6 +1060,7 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) {
                 .dataFetcher("form", getResolver(formType))
                 .dataFetcher("view", getResolver(dataHubViewType))
                 .dataFetcher("structuredProperty", getResolver(structuredPropertyType))
+                .dataFetcher("versionSet", getResolver(versionSetType))
                 .dataFetcher("listPolicies", new ListPoliciesResolver(this.entityClient))
                 .dataFetcher("getGrantedPrivileges", new GetGrantedPrivilegesResolver())
                 .dataFetcher("listUsers", new ListUsersResolver(this.entityClient))
@@ -2295,7 +2306,15 @@ private void configureTypeResolvers(final RuntimeWiring.Builder builder) {
         .type(
             "TimeSeriesAspect",
             typeWiring -> typeWiring.typeResolver(new TimeSeriesAspectInterfaceTypeResolver()))
-        .type("ResultsType", typeWiring -> typeWiring.typeResolver(new ResultsTypeResolver()));
+        .type("ResultsType", typeWiring -> typeWiring.typeResolver(new ResultsTypeResolver()))
+        .type(
+            "SupportsVersions",
+            typeWiring ->
+                typeWiring.typeResolver(
+                    new EntityInterfaceTypeResolver(
+                        loadableTypes.stream()
+                            .map(graphType -> (EntityType<?, ?>) graphType)
+                            .collect(Collectors.toList()))));
   }
 
   /** Configures custom type extensions leveraged within our GraphQL schema. */
@@ -3322,4 +3341,34 @@ private void configureMetadataAttributionResolver(final RuntimeWiring.Builder bu
                         entityTypes,
                         (env) -> ((MetadataAttribution) env.getSource()).getSource())));
   }
+
+  private void configureVersionPropertiesResolvers(final RuntimeWiring.Builder builder) {
+    builder.type(
+        "VersionProperties",
+        typeWiring ->
+            typeWiring.dataFetcher(
+                "versionSet",
+                new LoadableTypeResolver<>(
+                    versionSetType,
+                    (env) -> {
+                      final VersionProperties versionProperties = env.getSource();
+                      return versionProperties != null
+                          ? versionProperties.getVersionSet().getUrn()
+                          : null;
+                    })));
+  }
+
+  private void configureVersionSetResolvers(final RuntimeWiring.Builder builder) {
+    builder.type(
+        "VersionSet",
+        typeWiring ->
+            typeWiring
+                .dataFetcher(
+                    "latestVersion",
+                    new EntityTypeResolver(
+                        entityTypes, (env) -> ((VersionSet) env.getSource()).getLatestVersion()))
+                .dataFetcher(
+                    "versionsSearch",
+                    new VersionsSearchResolver(this.entityClient, this.viewService)));
+  }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/config/AppConfigResolver.java

@@ -189,6 +189,7 @@ public CompletableFuture<AppConfig> get(final DataFetchingEnvironment environmen
             .setEditableDatasetNameEnabled(_featureFlags.isEditableDatasetNameEnabled())
             .setShowSeparateSiblings(_featureFlags.isShowSeparateSiblings())
             .setShowManageStructuredProperties(_featureFlags.isShowManageStructuredProperties())
+            .setEntityVersioningEnabled(_featureFlags.isEntityVersioning())
             .build();
 
     appConfig.setFeatureFlags(featureFlagsConfig);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/LinkAssetVersionResolver.java

@@ -12,7 +12,9 @@
 import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.featureflags.FeatureFlags;
+import com.linkedin.datahub.graphql.generated.EntityType;
 import com.linkedin.datahub.graphql.generated.LinkVersionInput;
+import com.linkedin.datahub.graphql.generated.VersionSet;
 import com.linkedin.metadata.entity.IngestResult;
 import com.linkedin.metadata.entity.versioning.EntityVersioningService;
 import com.linkedin.metadata.entity.versioning.VersionPropertiesInput;
@@ -21,24 +23,22 @@
 import io.datahubproject.metadata.context.OperationContext;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang.StringUtils;
 
 /**
  * Currently only supports linking the latest version, but may be modified later to support inserts
  */
-public class LinkAssetVersionResolver implements DataFetcher<CompletableFuture<String>> {
+@Slf4j
+@RequiredArgsConstructor
+public class LinkAssetVersionResolver implements DataFetcher<CompletableFuture<VersionSet>> {
 
   private final EntityVersioningService entityVersioningService;
   private final FeatureFlags featureFlags;
 
-  public LinkAssetVersionResolver(
-      EntityVersioningService entityVersioningService, FeatureFlags featureFlags) {
-    this.entityVersioningService = entityVersioningService;
-    this.featureFlags = featureFlags;
-  }
-
   @Override
-  public CompletableFuture<String> get(DataFetchingEnvironment environment) throws Exception {
+  public CompletableFuture<VersionSet> get(DataFetchingEnvironment environment) throws Exception {
     final QueryContext context = environment.getContext();
     final LinkVersionInput input =
         bindArgument(environment.getArgument("input"), LinkVersionInput.class);
@@ -75,12 +75,22 @@ public CompletableFuture<String> get(DataFetchingEnvironment environment) throws
               entityVersioningService.linkLatestVersion(
                   opContext, versionSetUrn, entityUrn, versionPropertiesInput);
 
-          return linkResults.stream()
-              .filter(
-                  ingestResult -> input.getLinkedEntity().equals(ingestResult.getUrn().toString()))
-              .map(ingestResult -> ingestResult.getUrn().toString())
-              .findAny()
-              .orElse(StringUtils.EMPTY);
+          String successVersionSetUrn =
+              linkResults.stream()
+                  .filter(
+                      ingestResult ->
+                          input.getLinkedEntity().equals(ingestResult.getUrn().toString()))
+                  .map(ingestResult -> ingestResult.getUrn().toString())
+                  .findAny()
+                  .orElse(StringUtils.EMPTY);
+
+          if (StringUtils.isEmpty(successVersionSetUrn)) {
+            return null;
+          }
+          VersionSet versionSet = new VersionSet();
+          versionSet.setUrn(versionSetUrn.toString());
+          versionSet.setType(EntityType.VERSION_SET);
+          return versionSet;
         },
         this.getClass().getSimpleName(),
         "get");

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/entity/versioning/UnlinkAssetVersionResolver.java

@@ -12,14 +12,18 @@
 import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.featureflags.FeatureFlags;
+import com.linkedin.datahub.graphql.generated.EntityType;
 import com.linkedin.datahub.graphql.generated.UnlinkVersionInput;
+import com.linkedin.datahub.graphql.generated.VersionSet;
+import com.linkedin.metadata.entity.RollbackResult;
 import com.linkedin.metadata.entity.versioning.EntityVersioningService;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import io.datahubproject.metadata.context.OperationContext;
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
 
-public class UnlinkAssetVersionResolver implements DataFetcher<CompletableFuture<Boolean>> {
+public class UnlinkAssetVersionResolver implements DataFetcher<CompletableFuture<VersionSet>> {
 
   private final EntityVersioningService entityVersioningService;
   private final FeatureFlags featureFlags;
@@ -31,7 +35,7 @@ public UnlinkAssetVersionResolver(
   }
 
   @Override
-  public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throws Exception {
+  public CompletableFuture<VersionSet> get(DataFetchingEnvironment environment) throws Exception {
     if (!featureFlags.isEntityVersioning()) {
       throw new IllegalAccessError(
           "Entity Versioning is not configured, please enable before attempting to use this feature.");
@@ -58,8 +62,15 @@ public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throw
     }
     return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
-          entityVersioningService.unlinkVersion(opContext, versionSetUrn, entityUrn);
-          return true;
+          List<RollbackResult> results =
+              entityVersioningService.unlinkVersion(opContext, versionSetUrn, entityUrn);
+          if (results.isEmpty() || results.stream().allMatch(RollbackResult::isNoOp)) {
+            return null;
+          }
+          VersionSet versionSet = new VersionSet();
+          versionSet.setUrn(versionSetUrn.toString());
+          versionSet.setType(EntityType.VERSION_SET);
+          return versionSet;
         },
         this.getClass().getSimpleName(),
         "get");