Skip to content

Commit 731c29e

Browse files
david-leifkerdavid-leifker
authored
feat(graph-retriever): implement graph retriever (datahub-project#10241)
1 parent 90c1249 commit 731c29e

File tree

674 files changed

+10975
-7927
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

674 files changed

+10975
-7927
lines changed

‎datahub-frontend/app/auth/AuthModule.java

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import com.google.inject.name.Named;
1616
import com.linkedin.entity.client.SystemEntityClient;
1717
import com.linkedin.entity.client.SystemRestliEntityClient;
18+
import com.linkedin.metadata.models.registry.EmptyEntityRegistry;
1819
import com.linkedin.metadata.restli.DefaultRestliClientFactory;
1920
import com.linkedin.parseq.retry.backoff.ExponentialBackoff;
2021
import com.linkedin.util.Configuration;
@@ -112,7 +113,7 @@ protected void configure() {
112113
.toConstructor(
113114
SsoCallbackController.class.getConstructor(
114115
SsoManager.class,
115-
Authentication.class,
116+
OperationContext.class,
116117
SystemEntityClient.class,
117118
AuthServiceClient.class,
118119
org.pac4j.core.config.Config.class,
@@ -164,8 +165,9 @@ protected Authentication provideSystemAuthentication() {
164165
@Provides
165166
@Singleton
166167
@Named("systemOperationContext")
167-
protected OperationContext provideOperationContext(final Authentication systemAuthentication,
168-
final ConfigurationProvider configurationProvider) {
168+
protected OperationContext provideOperationContext(
169+
final Authentication systemAuthentication,
170+
final ConfigurationProvider configurationProvider) {
169171
ActorContext systemActorContext =
170172
ActorContext.builder()
171173
.systemAuth(true)
@@ -180,7 +182,7 @@ protected OperationContext provideOperationContext(final Authentication systemAu
180182
.operationContextConfig(systemConfig)
181183
.systemActorContext(systemActorContext)
182184
.searchContext(SearchContext.EMPTY)
183-
.entityRegistryContext(EntityRegistryContext.EMPTY)
185+
.entityRegistryContext(EntityRegistryContext.builder().build(EmptyEntityRegistry.EMPTY))
184186
// Authorizer.EMPTY doesn't actually apply to system auth
185187
.authorizerContext(AuthorizerContext.builder().authorizer(Authorizer.EMPTY).build())
186188
.build(systemAuthentication);
@@ -200,7 +202,6 @@ protected SystemEntityClient provideEntityClient(
200202
@Named("systemOperationContext") final OperationContext systemOperationContext,
201203
final ConfigurationProvider configurationProvider) {
202204
return new SystemRestliEntityClient(
203-
systemOperationContext,
204205
buildRestliClient(),
205206
new ExponentialBackoff(_configs.getInt(ENTITY_CLIENT_RETRY_INTERVAL)),
206207
_configs.getInt(ENTITY_CLIENT_NUM_RETRIES),

‎datahub-frontend/app/auth/sso/oidc/OidcCallbackLogic.java

Lines changed: 41 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
import auth.CookieConfigs;
1010
import auth.sso.SsoManager;
1111
import client.AuthServiceClient;
12-
import com.datahub.authentication.Authentication;
1312
import com.fasterxml.jackson.core.type.TypeReference;
1413
import com.fasterxml.jackson.databind.ObjectMapper;
1514
import com.linkedin.common.AuditStamp;
@@ -57,6 +56,8 @@
5756
import java.util.regex.Matcher;
5857
import java.util.regex.Pattern;
5958
import java.util.stream.Collectors;
59+
60+
import io.datahubproject.metadata.context.OperationContext;
6061
import lombok.extern.slf4j.Slf4j;
6162
import org.pac4j.core.config.Config;
6263
import org.pac4j.core.context.Cookie;
@@ -69,6 +70,8 @@
6970
import org.pac4j.play.PlayWebContext;
7071
import play.mvc.Result;
7172

73+
import javax.annotation.Nonnull;
74+
7275
/**
7376
* This class contains the logic that is executed when an OpenID Connect Identity Provider redirects
7477
* back to D DataHub after an authentication attempt.
@@ -82,23 +85,23 @@
8285
@Slf4j
8386
public class OidcCallbackLogic extends DefaultCallbackLogic<Result, PlayWebContext> {
8487

85-
private final SsoManager _ssoManager;
86-
private final SystemEntityClient _entityClient;
87-
private final Authentication _systemAuthentication;
88-
private final AuthServiceClient _authClient;
89-
private final CookieConfigs _cookieConfigs;
88+
private final SsoManager ssoManager;
89+
private final SystemEntityClient systemEntityClient;
90+
private final OperationContext systemOperationContext;
91+
private final AuthServiceClient authClient;
92+
private final CookieConfigs cookieConfigs;
9093

9194
public OidcCallbackLogic(
9295
final SsoManager ssoManager,
93-
final Authentication systemAuthentication,
96+
final OperationContext systemOperationContext,
9497
final SystemEntityClient entityClient,
9598
final AuthServiceClient authClient,
9699
final CookieConfigs cookieConfigs) {
97-
_ssoManager = ssoManager;
98-
_systemAuthentication = systemAuthentication;
99-
_entityClient = entityClient;
100-
_authClient = authClient;
101-
_cookieConfigs = cookieConfigs;
100+
this.ssoManager = ssoManager;
101+
this.systemOperationContext = systemOperationContext;
102+
systemEntityClient = entityClient;
103+
this.authClient = authClient;
104+
this.cookieConfigs = cookieConfigs;
102105
}
103106

104107
@Override
@@ -131,8 +134,8 @@ public Result perform(
131134
}
132135

133136
// By this point, we know that OIDC is the enabled provider.
134-
final OidcConfigs oidcConfigs = (OidcConfigs) _ssoManager.getSsoProvider().configs();
135-
return handleOidcCallback(oidcConfigs, result, context, getProfileManager(context));
137+
final OidcConfigs oidcConfigs = (OidcConfigs) ssoManager.getSsoProvider().configs();
138+
return handleOidcCallback(systemOperationContext, oidcConfigs, result, getProfileManager(context));
136139
}
137140

138141
@SuppressWarnings("unchecked")
@@ -153,9 +156,9 @@ private void setContextRedirectUrl(PlayWebContext context) {
153156
}
154157

155158
private Result handleOidcCallback(
159+
final OperationContext opContext,
156160
final OidcConfigs oidcConfigs,
157161
final Result result,
158-
final PlayWebContext context,
159162
final ProfileManager<UserProfile> profileManager) {
160163

161164
log.debug("Beginning OIDC Callback Handling...");
@@ -177,23 +180,23 @@ private Result handleOidcCallback(
177180
if (oidcConfigs.isJitProvisioningEnabled()) {
178181
log.debug("Just-in-time provisioning is enabled. Beginning provisioning process...");
179182
CorpUserSnapshot extractedUser = extractUser(corpUserUrn, profile);
180-
tryProvisionUser(extractedUser);
183+
tryProvisionUser(opContext, extractedUser);
181184
if (oidcConfigs.isExtractGroupsEnabled()) {
182185
// Extract groups & provision them.
183186
List<CorpGroupSnapshot> extractedGroups = extractGroups(profile);
184-
tryProvisionGroups(extractedGroups);
187+
tryProvisionGroups(opContext, extractedGroups);
185188
// Add users to groups on DataHub. Note that this clears existing group membership for a
186189
// user if it already exists.
187-
updateGroupMembership(corpUserUrn, createGroupMembership(extractedGroups));
190+
updateGroupMembership(opContext, corpUserUrn, createGroupMembership(extractedGroups));
188191
}
189192
} else if (oidcConfigs.isPreProvisioningRequired()) {
190193
// We should only allow logins for user accounts that have been pre-provisioned
191194
log.debug("Pre Provisioning is required. Beginning validation of extracted user...");
192-
verifyPreProvisionedUser(corpUserUrn);
195+
verifyPreProvisionedUser(opContext, corpUserUrn);
193196
}
194197
// Update user status to active on login.
195198
// If we want to prevent certain users from logging in, here's where we'll want to do it.
196-
setUserStatus(
199+
setUserStatus(opContext,
197200
corpUserUrn,
198201
new CorpUserStatus()
199202
.setStatus(Constants.CORP_USER_STATUS_ACTIVE)
@@ -209,15 +212,15 @@ private Result handleOidcCallback(
209212
}
210213

211214
// Successfully logged in - Generate GMS login token
212-
final String accessToken = _authClient.generateSessionTokenForUser(corpUserUrn.getId());
215+
final String accessToken = authClient.generateSessionTokenForUser(corpUserUrn.getId());
213216
return result
214217
.withSession(createSessionMap(corpUserUrn.toString(), accessToken))
215218
.withCookies(
216219
createActorCookie(
217220
corpUserUrn.toString(),
218-
_cookieConfigs.getTtlInHours(),
219-
_cookieConfigs.getAuthCookieSameSite(),
220-
_cookieConfigs.getAuthCookieSecure()));
221+
cookieConfigs.getTtlInHours(),
222+
cookieConfigs.getAuthCookieSameSite(),
223+
cookieConfigs.getAuthCookieSecure()));
221224
}
222225
return internalServerError(
223226
"Failed to authenticate current user. Cannot find valid identity provider profile in session.");
@@ -331,7 +334,7 @@ private List<CorpGroupSnapshot> extractGroups(CommonProfile profile) {
331334
String.format(
332335
"Attempting to extract groups from OIDC profile %s",
333336
profile.getAttributes().toString()));
334-
final OidcConfigs configs = (OidcConfigs) _ssoManager.getSsoProvider().configs();
337+
final OidcConfigs configs = (OidcConfigs) ssoManager.getSsoProvider().configs();
335338

336339
// First, attempt to extract a list of groups from the profile, using the group name attribute
337340
// config.
@@ -400,13 +403,13 @@ private GroupMembership createGroupMembership(final List<CorpGroupSnapshot> extr
400403
return groupMembershipAspect;
401404
}
402405

403-
private void tryProvisionUser(CorpUserSnapshot corpUserSnapshot) {
406+
private void tryProvisionUser(@Nonnull OperationContext opContext, CorpUserSnapshot corpUserSnapshot) {
404407

405408
log.debug(String.format("Attempting to provision user with urn %s", corpUserSnapshot.getUrn()));
406409

407410
// 1. Check if this user already exists.
408411
try {
409-
final Entity corpUser = _entityClient.get(corpUserSnapshot.getUrn(), _systemAuthentication);
412+
final Entity corpUser = systemEntityClient.get(opContext, corpUserSnapshot.getUrn());
410413
final CorpUserSnapshot existingCorpUserSnapshot = corpUser.getValue().getCorpUserSnapshot();
411414

412415
log.debug(String.format("Fetched GMS user with urn %s", corpUserSnapshot.getUrn()));
@@ -420,7 +423,7 @@ private void tryProvisionUser(CorpUserSnapshot corpUserSnapshot) {
420423
// 2. The user does not exist. Provision them.
421424
final Entity newEntity = new Entity();
422425
newEntity.setValue(Snapshot.create(corpUserSnapshot));
423-
_entityClient.update(newEntity, _systemAuthentication);
426+
systemEntityClient.update(opContext, newEntity);
424427
log.debug(String.format("Successfully provisioned user %s", corpUserSnapshot.getUrn()));
425428
}
426429
log.debug(
@@ -434,7 +437,7 @@ private void tryProvisionUser(CorpUserSnapshot corpUserSnapshot) {
434437
}
435438
}
436439

437-
private void tryProvisionGroups(List<CorpGroupSnapshot> corpGroups) {
440+
private void tryProvisionGroups(@Nonnull OperationContext opContext, List<CorpGroupSnapshot> corpGroups) {
438441

439442
log.debug(
440443
String.format(
@@ -446,7 +449,7 @@ private void tryProvisionGroups(List<CorpGroupSnapshot> corpGroups) {
446449
final Set<Urn> urnsToFetch =
447450
corpGroups.stream().map(CorpGroupSnapshot::getUrn).collect(Collectors.toSet());
448451
final Map<Urn, Entity> existingGroups =
449-
_entityClient.batchGet(urnsToFetch, _systemAuthentication);
452+
systemEntityClient.batchGet(opContext, urnsToFetch);
450453

451454
log.debug(String.format("Fetched GMS groups with urns %s", existingGroups.keySet()));
452455

@@ -484,11 +487,10 @@ private void tryProvisionGroups(List<CorpGroupSnapshot> corpGroups) {
484487
log.debug(String.format("Provisioning groups with urns %s", groupsToCreateUrns));
485488

486489
// Now batch create all entities identified to create.
487-
_entityClient.batchUpdate(
490+
systemEntityClient.batchUpdate(opContext,
488491
groupsToCreate.stream()
489492
.map(groupSnapshot -> new Entity().setValue(Snapshot.create(groupSnapshot)))
490-
.collect(Collectors.toSet()),
491-
_systemAuthentication);
493+
.collect(Collectors.toSet()));
492494

493495
log.debug(String.format("Successfully provisioned groups with urns %s", groupsToCreateUrns));
494496
} catch (RemoteInvocationException e) {
@@ -501,7 +503,7 @@ private void tryProvisionGroups(List<CorpGroupSnapshot> corpGroups) {
501503
}
502504
}
503505

504-
private void updateGroupMembership(Urn urn, GroupMembership groupMembership) {
506+
private void updateGroupMembership(@Nonnull OperationContext opContext, Urn urn, GroupMembership groupMembership) {
505507
log.debug(String.format("Updating group membership for user %s", urn));
506508
final MetadataChangeProposal proposal = new MetadataChangeProposal();
507509
proposal.setEntityUrn(urn);
@@ -510,18 +512,18 @@ private void updateGroupMembership(Urn urn, GroupMembership groupMembership) {
510512
proposal.setAspect(GenericRecordUtils.serializeAspect(groupMembership));
511513
proposal.setChangeType(ChangeType.UPSERT);
512514
try {
513-
_entityClient.ingestProposal(proposal, _systemAuthentication);
515+
systemEntityClient.ingestProposal(opContext, proposal);
514516
} catch (RemoteInvocationException e) {
515517
throw new RuntimeException(
516518
String.format("Failed to update group membership for user with urn %s", urn), e);
517519
}
518520
}
519521

520-
private void verifyPreProvisionedUser(CorpuserUrn urn) {
522+
private void verifyPreProvisionedUser(@Nonnull OperationContext opContext, CorpuserUrn urn) {
521523
// Validate that the user exists in the system (there is more than just a key aspect for them,
522524
// as of today).
523525
try {
524-
final Entity corpUser = _entityClient.get(urn, _systemAuthentication);
526+
final Entity corpUser = systemEntityClient.get(opContext, urn);
525527

526528
log.debug(String.format("Fetched GMS user with urn %s", urn));
527529

@@ -543,15 +545,15 @@ private void verifyPreProvisionedUser(CorpuserUrn urn) {
543545
}
544546
}
545547

546-
private void setUserStatus(final Urn urn, final CorpUserStatus newStatus) throws Exception {
548+
private void setUserStatus(@Nonnull OperationContext opContext, final Urn urn, final CorpUserStatus newStatus) throws Exception {
547549
// Update status aspect to be active.
548550
final MetadataChangeProposal proposal = new MetadataChangeProposal();
549551
proposal.setEntityUrn(urn);
550552
proposal.setEntityType(Constants.CORP_USER_ENTITY_NAME);
551553
proposal.setAspectName(Constants.CORP_USER_STATUS_ASPECT_NAME);
552554
proposal.setAspect(GenericRecordUtils.serializeAspect(newStatus));
553555
proposal.setChangeType(ChangeType.UPSERT);
554-
_entityClient.ingestProposal(proposal, _systemAuthentication);
556+
systemEntityClient.ingestProposal(opContext, proposal);
555557
}
556558

557559
private Optional<String> extractRegexGroup(final String patternStr, final String target) {

‎datahub-frontend/app/controllers/SsoCallbackController.java

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@
1515
import java.util.concurrent.CompletionStage;
1616
import javax.annotation.Nonnull;
1717
import javax.inject.Inject;
18+
import javax.inject.Named;
19+
20+
import io.datahubproject.metadata.context.OperationContext;
1821
import lombok.extern.slf4j.Slf4j;
1922
import org.pac4j.core.client.Client;
2023
import org.pac4j.core.client.Clients;
@@ -43,7 +46,7 @@ public class SsoCallbackController extends CallbackController {
4346
@Inject
4447
public SsoCallbackController(
4548
@Nonnull SsoManager ssoManager,
46-
@Nonnull Authentication systemAuthentication,
49+
@Named("systemOperationContext") @Nonnull OperationContext systemOperationContext,
4750
@Nonnull SystemEntityClient entityClient,
4851
@Nonnull AuthServiceClient authClient,
4952
@Nonnull Config config,
@@ -55,7 +58,7 @@ public SsoCallbackController(
5558
setCallbackLogic(
5659
new SsoCallbackLogic(
5760
ssoManager,
58-
systemAuthentication,
61+
systemOperationContext,
5962
entityClient,
6063
authClient,
6164
new CookieConfigs(configs)));
@@ -96,13 +99,13 @@ public class SsoCallbackLogic implements CallbackLogic<Result, PlayWebContext> {
9699

97100
SsoCallbackLogic(
98101
final SsoManager ssoManager,
99-
final Authentication systemAuthentication,
102+
final OperationContext systemOperationContext,
100103
final SystemEntityClient entityClient,
101104
final AuthServiceClient authClient,
102105
final CookieConfigs cookieConfigs) {
103106
_oidcCallbackLogic =
104107
new OidcCallbackLogic(
105-
ssoManager, systemAuthentication, entityClient, authClient, cookieConfigs);
108+
ssoManager, systemOperationContext, entityClient, authClient, cookieConfigs);
106109
}
107110

108111
@Override

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,7 @@
347347
import com.linkedin.datahub.graphql.types.view.DataHubViewType;
348348
import com.linkedin.entity.client.EntityClient;
349349
import com.linkedin.entity.client.SystemEntityClient;
350+
import com.linkedin.metadata.client.UsageStatsJavaClient;
350351
import com.linkedin.metadata.config.DataHubConfiguration;
351352
import com.linkedin.metadata.config.IngestionConfiguration;
352353
import com.linkedin.metadata.config.TestsConfiguration;
@@ -372,7 +373,6 @@
372373
import com.linkedin.metadata.timeline.TimelineService;
373374
import com.linkedin.metadata.timeseries.TimeseriesAspectService;
374375
import com.linkedin.metadata.version.GitVersion;
375-
import com.linkedin.usage.UsageClient;
376376
import graphql.execution.DataFetcherResult;
377377
import graphql.schema.DataFetcher;
378378
import graphql.schema.DataFetchingEnvironment;
@@ -411,7 +411,7 @@ public class GmsGraphQLEngine {
411411
private final EntityClient entityClient;
412412
private final SystemEntityClient systemEntityClient;
413413
private final GraphClient graphClient;
414-
private final UsageClient usageClient;
414+
private final UsageStatsJavaClient usageClient;
415415
private final SiblingGraphService siblingGraphService;
416416

417417
private final EntityService entityService;

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngineArgs.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import com.linkedin.datahub.graphql.featureflags.FeatureFlags;
1313
import com.linkedin.entity.client.EntityClient;
1414
import com.linkedin.entity.client.SystemEntityClient;
15+
import com.linkedin.metadata.client.UsageStatsJavaClient;
1516
import com.linkedin.metadata.config.DataHubConfiguration;
1617
import com.linkedin.metadata.config.IngestionConfiguration;
1718
import com.linkedin.metadata.config.TestsConfiguration;
@@ -35,7 +36,6 @@
3536
import com.linkedin.metadata.timeline.TimelineService;
3637
import com.linkedin.metadata.timeseries.TimeseriesAspectService;
3738
import com.linkedin.metadata.version.GitVersion;
38-
import com.linkedin.usage.UsageClient;
3939
import io.datahubproject.metadata.services.RestrictedService;
4040
import io.datahubproject.metadata.services.SecretService;
4141
import lombok.Data;
@@ -45,7 +45,7 @@ public class GmsGraphQLEngineArgs {
4545
EntityClient entityClient;
4646
SystemEntityClient systemEntityClient;
4747
GraphClient graphClient;
48-
UsageClient usageClient;
48+
UsageStatsJavaClient usageClient;
4949
AnalyticsService analyticsService;
5050
EntityService entityService;
5151
RecommendationsService recommendationsService;

‎datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/SubTypesResolver.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,10 @@ public CompletableFuture<SubTypes> get(DataFetchingEnvironment environment) thro
3636
EntityResponse entityResponse =
3737
_entityClient
3838
.batchGetV2(
39+
context.getOperationContext(),
3940
urn.getEntityType(),
4041
Collections.singleton(urn),
41-
Collections.singleton(_aspectName),
42-
context.getAuthentication())
42+
Collections.singleton(_aspectName))
4343
.get(urn);
4444
if (entityResponse != null && entityResponse.getAspects().containsKey(_aspectName)) {
4545
subType =

0 commit comments

Comments
 (0)