Skip to content

Commit

Permalink
feat(onboarding): adds framework and some steps for onboarding steps …
Browse files Browse the repository at this point in the history
…UI (#6462)

* feat(onboarding): adds models and API for onboarding steps feature

* feat(onboarding): adds backend for onboarding steps feature

* feat(onboarding): adds framework and some steps for onboarding steps UI
  • Loading branch information
aditya-radhakrishnan authored Dec 8, 2022
1 parent 9e3267a commit f0f0355
Show file tree
Hide file tree
Showing 68 changed files with 2,067 additions and 199 deletions.
1 change: 1 addition & 0 deletions datahub-graphql-core/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ graphqlCodegen {
"$projectDir/src/main/resources/auth.graphql".toString(),
"$projectDir/src/main/resources/timeline.graphql".toString(),
"$projectDir/src/main/resources/tests.graphql".toString(),
"$projectDir/src/main/resources/step.graphql".toString(),
]
outputDir = new File("$projectDir/src/mainGeneratedGraphQL/java")
packageName = "com.linkedin.datahub.graphql.generated"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ public class Constants {
public static final String INGESTION_SCHEMA_FILE = "ingestion.graphql";
public static final String TIMELINE_SCHEMA_FILE = "timeline.graphql";
public static final String TESTS_SCHEMA_FILE = "tests.graphql";
public static final String STEPS_SCHEMA_FILE = "step.graphql";
public static final String BROWSE_PATH_DELIMITER = "/";
public static final String VERSION_STAMP_FIELD_NAME = "versionStamp";
}
Original file line number Diff line number Diff line change
Expand Up @@ -192,6 +192,8 @@
import com.linkedin.datahub.graphql.resolvers.search.SearchAcrossEntitiesResolver;
import com.linkedin.datahub.graphql.resolvers.search.SearchAcrossLineageResolver;
import com.linkedin.datahub.graphql.resolvers.search.SearchResolver;
import com.linkedin.datahub.graphql.resolvers.step.BatchGetStepStatesResolver;
import com.linkedin.datahub.graphql.resolvers.step.BatchUpdateStepStatesResolver;
import com.linkedin.datahub.graphql.resolvers.tag.CreateTagResolver;
import com.linkedin.datahub.graphql.resolvers.tag.DeleteTagResolver;
import com.linkedin.datahub.graphql.resolvers.tag.SetTagColorResolver;
Expand Down Expand Up @@ -557,6 +559,7 @@ public GraphQLEngine.Builder builder() {
.addSchema(fileBasedSchema(INGESTION_SCHEMA_FILE))
.addSchema(fileBasedSchema(TIMELINE_SCHEMA_FILE))
.addSchema(fileBasedSchema(TESTS_SCHEMA_FILE))
.addSchema(fileBasedSchema(STEPS_SCHEMA_FILE))
.addDataLoaders(loaderSuppliers(loadableTypes))
.addDataLoader("Aspect", context -> createDataLoader(aspectType, context))
.configureRuntimeWiring(this::configureRuntimeWiring);
Expand Down Expand Up @@ -691,6 +694,7 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) {
.dataFetcher("listRoles", new ListRolesResolver(this.entityClient))
.dataFetcher("getInviteToken", new GetInviteTokenResolver(this.inviteTokenService))
.dataFetcher("listPosts", new ListPostsResolver(this.entityClient))
.dataFetcher("batchGetStepStates", new BatchGetStepStatesResolver(this.entityClient))
);
}

Expand Down Expand Up @@ -814,6 +818,7 @@ private void configureMutationResolvers(final RuntimeWiring.Builder builder) {
.dataFetcher("createInviteToken", new CreateInviteTokenResolver(this.inviteTokenService))
.dataFetcher("acceptRole", new AcceptRoleResolver(this.roleService, this.inviteTokenService))
.dataFetcher("createPost", new CreatePostResolver(this.postService))
.dataFetcher("batchUpdateStepStates", new BatchUpdateStepStatesResolver(this.entityClient))
);
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,135 @@
package com.linkedin.datahub.graphql.resolvers.step;

import com.datahub.authentication.Authentication;
import com.google.common.collect.ImmutableSet;
import com.linkedin.common.urn.Urn;
import com.linkedin.datahub.graphql.QueryContext;
import com.linkedin.datahub.graphql.generated.BatchGetStepStatesInput;
import com.linkedin.datahub.graphql.generated.BatchGetStepStatesResult;
import com.linkedin.datahub.graphql.generated.StepStateResult;
import com.linkedin.datahub.graphql.generated.StringMapEntry;
import com.linkedin.entity.EntityResponse;
import com.linkedin.entity.EnvelopedAspectMap;
import com.linkedin.entity.client.EntityClient;
import com.linkedin.metadata.key.DataHubStepStateKey;
import com.linkedin.r2.RemoteInvocationException;
import com.linkedin.step.DataHubStepStateProperties;
import graphql.schema.DataFetcher;
import graphql.schema.DataFetchingEnvironment;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*;
import static com.linkedin.metadata.Constants.*;
import static com.linkedin.metadata.utils.EntityKeyUtils.*;


