Skip to content

Commit

Permalink
fix(graphql): add data process instance entity type (datahub-project#…
Browse files Browse the repository at this point in the history
  • Loading branch information
RyanHolstien committed Jan 8, 2025
1 parent 92f013e commit accc254
Show file tree
Hide file tree
Showing 5 changed files with 110 additions and 18 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -346,6 +346,7 @@
import com.linkedin.datahub.graphql.types.datajob.DataJobType;
import com.linkedin.datahub.graphql.types.dataplatform.DataPlatformType;
import com.linkedin.datahub.graphql.types.dataplatforminstance.DataPlatformInstanceType;
import com.linkedin.datahub.graphql.types.dataprocessinst.DataProcessInstanceType;
import com.linkedin.datahub.graphql.types.dataprocessinst.mappers.DataProcessInstanceRunEventMapper;
import com.linkedin.datahub.graphql.types.dataproduct.DataProductType;
import com.linkedin.datahub.graphql.types.dataset.DatasetType;
Expand Down Expand Up @@ -530,6 +531,7 @@ public class GmsGraphQLEngine {
private final FormType formType;
private final IncidentType incidentType;
private final RestrictedType restrictedType;
private final DataProcessInstanceType dataProcessInstanceType;

private final int graphQLQueryComplexityLimit;
private final int graphQLQueryDepthLimit;
Expand Down Expand Up @@ -649,6 +651,7 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
this.formType = new FormType(entityClient);
this.incidentType = new IncidentType(entityClient);
this.restrictedType = new RestrictedType(entityClient, restrictedService);
this.dataProcessInstanceType = new DataProcessInstanceType(entityClient);

this.graphQLQueryComplexityLimit = args.graphQLQueryComplexityLimit;
this.graphQLQueryDepthLimit = args.graphQLQueryDepthLimit;
Expand Down Expand Up @@ -699,7 +702,8 @@ public GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
formType,
incidentType,
restrictedType,
businessAttributeType));
businessAttributeType,
dataProcessInstanceType));
this.loadableTypes = new ArrayList<>(entityTypes);
// Extend loadable types with types from the plugins
// This allows us to offer search and browse capabilities out of the box for
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -225,6 +225,11 @@ public Entity apply(@Nullable QueryContext context, Urn input) {
((BusinessAttribute) partialEntity).setUrn(input.toString());
((BusinessAttribute) partialEntity).setType(EntityType.BUSINESS_ATTRIBUTE);
}
if (input.getEntityType().equals(DATA_PROCESS_INSTANCE_ENTITY_NAME)) {
partialEntity = new DataProcessInstance();
((DataProcessInstance) partialEntity).setUrn(input.toString());
((DataProcessInstance) partialEntity).setType(EntityType.DATA_PROCESS_INSTANCE);
}
return partialEntity;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
package com.linkedin.datahub.graphql.types.dataprocessinst;

import static com.linkedin.datahub.graphql.authorization.AuthorizationUtils.canView;
import static com.linkedin.metadata.Constants.*;

import com.google.common.collect.ImmutableSet;
import com.linkedin.common.urn.Urn;
import com.linkedin.common.urn.UrnUtils;
import com.linkedin.datahub.graphql.QueryContext;
import com.linkedin.datahub.graphql.generated.DataProcessInstance;
import com.linkedin.datahub.graphql.generated.Entity;
import com.linkedin.datahub.graphql.generated.EntityType;
import com.linkedin.datahub.graphql.types.dataprocessinst.mappers.DataProcessInstanceMapper;
import com.linkedin.entity.EntityResponse;
import com.linkedin.entity.client.EntityClient;
import graphql.execution.DataFetcherResult;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;

