diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/AspectMappingRegistry.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/AspectMappingRegistry.java new file mode 100644 index 0000000000000..f0a3c65207580 --- /dev/null +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/AspectMappingRegistry.java @@ -0,0 +1,130 @@ +package com.linkedin.datahub.graphql; + +import graphql.language.ArrayValue; +import graphql.language.StringValue; +import graphql.schema.GraphQLArgument; +import graphql.schema.GraphQLDirective; +import graphql.schema.GraphQLObjectType; +import graphql.schema.GraphQLSchema; +import java.util.*; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import lombok.extern.slf4j.Slf4j; + +/** + * Maps GraphQL field selections to the minimum set of aspects needed to resolve them, enabling + * performance optimization by fetching only required aspects instead of all aspects. + * + *

This class scans the GraphQL schema for two directives: - @aspectMapping(aspects: + * ["aspectName"]) - declares which aspects a field needs - @noAspects - indicates a field needs no + * aspects (computed fields, custom resolvers) + * + *

To use in entity types, add one line to batchLoad: Set aspects = + * AspectUtils.getOptimizedAspects(context, "Dataset", ALL_ASPECTS, "datasetKey"); + * + *

If any field lacks a mapping directive, getRequiredAspects returns null and the entity type + * falls back to fetching all aspects for safety. + */ +@Slf4j +public class AspectMappingRegistry { + private final Map> fieldToAspects = new HashMap<>(); + + public AspectMappingRegistry(GraphQLSchema schema) { + buildMappingFromSchema(schema); + } + + private void buildMappingFromSchema(GraphQLSchema schema) { + schema + .getTypeMap() + .values() + .forEach( + type -> { + if (type instanceof GraphQLObjectType) { + GraphQLObjectType objectType = (GraphQLObjectType) type; + String typeName = objectType.getName(); + + objectType + .getFieldDefinitions() + .forEach( + field -> { + String fieldName = field.getName(); + GraphQLDirective aspectsDirective = field.getDirective("aspectMapping"); + GraphQLDirective noAspectsDirective = field.getDirective("noAspects"); + + if (aspectsDirective != null) { + GraphQLArgument aspectsArg = aspectsDirective.getArgument("aspects"); + if (aspectsArg != null + && aspectsArg.getArgumentValue().getValue() instanceof ArrayValue) { + ArrayValue aspectsArray = + (ArrayValue) aspectsArg.getArgumentValue().getValue(); + Set aspects = + aspectsArray.getValues().stream() + .map(value -> ((StringValue) value).getValue()) + .collect(Collectors.toSet()); + + String key = typeName + "." + fieldName; + fieldToAspects.put(key, aspects); + log.debug( + "Mapped {}.{} to aspects: {}", typeName, fieldName, aspects); + } + } else if (noAspectsDirective != null) { + String key = typeName + "." + fieldName; + fieldToAspects.put(key, new HashSet<>()); + log.debug( + "Mapped {}.{} to to request no specific aspects.", + typeName, + fieldName); + } + }); + } + }); + + log.info("Built aspect mapping registry with {} field mappings", fieldToAspects.size()); + } + + /** + * Get required aspects for the given fields on a type. Returns null if any field is unmapped + * (fallback to all aspects). + * + *

This method filters the selection set to only include fields that directly belong to the + * specified type, regardless of where that type appears in the query tree. This allows it to work + * correctly for both top-level queries and nested entities (e.g., Dataset inside SearchResult). + */ + @Nullable + public Set getRequiredAspects( + String typeName, List requestedFields) { + Set aspects = new HashSet<>(); + + for (graphql.schema.SelectedField field : requestedFields) { + String fieldName = field.getName(); + + // Skip introspection fields + if (fieldName.startsWith("__")) { + continue; + } + + // Only process fields that belong to the target type + // getObjectTypeNames() returns the set of types this field belongs to (accounting for + // interfaces/unions) + if (!field.getObjectTypeNames().contains(typeName)) { + continue; + } + + String key = typeName + "." + fieldName; + Set fieldAspects = fieldToAspects.get(key); + + if (fieldAspects != null) { + aspects.addAll(fieldAspects); + log.debug("Field {} mapped to aspects: {}", key, fieldAspects); + } else { + // Unmapped field - fallback to all aspects to be conservative + log.debug( + "Field {} has no @aspectMapping or @noAspects directives, will fetch all aspects", key); + return null; + } + } + + log.debug("Computed required aspects for {}: {}", typeName, aspects); + return aspects.isEmpty() ? Collections.emptySet() : aspects; + } +} diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/QueryContext.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/QueryContext.java index 5ad82b5d70375..aaa14581b2d09 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/QueryContext.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/QueryContext.java @@ -4,7 +4,9 @@ import com.datahub.authentication.Authentication; import com.datahub.plugins.auth.authorization.Authorizer; import com.linkedin.metadata.config.DataHubAppConfiguration; +import graphql.schema.DataFetchingEnvironment; import io.datahubproject.metadata.context.OperationContext; +import javax.annotation.Nullable; /** Provided as input to GraphQL resolvers; used to carry information about GQL request context. */ public interface QueryContext { @@ -34,4 +36,38 @@ default String getActorUrn() { OperationContext getOperationContext(); DataHubAppConfiguration getDataHubAppConfig(); + + /** + * Returns the {@link DataFetchingEnvironment} associated with the current GraphQL request. This + * provides access to the GraphQL query structure, requested fields, and other execution context. + * + * @return the DataFetchingEnvironment, or null if not available + */ + @Nullable + DataFetchingEnvironment getDataFetchingEnvironment(); + + /** + * Sets the {@link DataFetchingEnvironment} for the current GraphQL request. This is typically + * called by GraphQL resolvers to provide access to the execution context. + * + * @param environment the DataFetchingEnvironment to associate with this context + */ + void setDataFetchingEnvironment(@Nullable DataFetchingEnvironment environment); + + /** + * Returns the {@link AspectMappingRegistry} for optimizing aspect fetching based on GraphQL field + * selections. + * + * @return the AspectMappingRegistry, or null if not available + */ + @Nullable + AspectMappingRegistry getAspectMappingRegistry(); + + /** + * Sets the {@link AspectMappingRegistry} for this context. This is typically called during + * context initialization. + * + * @param aspectMappingRegistry the AspectMappingRegistry to use + */ + void setAspectMappingRegistry(@Nullable AspectMappingRegistry aspectMappingRegistry); } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityTypeResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityTypeResolver.java index 3c285f30661bc..6b799ff36d7bb 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityTypeResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityTypeResolver.java @@ -2,6 +2,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterables; +import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.generated.Entity; import graphql.schema.DataFetcher; import graphql.schema.DataFetchingEnvironment; @@ -44,6 +45,12 @@ private boolean isOnlySelectingIdentityFields(DataFetchingEnvironment environmen @Override public CompletableFuture get(DataFetchingEnvironment environment) { + // Set the DataFetchingEnvironment in the QueryContext for access in batchLoad methods + QueryContext context = environment.getContext(); + if (context != null) { + context.setDataFetchingEnvironment(environment); + } + final Entity resolvedEntity = _entityProvider.apply(environment); if (resolvedEntity == null) { return CompletableFuture.completedFuture(null); diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/LoadableTypeResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/LoadableTypeResolver.java index 3868b1a35b64f..1b0f41ce1ef11 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/LoadableTypeResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/LoadableTypeResolver.java @@ -1,5 +1,6 @@ package com.linkedin.datahub.graphql.resolvers.load; +import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.datahub.graphql.types.LoadableType; import graphql.schema.DataFetcher; import graphql.schema.DataFetchingEnvironment; @@ -32,6 +33,12 @@ public LoadableTypeResolver( @Override public CompletableFuture get(DataFetchingEnvironment environment) { + // Set the DataFetchingEnvironment in the QueryContext for access in batchLoad methods + QueryContext context = environment.getContext(); + if (context != null) { + context.setDataFetchingEnvironment(environment); + } + final K key = _keyProvider.apply(environment); if (key == null) { return null; diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java index a0579d4f2b75e..4ff992095e9da 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java @@ -34,6 +34,7 @@ import com.linkedin.datahub.graphql.types.mappers.BrowsePathsMapper; import com.linkedin.datahub.graphql.types.mappers.BrowseResultMapper; import com.linkedin.datahub.graphql.types.mappers.UrnSearchResultsMapper; +import com.linkedin.datahub.graphql.util.AspectUtils; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.client.EntityClient; import com.linkedin.metadata.Constants; @@ -56,7 +57,9 @@ import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; +import lombok.extern.slf4j.Slf4j; +@Slf4j public class DatasetType implements SearchableEntityType, BrowsableEntityType, @@ -97,6 +100,7 @@ public class DatasetType private static final Set FACET_FIELDS = ImmutableSet.of("origin", "platform"); private static final String ENTITY_NAME = "dataset"; + private static final String KEY_ASPECT = "datasetKey"; private final EntityClient entityClient; @@ -135,12 +139,16 @@ public List> batchLoad( try { final List urns = urnStrs.stream().map(UrnUtils::getUrn).collect(Collectors.toList()); + // Determine optimal aspects to fetch based on GraphQL field selections + Set aspectsToResolve = + AspectUtils.getOptimizedAspects(context, "Dataset", ASPECTS_TO_RESOLVE, KEY_ASPECT); + final Map datasetMap = entityClient.batchGetV2( context.getOperationContext(), Constants.DATASET_ENTITY_NAME, new HashSet<>(urns), - ASPECTS_TO_RESOLVE); + aspectsToResolve); final List gmsResults = new ArrayList<>(urnStrs.size()); for (Urn urn : urns) { diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/util/AspectUtils.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/util/AspectUtils.java new file mode 100644 index 0000000000000..c6d728ae1d2fa --- /dev/null +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/util/AspectUtils.java @@ -0,0 +1,74 @@ +package com.linkedin.datahub.graphql.util; + +import com.linkedin.datahub.graphql.QueryContext; +import java.util.Collections; +import java.util.HashSet; +import java.util.Set; +import javax.annotation.Nonnull; +import lombok.extern.slf4j.Slf4j; + +/** + * Utility methods for optimizing aspect fetching in GraphQL entity types by determining which + * aspects need to be fetched based on requested fields. + */ +@Slf4j +public class AspectUtils { + + private AspectUtils() {} + + /** + * Determines optimal aspects to fetch based on GraphQL field selections. Falls back to + * defaultAspects if optimization isn't possible (missing registry, unmapped fields, etc). + * + *

Usage in entity type batchLoad: Set aspects = + * AspectUtils.getOptimizedAspects(context, "Dataset", ALL_ASPECTS, "datasetKey"); + * + * @param context the QueryContext containing AspectMappingRegistry and DataFetchingEnvironment + * @param entityTypeName the GraphQL type name (e.g., "Dataset", "CorpUser") + * @param defaultAspects the full set of aspects to use as fallback + * @param alwaysIncludeAspects aspects to always include (e.g., key aspects) + * @return optimized aspect set, or defaultAspects if optimization isn't possible + */ + @Nonnull + public static Set getOptimizedAspects( + @Nonnull final QueryContext context, + @Nonnull final String entityTypeName, + @Nonnull final Set defaultAspects, + @Nonnull final String... alwaysIncludeAspects) { + + // Check if we have the necessary context for optimization + if (context.getDataFetchingEnvironment() == null + || context.getAspectMappingRegistry() == null) { + log.debug( + "DataFetchingEnvironment or AspectMappingRegistry not available for {}, fetching all aspects", + entityTypeName); + return defaultAspects; + } + + // Attempt to determine required aspects from GraphQL field selections + Set requiredAspects = + context + .getAspectMappingRegistry() + .getRequiredAspects( + entityTypeName, context.getDataFetchingEnvironment().getSelectionSet().getFields()); + + // If we couldn't determine required aspects (e.g., unmapped field), fall back to all aspects + if (requiredAspects == null) { + log.debug( + "Could not determine required aspects for {}, falling back to fetching all aspects", + entityTypeName); + return defaultAspects; + } + + // Successfully optimized - build the minimal aspect set + Set optimizedAspects = new HashSet<>(requiredAspects); + + // Add any aspects that should always be included + if (alwaysIncludeAspects != null && alwaysIncludeAspects.length > 0) { + Collections.addAll(optimizedAspects, alwaysIncludeAspects); + } + + log.info("Fetching optimized aspect set for {}: {}", entityTypeName, optimizedAspects); + return optimizedAspects; + } +} diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql index 47384cece3406..43368cc87de82 100644 --- a/datahub-graphql-core/src/main/resources/entity.graphql +++ b/datahub-graphql-core/src/main/resources/entity.graphql @@ -1,6 +1,12 @@ # Extending the GQL type system to include Long type used for dates scalar Long +# Used for performance reasons to map fields to aspects so we only fetch aspects that we need +directive @aspectMapping(aspects: [String!]!) on FIELD_DEFINITION + +# Optional: directive to explicitly mark fields that don't need aspects +directive @noAspects on FIELD_DEFINITION + """ Root GraphQL API Schema """ @@ -1655,118 +1661,124 @@ type Dataset implements EntityWithRelationships & Entity & BrowsableEntity { """ The primary key of the Dataset """ - urn: String! + urn: String! @noAspects """ The standard Entity Type """ - type: EntityType! + type: EntityType! @noAspects """ The timestamp for the last time this entity was ingested + In order to fetch this we need to fetch other aspects """ - lastIngested: Long + lastIngested: Long @noAspects """ Standardized platform urn where the dataset is defined """ - platform: DataPlatform! - + platform: DataPlatform! @aspectMapping(aspects: ["datasetKey"]) # we also have custom graphql resolver for this """ The parent container in which the entity resides """ - container: Container + container: Container @aspectMapping(aspects: ["container"]) """ Recursively get the lineage of containers for this entity """ - parentContainers: ParentContainersResult - + parentContainers: ParentContainersResult @noAspects # parentContainers has own resolver """ Unique guid for dataset No longer to be used as the Dataset display name. Use properties.name instead """ name: String! + @aspectMapping( + aspects: ["datasetProperties", "editableDatasetProperties", "datasetKey"] + ) """ An additional set of read only properties """ - properties: DatasetProperties + properties: DatasetProperties @aspectMapping(aspects: ["datasetProperties"]) """ An additional set of of read write properties """ editableProperties: DatasetEditableProperties + @aspectMapping(aspects: ["editableDatasetProperties"]) """ Ownership metadata of the dataset """ - ownership: Ownership + ownership: Ownership @aspectMapping(aspects: ["ownership"]) """ The deprecation status of the dataset """ deprecation: Deprecation + @aspectMapping(aspects: ["datasetDeprecation", "deprecation"]) """ References to internal resources related to the dataset """ institutionalMemory: InstitutionalMemory + @aspectMapping(aspects: ["institutionalMemory"]) """ Schema metadata of the dataset, available by version number """ - schemaMetadata(version: Long): SchemaMetadata - + schemaMetadata(version: Long): SchemaMetadata @noAspects # schemaMetadata uses separate AspectResolver """ Editable schema metadata of the dataset """ editableSchemaMetadata: EditableSchemaMetadata + @aspectMapping(aspects: ["editableSchemaMetadata"]) """ Status of the Dataset """ - status: Status + status: Status @aspectMapping(aspects: ["status"]) """ Embed information about the Dataset """ - embed: Embed + embed: Embed @aspectMapping(aspects: ["embed"]) """ Tags used for searching dataset """ - tags: GlobalTags + tags: GlobalTags @aspectMapping(aspects: ["globalTags"]) """ The structured glossary terms associated with the dataset """ - glossaryTerms: GlossaryTerms + glossaryTerms: GlossaryTerms @aspectMapping(aspects: ["glossaryTerms"]) """ The specific instance of the data platform that this entity belongs to """ dataPlatformInstance: DataPlatformInstance + @aspectMapping(aspects: ["dataPlatformInstance"]) """ The Domain associated with the Dataset """ - domain: DomainAssociation + domain: DomainAssociation @aspectMapping(aspects: ["domains"]) """ The application associated with the dataset """ - application: ApplicationAssociation + application: ApplicationAssociation @aspectMapping(aspects: ["applications"]) """ The forms associated with the Dataset """ - forms: Forms + forms: Forms @aspectMapping(aspects: ["forms"]) """ The Roles and the properties to access the dataset """ - access: Access + access: Access @aspectMapping(aspects: ["access"]) """ Statistics about how this Dataset is used @@ -1778,13 +1790,11 @@ type Dataset implements EntityWithRelationships & Entity & BrowsableEntity { range: TimeRange startTimeMillis: Long timeZone: String - ): UsageQueryResult - + ): UsageQueryResult @noAspects # has own custom resolver """ Experimental - Summary operational & usage statistics about a Dataset """ - statsSummary: DatasetStatsSummary - + statsSummary: DatasetStatsSummary @noAspects # has own custom resolver """ Profile Stats resource that retrieves the events in a previous unit of time in descending order If no start or end time are provided, the most recent events will be returned @@ -1794,8 +1804,7 @@ type Dataset implements EntityWithRelationships & Entity & BrowsableEntity { endTimeMillis: Long filter: FilterInput limit: Int - ): [DatasetProfile!] - + ): [DatasetProfile!] @noAspects # has own custom resolver """ Operational events for an entity. """ @@ -1804,8 +1813,7 @@ type Dataset implements EntityWithRelationships & Entity & BrowsableEntity { endTimeMillis: Long filter: FilterInput limit: Int - ): [Operation!] - + ): [Operation!] @noAspects # has own custom resolver """ Assertions associated with the Dataset """ @@ -1813,92 +1821,91 @@ type Dataset implements EntityWithRelationships & Entity & BrowsableEntity { start: Int count: Int includeSoftDeleted: Boolean - ): EntityAssertionsResult - + ): EntityAssertionsResult @noAspects # has own custom resolver """ Edges extending from this entity """ relationships(input: RelationshipsInput!): EntityRelationshipsResult - + @noAspects # has own custom resolver """ Edges extending from this entity grouped by direction in the lineage graph """ - lineage(input: LineageInput!): EntityLineageResult - + lineage(input: LineageInput!): EntityLineageResult @noAspects # has own custom resolver """ The browse paths corresponding to the dataset. If no Browse Paths have been generated before, this will be null. """ - browsePaths: [BrowsePath!] - + browsePaths: [BrowsePath!] @noAspects # has own custom resolver """ The browse path V2 corresponding to an entity. If no Browse Paths V2 have been generated before, this will be null. """ - browsePathV2: BrowsePathV2 + browsePathV2: BrowsePathV2 @aspectMapping(aspects: ["browsePathsV2"]) """ Experimental! The resolved health statuses of the Dataset """ - health: [Health!] - + health: [Health!] @noAspects # has own custom resolver """ Schema metadata of the dataset """ - schema: Schema @deprecated(reason: "Use `schemaMetadata`") + schema: Schema + @deprecated(reason: "Use `schemaMetadata`") + @aspectMapping(aspects: ["schemaMetadata"]) """ Deprecated, use properties field instead External URL associated with the Dataset """ - externalUrl: String @deprecated + externalUrl: String @deprecated @aspectMapping(aspects: ["datasetProperties"]) """ Deprecated, see the properties field instead Environment in which the dataset belongs to or where it was generated Note that this field will soon be deprecated in favor of a more standardized concept of Environment """ - origin: FabricType! @deprecated + origin: FabricType! + @deprecated + @aspectMapping(aspects: ["origin", "datasetKey"]) """ Deprecated, use the properties field instead Read only technical description for dataset """ - description: String @deprecated + description: String @deprecated @aspectMapping(aspects: ["datasetProperties"]) """ Deprecated, do not use this field The logical type of the dataset ie table, stream, etc """ - platformNativeType: PlatformNativeType @deprecated + platformNativeType: PlatformNativeType @deprecated @noAspects """ Deprecated, use properties instead Native Dataset Uri Uri should not include any environment specific properties """ - uri: String @deprecated + uri: String @deprecated @aspectMapping(aspects: ["datasetProperties"]) """ Deprecated, use tags field instead The structured tags associated with the dataset """ - globalTags: GlobalTags @deprecated + globalTags: GlobalTags @deprecated @aspectMapping(aspects: ["globalTags"]) """ Sub Types that this entity implements """ - subTypes: SubTypes + subTypes: SubTypes @aspectMapping(aspects: ["subTypes"]) """ View related properties. Only relevant if subtypes field contains view. """ - viewProperties: ViewProperties + viewProperties: ViewProperties @aspectMapping(aspects: ["viewProperties"]) """ Experimental API. For fetching extra entities that do not have custom UI code yet """ - aspects(input: AspectParams): [RawAspect!] - + aspects(input: AspectParams): [RawAspect!] @noAspects # uses custom resolver """ History of datajob runs that either produced or consumed this dataset """ @@ -1906,44 +1913,42 @@ type Dataset implements EntityWithRelationships & Entity & BrowsableEntity { start: Int count: Int direction: RelationshipDirection! - ): DataProcessInstanceResult - + ): DataProcessInstanceResult @noAspects # uses custom resolver """ Metadata about the datasets siblings """ - siblings: SiblingProperties + siblings: SiblingProperties @aspectMapping(aspects: ["siblings"]) """ Executes a search on only the siblings of an entity """ - siblingsSearch(input: ScrollAcrossEntitiesInput!): ScrollResults - + siblingsSearch(input: ScrollAcrossEntitiesInput!): ScrollResults @noAspects # uses custom resolver """ Lineage information for the column-level. Includes a list of objects detailing which columns are upstream and which are downstream of each other. The upstream and downstream columns are from datasets. """ fineGrainedLineages: [FineGrainedLineage!] + @aspectMapping(aspects: ["upstreamLineage"]) """ Privileges given to a user relevant to this entity """ - privileges: EntityPrivileges - + privileges: EntityPrivileges @noAspects # uses custom resolver """ Whether or not this entity exists on DataHub """ - exists: Boolean - + exists: Boolean @noAspects # uses custom resolver """ Structured properties about this Dataset """ structuredProperties: StructuredProperties + @aspectMapping(aspects: ["structuredProperties"]) """ Statistics about how this Dataset has been operated on """ - operationsStats(input: OperationsStatsInput): OperationsQueryResult + operationsStats(input: OperationsStatsInput): OperationsQueryResult @noAspects # uses custom resolver } type RoleAssociation { diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/graphql/GraphQLEngineFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/graphql/GraphQLEngineFactory.java index 9f1d9225ade93..3b076eb5813ae 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/graphql/GraphQLEngineFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/graphql/GraphQLEngineFactory.java @@ -6,6 +6,7 @@ import com.datahub.authentication.token.StatefulTokenService; import com.datahub.authentication.user.NativeUserService; import com.datahub.authorization.role.RoleService; +import com.linkedin.datahub.graphql.AspectMappingRegistry; import com.linkedin.datahub.graphql.GmsGraphQLEngine; import com.linkedin.datahub.graphql.GmsGraphQLEngineArgs; import com.linkedin.datahub.graphql.GraphQLEngine; @@ -296,7 +297,25 @@ protected GraphQLEngine graphQLEngine( args.setMetricUtils(metricUtils); args.setS3Util(s3Util); - return new GmsGraphQLEngine(args).builder().build(); + // Create the GmsGraphQLEngine and build the GraphQL schema + GmsGraphQLEngine gmsGraphQLEngine = new GmsGraphQLEngine(args); + GraphQLEngine graphQLEngine = gmsGraphQLEngine.builder().build(); + + // Create the AspectMappingRegistry with the built schema + // This enables entity types to optimize aspect fetching based on GraphQL field selections + this.aspectMappingRegistry = + new AspectMappingRegistry(graphQLEngine.getGraphQL().getGraphQLSchema()); + + return graphQLEngine; + } + + // Store the AspectMappingRegistry to expose it as a bean + private AspectMappingRegistry aspectMappingRegistry; + + /** Provides the AspectMappingRegistry bean for use in resolvers and entity types. */ + @Bean(name = "aspectMappingRegistry") + protected AspectMappingRegistry aspectMappingRegistry() { + return this.aspectMappingRegistry; } @Bean(name = "graphQLWorkerPool") diff --git a/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java b/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java index e1e233f31aefd..31ef8b7ac310a 100644 --- a/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java +++ b/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/GraphQLController.java @@ -54,6 +54,8 @@ public class GraphQLController { @Inject MetricUtils metricUtils; + @Inject com.linkedin.datahub.graphql.AspectMappingRegistry aspectMappingRegistry; + @Nonnull @Inject @Named("systemOperationContext") @@ -130,6 +132,8 @@ CompletableFuture> postGraphQL( operationName, query, variables); + // Set the AspectMappingRegistry for aspect-level optimizations + context.setAspectMappingRegistry(aspectMappingRegistry); Span.current().setAttribute(ACTOR_URN_ATTR, context.getActorUrn()); final String threadName = Thread.currentThread().getName(); diff --git a/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/SpringQueryContext.java b/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/SpringQueryContext.java index cd8ce56bf36f9..d33312faa1c0d 100644 --- a/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/SpringQueryContext.java +++ b/metadata-service/graphql-servlet-impl/src/main/java/com/datahub/graphql/SpringQueryContext.java @@ -2,10 +2,12 @@ import com.datahub.authentication.Authentication; import com.datahub.plugins.auth.authorization.Authorizer; +import com.linkedin.datahub.graphql.AspectMappingRegistry; import com.linkedin.datahub.graphql.QueryContext; import com.linkedin.metadata.config.DataHubAppConfiguration; import graphql.language.OperationDefinition; import graphql.parser.Parser; +import graphql.schema.DataFetchingEnvironment; import io.datahubproject.metadata.context.OperationContext; import io.datahubproject.metadata.context.RequestContext; import jakarta.servlet.http.HttpServletRequest; @@ -24,6 +26,10 @@ public class SpringQueryContext implements QueryContext { @Nonnull private final OperationContext operationContext; @Nonnull private final DataHubAppConfiguration dataHubAppConfig; + // Mutable fields for request-scoped data + @Nullable private DataFetchingEnvironment dataFetchingEnvironment; + @Nullable private AspectMappingRegistry aspectMappingRegistry; + public SpringQueryContext( final boolean isAuthenticated, final Authentication authentication, @@ -63,4 +69,26 @@ public SpringQueryContext( this.dataHubAppConfig = dataHubAppConfig; } + + @Override + @Nullable + public DataFetchingEnvironment getDataFetchingEnvironment() { + return dataFetchingEnvironment; + } + + @Override + public void setDataFetchingEnvironment(@Nullable DataFetchingEnvironment environment) { + this.dataFetchingEnvironment = environment; + } + + @Override + @Nullable + public AspectMappingRegistry getAspectMappingRegistry() { + return aspectMappingRegistry; + } + + @Override + public void setAspectMappingRegistry(@Nullable AspectMappingRegistry aspectMappingRegistry) { + this.aspectMappingRegistry = aspectMappingRegistry; + } }