@Slf4j
@RequiredArgsConstructor
public class BatchGetStepStatesResolver implements DataFetcher<CompletableFuture<BatchGetStepStatesResult>> {
private final EntityClient _entityClient;

@Override
public CompletableFuture<BatchGetStepStatesResult> get(@Nonnull final DataFetchingEnvironment environment)
throws Exception {
final QueryContext context = environment.getContext();
final Authentication authentication = context.getAuthentication();
final BatchGetStepStatesInput input =
bindArgument(environment.getArgument("input"), BatchGetStepStatesInput.class);

return CompletableFuture.supplyAsync(() -> {
Map<Urn, String> urnsToIdsMap;
Set<Urn> urns;
Map<Urn, EntityResponse> entityResponseMap;

try {
urnsToIdsMap = buildUrnToIdMap(input.getIds(), authentication);
urns = urnsToIdsMap.keySet();
entityResponseMap = _entityClient.batchGetV2(DATAHUB_STEP_STATE_ENTITY_NAME, urns,
ImmutableSet.of(DATAHUB_STEP_STATE_PROPERTIES_ASPECT_NAME), authentication);
} catch (Exception e) {
throw new RuntimeException(e);
}

final Map<Urn, DataHubStepStateProperties> stepStatePropertiesMap = new HashMap<>();
for (Map.Entry<Urn, EntityResponse> entry : entityResponseMap.entrySet()) {
final Urn urn = entry.getKey();
final DataHubStepStateProperties stepStateProperties = getStepStateProperties(urn, entry.getValue());
if (stepStateProperties != null) {
stepStatePropertiesMap.put(urn, stepStateProperties);
}
}

final List<StepStateResult> results = stepStatePropertiesMap.entrySet()
.stream()
.map(entry -> buildStepStateResult(urnsToIdsMap.get(entry.getKey()), entry.getValue()))
.collect(Collectors.toList());
final BatchGetStepStatesResult result = new BatchGetStepStatesResult();
result.setResults(results);
return result;
});
}

@Nonnull
private Map<Urn, String> buildUrnToIdMap(@Nonnull final List<String> ids, @Nonnull final Authentication authentication)
throws RemoteInvocationException {
final Map<Urn, String> urnToIdMap = new HashMap<>();
for (final String id : ids) {
final Urn urn = getStepStateUrn(id);
if (_entityClient.exists(urn, authentication)) {
urnToIdMap.put(urn, id);
}
}

return urnToIdMap;
}

@Nonnull
private Urn getStepStateUrn(@Nonnull final String id) {
final DataHubStepStateKey stepStateKey = new DataHubStepStateKey().setId(id);
return convertEntityKeyToUrn(stepStateKey, DATAHUB_STEP_STATE_ENTITY_NAME);
}

@Nullable
private DataHubStepStateProperties getStepStateProperties(@Nonnull final Urn urn,
@Nonnull final EntityResponse entityResponse) {
final EnvelopedAspectMap aspectMap = entityResponse.getAspects();
// If aspect is not present, log the error and return null.
if (!aspectMap.containsKey(DATAHUB_STEP_STATE_PROPERTIES_ASPECT_NAME)) {
log.error("Failed to find step state properties for urn: " + urn);
return null;
}
return new DataHubStepStateProperties(aspectMap.get(DATAHUB_STEP_STATE_PROPERTIES_ASPECT_NAME).getValue().data());
}

@Nonnull
private StepStateResult buildStepStateResult(@Nonnull final String id,
@Nonnull final DataHubStepStateProperties stepStateProperties) {
final StepStateResult result = new StepStateResult();
result.setId(id);
final List<StringMapEntry> mappedProperties = stepStateProperties
.getProperties()
.entrySet()
.stream()
.map(entry -> buildStringMapEntry(entry.getKey(), entry.getValue()))
.collect(Collectors.toList());
result.setProperties(mappedProperties);
return result;
}

@Nonnull
private StringMapEntry buildStringMapEntry(@Nonnull final String key, @Nonnull final String value) {
final StringMapEntry entry = new StringMapEntry();
entry.setKey(key);
entry.setValue(value);
return entry;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
package com.linkedin.datahub.graphql.resolvers.step;

import com.datahub.authentication.Authentication;
import com.linkedin.common.AuditStamp;
import com.linkedin.common.urn.Urn;
import com.linkedin.common.urn.UrnUtils;
import com.linkedin.data.template.StringMap;
import com.linkedin.datahub.graphql.QueryContext;
import com.linkedin.datahub.graphql.generated.BatchUpdateStepStatesInput;
import com.linkedin.datahub.graphql.generated.BatchUpdateStepStatesResult;
import com.linkedin.datahub.graphql.generated.StepStateInput;
import com.linkedin.datahub.graphql.generated.StringMapEntryInput;
import com.linkedin.datahub.graphql.generated.UpdateStepStateResult;
import com.linkedin.entity.client.EntityClient;
import com.linkedin.metadata.key.DataHubStepStateKey;
import com.linkedin.mxe.MetadataChangeProposal;
import com.linkedin.step.DataHubStepStateProperties;
import graphql.schema.DataFetcher;
import graphql.schema.DataFetchingEnvironment;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*;
import static com.linkedin.metadata.Constants.*;
import static com.linkedin.metadata.entity.AspectUtils.*;


@Slf4j
@RequiredArgsConstructor
public class BatchUpdateStepStatesResolver implements DataFetcher<CompletableFuture<BatchUpdateStepStatesResult>> {
private final EntityClient _entityClient;

@Override
public CompletableFuture<BatchUpdateStepStatesResult> get(@Nonnull final DataFetchingEnvironment environment)
throws Exception {
final QueryContext context = environment.getContext();
final Authentication authentication = context.getAuthentication();

final BatchUpdateStepStatesInput input =
bindArgument(environment.getArgument("input"), BatchUpdateStepStatesInput.class);
final List<StepStateInput> states = input.getStates();
final String actorUrnStr = authentication.getActor().toUrnStr();

return CompletableFuture.supplyAsync(() -> {
final Urn actorUrn = UrnUtils.getUrn(actorUrnStr);
final AuditStamp auditStamp = new AuditStamp().setActor(actorUrn).setTime(System.currentTimeMillis());
final List<UpdateStepStateResult> results = states
.stream()
.map(state -> buildUpdateStepStateResult(state, auditStamp, authentication))
.collect(Collectors.toList());
final BatchUpdateStepStatesResult result = new BatchUpdateStepStatesResult();
result.setResults(results);
return result;
});
}

private UpdateStepStateResult buildUpdateStepStateResult(@Nonnull final StepStateInput state,
@Nonnull final AuditStamp auditStamp,
@Nonnull final Authentication authentication) {
final String id = state.getId();
final UpdateStepStateResult updateStepStateResult = new UpdateStepStateResult();
updateStepStateResult.setId(id);
final boolean success = updateStepState(id, state.getProperties(), auditStamp, authentication);
updateStepStateResult.setSucceeded(success);
return updateStepStateResult;
}

private boolean updateStepState(@Nonnull final String id,
@Nonnull final List<StringMapEntryInput> inputProperties, @Nonnull final AuditStamp auditStamp,
@Nonnull final Authentication authentication) {
final Map<String, String> properties =
inputProperties.stream().collect(Collectors.toMap(StringMapEntryInput::getKey, StringMapEntryInput::getValue));
try {
final DataHubStepStateKey stepStateKey = new DataHubStepStateKey().setId(id);
final DataHubStepStateProperties stepStateProperties =
new DataHubStepStateProperties().setProperties(new StringMap(properties)).setLastModified(auditStamp);

final MetadataChangeProposal proposal =
buildMetadataChangeProposal(DATAHUB_STEP_STATE_ENTITY_NAME, stepStateKey,
DATAHUB_STEP_STATE_PROPERTIES_ASPECT_NAME, stepStateProperties);
_entityClient.ingestProposal(proposal, authentication, false);
return true;
} catch (Exception e) {
log.error("Could not update step state for id {}", id, e);
return false;
}
}
}
Original file line number Diff line number Diff line change
@@ -1,11 +1,17 @@
package com.linkedin.datahub.graphql.resolvers.test;

import com.linkedin.data.template.RecordTemplate;
import com.linkedin.data.template.SetMode;
import com.linkedin.datahub.graphql.QueryContext;
import com.linkedin.datahub.graphql.generated.TestDefinitionInput;
import com.linkedin.entity.Aspect;
import com.linkedin.entity.EntityResponse;
import com.linkedin.entity.EnvelopedAspect;
import com.linkedin.entity.EnvelopedAspectMap;
import com.linkedin.metadata.authorization.PoliciesConfig;
import com.linkedin.test.TestDefinition;
import com.linkedin.test.TestDefinitionType;
import java.util.Map;
import java.util.Optional;
import javax.annotation.Nonnull;

Expand All @@ -28,5 +34,15 @@ public static TestDefinition mapDefinition(final TestDefinitionInput testDefInpu
return result;
}

public static EntityResponse buildEntityResponse(Map<String, RecordTemplate> aspects) {
final EntityResponse entityResponse = new EntityResponse();
final EnvelopedAspectMap aspectMap = new EnvelopedAspectMap();
for (Map.Entry<String, RecordTemplate> entry : aspects.entrySet()) {
aspectMap.put(entry.getKey(), new EnvelopedAspect().setValue(new Aspect(entry.getValue().data())));
}
entityResponse.setAspects(aspectMap);
return entityResponse;
}

private TestUtils() { }
}
Loading

0 comments on commit f0f0355

Please sign in to comment.