public class DataProcessInstanceType
implements com.linkedin.datahub.graphql.types.EntityType<DataProcessInstance, String> {

private static final Set<String> ASPECTS_TO_RESOLVE =
ImmutableSet.of(
DATA_PROCESS_INSTANCE_KEY_ASPECT_NAME,
DATA_PROCESS_INSTANCE_INPUT_ASPECT_NAME,
DATA_PROCESS_INSTANCE_OUTPUT_ASPECT_NAME,
DATA_PROCESS_INSTANCE_RELATIONSHIPS_ASPECT_NAME,
STATUS_ASPECT_NAME,
TEST_RESULTS_ASPECT_NAME,
LINEAGE_FEATURES_ASPECT_NAME);

private static final Set<String> FACET_FIELDS = ImmutableSet.of("origin", "platform");
private static final String ENTITY_NAME = "dataProcessInstance";

private final EntityClient entityClient;

public DataProcessInstanceType(final EntityClient entityClient) {
this.entityClient = entityClient;
}

@Override
public Class<DataProcessInstance> objectClass() {
return DataProcessInstance.class;
}

@Override
public EntityType type() {
return EntityType.DATA_PROCESS_INSTANCE;
}

@Override
public Function<Entity, String> getKeyProvider() {
return Entity::getUrn;
}

@Override
public List<DataFetcherResult<DataProcessInstance>> batchLoad(
@Nonnull final List<String> urnStrs, @Nonnull final QueryContext context) {
try {
final List<Urn> urns = urnStrs.stream().map(UrnUtils::getUrn).collect(Collectors.toList());

final Map<Urn, EntityResponse> dataProcessInstanceMap =
entityClient.batchGetV2(
context.getOperationContext(),
DATA_PROCESS_INSTANCE_ENTITY_NAME,
urns.stream()
.filter(urn -> canView(context.getOperationContext(), urn))
.collect(Collectors.toSet()),
ASPECTS_TO_RESOLVE);

final List<EntityResponse> gmsResults = new ArrayList<>(urnStrs.size());
for (Urn urn : urns) {
gmsResults.add(dataProcessInstanceMap.getOrDefault(urn, null));
}
return gmsResults.stream()
.map(
gmsDataProcessInstance ->
gmsDataProcessInstance == null
? null
: DataFetcherResult.<DataProcessInstance>newResult()
.data(DataProcessInstanceMapper.map(context, gmsDataProcessInstance))
.build())
.collect(Collectors.toList());
} catch (Exception e) {
throw new RuntimeException("Failed to batch load Datasets", e);
}
}
}
3 changes: 3 additions & 0 deletions li-utils/src/main/java/com/linkedin/metadata/Constants.java
Original file line number Diff line number Diff line change
Expand Up @@ -421,6 +421,9 @@ public class Constants {
"dataProcessInstanceRunEvent";
public static final String DATA_PROCESS_INSTANCE_RELATIONSHIPS_ASPECT_NAME =
"dataProcessInstanceRelationships";
public static final String DATA_PROCESS_INSTANCE_INPUT_ASPECT_NAME = "dataProcessInstanceInput";
public static final String DATA_PROCESS_INSTANCE_OUTPUT_ASPECT_NAME = "dataProcessInstanceOutput";
public static final String DATA_PROCESS_INSTANCE_KEY_ASPECT_NAME = "dataProcessInstanceKey";

// Business Attribute
public static final String BUSINESS_ATTRIBUTE_KEY_ASPECT_NAME = "businessAttributeKey";
Expand Down
Original file line number Diff line number Diff line change
@@ -1,21 +1,6 @@
package com.datahub.authorization;

import static com.linkedin.metadata.Constants.CHART_ENTITY_NAME;
import static com.linkedin.metadata.Constants.DASHBOARD_ENTITY_NAME;
import static com.linkedin.metadata.Constants.DATASET_ENTITY_NAME;
import static com.linkedin.metadata.Constants.DATA_FLOW_ENTITY_NAME;
import static com.linkedin.metadata.Constants.DATA_JOB_ENTITY_NAME;
import static com.linkedin.metadata.Constants.DATA_PRODUCT_ENTITY_NAME;
import static com.linkedin.metadata.Constants.DOMAIN_ENTITY_NAME;
import static com.linkedin.metadata.Constants.GLOSSARY_NODE_ENTITY_NAME;
import static com.linkedin.metadata.Constants.GLOSSARY_TERM_ENTITY_NAME;
import static com.linkedin.metadata.Constants.ML_FEATURE_ENTITY_NAME;
import static com.linkedin.metadata.Constants.ML_FEATURE_TABLE_ENTITY_NAME;
import static com.linkedin.metadata.Constants.ML_MODEL_ENTITY_NAME;
import static com.linkedin.metadata.Constants.ML_MODEL_GROUP_ENTITY_NAME;
import static com.linkedin.metadata.Constants.ML_PRIMARY_KEY_ENTITY_NAME;
import static com.linkedin.metadata.Constants.NOTEBOOK_ENTITY_NAME;
import static com.linkedin.metadata.Constants.REST_API_AUTHORIZATION_ENABLED_ENV;
import static com.linkedin.metadata.Constants.*;
import static com.linkedin.metadata.authorization.ApiGroup.ENTITY;
import static com.linkedin.metadata.authorization.ApiOperation.CREATE;
import static com.linkedin.metadata.authorization.ApiOperation.DELETE;
Expand Down Expand Up @@ -89,7 +74,8 @@ public class AuthUtil {
GLOSSARY_NODE_ENTITY_NAME,
DOMAIN_ENTITY_NAME,
DATA_PRODUCT_ENTITY_NAME,
NOTEBOOK_ENTITY_NAME);
NOTEBOOK_ENTITY_NAME,
DATA_PROCESS_INSTANCE_ENTITY_NAME);

/** OpenAPI/Rest.li Methods */
public static List<Pair<MetadataChangeProposal, Integer>> isAPIAuthorized(
Expand Down

0 comments on commit accc254

Please sign in to comment.