From aef49b8fb2478f8a1b902aaee16fee9c07c7beab Mon Sep 17 00:00:00 2001 From: david-leifker <114954101+david-leifker@users.noreply.github.com> Date: Thu, 21 Sep 2023 22:00:14 -0500 Subject: [PATCH] feat(entity-client): enable client side cache for entity-client and usage-client (#8877) --- datahub-frontend/app/auth/AuthModule.java | 25 +++- .../app/auth/sso/oidc/OidcCallbackLogic.java | 7 +- .../app/config/ConfigurationProvider.java | 27 ++++ .../controllers/SsoCallbackController.java | 6 +- datahub-frontend/play.gradle | 8 +- .../datahub/graphql/GmsGraphQLEngine.java | 3 + .../datahub/graphql/GmsGraphQLEngineArgs.java | 2 + .../dataset/DatasetStatsSummaryResolver.java | 19 ++- .../dataset/DatasetUsageStatsResolver.java | 8 +- .../dashboard/DashboardStatsSummaryTest.java | 3 +- .../DatasetStatsSummaryResolverTest.java | 17 ++- .../common/steps/GMSDisableWriteModeStep.java | 8 +- .../common/steps/GMSEnableWriteModeStep.java | 9 +- .../upgrade/config/NoCodeUpgradeConfig.java | 10 +- .../upgrade/config/RestoreBackupConfig.java | 10 +- .../datahub/upgrade/nocode/NoCodeUpgrade.java | 12 +- .../upgrade/restorebackup/RestoreBackup.java | 15 +- .../client/SystemJavaEntityClient.java | 39 +++++ ...sInstanceRunEventChangeEventGenerator.java | 12 +- .../EntityChangeEventGenerator.java | 8 +- ...eConsumerApplicationTestConfiguration.java | 4 +- .../kafka/config/EntityHydratorConfig.java | 24 +-- .../event/EntityChangeEventGeneratorHook.java | 17 +-- .../hook/siblings/SiblingAssociationHook.java | 38 ++--- .../kafka/hydrator/EntityHydrator.java | 22 ++- .../EntityChangeEventGeneratorHookTest.java | 28 ++-- .../siblings/SiblingAssociationHookTest.java | 78 +++------- .../spring/MCLSpringTestConfiguration.java | 6 +- .../kafka/MceConsumerApplication.java | 2 +- .../kafka/MetadataChangeEventsProcessor.java | 8 +- .../MetadataChangeProposalsProcessor.java | 12 +- metadata-service/configuration/build.gradle | 1 + .../config/cache/CacheConfiguration.java | 2 + .../cache/client/ClientCacheConfig.java | 10 ++ .../client/ClientCacheConfiguration.java | 9 ++ .../cache/client/EntityClientCacheConfig.java | 17 +++ .../cache/client/UsageClientCacheConfig.java | 12 ++ .../spring/YamlPropertySourceFactory.java | 10 +- .../src/main/resources/application.yml | 24 +++ .../factory/auth/AuthorizerChainFactory.java | 2 +- .../auth/DataHubAuthorizerFactory.java | 2 +- .../auth/DataHubTokenServiceFactory.java | 3 +- .../gms/factory/auth/GroupServiceFactory.java | 2 +- .../auth/InviteTokenServiceFactory.java | 2 +- .../auth/NativeUserServiceFactory.java | 2 +- .../gms/factory/auth/PostServiceFactory.java | 2 +- .../gms/factory/auth/RoleServiceFactory.java | 2 +- .../auth/SystemAuthenticationFactory.java | 2 +- .../ElasticSearchGraphServiceFactory.java | 2 +- ...ticSearchSystemMetadataServiceFactory.java | 2 +- .../ElasticsearchSSLContextFactory.java | 2 +- .../factory/common/GraphServiceFactory.java | 2 +- .../common/IndexConventionFactory.java | 2 +- .../common/LocalEbeanServerConfigFactory.java | 2 +- .../factory/common/Neo4jDriverFactory.java | 2 +- .../common/RestHighLevelClientFactory.java | 2 +- .../factory/config/ConfigurationProvider.java | 2 +- .../DataProductServiceFactory.java | 2 +- .../entity/JavaEntityClientFactory.java | 29 +++- .../entity/RestliEntityClientFactory.java | 19 ++- .../entity/RetentionServiceFactory.java | 2 +- .../ConfigEntityRegistryFactory.java | 2 +- .../PluginEntityRegistryFactory.java | 2 +- .../factory/graphql/GraphQLEngineFactory.java | 6 + .../ingestion/IngestionSchedulerFactory.java | 2 +- .../DataHubKafkaEventProducerFactory.java | 2 +- .../kafka/DataHubKafkaProducerFactory.java | 2 +- .../AwsGlueSchemaRegistryFactory.java | 2 +- .../KafkaSchemaRegistryFactory.java | 2 +- .../lineage/LineageServiceFactory.java | 2 +- .../OwnershipTypeServiceFactory.java | 2 +- .../factory/query/QueryServiceFactory.java | 2 +- .../BaseElasticSearchComponentsFactory.java | 2 +- .../CachingEntitySearchServiceFactory.java | 2 +- .../ElasticSearchBulkProcessorFactory.java | 2 +- .../ElasticSearchIndexBuilderFactory.java | 2 +- .../search/ElasticSearchServiceFactory.java | 2 +- .../search/LineageSearchServiceFactory.java | 2 +- .../SearchDocumentTransformerFactory.java | 2 +- .../factory/search/SearchServiceFactory.java | 2 +- .../search/SettingsBuilderFactory.java | 2 +- .../search/views/ViewServiceFactory.java | 2 +- .../settings/SettingsServiceFactory.java | 2 +- .../factory/telemetry/MixpanelApiFactory.java | 2 +- .../MixpanelMessageBuilderFactory.java | 2 +- .../telemetry/TrackingServiceFactory.java | 2 +- ...tyChangeEventGeneratorRegistryFactory.java | 6 +- .../timeline/TimelineServiceFactory.java | 2 +- ...cSearchTimeseriesAspectServiceFactory.java | 2 +- .../gms/factory/usage/UsageClientFactory.java | 15 +- .../IngestRetentionPoliciesStepFactory.java | 2 +- .../openapi/util/OpenApiEntitiesUtilTest.java | 2 +- metadata-service/restli-client/build.gradle | 1 + .../linkedin/common/client/ClientCache.java | 134 +++++++++++++++++ .../entity/client/EntityClientCache.java | 141 ++++++++++++++++++ .../entity/client/SystemEntityClient.java | 91 +++++++++++ .../client/SystemRestliEntityClient.java | 25 ++++ .../java/com/linkedin/usage/UsageClient.java | 33 +++- .../com/linkedin/usage/UsageClientCache.java | 75 ++++++++++ .../metadata/utils/metrics/MetricUtils.java | 5 + 100 files changed, 951 insertions(+), 298 deletions(-) create mode 100644 datahub-frontend/app/config/ConfigurationProvider.java create mode 100644 metadata-io/src/main/java/com/linkedin/metadata/client/SystemJavaEntityClient.java create mode 100644 metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/ClientCacheConfig.java create mode 100644 metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/ClientCacheConfiguration.java create mode 100644 metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/EntityClientCacheConfig.java create mode 100644 metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/UsageClientCacheConfig.java rename metadata-service/{factories/src/main/java/com/linkedin/gms/factory => configuration/src/main/java/com/linkedin/metadata}/spring/YamlPropertySourceFactory.java (87%) create mode 100644 metadata-service/restli-client/src/main/java/com/linkedin/common/client/ClientCache.java create mode 100644 metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClientCache.java create mode 100644 metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemEntityClient.java create mode 100644 metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemRestliEntityClient.java create mode 100644 metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClientCache.java diff --git a/datahub-frontend/app/auth/AuthModule.java b/datahub-frontend/app/auth/AuthModule.java index eb95078b1a..98f3b82285 100644 --- a/datahub-frontend/app/auth/AuthModule.java +++ b/datahub-frontend/app/auth/AuthModule.java @@ -11,16 +11,19 @@ import com.datahub.authentication.Authentication; import com.google.inject.AbstractModule; import com.google.inject.Provides; import com.google.inject.Singleton; -import com.linkedin.entity.client.EntityClient; -import com.linkedin.entity.client.RestliEntityClient; +import com.linkedin.entity.client.SystemEntityClient; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.metadata.restli.DefaultRestliClientFactory; import com.linkedin.parseq.retry.backoff.ExponentialBackoff; import com.linkedin.util.Configuration; +import config.ConfigurationProvider; import controllers.SsoCallbackController; + import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Collections; import java.util.List; + import org.apache.commons.codec.digest.DigestUtils; import org.apache.http.impl.client.CloseableHttpClient; import org.apache.http.impl.client.HttpClients; @@ -34,6 +37,7 @@ import org.pac4j.play.store.PlayCacheSessionStore; import org.pac4j.play.store.PlayCookieSessionStore; import org.pac4j.play.store.PlaySessionStore; import org.pac4j.play.store.ShiroAesDataEncrypter; +import org.springframework.context.annotation.AnnotationConfigApplicationContext; import play.Environment; import play.cache.SyncCacheApi; import utils.ConfigUtil; @@ -104,7 +108,7 @@ public class AuthModule extends AbstractModule { bind(SsoCallbackController.class).toConstructor(SsoCallbackController.class.getConstructor( SsoManager.class, Authentication.class, - EntityClient.class, + SystemEntityClient.class, AuthServiceClient.class, com.typesafe.config.Config.class)); } catch (NoSuchMethodException | SecurityException e) { @@ -161,10 +165,19 @@ public class AuthModule extends AbstractModule { @Provides @Singleton - protected EntityClient provideEntityClient() { - return new RestliEntityClient(buildRestliClient(), + protected ConfigurationProvider provideConfigurationProvider() { + AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(ConfigurationProvider.class); + return context.getBean(ConfigurationProvider.class); + } + + @Provides + @Singleton + protected SystemEntityClient provideEntityClient(final Authentication systemAuthentication, + final ConfigurationProvider configurationProvider) { + return new SystemRestliEntityClient(buildRestliClient(), new ExponentialBackoff(_configs.getInt(ENTITY_CLIENT_RETRY_INTERVAL)), - _configs.getInt(ENTITY_CLIENT_NUM_RETRIES)); + _configs.getInt(ENTITY_CLIENT_NUM_RETRIES), systemAuthentication, + configurationProvider.getCache().getClient().getEntityClient()); } @Provides diff --git a/datahub-frontend/app/auth/sso/oidc/OidcCallbackLogic.java b/datahub-frontend/app/auth/sso/oidc/OidcCallbackLogic.java index 85139d1db0..4bde0872fc 100644 --- a/datahub-frontend/app/auth/sso/oidc/OidcCallbackLogic.java +++ b/datahub-frontend/app/auth/sso/oidc/OidcCallbackLogic.java @@ -13,7 +13,7 @@ import com.linkedin.common.urn.CorpuserUrn; import com.linkedin.common.urn.Urn; import com.linkedin.data.template.SetMode; import com.linkedin.entity.Entity; -import com.linkedin.entity.client.EntityClient; +import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.events.metadata.ChangeType; import com.linkedin.identity.CorpGroupInfo; import com.linkedin.identity.CorpUserEditableInfo; @@ -78,13 +78,14 @@ import static play.mvc.Results.internalServerError; public class OidcCallbackLogic extends DefaultCallbackLogic { private final SsoManager _ssoManager; - private final EntityClient _entityClient; + private final SystemEntityClient _entityClient; private final Authentication _systemAuthentication; private final AuthServiceClient _authClient; private final CookieConfigs _cookieConfigs; public OidcCallbackLogic(final SsoManager ssoManager, final Authentication systemAuthentication, - final EntityClient entityClient, final AuthServiceClient authClient, final CookieConfigs cookieConfigs) { + final SystemEntityClient entityClient, final AuthServiceClient authClient, + final CookieConfigs cookieConfigs) { _ssoManager = ssoManager; _systemAuthentication = systemAuthentication; _entityClient = entityClient; diff --git a/datahub-frontend/app/config/ConfigurationProvider.java b/datahub-frontend/app/config/ConfigurationProvider.java new file mode 100644 index 0000000000..00a5472ec3 --- /dev/null +++ b/datahub-frontend/app/config/ConfigurationProvider.java @@ -0,0 +1,27 @@ +package config; + +import com.linkedin.metadata.config.cache.CacheConfiguration; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; +import lombok.Data; + +import org.springframework.boot.context.properties.ConfigurationProperties; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.PropertySource; + + +/** + * Minimal sharing between metadata-service and frontend + * Initially for use of client caching configuration. + * Does not use the factories module to avoid transitive dependencies. + */ +@EnableConfigurationProperties +@PropertySource(value = "application.yml", factory = YamlPropertySourceFactory.class) +@ConfigurationProperties +@Data +public class ConfigurationProvider { + + /** + * Configuration for caching + */ + private CacheConfiguration cache; +} diff --git a/datahub-frontend/app/controllers/SsoCallbackController.java b/datahub-frontend/app/controllers/SsoCallbackController.java index 5a36d833de..7a4b5585cc 100644 --- a/datahub-frontend/app/controllers/SsoCallbackController.java +++ b/datahub-frontend/app/controllers/SsoCallbackController.java @@ -3,7 +3,7 @@ package controllers; import auth.CookieConfigs; import client.AuthServiceClient; import com.datahub.authentication.Authentication; -import com.linkedin.entity.client.EntityClient; +import com.linkedin.entity.client.SystemEntityClient; import java.net.URLEncoder; import java.nio.charset.StandardCharsets; import java.util.concurrent.CompletableFuture; @@ -40,7 +40,7 @@ public class SsoCallbackController extends CallbackController { public SsoCallbackController( @Nonnull SsoManager ssoManager, @Nonnull Authentication systemAuthentication, - @Nonnull EntityClient entityClient, + @Nonnull SystemEntityClient entityClient, @Nonnull AuthServiceClient authClient, @Nonnull com.typesafe.config.Config configs) { _ssoManager = ssoManager; @@ -79,7 +79,7 @@ public class SsoCallbackController extends CallbackController { private final OidcCallbackLogic _oidcCallbackLogic; SsoCallbackLogic(final SsoManager ssoManager, final Authentication systemAuthentication, - final EntityClient entityClient, final AuthServiceClient authClient, final CookieConfigs cookieConfigs) { + final SystemEntityClient entityClient, final AuthServiceClient authClient, final CookieConfigs cookieConfigs) { _oidcCallbackLogic = new OidcCallbackLogic(ssoManager, systemAuthentication, entityClient, authClient, cookieConfigs); } diff --git a/datahub-frontend/play.gradle b/datahub-frontend/play.gradle index e40f8e3eeb..daecba16cb 100644 --- a/datahub-frontend/play.gradle +++ b/datahub-frontend/play.gradle @@ -16,9 +16,6 @@ dependencies { implementation project(':datahub-web-react') constraints { - play(externalDependency.springCore) - play(externalDependency.springBeans) - play(externalDependency.springContext) play(externalDependency.jacksonDataBind) play('com.nimbusds:oauth2-oidc-sdk:8.36.2') play('com.nimbusds:nimbus-jose-jwt:8.18') @@ -35,7 +32,12 @@ dependencies { implementation project(":metadata-service:restli-client") implementation project(":metadata-service:auth-config") + implementation project(":metadata-service:configuration") + implementation externalDependency.springCore + implementation externalDependency.springBeans + implementation externalDependency.springContext + implementation externalDependency.springBootAutoconfigure implementation externalDependency.jettyJaas implementation externalDependency.graphqlJava implementation externalDependency.antlr4Runtime diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java index d86234cf59..3ba0cc1f74 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java @@ -302,6 +302,7 @@ import com.linkedin.datahub.graphql.types.tag.TagType; import com.linkedin.datahub.graphql.types.test.TestType; import com.linkedin.datahub.graphql.types.view.DataHubViewType; import com.linkedin.entity.client.EntityClient; +import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.metadata.config.DataHubConfiguration; import com.linkedin.metadata.config.IngestionConfiguration; import com.linkedin.metadata.config.TestsConfiguration; @@ -364,6 +365,7 @@ import static graphql.scalars.ExtendedScalars.*; public class GmsGraphQLEngine { private final EntityClient entityClient; + private final SystemEntityClient systemEntityClient; private final GraphClient graphClient; private final UsageClient usageClient; private final SiblingGraphService siblingGraphService; @@ -476,6 +478,7 @@ public class GmsGraphQLEngine { this.graphQLPlugins.forEach(plugin -> plugin.init(args)); this.entityClient = args.entityClient; + this.systemEntityClient = args.systemEntityClient; this.graphClient = args.graphClient; this.usageClient = args.usageClient; this.siblingGraphService = args.siblingGraphService; diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngineArgs.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngineArgs.java index cbcf42c4f9..157fb10ce7 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngineArgs.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngineArgs.java @@ -11,6 +11,7 @@ import com.datahub.authorization.role.RoleService; import com.linkedin.datahub.graphql.analytics.service.AnalyticsService; import com.linkedin.datahub.graphql.featureflags.FeatureFlags; import com.linkedin.entity.client.EntityClient; +import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.metadata.config.DataHubConfiguration; import com.linkedin.metadata.config.IngestionConfiguration; import com.linkedin.metadata.config.TestsConfiguration; @@ -38,6 +39,7 @@ import lombok.Data; @Data public class GmsGraphQLEngineArgs { EntityClient entityClient; + SystemEntityClient systemEntityClient; GraphClient graphClient; UsageClient usageClient; AnalyticsService analyticsService; diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetStatsSummaryResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetStatsSummaryResolver.java index f27fd604a7..23be49c7e7 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetStatsSummaryResolver.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetStatsSummaryResolver.java @@ -1,13 +1,16 @@ package com.linkedin.datahub.graphql.resolvers.dataset; +import com.datahub.authorization.ResourceSpec; import com.google.common.cache.Cache; import com.google.common.cache.CacheBuilder; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; import com.linkedin.datahub.graphql.QueryContext; +import com.linkedin.datahub.graphql.authorization.AuthorizationUtils; import com.linkedin.datahub.graphql.generated.CorpUser; import com.linkedin.datahub.graphql.generated.DatasetStatsSummary; import com.linkedin.datahub.graphql.generated.Entity; +import com.linkedin.metadata.authorization.PoliciesConfig; import com.linkedin.usage.UsageClient; import com.linkedin.usage.UsageTimeRange; import com.linkedin.usage.UserUsageCounts; @@ -15,6 +18,7 @@ import graphql.schema.DataFetcher; import graphql.schema.DataFetchingEnvironment; import java.util.List; import java.util.Objects; +import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; @@ -55,8 +59,15 @@ public class DatasetStatsSummaryResolver implements DataFetcher executable() { return (context) -> { try { - _entityClient.setWritable(false, _systemAuthentication); + _entityClient.setWritable(false); } catch (Exception e) { e.printStackTrace(); context.report().addLine("Failed to turn write mode off in GMS"); diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/GMSEnableWriteModeStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/GMSEnableWriteModeStep.java index 270eff8df2..8df0212398 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/GMSEnableWriteModeStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/GMSEnableWriteModeStep.java @@ -1,20 +1,17 @@ package com.linkedin.datahub.upgrade.common.steps; -import com.datahub.authentication.Authentication; import com.linkedin.datahub.upgrade.UpgradeContext; import com.linkedin.datahub.upgrade.UpgradeStep; import com.linkedin.datahub.upgrade.UpgradeStepResult; import com.linkedin.datahub.upgrade.impl.DefaultUpgradeStepResult; -import com.linkedin.entity.client.RestliEntityClient; +import com.linkedin.entity.client.SystemRestliEntityClient; import java.util.function.Function; import lombok.RequiredArgsConstructor; @RequiredArgsConstructor public class GMSEnableWriteModeStep implements UpgradeStep { - - private final Authentication _systemAuthentication; - private final RestliEntityClient _entityClient; + private final SystemRestliEntityClient _entityClient; @Override public String id() { @@ -30,7 +27,7 @@ public class GMSEnableWriteModeStep implements UpgradeStep { public Function executable() { return (context) -> { try { - _entityClient.setWritable(true, _systemAuthentication); + _entityClient.setWritable(true); } catch (Exception e) { e.printStackTrace(); context.report().addLine("Failed to turn write mode back on in GMS"); diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/NoCodeUpgradeConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/NoCodeUpgradeConfig.java index 30175c6fa7..cd264e529e 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/NoCodeUpgradeConfig.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/NoCodeUpgradeConfig.java @@ -1,8 +1,7 @@ package com.linkedin.datahub.upgrade.config; -import com.datahub.authentication.Authentication; import com.linkedin.datahub.upgrade.nocode.NoCodeUpgrade; -import com.linkedin.entity.client.RestliEntityClient; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.models.registry.EntityRegistry; import io.ebean.Database; @@ -21,15 +20,14 @@ public class NoCodeUpgradeConfig { ApplicationContext applicationContext; @Bean(name = "noCodeUpgrade") - @DependsOn({"ebeanServer", "entityService", "systemAuthentication", "restliEntityClient", "entityRegistry"}) + @DependsOn({"ebeanServer", "entityService", "systemRestliEntityClient", "entityRegistry"}) @Nonnull public NoCodeUpgrade createInstance() { final Database ebeanServer = applicationContext.getBean(Database.class); final EntityService entityService = applicationContext.getBean(EntityService.class); - final Authentication systemAuthentication = applicationContext.getBean(Authentication.class); - final RestliEntityClient entityClient = applicationContext.getBean(RestliEntityClient.class); + final SystemRestliEntityClient entityClient = applicationContext.getBean(SystemRestliEntityClient.class); final EntityRegistry entityRegistry = applicationContext.getBean(EntityRegistry.class); - return new NoCodeUpgrade(ebeanServer, entityService, entityRegistry, systemAuthentication, entityClient); + return new NoCodeUpgrade(ebeanServer, entityService, entityRegistry, entityClient); } } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreBackupConfig.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreBackupConfig.java index 9b0fcf279a..97a0880053 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreBackupConfig.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/RestoreBackupConfig.java @@ -1,8 +1,7 @@ package com.linkedin.datahub.upgrade.config; -import com.datahub.authentication.Authentication; import com.linkedin.datahub.upgrade.restorebackup.RestoreBackup; -import com.linkedin.entity.client.RestliEntityClient; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.models.registry.EntityRegistry; @@ -22,19 +21,18 @@ public class RestoreBackupConfig { ApplicationContext applicationContext; @Bean(name = "restoreBackup") - @DependsOn({"ebeanServer", "entityService", "systemAuthentication", "restliEntityClient", "graphService", + @DependsOn({"ebeanServer", "entityService", "systemRestliEntityClient", "graphService", "searchService", "entityRegistry"}) @Nonnull public RestoreBackup createInstance() { final Database ebeanServer = applicationContext.getBean(Database.class); final EntityService entityService = applicationContext.getBean(EntityService.class); - final Authentication systemAuthentication = applicationContext.getBean(Authentication.class); - final RestliEntityClient entityClient = applicationContext.getBean(RestliEntityClient.class); + final SystemRestliEntityClient entityClient = applicationContext.getBean(SystemRestliEntityClient.class); final GraphService graphClient = applicationContext.getBean(GraphService.class); final EntitySearchService searchClient = applicationContext.getBean(EntitySearchService.class); final EntityRegistry entityRegistry = applicationContext.getBean(EntityRegistry.class); - return new RestoreBackup(ebeanServer, entityService, entityRegistry, systemAuthentication, entityClient, + return new RestoreBackup(ebeanServer, entityService, entityRegistry, entityClient, graphClient, searchClient); } } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/NoCodeUpgrade.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/NoCodeUpgrade.java index ee4a3bc504..a299deb874 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/NoCodeUpgrade.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/NoCodeUpgrade.java @@ -1,13 +1,12 @@ package com.linkedin.datahub.upgrade.nocode; -import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableMap; import com.linkedin.datahub.upgrade.Upgrade; import com.linkedin.datahub.upgrade.UpgradeCleanupStep; import com.linkedin.datahub.upgrade.UpgradeStep; import com.linkedin.datahub.upgrade.common.steps.GMSEnableWriteModeStep; import com.linkedin.datahub.upgrade.common.steps.GMSQualificationStep; -import com.linkedin.entity.client.RestliEntityClient; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.models.registry.EntityRegistry; import io.ebean.Database; @@ -30,12 +29,10 @@ public class NoCodeUpgrade implements Upgrade { final Database server, final EntityService entityService, final EntityRegistry entityRegistry, - final Authentication systemAuthentication, - final RestliEntityClient entityClient) { + final SystemRestliEntityClient entityClient) { _steps = buildUpgradeSteps( server, entityService, entityRegistry, - systemAuthentication, entityClient); _cleanupSteps = buildCleanupSteps(); } @@ -63,15 +60,14 @@ public class NoCodeUpgrade implements Upgrade { final Database server, final EntityService entityService, final EntityRegistry entityRegistry, - final Authentication systemAuthentication, - final RestliEntityClient entityClient) { + final SystemRestliEntityClient entityClient) { final List steps = new ArrayList<>(); steps.add(new RemoveAspectV2TableStep(server)); steps.add(new GMSQualificationStep(ImmutableMap.of("noCode", "true"))); steps.add(new UpgradeQualificationStep(server)); steps.add(new CreateAspectTableStep(server)); steps.add(new DataMigrationStep(server, entityService, entityRegistry)); - steps.add(new GMSEnableWriteModeStep(systemAuthentication, entityClient)); + steps.add(new GMSEnableWriteModeStep(entityClient)); return steps; } } diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreBackup.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreBackup.java index 67718a6739..9175ad606e 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreBackup.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreBackup.java @@ -1,6 +1,5 @@ package com.linkedin.datahub.upgrade.restorebackup; -import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableList; import com.linkedin.datahub.upgrade.Upgrade; import com.linkedin.datahub.upgrade.UpgradeCleanupStep; @@ -9,7 +8,7 @@ import com.linkedin.datahub.upgrade.common.steps.ClearGraphServiceStep; import com.linkedin.datahub.upgrade.common.steps.ClearSearchServiceStep; import com.linkedin.datahub.upgrade.common.steps.GMSDisableWriteModeStep; import com.linkedin.datahub.upgrade.common.steps.GMSEnableWriteModeStep; -import com.linkedin.entity.client.RestliEntityClient; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.models.registry.EntityRegistry; @@ -27,11 +26,10 @@ public class RestoreBackup implements Upgrade { final Database server, final EntityService entityService, final EntityRegistry entityRegistry, - final Authentication systemAuthentication, - final RestliEntityClient entityClient, + final SystemRestliEntityClient entityClient, final GraphService graphClient, final EntitySearchService searchClient) { - _steps = buildSteps(server, entityService, entityRegistry, systemAuthentication, entityClient, graphClient, searchClient); + _steps = buildSteps(server, entityService, entityRegistry, entityClient, graphClient, searchClient); } @Override @@ -48,17 +46,16 @@ public class RestoreBackup implements Upgrade { final Database server, final EntityService entityService, final EntityRegistry entityRegistry, - final Authentication systemAuthentication, - final RestliEntityClient entityClient, + final SystemRestliEntityClient entityClient, final GraphService graphClient, final EntitySearchService searchClient) { final List steps = new ArrayList<>(); - steps.add(new GMSDisableWriteModeStep(systemAuthentication, entityClient)); + steps.add(new GMSDisableWriteModeStep(entityClient)); steps.add(new ClearSearchServiceStep(searchClient, true)); steps.add(new ClearGraphServiceStep(graphClient, true)); steps.add(new ClearAspectV2TableStep(server)); steps.add(new RestoreStorageStep(entityService, entityRegistry)); - steps.add(new GMSEnableWriteModeStep(systemAuthentication, entityClient)); + steps.add(new GMSEnableWriteModeStep(entityClient)); return steps; } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/client/SystemJavaEntityClient.java b/metadata-io/src/main/java/com/linkedin/metadata/client/SystemJavaEntityClient.java new file mode 100644 index 0000000000..6b5a3d5bfb --- /dev/null +++ b/metadata-io/src/main/java/com/linkedin/metadata/client/SystemJavaEntityClient.java @@ -0,0 +1,39 @@ +package com.linkedin.metadata.client; + +import com.datahub.authentication.Authentication; +import com.linkedin.entity.client.EntityClientCache; +import com.linkedin.metadata.config.cache.client.EntityClientCacheConfig; +import com.linkedin.entity.client.RestliEntityClient; +import com.linkedin.entity.client.SystemEntityClient; +import com.linkedin.metadata.entity.DeleteEntityService; +import com.linkedin.metadata.entity.EntityService; +import com.linkedin.metadata.event.EventProducer; +import com.linkedin.metadata.search.EntitySearchService; +import com.linkedin.metadata.search.LineageSearchService; +import com.linkedin.metadata.search.SearchService; +import com.linkedin.metadata.search.client.CachingEntitySearchService; +import com.linkedin.metadata.timeseries.TimeseriesAspectService; +import lombok.Getter; + + +/** + * Java backed SystemEntityClient + */ +@Getter +public class SystemJavaEntityClient extends JavaEntityClient implements SystemEntityClient { + + private final EntityClientCache entityClientCache; + private final Authentication systemAuthentication; + + public SystemJavaEntityClient(EntityService entityService, DeleteEntityService deleteEntityService, + EntitySearchService entitySearchService, CachingEntitySearchService cachingEntitySearchService, + SearchService searchService, LineageSearchService lineageSearchService, + TimeseriesAspectService timeseriesAspectService, EventProducer eventProducer, + RestliEntityClient restliEntityClient, Authentication systemAuthentication, + EntityClientCacheConfig cacheConfig) { + super(entityService, deleteEntityService, entitySearchService, cachingEntitySearchService, searchService, + lineageSearchService, timeseriesAspectService, eventProducer, restliEntityClient); + this.systemAuthentication = systemAuthentication; + this.entityClientCache = buildEntityClientCache(SystemJavaEntityClient.class, systemAuthentication, cacheConfig); + } +} diff --git a/metadata-io/src/main/java/com/linkedin/metadata/timeline/eventgenerator/DataProcessInstanceRunEventChangeEventGenerator.java b/metadata-io/src/main/java/com/linkedin/metadata/timeline/eventgenerator/DataProcessInstanceRunEventChangeEventGenerator.java index fee9cd9bca..a3e5a051a4 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/timeline/eventgenerator/DataProcessInstanceRunEventChangeEventGenerator.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/timeline/eventgenerator/DataProcessInstanceRunEventChangeEventGenerator.java @@ -1,6 +1,5 @@ package com.linkedin.metadata.timeline.eventgenerator; -import com.datahub.authentication.Authentication; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; import com.linkedin.dataprocess.DataProcessInstanceRelationships; @@ -8,7 +7,7 @@ import com.linkedin.dataprocess.DataProcessInstanceRunEvent; import com.linkedin.dataprocess.DataProcessRunStatus; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.EnvelopedAspectMap; -import com.linkedin.entity.client.EntityClient; +import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.metadata.timeline.data.ChangeCategory; import com.linkedin.metadata.timeline.data.ChangeEvent; import com.linkedin.metadata.timeline.data.ChangeOperation; @@ -27,9 +26,8 @@ public class DataProcessInstanceRunEventChangeEventGenerator private static final String COMPLETED_STATUS = "COMPLETED"; private static final String STARTED_STATUS = "STARTED"; - public DataProcessInstanceRunEventChangeEventGenerator(@Nonnull final EntityClient entityClient, @Nonnull final - Authentication authentication) { - super(entityClient, authentication); + public DataProcessInstanceRunEventChangeEventGenerator(@Nonnull final SystemEntityClient entityClient) { + super(entityClient); } @Override @@ -108,8 +106,8 @@ public class DataProcessInstanceRunEventChangeEventGenerator EntityResponse entityResponse; try { entityUrn = Urn.createFromString(entityUrnString); - entityResponse = _entityClient.getV2(DATA_PROCESS_INSTANCE_ENTITY_NAME, entityUrn, - Collections.singleton(DATA_PROCESS_INSTANCE_RELATIONSHIPS_ASPECT_NAME), _authentication); + entityResponse = _entityClient.getV2(entityUrn, + Collections.singleton(DATA_PROCESS_INSTANCE_RELATIONSHIPS_ASPECT_NAME)); } catch (Exception e) { return null; } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/timeline/eventgenerator/EntityChangeEventGenerator.java b/metadata-io/src/main/java/com/linkedin/metadata/timeline/eventgenerator/EntityChangeEventGenerator.java index 7f6aa5e532..d5539ec3d3 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/timeline/eventgenerator/EntityChangeEventGenerator.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/timeline/eventgenerator/EntityChangeEventGenerator.java @@ -5,7 +5,7 @@ import com.github.fge.jsonpatch.JsonPatch; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; import com.linkedin.data.template.RecordTemplate; -import com.linkedin.entity.client.EntityClient; +import com.linkedin.entity.client.SystemEntityClient; import com.linkedin.metadata.entity.EntityAspect; import com.linkedin.metadata.timeline.data.ChangeCategory; import com.linkedin.metadata.timeline.data.ChangeEvent; @@ -19,16 +19,14 @@ import javax.annotation.Nonnull; */ public abstract class EntityChangeEventGenerator { // TODO: Add a check for supported aspects - protected EntityClient _entityClient; + protected SystemEntityClient _entityClient; protected Authentication _authentication; public EntityChangeEventGenerator() { } - public EntityChangeEventGenerator(@Nonnull final EntityClient entityClient, - @Nonnull final Authentication authentication) { + public EntityChangeEventGenerator(@Nonnull final SystemEntityClient entityClient) { _entityClient = entityClient; - _authentication = authentication; } @Deprecated diff --git a/metadata-jobs/mae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTestConfiguration.java b/metadata-jobs/mae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTestConfiguration.java index 3b44ede0f1..a214117f4e 100644 --- a/metadata-jobs/mae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTestConfiguration.java +++ b/metadata-jobs/mae-consumer-job/src/test/java/com/linkedin/metadata/kafka/MaeConsumerApplicationTestConfiguration.java @@ -1,6 +1,6 @@ package com.linkedin.metadata.kafka; -import com.linkedin.entity.client.RestliEntityClient; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.gms.factory.auth.SystemAuthenticationFactory; import com.linkedin.metadata.dao.producer.KafkaHealthChecker; import com.linkedin.metadata.entity.EntityServiceImpl; @@ -24,7 +24,7 @@ public class MaeConsumerApplicationTestConfiguration { private EntityServiceImpl _entityServiceImpl; @MockBean - private RestliEntityClient restliEntityClient; + private SystemRestliEntityClient restliEntityClient; @MockBean private Database ebeanServer; diff --git a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/config/EntityHydratorConfig.java b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/config/EntityHydratorConfig.java index 2d8c52566e..a9e54e5354 100644 --- a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/config/EntityHydratorConfig.java +++ b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/config/EntityHydratorConfig.java @@ -1,10 +1,10 @@ package com.linkedin.metadata.kafka.config; -import com.datahub.authentication.Authentication; -import com.linkedin.entity.client.RestliEntityClient; -import com.linkedin.gms.factory.auth.SystemAuthenticationFactory; +import com.google.common.collect.ImmutableSet; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.gms.factory.entity.RestliEntityClientFactory; import com.linkedin.metadata.kafka.hydrator.EntityHydrator; +import com.linkedin.metadata.models.registry.EntityRegistry; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.context.annotation.Bean; @@ -13,19 +13,25 @@ import org.springframework.context.annotation.Import; @Configuration -@Import({RestliEntityClientFactory.class, SystemAuthenticationFactory.class}) +@Import({RestliEntityClientFactory.class}) public class EntityHydratorConfig { @Autowired - @Qualifier("systemAuthentication") - private Authentication _systemAuthentication; + @Qualifier("systemRestliEntityClient") + private SystemRestliEntityClient _entityClient; @Autowired - @Qualifier("restliEntityClient") - private RestliEntityClient _entityClient; + private EntityRegistry _entityRegistry; + + public final static ImmutableSet EXCLUDED_ASPECTS = ImmutableSet.builder() + .add("datasetUpstreamLineage", "upstreamLineage") + .add("dataJobInputOutput") + .add("dataProcessInstanceRelationships", "dataProcessInstanceInput", "dataProcessInstanceOutput") + .add("inputFields") + .build(); @Bean public EntityHydrator getEntityHydrator() { - return new EntityHydrator(_systemAuthentication, _entityClient); + return new EntityHydrator(_entityRegistry, _entityClient); } } diff --git a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHook.java b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHook.java index 55077c46a1..3b65ecccad 100644 --- a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHook.java +++ b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHook.java @@ -1,15 +1,12 @@ package com.linkedin.metadata.kafka.hook.event; -import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableSet; import com.linkedin.common.AuditStamp; import com.linkedin.common.urn.Urn; import com.linkedin.data.DataMap; import com.linkedin.data.template.RecordTemplate; import com.linkedin.data.template.SetMode; -import com.linkedin.entity.client.EntityClient; -import com.linkedin.entity.client.RestliEntityClient; -import com.linkedin.gms.factory.auth.SystemAuthenticationFactory; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.gms.factory.entity.RestliEntityClientFactory; import com.linkedin.gms.factory.entityregistry.EntityRegistryFactory; import com.linkedin.metadata.Constants; @@ -46,8 +43,7 @@ import org.springframework.stereotype.Component; */ @Slf4j @Component -@Import({EntityChangeEventGeneratorRegistry.class, EntityRegistryFactory.class, RestliEntityClientFactory.class, - SystemAuthenticationFactory.class}) +@Import({EntityChangeEventGeneratorRegistry.class, EntityRegistryFactory.class, RestliEntityClientFactory.class}) public class EntityChangeEventGeneratorHook implements MetadataChangeLogHook { /** @@ -83,20 +79,18 @@ public class EntityChangeEventGeneratorHook implements MetadataChangeLogHook { */ private static final Set SUPPORTED_OPERATIONS = ImmutableSet.of("CREATE", "UPSERT", "DELETE"); private final EntityChangeEventGeneratorRegistry _entityChangeEventGeneratorRegistry; - private final EntityClient _entityClient; - private final Authentication _systemAuthentication; + private final SystemRestliEntityClient _entityClient; private final EntityRegistry _entityRegistry; private final Boolean _isEnabled; @Autowired public EntityChangeEventGeneratorHook( @Nonnull final EntityChangeEventGeneratorRegistry entityChangeEventGeneratorRegistry, - @Nonnull final RestliEntityClient entityClient, @Nonnull final Authentication systemAuthentication, + @Nonnull final SystemRestliEntityClient entityClient, @Nonnull final EntityRegistry entityRegistry, @Nonnull @Value("${entityChangeEvents.enabled:true}") Boolean isEnabled) { _entityChangeEventGeneratorRegistry = Objects.requireNonNull(entityChangeEventGeneratorRegistry); _entityClient = Objects.requireNonNull(entityClient); - _systemAuthentication = Objects.requireNonNull(systemAuthentication); _entityRegistry = Objects.requireNonNull(entityRegistry); _isEnabled = isEnabled; } @@ -189,8 +183,7 @@ public class EntityChangeEventGeneratorHook implements MetadataChangeLogHook { _entityClient.producePlatformEvent( Constants.CHANGE_EVENT_PLATFORM_EVENT_NAME, partitioningKey, - event, - _systemAuthentication + event ); } diff --git a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java index 06545ef352..7cbe53dee9 100644 --- a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java +++ b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHook.java @@ -1,6 +1,5 @@ package com.linkedin.metadata.kafka.hook.siblings; -import com.datahub.authentication.Authentication; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; @@ -13,9 +12,8 @@ import com.linkedin.common.urn.Urn; import com.linkedin.dataset.UpstreamArray; import com.linkedin.dataset.UpstreamLineage; import com.linkedin.entity.EntityResponse; -import com.linkedin.entity.client.RestliEntityClient; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.events.metadata.ChangeType; -import com.linkedin.gms.factory.auth.SystemAuthenticationFactory; import com.linkedin.gms.factory.entity.RestliEntityClientFactory; import com.linkedin.gms.factory.entityregistry.EntityRegistryFactory; import com.linkedin.gms.factory.search.EntitySearchServiceFactory; @@ -60,7 +58,7 @@ import static com.linkedin.metadata.Constants.*; @Slf4j @Component @Singleton -@Import({EntityRegistryFactory.class, RestliEntityClientFactory.class, EntitySearchServiceFactory.class, SystemAuthenticationFactory.class}) +@Import({EntityRegistryFactory.class, RestliEntityClientFactory.class, EntitySearchServiceFactory.class}) public class SiblingAssociationHook implements MetadataChangeLogHook { public static final String SIBLING_ASSOCIATION_SYSTEM_ACTOR = "urn:li:corpuser:__datahub_system_sibling_hook"; @@ -73,23 +71,20 @@ public class SiblingAssociationHook implements MetadataChangeLogHook { public static final String SOURCE_SUBTYPE_V2 = "Source"; private final EntityRegistry _entityRegistry; - private final RestliEntityClient _entityClient; + private final SystemRestliEntityClient _entityClient; private final EntitySearchService _searchService; - private final Authentication _systemAuthentication; private final boolean _isEnabled; @Autowired public SiblingAssociationHook( @Nonnull final EntityRegistry entityRegistry, - @Nonnull final RestliEntityClient entityClient, + @Nonnull final SystemRestliEntityClient entityClient, @Nonnull final EntitySearchService searchService, - @Nonnull final Authentication systemAuthentication, @Nonnull @Value("${siblings.enabled:true}") Boolean isEnabled ) { _entityRegistry = entityRegistry; _entityClient = entityClient; _searchService = searchService; - _systemAuthentication = systemAuthentication; _isEnabled = isEnabled; } @@ -251,9 +246,9 @@ public class SiblingAssociationHook implements MetadataChangeLogHook { dbtSiblingProposal.setEntityUrn(dbtUrn); try { - _entityClient.ingestProposal(dbtSiblingProposal, _systemAuthentication); + _entityClient.ingestProposal(dbtSiblingProposal, true); } catch (RemoteInvocationException e) { - log.error("Error while associating {} with {}: {}", dbtUrn.toString(), sourceUrn.toString(), e.toString()); + log.error("Error while associating {} with {}: {}", dbtUrn, sourceUrn, e.toString()); throw new RuntimeException("Error ingesting sibling proposal. Skipping processing.", e); } @@ -274,9 +269,9 @@ public class SiblingAssociationHook implements MetadataChangeLogHook { List filteredNewSiblingsArray = newSiblingsUrnArray.stream().filter(urn -> { try { - return _entityClient.exists(urn, _systemAuthentication); + return _entityClient.exists(urn); } catch (RemoteInvocationException e) { - log.error("Error while checking existence of {}: {}", urn.toString(), e.toString()); + log.error("Error while checking existence of {}: {}", urn, e.toString()); throw new RuntimeException("Error checking existence. Skipping processing.", e); } }).collect(Collectors.toList()); @@ -294,9 +289,9 @@ public class SiblingAssociationHook implements MetadataChangeLogHook { sourceSiblingProposal.setEntityUrn(sourceUrn); try { - _entityClient.ingestProposal(sourceSiblingProposal, _systemAuthentication); + _entityClient.ingestProposal(sourceSiblingProposal, true); } catch (RemoteInvocationException e) { - log.error("Error while associating {} with {}: {}", dbtUrn.toString(), sourceUrn.toString(), e.toString()); + log.error("Error while associating {} with {}: {}", dbtUrn, sourceUrn, e.toString()); throw new RuntimeException("Error ingesting sibling proposal. Skipping processing.", e); } } @@ -406,11 +401,8 @@ public class SiblingAssociationHook implements MetadataChangeLogHook { ) { try { EntityResponse entityResponse = _entityClient.getV2( - DATASET_ENTITY_NAME, urn, - ImmutableSet.of(SUB_TYPES_ASPECT_NAME), - _systemAuthentication - ); + ImmutableSet.of(SUB_TYPES_ASPECT_NAME)); if (entityResponse != null && entityResponse.hasAspects() && entityResponse.getAspects().containsKey(Constants.SUB_TYPES_ASPECT_NAME)) { return new SubTypes(entityResponse.getAspects().get(Constants.SUB_TYPES_ASPECT_NAME).getValue().data()); @@ -427,10 +419,8 @@ public class SiblingAssociationHook implements MetadataChangeLogHook { ) { try { EntityResponse entityResponse = _entityClient.getV2( - DATASET_ENTITY_NAME, urn, - ImmutableSet.of(UPSTREAM_LINEAGE_ASPECT_NAME), - _systemAuthentication + ImmutableSet.of(UPSTREAM_LINEAGE_ASPECT_NAME) ); if (entityResponse != null && entityResponse.hasAspects() && entityResponse.getAspects().containsKey(Constants.UPSTREAM_LINEAGE_ASPECT_NAME)) { @@ -448,10 +438,8 @@ public class SiblingAssociationHook implements MetadataChangeLogHook { ) { try { EntityResponse entityResponse = _entityClient.getV2( - DATASET_ENTITY_NAME, urn, - ImmutableSet.of(SIBLINGS_ASPECT_NAME), - _systemAuthentication + ImmutableSet.of(SIBLINGS_ASPECT_NAME) ); if (entityResponse != null && entityResponse.hasAspects() && entityResponse.getAspects().containsKey(Constants.SIBLINGS_ASPECT_NAME)) { diff --git a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hydrator/EntityHydrator.java b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hydrator/EntityHydrator.java index d768ada176..0a3b38517e 100644 --- a/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hydrator/EntityHydrator.java +++ b/metadata-jobs/mae-consumer/src/main/java/com/linkedin/metadata/kafka/hydrator/EntityHydrator.java @@ -1,28 +1,32 @@ package com.linkedin.metadata.kafka.hydrator; -import com.datahub.authentication.Authentication; import com.fasterxml.jackson.databind.node.JsonNodeFactory; import com.fasterxml.jackson.databind.node.ObjectNode; import com.linkedin.common.urn.Urn; import com.linkedin.entity.EntityResponse; -import com.linkedin.entity.client.EntityClient; +import com.linkedin.entity.client.SystemRestliEntityClient; +import com.linkedin.metadata.models.AspectSpec; +import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.r2.RemoteInvocationException; import java.net.URISyntaxException; import java.util.Collections; import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import static com.linkedin.metadata.Constants.*; +import static com.linkedin.metadata.kafka.config.EntityHydratorConfig.EXCLUDED_ASPECTS; @Slf4j @RequiredArgsConstructor public class EntityHydrator { - private final Authentication _systemAuthentication; - private final EntityClient _entityClient; - + private final EntityRegistry _entityRegistry; + private final SystemRestliEntityClient _entityClient; private final ChartHydrator _chartHydrator = new ChartHydrator(); private final CorpUserHydrator _corpUserHydrator = new CorpUserHydrator(); private final DashboardHydrator _dashboardHydrator = new DashboardHydrator(); @@ -43,8 +47,12 @@ public class EntityHydrator { // Hydrate fields from snapshot EntityResponse entityResponse; try { - entityResponse = _entityClient.batchGetV2(entityTypeName, Collections.singleton(urnObj), null, - this._systemAuthentication).get(urnObj); + Set aspectNames = Optional.ofNullable(_entityRegistry.getEntitySpecs().get(urnObj.getEntityType())) + .map(spec -> spec.getAspectSpecs().stream().map(AspectSpec::getName) + .filter(aspectName -> !EXCLUDED_ASPECTS.contains(aspectName)) + .collect(Collectors.toSet())) + .orElse(Set.of()); + entityResponse = _entityClient.batchGetV2(Collections.singleton(urnObj), aspectNames).get(urnObj); } catch (RemoteInvocationException | URISyntaxException e) { log.error("Error while calling GMS to hydrate entity for urn {}", urn); return Optional.empty(); diff --git a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHookTest.java b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHookTest.java index d8759da0fe..7d9619f3e2 100644 --- a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHookTest.java +++ b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/event/EntityChangeEventGeneratorHookTest.java @@ -1,6 +1,5 @@ package com.linkedin.metadata.kafka.hook.event; -import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.linkedin.assertion.AssertionResult; @@ -38,8 +37,7 @@ import com.linkedin.entity.Aspect; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.EnvelopedAspect; import com.linkedin.entity.EnvelopedAspectMap; -import com.linkedin.entity.client.EntityClient; -import com.linkedin.entity.client.RestliEntityClient; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.key.DatasetKey; @@ -66,6 +64,7 @@ import com.linkedin.platform.event.v1.EntityChangeEvent; import com.linkedin.platform.event.v1.Parameters; import java.net.URISyntaxException; import java.util.Map; + import org.mockito.Mockito; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -92,22 +91,19 @@ public class EntityChangeEventGeneratorHookTest { private static final String TEST_DATA_FLOW_URN = "urn:li:dataFlow:flow"; private static final String TEST_DATA_JOB_URN = "urn:li:dataJob:job"; private Urn actorUrn; - private Authentication _mockAuthentication; - private RestliEntityClient _mockClient; + private SystemRestliEntityClient _mockClient; private EntityService _mockEntityService; private EntityChangeEventGeneratorHook _entityChangeEventHook; @BeforeMethod public void setupTest() throws URISyntaxException { actorUrn = Urn.createFromString(TEST_ACTOR_URN); - _mockAuthentication = Mockito.mock(Authentication.class); - _mockClient = Mockito.mock(RestliEntityClient.class); + _mockClient = Mockito.mock(SystemRestliEntityClient.class); _mockEntityService = Mockito.mock(EntityService.class); EntityChangeEventGeneratorRegistry entityChangeEventGeneratorRegistry = createEntityChangeEventGeneratorRegistry(); _entityChangeEventHook = - new EntityChangeEventGeneratorHook(entityChangeEventGeneratorRegistry, _mockClient, _mockAuthentication, - createMockEntityRegistry(), true); + new EntityChangeEventGeneratorHook(entityChangeEventGeneratorRegistry, _mockClient, createMockEntityRegistry(), true); } @Test @@ -498,8 +494,7 @@ public class EntityChangeEventGeneratorHookTest { final EntityResponse entityResponse = buildEntityResponse(ImmutableMap.of(DATA_PROCESS_INSTANCE_RELATIONSHIPS_ASPECT_NAME, relationships)); - Mockito.when(_mockClient.getV2(eq(DATA_PROCESS_INSTANCE_ENTITY_NAME), eq(dataProcessInstanceUrn), - any(), eq(_mockAuthentication))).thenReturn(entityResponse); + Mockito.when(_mockClient.getV2(eq(dataProcessInstanceUrn), any())).thenReturn(entityResponse); _entityChangeEventHook.invoke(event); @@ -540,8 +535,7 @@ public class EntityChangeEventGeneratorHookTest { final EntityResponse entityResponse = buildEntityResponse(ImmutableMap.of(DATA_PROCESS_INSTANCE_RELATIONSHIPS_ASPECT_NAME, relationships)); - Mockito.when(_mockClient.getV2(eq(DATA_PROCESS_INSTANCE_ENTITY_NAME), eq(dataProcessInstanceUrn), - any(), eq(_mockAuthentication))).thenReturn(entityResponse); + Mockito.when(_mockClient.getV2(eq(dataProcessInstanceUrn), any())).thenReturn(entityResponse); _entityChangeEventHook.invoke(event); @@ -618,7 +612,7 @@ public class EntityChangeEventGeneratorHookTest { // Run change event generators registry.register(ASSERTION_RUN_EVENT_ASPECT_NAME, new AssertionRunEventChangeEventGenerator()); registry.register(DATA_PROCESS_INSTANCE_RUN_EVENT_ASPECT_NAME, - new DataProcessInstanceRunEventChangeEventGenerator(_mockClient, _mockAuthentication)); + new DataProcessInstanceRunEventChangeEventGenerator(_mockClient)); return registry; } @@ -668,14 +662,14 @@ public class EntityChangeEventGeneratorHookTest { return registry; } - private void verifyProducePlatformEvent(EntityClient mockClient, PlatformEvent platformEvent) throws Exception { + private void verifyProducePlatformEvent(SystemRestliEntityClient mockClient, PlatformEvent platformEvent) throws Exception { verifyProducePlatformEvent(mockClient, platformEvent, true); } - private void verifyProducePlatformEvent(EntityClient mockClient, PlatformEvent platformEvent, boolean noMoreInteractions) throws Exception { + private void verifyProducePlatformEvent(SystemRestliEntityClient mockClient, PlatformEvent platformEvent, boolean noMoreInteractions) throws Exception { // Verify event has been emitted. verify(mockClient, Mockito.times(1)).producePlatformEvent(eq(CHANGE_EVENT_PLATFORM_EVENT_NAME), Mockito.anyString(), - argThat(new PlatformEventMatcher(platformEvent)), Mockito.any(Authentication.class)); + argThat(new PlatformEventMatcher(platformEvent))); if (noMoreInteractions) { Mockito.verifyNoMoreInteractions(_mockClient); diff --git a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java index 78d304d67b..6a2a05aa4b 100644 --- a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java +++ b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/siblings/SiblingAssociationHookTest.java @@ -1,6 +1,5 @@ package com.linkedin.metadata.kafka.hook.siblings; -import com.datahub.authentication.Authentication; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.linkedin.common.FabricType; @@ -19,7 +18,7 @@ import com.linkedin.entity.Aspect; import com.linkedin.entity.EntityResponse; import com.linkedin.entity.EnvelopedAspect; import com.linkedin.entity.EnvelopedAspectMap; -import com.linkedin.entity.client.RestliEntityClient; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.events.metadata.ChangeType; import com.linkedin.metadata.key.DatasetKey; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; @@ -44,19 +43,16 @@ import static org.mockito.ArgumentMatchers.*; public class SiblingAssociationHookTest { private SiblingAssociationHook _siblingAssociationHook; - RestliEntityClient _mockEntityClient; + SystemRestliEntityClient _mockEntityClient; EntitySearchService _mockSearchService; - Authentication _mockAuthentication; @BeforeMethod public void setupTest() { EntityRegistry registry = new ConfigEntityRegistry( SiblingAssociationHookTest.class.getClassLoader().getResourceAsStream("test-entity-registry-siblings.yml")); - _mockEntityClient = Mockito.mock(RestliEntityClient.class); + _mockEntityClient = Mockito.mock(SystemRestliEntityClient.class); _mockSearchService = Mockito.mock(EntitySearchService.class); - _mockAuthentication = Mockito.mock(Authentication.class); - _siblingAssociationHook = new SiblingAssociationHook(registry, _mockEntityClient, _mockSearchService, _mockAuthentication, - true); + _siblingAssociationHook = new SiblingAssociationHook(registry, _mockEntityClient, _mockSearchService, true); _siblingAssociationHook.setEnabled(true); } @@ -69,15 +65,13 @@ public class SiblingAssociationHookTest { EntityResponse mockResponse = new EntityResponse(); mockResponse.setAspects(mockResponseMap); - Mockito.when(_mockEntityClient.exists(Mockito.any(), Mockito.any())).thenReturn(true); + Mockito.when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); Mockito.when( _mockEntityClient.getV2( - DATASET_ENTITY_NAME, Urn.createFromString("urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.jaffle_shop.customers,PROD)"), - ImmutableSet.of(SUB_TYPES_ASPECT_NAME), - _mockAuthentication + ImmutableSet.of(SUB_TYPES_ASPECT_NAME) )).thenReturn(mockResponse); @@ -105,10 +99,7 @@ public class SiblingAssociationHookTest { proposal.setAspect(GenericRecordUtils.serializeAspect(dbtSiblingsAspect)); proposal.setChangeType(ChangeType.UPSERT); - Mockito.verify(_mockEntityClient, Mockito.times(1)).ingestProposal( - Mockito.eq(proposal), - Mockito.eq(_mockAuthentication) - ); + Mockito.verify(_mockEntityClient, Mockito.times(1)).ingestProposal(Mockito.eq(proposal), eq(true)); final Siblings sourceSiblingsAspect = new Siblings() .setSiblings(new UrnArray(ImmutableList.of(Urn.createFromString("urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.jaffle_shop.customers,PROD)")))) @@ -121,10 +112,7 @@ public class SiblingAssociationHookTest { proposal2.setAspect(GenericRecordUtils.serializeAspect(sourceSiblingsAspect)); proposal2.setChangeType(ChangeType.UPSERT); - Mockito.verify(_mockEntityClient, Mockito.times(1)).ingestProposal( - Mockito.eq(proposal2), - Mockito.eq(_mockAuthentication) - ); + Mockito.verify(_mockEntityClient, Mockito.times(1)).ingestProposal(Mockito.eq(proposal2), eq(true)); } @Test @@ -132,23 +120,20 @@ public class SiblingAssociationHookTest { SubTypes mockSourceSubtypesAspect = new SubTypes(); mockSourceSubtypesAspect.setTypeNames(new StringArray(ImmutableList.of("model"))); - Mockito.when(_mockEntityClient.exists(Mockito.any(), Mockito.any())).thenReturn(true); + Mockito.when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); EnvelopedAspectMap mockResponseMap = new EnvelopedAspectMap(); mockResponseMap.put(SUB_TYPES_ASPECT_NAME, new EnvelopedAspect().setValue(new Aspect(mockSourceSubtypesAspect.data()))); EntityResponse mockResponse = new EntityResponse(); mockResponse.setAspects(mockResponseMap); - Mockito.when(_mockEntityClient.exists(Mockito.any(), Mockito.any())).thenReturn(true); + Mockito.when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); Mockito.when( _mockEntityClient.getV2( - DATASET_ENTITY_NAME, Urn.createFromString("urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.jaffle_shop.customers,PROD)"), - ImmutableSet.of(SUB_TYPES_ASPECT_NAME), - _mockAuthentication - )).thenReturn(mockResponse); + ImmutableSet.of(SUB_TYPES_ASPECT_NAME))).thenReturn(mockResponse); MetadataChangeLog event = createEvent(DATASET_ENTITY_NAME, UPSTREAM_LINEAGE_ASPECT_NAME, ChangeType.UPSERT); Upstream upstream = createUpstream("urn:li:dataset:(urn:li:dataPlatform:bigquery,my-proj.jaffle_shop.customers,PROD)", DatasetLineageType.TRANSFORMED); @@ -174,15 +159,12 @@ public class SiblingAssociationHookTest { proposal.setAspect(GenericRecordUtils.serializeAspect(dbtSiblingsAspect)); proposal.setChangeType(ChangeType.UPSERT); - Mockito.verify(_mockEntityClient, Mockito.times(0)).ingestProposal( - Mockito.eq(proposal), - Mockito.eq(_mockAuthentication) - ); + Mockito.verify(_mockEntityClient, Mockito.times(0)).ingestProposal(Mockito.eq(proposal), eq(true)); } @Test public void testInvokeWhenThereIsAPairWithBigqueryDownstreamNode() throws Exception { - Mockito.when(_mockEntityClient.exists(Mockito.any(), Mockito.any())).thenReturn(true); + Mockito.when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); MetadataChangeLog event = createEvent(DATASET_ENTITY_NAME, UPSTREAM_LINEAGE_ASPECT_NAME, ChangeType.UPSERT); @@ -208,10 +190,7 @@ public class SiblingAssociationHookTest { proposal.setAspect(GenericRecordUtils.serializeAspect(dbtSiblingsAspect)); proposal.setChangeType(ChangeType.UPSERT); - Mockito.verify(_mockEntityClient, Mockito.times(1)).ingestProposal( - Mockito.eq(proposal), - Mockito.eq(_mockAuthentication) - ); + Mockito.verify(_mockEntityClient, Mockito.times(1)).ingestProposal(Mockito.eq(proposal), eq(true)); final Siblings sourceSiblingsAspect = new Siblings() .setSiblings(new UrnArray(ImmutableList.of(Urn.createFromString("urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.jaffle_shop.customers,PROD)")))) @@ -224,15 +203,12 @@ public class SiblingAssociationHookTest { proposal2.setAspect(GenericRecordUtils.serializeAspect(sourceSiblingsAspect)); proposal2.setChangeType(ChangeType.UPSERT); - Mockito.verify(_mockEntityClient, Mockito.times(1)).ingestProposal( - Mockito.eq(proposal2), - Mockito.eq(_mockAuthentication) - ); + Mockito.verify(_mockEntityClient, Mockito.times(1)).ingestProposal(Mockito.eq(proposal2), eq(true)); } @Test public void testInvokeWhenThereIsAKeyBeingReingested() throws Exception { - Mockito.when(_mockEntityClient.exists(Mockito.any(), Mockito.any())).thenReturn(true); + Mockito.when(_mockEntityClient.exists(Mockito.any())).thenReturn(true); SearchResult returnSearchResult = new SearchResult(); SearchEntityArray returnEntityArray = new SearchEntityArray(); @@ -271,10 +247,7 @@ public class SiblingAssociationHookTest { proposal.setAspect(GenericRecordUtils.serializeAspect(dbtSiblingsAspect)); proposal.setChangeType(ChangeType.UPSERT); - Mockito.verify(_mockEntityClient, Mockito.times(1)).ingestProposal( - Mockito.eq(proposal), - Mockito.eq(_mockAuthentication) - ); + Mockito.verify(_mockEntityClient, Mockito.times(1)).ingestProposal(Mockito.eq(proposal), eq(true)); final Siblings sourceSiblingsAspect = new Siblings() .setSiblings(new UrnArray(ImmutableList.of(Urn.createFromString("urn:li:dataset:(urn:li:dataPlatform:dbt,my-proj.jaffle_shop.customers,PROD)")))) @@ -287,10 +260,7 @@ public class SiblingAssociationHookTest { proposal2.setAspect(GenericRecordUtils.serializeAspect(sourceSiblingsAspect)); proposal2.setChangeType(ChangeType.UPSERT); - Mockito.verify(_mockEntityClient, Mockito.times(1)).ingestProposal( - Mockito.eq(proposal2), - Mockito.eq(_mockAuthentication) - ); + Mockito.verify(_mockEntityClient, Mockito.times(1)).ingestProposal(Mockito.eq(proposal2), eq(true)); } @Test public void testInvokeWhenSourceUrnHasTwoDbtUpstreams() throws Exception { @@ -309,10 +279,7 @@ public class SiblingAssociationHookTest { _siblingAssociationHook.invoke(event); - Mockito.verify(_mockEntityClient, Mockito.times(0)).ingestProposal( - Mockito.any(), - Mockito.eq(_mockAuthentication) - ); + Mockito.verify(_mockEntityClient, Mockito.times(0)).ingestProposal(Mockito.any(), eq(true)); } @@ -335,12 +302,7 @@ public class SiblingAssociationHookTest { _siblingAssociationHook.invoke(event); - Mockito.verify(_mockEntityClient, Mockito.times(2)).ingestProposal( - Mockito.any(), - Mockito.eq(_mockAuthentication) - ); - - + Mockito.verify(_mockEntityClient, Mockito.times(2)).ingestProposal(Mockito.any(), eq(true)); } private MetadataChangeLog createEvent(String entityType, String aspectName, ChangeType changeType) { diff --git a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/spring/MCLSpringTestConfiguration.java b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/spring/MCLSpringTestConfiguration.java index ef80c49ec4..dc5a6cd232 100644 --- a/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/spring/MCLSpringTestConfiguration.java +++ b/metadata-jobs/mae-consumer/src/test/java/com/linkedin/metadata/kafka/hook/spring/MCLSpringTestConfiguration.java @@ -2,7 +2,7 @@ package com.linkedin.metadata.kafka.hook.spring; import com.datahub.authentication.Authentication; import com.datahub.metadata.ingestion.IngestionScheduler; -import com.linkedin.entity.client.RestliEntityClient; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.gms.factory.kafka.schemaregistry.SchemaRegistryConfig; import com.linkedin.metadata.boot.kafka.DataHubUpgradeKafkaListener; import com.linkedin.metadata.graph.elastic.ElasticSearchGraphService; @@ -44,8 +44,8 @@ public class MCLSpringTestConfiguration { @MockBean public IngestionScheduler ingestionScheduler; - @MockBean - public RestliEntityClient entityClient; + @MockBean(name = "systemRestliEntityClient") + public SystemRestliEntityClient entityClient; @MockBean public ElasticSearchService searchService; diff --git a/metadata-jobs/mce-consumer-job/src/main/java/com/linkedin/metadata/kafka/MceConsumerApplication.java b/metadata-jobs/mce-consumer-job/src/main/java/com/linkedin/metadata/kafka/MceConsumerApplication.java index 9b4fe15c11..f0c59240a9 100644 --- a/metadata-jobs/mce-consumer-job/src/main/java/com/linkedin/metadata/kafka/MceConsumerApplication.java +++ b/metadata-jobs/mce-consumer-job/src/main/java/com/linkedin/metadata/kafka/MceConsumerApplication.java @@ -1,8 +1,8 @@ package com.linkedin.metadata.kafka; import com.linkedin.gms.factory.entity.RestliEntityClientFactory; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; import com.linkedin.gms.factory.telemetry.ScheduledAnalyticsFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import org.springframework.boot.SpringApplication; import org.springframework.boot.actuate.autoconfigure.solr.SolrHealthContributorAutoConfiguration; import org.springframework.boot.autoconfigure.SpringBootApplication; diff --git a/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeEventsProcessor.java b/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeEventsProcessor.java index 74679d30b2..c30dd6e6f9 100644 --- a/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeEventsProcessor.java +++ b/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeEventsProcessor.java @@ -4,8 +4,7 @@ import com.codahale.metrics.Histogram; import com.codahale.metrics.MetricRegistry; import com.datahub.authentication.Authentication; import com.linkedin.entity.Entity; -import com.linkedin.entity.client.RestliEntityClient; -import com.linkedin.gms.factory.auth.SystemAuthenticationFactory; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.gms.factory.entity.RestliEntityClientFactory; import com.linkedin.gms.factory.kafka.KafkaEventConsumerFactory; import com.linkedin.gms.factory.kafka.DataHubKafkaProducerFactory; @@ -40,15 +39,14 @@ import org.springframework.stereotype.Component; @Slf4j @Component @Conditional(MetadataChangeProposalProcessorCondition.class) -@Import({RestliEntityClientFactory.class, SystemAuthenticationFactory.class, KafkaEventConsumerFactory.class, - DataHubKafkaProducerFactory.class}) +@Import({RestliEntityClientFactory.class, KafkaEventConsumerFactory.class, DataHubKafkaProducerFactory.class}) @EnableKafka @RequiredArgsConstructor public class MetadataChangeEventsProcessor { @NonNull private final Authentication systemAuthentication; - private final RestliEntityClient entityClient; + private final SystemRestliEntityClient entityClient; private final Producer kafkaProducer; private final Histogram kafkaLagStats = MetricUtils.get().histogram(MetricRegistry.name(this.getClass(), "kafkaLag")); diff --git a/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeProposalsProcessor.java b/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeProposalsProcessor.java index 289d70ef8c..79f8c90af8 100644 --- a/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeProposalsProcessor.java +++ b/metadata-jobs/mce-consumer/src/main/java/com/linkedin/metadata/kafka/MetadataChangeProposalsProcessor.java @@ -2,9 +2,7 @@ package com.linkedin.metadata.kafka; import com.codahale.metrics.Histogram; import com.codahale.metrics.MetricRegistry; -import com.datahub.authentication.Authentication; -import com.linkedin.entity.client.RestliEntityClient; -import com.linkedin.gms.factory.auth.SystemAuthenticationFactory; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.gms.factory.entity.RestliEntityClientFactory; import com.linkedin.gms.factory.kafka.KafkaEventConsumerFactory; import com.linkedin.gms.factory.kafka.DataHubKafkaProducerFactory; @@ -35,15 +33,13 @@ import org.springframework.stereotype.Component; @Slf4j @Component -@Import({RestliEntityClientFactory.class, SystemAuthenticationFactory.class, KafkaEventConsumerFactory.class, - DataHubKafkaProducerFactory.class}) +@Import({RestliEntityClientFactory.class, KafkaEventConsumerFactory.class, DataHubKafkaProducerFactory.class}) @Conditional(MetadataChangeProposalProcessorCondition.class) @EnableKafka @RequiredArgsConstructor public class MetadataChangeProposalsProcessor { - private final Authentication systemAuthentication; - private final RestliEntityClient entityClient; + private final SystemRestliEntityClient entityClient; private final Producer kafkaProducer; private final Histogram kafkaLagStats = MetricUtils.get().histogram(MetricRegistry.name(this.getClass(), "kafkaLag")); @@ -64,7 +60,7 @@ public class MetadataChangeProposalsProcessor { event = EventUtils.avroToPegasusMCP(record); log.debug("MetadataChangeProposal {}", event); // TODO: Get this from the event itself. - entityClient.ingestProposal(event, this.systemAuthentication, false); + entityClient.ingestProposal(event, false); } catch (Throwable throwable) { log.error("MCP Processor Error", throwable); log.error("Message: {}", record); diff --git a/metadata-service/configuration/build.gradle b/metadata-service/configuration/build.gradle index 30fa3079d2..bf79469633 100644 --- a/metadata-service/configuration/build.gradle +++ b/metadata-service/configuration/build.gradle @@ -7,6 +7,7 @@ dependencies { implementation externalDependency.slf4jApi implementation externalDependency.springCore + implementation externalDependency.springBeans compileOnly externalDependency.lombok diff --git a/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/CacheConfiguration.java b/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/CacheConfiguration.java index 38934cb9a3..aff0e23e3b 100644 --- a/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/CacheConfiguration.java +++ b/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/CacheConfiguration.java @@ -1,5 +1,6 @@ package com.linkedin.metadata.config.cache; +import com.linkedin.metadata.config.cache.client.ClientCacheConfiguration; import lombok.Data; @@ -8,4 +9,5 @@ public class CacheConfiguration { PrimaryCacheConfiguration primary; HomepageCacheConfiguration homepage; SearchCacheConfiguration search; + ClientCacheConfiguration client; } diff --git a/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/ClientCacheConfig.java b/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/ClientCacheConfig.java new file mode 100644 index 0000000000..3cf7ef2079 --- /dev/null +++ b/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/ClientCacheConfig.java @@ -0,0 +1,10 @@ +package com.linkedin.metadata.config.cache.client; + + +public interface ClientCacheConfig { + boolean isEnabled(); + boolean isStatsEnabled(); + int getStatsIntervalSeconds(); + int getDefaultTTLSeconds(); + int getMaxBytes(); +} diff --git a/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/ClientCacheConfiguration.java b/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/ClientCacheConfiguration.java new file mode 100644 index 0000000000..d940bbe135 --- /dev/null +++ b/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/ClientCacheConfiguration.java @@ -0,0 +1,9 @@ +package com.linkedin.metadata.config.cache.client; + +import lombok.Data; + +@Data +public class ClientCacheConfiguration { + EntityClientCacheConfig entityClient; + UsageClientCacheConfig usageClient; +} diff --git a/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/EntityClientCacheConfig.java b/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/EntityClientCacheConfig.java new file mode 100644 index 0000000000..595b614f2f --- /dev/null +++ b/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/EntityClientCacheConfig.java @@ -0,0 +1,17 @@ +package com.linkedin.metadata.config.cache.client; + +import lombok.Data; + +import java.util.Map; + +@Data +public class EntityClientCacheConfig implements ClientCacheConfig { + private boolean enabled; + private boolean statsEnabled; + private int statsIntervalSeconds; + private int defaultTTLSeconds; + private int maxBytes; + + // entityName -> aspectName -> cache ttl override + private Map> entityAspectTTLSeconds; +} diff --git a/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/UsageClientCacheConfig.java b/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/UsageClientCacheConfig.java new file mode 100644 index 0000000000..3aebec9422 --- /dev/null +++ b/metadata-service/configuration/src/main/java/com/linkedin/metadata/config/cache/client/UsageClientCacheConfig.java @@ -0,0 +1,12 @@ +package com.linkedin.metadata.config.cache.client; + +import lombok.Data; + +@Data +public class UsageClientCacheConfig implements ClientCacheConfig { + private boolean enabled; + private boolean statsEnabled; + private int statsIntervalSeconds; + private int defaultTTLSeconds; + private int maxBytes; +} diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/spring/YamlPropertySourceFactory.java b/metadata-service/configuration/src/main/java/com/linkedin/metadata/spring/YamlPropertySourceFactory.java similarity index 87% rename from metadata-service/factories/src/main/java/com/linkedin/gms/factory/spring/YamlPropertySourceFactory.java rename to metadata-service/configuration/src/main/java/com/linkedin/metadata/spring/YamlPropertySourceFactory.java index 1542407697..c10399c4f3 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/spring/YamlPropertySourceFactory.java +++ b/metadata-service/configuration/src/main/java/com/linkedin/metadata/spring/YamlPropertySourceFactory.java @@ -1,14 +1,18 @@ -package com.linkedin.gms.factory.spring; +package com.linkedin.metadata.spring; -import java.io.IOException; -import java.util.Properties; import org.springframework.beans.factory.config.YamlPropertiesFactoryBean; import org.springframework.core.env.PropertiesPropertySource; import org.springframework.core.env.PropertySource; import org.springframework.core.io.support.EncodedResource; import org.springframework.core.io.support.PropertySourceFactory; +import java.io.IOException; +import java.util.Properties; + +/** + * Required for Spring to parse the application.yml provided by this module + */ public class YamlPropertySourceFactory implements PropertySourceFactory { @Override diff --git a/metadata-service/configuration/src/main/resources/application.yml b/metadata-service/configuration/src/main/resources/application.yml index ea959bebf2..42749d8205 100644 --- a/metadata-service/configuration/src/main/resources/application.yml +++ b/metadata-service/configuration/src/main/resources/application.yml @@ -327,3 +327,27 @@ cache: lineage: ttlSeconds: ${CACHE_SEARCH_LINEAGE_TTL_SECONDS:86400} # 1 day lightningThreshold: ${CACHE_SEARCH_LINEAGE_LIGHTNING_THRESHOLD:300} + client: + usageClient: + enabled: ${CACHE_CLIENT_USAGE_CLIENT_ENABLED:true} + statsEnabled: ${CACHE_CLIENT_USAGE_CLIENT_STATS_ENABLED:true} + statsIntervalSeconds: ${CACHE_CLIENT_USAGE_CLIENT_STATS_INTERVAL_SECONDS:120} + defaultTTLSeconds: ${CACHE_CLIENT_USAGE_CLIENT_TTL_SECONDS:86400} # 1 day + maxBytes: ${CACHE_CLIENT_USAGE_CLIENT_MAX_BYTES:52428800} # 50MB + entityClient: + enabled: ${CACHE_CLIENT_ENTITY_CLIENT_ENABLED:true} + statsEnabled: ${CACHE_CLIENT_ENTITY_CLIENT_STATS_ENABLED:true} + statsIntervalSeconds: ${CACHE_CLIENT_ENTITY_CLIENT_STATS_INTERVAL_SECONDS:120} + defaultTTLSeconds: ${CACHE_CLIENT_ENTITY_CLIENT_TTL_SECONDS:0} # do not cache entity/aspects by default + maxBytes: ${CACHE_CLIENT_USAGE_ENTITY_MAX_BYTES:104857600} # 100MB + entityAspectTTLSeconds: + # cache user aspects for 20s + corpuser: + corpUserKey: 20 + corpUserInfo: 20 + corpUserEditableInfo: 20 + corpUserStatus: 20 + globalTags: 20 + status: 20 + corpUserCredentials: 20 + corpUserSettings: 20 diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/AuthorizerChainFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/AuthorizerChainFactory.java index ed07239817..bf50a0c7b6 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/AuthorizerChainFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/AuthorizerChainFactory.java @@ -19,7 +19,7 @@ import com.datahub.plugins.loader.IsolatedClassLoader; import com.datahub.plugins.loader.PluginPermissionManagerImpl; import com.google.common.collect.ImmutableMap; import com.linkedin.gms.factory.config.ConfigurationProvider; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.client.JavaEntityClient; import java.nio.file.Path; import java.nio.file.Paths; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/DataHubAuthorizerFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/DataHubAuthorizerFactory.java index 30e03d87a8..5b298a4535 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/DataHubAuthorizerFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/DataHubAuthorizerFactory.java @@ -4,7 +4,7 @@ import com.datahub.authentication.Authentication; import com.datahub.authorization.DataHubAuthorizer; import com.linkedin.metadata.client.JavaEntityClient; import com.linkedin.gms.factory.entity.RestliEntityClientFactory; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import javax.annotation.Nonnull; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Qualifier; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/DataHubTokenServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/DataHubTokenServiceFactory.java index fc010a1aa2..6b2a61882b 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/DataHubTokenServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/DataHubTokenServiceFactory.java @@ -1,9 +1,10 @@ package com.linkedin.gms.factory.auth; import com.datahub.authentication.token.StatefulTokenService; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; import com.linkedin.metadata.entity.EntityService; import javax.annotation.Nonnull; + +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.beans.factory.annotation.Value; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/GroupServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/GroupServiceFactory.java index 9d29b8e77d..57598abf80 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/GroupServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/GroupServiceFactory.java @@ -4,7 +4,7 @@ package com.linkedin.gms.factory.auth; import com.datahub.authentication.group.GroupService; import com.linkedin.metadata.client.JavaEntityClient; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.graph.GraphClient; import javax.annotation.Nonnull; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/InviteTokenServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/InviteTokenServiceFactory.java index 47f7ef0e0c..105f4c677a 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/InviteTokenServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/InviteTokenServiceFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.auth; import com.datahub.authentication.invite.InviteTokenService; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.client.JavaEntityClient; import com.linkedin.metadata.secret.SecretService; import javax.annotation.Nonnull; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/NativeUserServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/NativeUserServiceFactory.java index ca52420b44..3df499ea93 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/NativeUserServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/NativeUserServiceFactory.java @@ -4,7 +4,7 @@ package com.linkedin.gms.factory.auth; import com.datahub.authentication.user.NativeUserService; import com.linkedin.metadata.client.JavaEntityClient; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.secret.SecretService; import javax.annotation.Nonnull; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/PostServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/PostServiceFactory.java index 8e5e5e5cfc..cc6f5c8272 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/PostServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/PostServiceFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.auth; import com.datahub.authentication.post.PostService; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.client.JavaEntityClient; import javax.annotation.Nonnull; import org.springframework.beans.factory.annotation.Autowired; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/RoleServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/RoleServiceFactory.java index 42f3e797c3..8a85f63cdd 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/RoleServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/RoleServiceFactory.java @@ -3,7 +3,7 @@ package com.linkedin.gms.factory.auth; import com.datahub.authorization.role.RoleService; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.client.JavaEntityClient; import javax.annotation.Nonnull; import org.springframework.beans.factory.annotation.Autowired; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/SystemAuthenticationFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/SystemAuthenticationFactory.java index d6c171dc74..5bdd8cbf83 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/SystemAuthenticationFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/SystemAuthenticationFactory.java @@ -3,7 +3,7 @@ package com.linkedin.gms.factory.auth; import com.datahub.authentication.Actor; import com.datahub.authentication.ActorType; import com.datahub.authentication.Authentication; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import javax.annotation.Nonnull; import lombok.Data; import org.springframework.beans.factory.annotation.Value; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticSearchGraphServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticSearchGraphServiceFactory.java index c1c5acbc1f..51c7db5e37 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticSearchGraphServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticSearchGraphServiceFactory.java @@ -3,7 +3,7 @@ package com.linkedin.gms.factory.common; import com.linkedin.gms.factory.config.ConfigurationProvider; import com.linkedin.gms.factory.entityregistry.EntityRegistryFactory; import com.linkedin.gms.factory.search.BaseElasticSearchComponentsFactory; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.models.registry.LineageRegistry; import com.linkedin.metadata.graph.elastic.ESGraphQueryDAO; import com.linkedin.metadata.graph.elastic.ESGraphWriteDAO; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticSearchSystemMetadataServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticSearchSystemMetadataServiceFactory.java index 89f196b056..504618ba9c 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticSearchSystemMetadataServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticSearchSystemMetadataServiceFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.common; import com.linkedin.gms.factory.search.BaseElasticSearchComponentsFactory; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.systemmetadata.ESSystemMetadataDAO; import com.linkedin.metadata.systemmetadata.ElasticSearchSystemMetadataService; import javax.annotation.Nonnull; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticsearchSSLContextFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticsearchSSLContextFactory.java index d57da33642..0dce80b989 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticsearchSSLContextFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/ElasticsearchSSLContextFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.common; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import org.apache.http.ssl.SSLContextBuilder; import org.springframework.beans.factory.annotation.Value; import org.springframework.context.annotation.Bean; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/GraphServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/GraphServiceFactory.java index 02e31c7dc4..94593eb1fb 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/GraphServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/GraphServiceFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.common; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.graph.neo4j.Neo4jGraphService; import com.linkedin.metadata.graph.elastic.ElasticSearchGraphService; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/IndexConventionFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/IndexConventionFactory.java index a2816830f3..ada8466d30 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/IndexConventionFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/IndexConventionFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.common; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; import com.linkedin.metadata.utils.elasticsearch.IndexConventionImpl; import org.springframework.beans.factory.annotation.Value; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/LocalEbeanServerConfigFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/LocalEbeanServerConfigFactory.java index 5ab5b14160..6bf8ff123b 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/LocalEbeanServerConfigFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/LocalEbeanServerConfigFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.common; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.utils.metrics.MetricUtils; import io.ebean.config.ServerConfig; import io.ebean.datasource.DataSourceConfig; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/Neo4jDriverFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/Neo4jDriverFactory.java index a364504d44..65b6115d66 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/Neo4jDriverFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/Neo4jDriverFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.common; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import java.util.concurrent.TimeUnit; import org.neo4j.driver.AuthTokens; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/RestHighLevelClientFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/RestHighLevelClientFactory.java index 5f50b8f7f0..3c40b30bfc 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/RestHighLevelClientFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/common/RestHighLevelClientFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.common; import com.linkedin.gms.factory.auth.AwsRequestSigningApacheInterceptor; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import java.io.IOException; import javax.annotation.Nonnull; import javax.net.ssl.HostnameVerifier; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/config/ConfigurationProvider.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/config/ConfigurationProvider.java index e07630111a..465480be34 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/config/ConfigurationProvider.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/config/ConfigurationProvider.java @@ -12,7 +12,7 @@ import com.linkedin.metadata.config.cache.CacheConfiguration; import com.linkedin.metadata.config.kafka.KafkaConfiguration; import com.linkedin.metadata.config.search.ElasticSearchConfiguration; import com.linkedin.datahub.graphql.featureflags.FeatureFlags; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.config.telemetry.TelemetryConfiguration; import lombok.Data; import org.springframework.boot.context.properties.ConfigurationProperties; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/dataproduct/DataProductServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/dataproduct/DataProductServiceFactory.java index c0f2c8e1f1..6eab711603 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/dataproduct/DataProductServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/dataproduct/DataProductServiceFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.dataproduct; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.client.JavaEntityClient; import com.linkedin.metadata.graph.GraphClient; import com.linkedin.metadata.service.DataProductService; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/JavaEntityClientFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/JavaEntityClientFactory.java index c9c3953f4d..e1c24b8054 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/JavaEntityClientFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/JavaEntityClientFactory.java @@ -1,8 +1,11 @@ package com.linkedin.gms.factory.entity; +import com.datahub.authentication.Authentication; +import com.linkedin.gms.factory.config.ConfigurationProvider; import com.linkedin.metadata.client.JavaEntityClient; import com.linkedin.entity.client.RestliEntityClient; import com.linkedin.gms.factory.kafka.DataHubKafkaProducerFactory; +import com.linkedin.metadata.client.SystemJavaEntityClient; import com.linkedin.metadata.entity.DeleteEntityService; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.event.EventProducer; @@ -53,12 +56,8 @@ public class JavaEntityClientFactory { @Qualifier("kafkaEventProducer") private EventProducer _eventProducer; - @Autowired - @Qualifier("restliEntityClient") - private RestliEntityClient _restliEntityClient; - @Bean("javaEntityClient") - public JavaEntityClient getJavaEntityClient() { + public JavaEntityClient getJavaEntityClient(@Qualifier("restliEntityClient") final RestliEntityClient restliEntityClient) { return new JavaEntityClient( _entityService, _deleteEntityService, @@ -68,6 +67,24 @@ public class JavaEntityClientFactory { _lineageSearchService, _timeseriesAspectService, _eventProducer, - _restliEntityClient); + restliEntityClient); + } + + @Bean("systemJavaEntityClient") + public SystemJavaEntityClient systemJavaEntityClient(@Qualifier("configurationProvider") final ConfigurationProvider configurationProvider, + @Qualifier("systemAuthentication") final Authentication systemAuthentication, + @Qualifier("systemRestliEntityClient") final RestliEntityClient restliEntityClient) { + return new SystemJavaEntityClient( + _entityService, + _deleteEntityService, + _entitySearchService, + _cachingEntitySearchService, + _searchService, + _lineageSearchService, + _timeseriesAspectService, + _eventProducer, + restliEntityClient, + systemAuthentication, + configurationProvider.getCache().getClient().getEntityClient()); } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/RestliEntityClientFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/RestliEntityClientFactory.java index e149ecedfa..dfc5e83539 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/RestliEntityClientFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/RestliEntityClientFactory.java @@ -1,10 +1,14 @@ package com.linkedin.gms.factory.entity; +import com.datahub.authentication.Authentication; import com.linkedin.entity.client.RestliEntityClient; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.entity.client.SystemRestliEntityClient; +import com.linkedin.gms.factory.config.ConfigurationProvider; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.restli.DefaultRestliClientFactory; import com.linkedin.parseq.retry.backoff.ExponentialBackoff; import com.linkedin.restli.client.Client; +import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.beans.factory.annotation.Value; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; @@ -48,4 +52,17 @@ public class RestliEntityClientFactory { } return new RestliEntityClient(restClient, new ExponentialBackoff(retryInterval), numRetries); } + + @Bean("systemRestliEntityClient") + public SystemRestliEntityClient systemRestliEntityClient(@Qualifier("configurationProvider") final ConfigurationProvider configurationProvider, + @Qualifier("systemAuthentication") final Authentication systemAuthentication) { + final Client restClient; + if (gmsUri != null) { + restClient = DefaultRestliClientFactory.getRestLiClient(URI.create(gmsUri), gmsSslProtocol); + } else { + restClient = DefaultRestliClientFactory.getRestLiClient(gmsHost, gmsPort, gmsUseSSL, gmsSslProtocol); + } + return new SystemRestliEntityClient(restClient, new ExponentialBackoff(retryInterval), numRetries, + systemAuthentication, configurationProvider.getCache().getClient().getEntityClient()); + } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/RetentionServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/RetentionServiceFactory.java index b13bf5813d..ff56f19e4f 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/RetentionServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entity/RetentionServiceFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.entity; import com.datastax.oss.driver.api.core.CqlSession; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.RetentionService; import com.linkedin.metadata.entity.cassandra.CassandraRetentionService; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityregistry/ConfigEntityRegistryFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityregistry/ConfigEntityRegistryFactory.java index 471f079683..cda21f8907 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityregistry/ConfigEntityRegistryFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityregistry/ConfigEntityRegistryFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.entityregistry; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.models.registry.ConfigEntityRegistry; import com.linkedin.metadata.models.registry.EntityRegistryException; import java.io.IOException; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityregistry/PluginEntityRegistryFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityregistry/PluginEntityRegistryFactory.java index 150e1e48f3..6dbb07309c 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityregistry/PluginEntityRegistryFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/entityregistry/PluginEntityRegistryFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.entityregistry; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.models.registry.PluginEntityRegistryLoader; import java.io.FileNotFoundException; import java.net.MalformedURLException; 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 d7aee59ca6..c50b4c9088 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 @@ -20,6 +20,7 @@ import com.linkedin.gms.factory.config.ConfigurationProvider; import com.linkedin.gms.factory.entityregistry.EntityRegistryFactory; import com.linkedin.gms.factory.entity.RestliEntityClientFactory; import com.linkedin.gms.factory.recommendation.RecommendationServiceFactory; +import com.linkedin.metadata.client.SystemJavaEntityClient; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.graph.GraphClient; import com.linkedin.metadata.graph.GraphService; @@ -65,6 +66,10 @@ public class GraphQLEngineFactory { @Qualifier("javaEntityClient") private JavaEntityClient _entityClient; + @Autowired + @Qualifier("systemJavaEntityClient") + private SystemJavaEntityClient _systemEntityClient; + @Autowired @Qualifier("graphClient") private GraphClient _graphClient; @@ -170,6 +175,7 @@ public class GraphQLEngineFactory { protected GraphQLEngine getInstance() { GmsGraphQLEngineArgs args = new GmsGraphQLEngineArgs(); args.setEntityClient(_entityClient); + args.setSystemEntityClient(_systemEntityClient); args.setGraphClient(_graphClient); args.setUsageClient(_usageClient); if (isAnalyticsEnabled) { diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/ingestion/IngestionSchedulerFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/ingestion/IngestionSchedulerFactory.java index b310ee25cb..9beb617c4f 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/ingestion/IngestionSchedulerFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/ingestion/IngestionSchedulerFactory.java @@ -6,7 +6,7 @@ import com.linkedin.entity.client.RestliEntityClient; import com.linkedin.gms.factory.auth.SystemAuthenticationFactory; import com.linkedin.gms.factory.config.ConfigurationProvider; import com.linkedin.gms.factory.entity.RestliEntityClientFactory; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import javax.annotation.Nonnull; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Qualifier; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/DataHubKafkaEventProducerFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/DataHubKafkaEventProducerFactory.java index 66f5560664..675f015d9e 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/DataHubKafkaEventProducerFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/DataHubKafkaEventProducerFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.kafka; import com.linkedin.gms.factory.common.TopicConventionFactory; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.dao.producer.KafkaEventProducer; import com.linkedin.metadata.dao.producer.KafkaHealthChecker; import com.linkedin.mxe.TopicConvention; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/DataHubKafkaProducerFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/DataHubKafkaProducerFactory.java index e58661b357..c67a2e7046 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/DataHubKafkaProducerFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/DataHubKafkaProducerFactory.java @@ -6,7 +6,7 @@ import com.linkedin.gms.factory.kafka.schemaregistry.AwsGlueSchemaRegistryFactor import com.linkedin.gms.factory.kafka.schemaregistry.InternalSchemaRegistryFactory; import com.linkedin.gms.factory.kafka.schemaregistry.KafkaSchemaRegistryFactory; import com.linkedin.gms.factory.kafka.schemaregistry.SchemaRegistryConfig; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import java.util.Arrays; import java.util.Map; import org.apache.avro.generic.IndexedRecord; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/schemaregistry/AwsGlueSchemaRegistryFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/schemaregistry/AwsGlueSchemaRegistryFactory.java index 59f08e3733..ac1cbbc5cc 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/schemaregistry/AwsGlueSchemaRegistryFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/schemaregistry/AwsGlueSchemaRegistryFactory.java @@ -5,7 +5,7 @@ import com.amazonaws.services.schemaregistry.serializers.GlueSchemaRegistryKafka import com.amazonaws.services.schemaregistry.utils.AWSSchemaRegistryConstants; import com.amazonaws.services.schemaregistry.utils.AvroRecordType; import com.linkedin.gms.factory.config.ConfigurationProvider; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import java.util.HashMap; import java.util.Map; import java.util.Optional; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/schemaregistry/KafkaSchemaRegistryFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/schemaregistry/KafkaSchemaRegistryFactory.java index d0e11baab9..7b72ba3f3b 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/schemaregistry/KafkaSchemaRegistryFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/kafka/schemaregistry/KafkaSchemaRegistryFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.kafka.schemaregistry; import com.linkedin.gms.factory.config.ConfigurationProvider; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import io.confluent.kafka.schemaregistry.client.SchemaRegistryClientConfig; import io.confluent.kafka.serializers.AbstractKafkaSchemaSerDeConfig; import io.confluent.kafka.serializers.KafkaAvroDeserializer; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/lineage/LineageServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/lineage/LineageServiceFactory.java index f76549c90a..8596a14b7f 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/lineage/LineageServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/lineage/LineageServiceFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.lineage; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.client.JavaEntityClient; import javax.annotation.Nonnull; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/ownership/OwnershipTypeServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/ownership/OwnershipTypeServiceFactory.java index 512a0a1fa4..3a1f18692f 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/ownership/OwnershipTypeServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/ownership/OwnershipTypeServiceFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.ownership; import com.datahub.authentication.Authentication; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.client.JavaEntityClient; import com.linkedin.metadata.service.OwnershipTypeService; import javax.annotation.Nonnull; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/query/QueryServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/query/QueryServiceFactory.java index f2bdce9083..f98c5bd504 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/query/QueryServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/query/QueryServiceFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.query; import com.datahub.authentication.Authentication; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.client.JavaEntityClient; import com.linkedin.metadata.service.QueryService; import javax.annotation.Nonnull; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/BaseElasticSearchComponentsFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/BaseElasticSearchComponentsFactory.java index 620af80372..c99d429e98 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/BaseElasticSearchComponentsFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/BaseElasticSearchComponentsFactory.java @@ -2,7 +2,7 @@ package com.linkedin.gms.factory.search; import com.linkedin.gms.factory.common.IndexConventionFactory; import com.linkedin.gms.factory.common.RestHighLevelClientFactory; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.search.elasticsearch.indexbuilder.ESIndexBuilder; import com.linkedin.metadata.search.elasticsearch.update.ESBulkProcessor; import com.linkedin.metadata.utils.elasticsearch.IndexConvention; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/CachingEntitySearchServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/CachingEntitySearchServiceFactory.java index 7b20e798b7..845c63c32e 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/CachingEntitySearchServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/CachingEntitySearchServiceFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.search; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.search.client.CachingEntitySearchService; import javax.annotation.Nonnull; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchBulkProcessorFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchBulkProcessorFactory.java index fc6f92b267..5deffdb01d 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchBulkProcessorFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchBulkProcessorFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.search; import com.linkedin.gms.factory.common.RestHighLevelClientFactory; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import javax.annotation.Nonnull; import com.linkedin.metadata.search.elasticsearch.update.ESBulkProcessor; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchIndexBuilderFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchIndexBuilderFactory.java index 495d77ccbb..b619ee9516 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchIndexBuilderFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchIndexBuilderFactory.java @@ -6,7 +6,7 @@ import com.linkedin.gms.factory.common.GitVersionFactory; import com.linkedin.gms.factory.common.IndexConventionFactory; import com.linkedin.gms.factory.common.RestHighLevelClientFactory; import com.linkedin.gms.factory.config.ConfigurationProvider; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.search.elasticsearch.indexbuilder.ESIndexBuilder; import com.linkedin.metadata.version.GitVersion; import javax.annotation.Nonnull; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchServiceFactory.java index 03dd2d072b..a2a0dbaf89 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/ElasticSearchServiceFactory.java @@ -7,7 +7,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.dataformat.yaml.YAMLMapper; import com.linkedin.gms.factory.config.ConfigurationProvider; import com.linkedin.gms.factory.entityregistry.EntityRegistryFactory; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.elasticsearch.ElasticSearchService; import com.linkedin.metadata.search.elasticsearch.indexbuilder.EntityIndexBuilders; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/LineageSearchServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/LineageSearchServiceFactory.java index 94b3f40849..e2eef83bc6 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/LineageSearchServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/LineageSearchServiceFactory.java @@ -2,7 +2,7 @@ package com.linkedin.gms.factory.search; import com.linkedin.gms.factory.common.GraphServiceFactory; import com.linkedin.gms.factory.config.ConfigurationProvider; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.graph.GraphService; import com.linkedin.metadata.search.LineageSearchService; import com.linkedin.metadata.search.SearchService; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SearchDocumentTransformerFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SearchDocumentTransformerFactory.java index e1fe0399cb..a186d2de77 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SearchDocumentTransformerFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SearchDocumentTransformerFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.search; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.search.transformer.SearchDocumentTransformer; import org.springframework.beans.factory.annotation.Value; import org.springframework.context.annotation.Bean; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SearchServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SearchServiceFactory.java index 70307e51f3..64bb0218a0 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SearchServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SearchServiceFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.search; import com.linkedin.gms.factory.config.ConfigurationProvider; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.EntitySearchService; import com.linkedin.metadata.search.SearchService; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SettingsBuilderFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SettingsBuilderFactory.java index b6bfef6ed8..840a370957 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SettingsBuilderFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/SettingsBuilderFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.search; import com.linkedin.gms.factory.entityregistry.EntityRegistryFactory; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.search.elasticsearch.indexbuilder.SettingsBuilder; import org.springframework.beans.factory.annotation.Autowired; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/views/ViewServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/views/ViewServiceFactory.java index 006b992191..60bcd9ea22 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/views/ViewServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/search/views/ViewServiceFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.search.views; import com.datahub.authentication.Authentication; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.client.JavaEntityClient; import com.linkedin.metadata.service.ViewService; import javax.annotation.Nonnull; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/settings/SettingsServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/settings/SettingsServiceFactory.java index 73ec79fa7e..2e22d43913 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/settings/SettingsServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/settings/SettingsServiceFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.settings; import com.datahub.authentication.Authentication; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.client.JavaEntityClient; import com.linkedin.metadata.service.SettingsService; import javax.annotation.Nonnull; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/telemetry/MixpanelApiFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/telemetry/MixpanelApiFactory.java index b2982d1f8e..8178ce1399 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/telemetry/MixpanelApiFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/telemetry/MixpanelApiFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.telemetry; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.mixpanel.mixpanelapi.MixpanelAPI; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.context.annotation.Bean; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/telemetry/MixpanelMessageBuilderFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/telemetry/MixpanelMessageBuilderFactory.java index aa8596786c..5385c5e81f 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/telemetry/MixpanelMessageBuilderFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/telemetry/MixpanelMessageBuilderFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.telemetry; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.mixpanel.mixpanelapi.MessageBuilder; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.context.annotation.Bean; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/telemetry/TrackingServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/telemetry/TrackingServiceFactory.java index 3b53a6fe92..bb166af550 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/telemetry/TrackingServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/telemetry/TrackingServiceFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.telemetry; import com.datahub.telemetry.TrackingService; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.secret.SecretService; import com.linkedin.metadata.version.GitVersion; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/timeline/EntityChangeEventGeneratorRegistryFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/timeline/EntityChangeEventGeneratorRegistryFactory.java index e9b9850c01..89a7e7dd8d 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/timeline/EntityChangeEventGeneratorRegistryFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/timeline/EntityChangeEventGeneratorRegistryFactory.java @@ -1,7 +1,7 @@ package com.linkedin.gms.factory.timeline; import com.datahub.authentication.Authentication; -import com.linkedin.entity.client.RestliEntityClient; +import com.linkedin.entity.client.SystemRestliEntityClient; import com.linkedin.metadata.timeline.eventgenerator.AssertionRunEventChangeEventGenerator; import com.linkedin.metadata.timeline.eventgenerator.DataProcessInstanceRunEventChangeEventGenerator; import com.linkedin.metadata.timeline.eventgenerator.DatasetPropertiesChangeEventGenerator; @@ -38,7 +38,7 @@ public class EntityChangeEventGeneratorRegistryFactory { @Singleton @Nonnull protected com.linkedin.metadata.timeline.eventgenerator.EntityChangeEventGeneratorRegistry entityChangeEventGeneratorRegistry() { - final RestliEntityClient entityClient = applicationContext.getBean(RestliEntityClient.class); + final SystemRestliEntityClient entityClient = applicationContext.getBean(SystemRestliEntityClient.class); final Authentication systemAuthentication = applicationContext.getBean(Authentication.class); final com.linkedin.metadata.timeline.eventgenerator.EntityChangeEventGeneratorRegistry registry = @@ -74,7 +74,7 @@ public class EntityChangeEventGeneratorRegistryFactory { // Data Process Instance differs registry.register(DATA_PROCESS_INSTANCE_RUN_EVENT_ASPECT_NAME, - new DataProcessInstanceRunEventChangeEventGenerator(entityClient, systemAuthentication)); + new DataProcessInstanceRunEventChangeEventGenerator(entityClient)); // TODO: Add ML models. diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/timeline/TimelineServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/timeline/TimelineServiceFactory.java index df9d80eb63..baa22d4013 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/timeline/TimelineServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/timeline/TimelineServiceFactory.java @@ -1,6 +1,6 @@ package com.linkedin.gms.factory.timeline; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.entity.AspectDao; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.timeline.TimelineService; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/timeseries/ElasticSearchTimeseriesAspectServiceFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/timeseries/ElasticSearchTimeseriesAspectServiceFactory.java index 717adf7d55..e3cc772f21 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/timeseries/ElasticSearchTimeseriesAspectServiceFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/timeseries/ElasticSearchTimeseriesAspectServiceFactory.java @@ -2,7 +2,7 @@ package com.linkedin.gms.factory.timeseries; import com.linkedin.gms.factory.entityregistry.EntityRegistryFactory; import com.linkedin.gms.factory.search.BaseElasticSearchComponentsFactory; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.timeseries.elastic.ElasticSearchTimeseriesAspectService; import com.linkedin.metadata.timeseries.elastic.indexbuilder.TimeseriesAspectIndexBuilders; diff --git a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/usage/UsageClientFactory.java b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/usage/UsageClientFactory.java index e4cbb92ceb..e83cbc82d8 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/gms/factory/usage/UsageClientFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/gms/factory/usage/UsageClientFactory.java @@ -1,10 +1,14 @@ package com.linkedin.gms.factory.usage; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.datahub.authentication.Authentication; +import com.linkedin.gms.factory.config.ConfigurationProvider; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.restli.DefaultRestliClientFactory; import com.linkedin.parseq.retry.backoff.ExponentialBackoff; import com.linkedin.restli.client.Client; import com.linkedin.usage.UsageClient; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.beans.factory.annotation.Value; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; @@ -33,10 +37,15 @@ public class UsageClientFactory { @Value("${usageClient.numRetries:3}") private int numRetries; + @Autowired + @Qualifier("configurationProvider") + private ConfigurationProvider configurationProvider; + @Bean("usageClient") - public UsageClient getUsageClient() { + public UsageClient getUsageClient(@Qualifier("systemAuthentication") final Authentication systemAuthentication) { Client restClient = DefaultRestliClientFactory.getRestLiClient(gmsHost, gmsPort, gmsUseSSL, gmsSslProtocol); - return new UsageClient(restClient, new ExponentialBackoff(retryInterval), numRetries); + return new UsageClient(restClient, new ExponentialBackoff(retryInterval), numRetries, systemAuthentication, + configurationProvider.getCache().getClient().getUsageClient()); } } diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/IngestRetentionPoliciesStepFactory.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/IngestRetentionPoliciesStepFactory.java index 91fc58d074..e038cb230c 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/IngestRetentionPoliciesStepFactory.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/factories/IngestRetentionPoliciesStepFactory.java @@ -1,7 +1,7 @@ package com.linkedin.metadata.boot.factories; import com.linkedin.gms.factory.entity.RetentionServiceFactory; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.boot.steps.IngestRetentionPoliciesStep; import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.entity.RetentionService; diff --git a/metadata-service/openapi-entity-servlet/src/test/java/io/datahubproject/openapi/util/OpenApiEntitiesUtilTest.java b/metadata-service/openapi-entity-servlet/src/test/java/io/datahubproject/openapi/util/OpenApiEntitiesUtilTest.java index e0fec07452..b4e87eedea 100644 --- a/metadata-service/openapi-entity-servlet/src/test/java/io/datahubproject/openapi/util/OpenApiEntitiesUtilTest.java +++ b/metadata-service/openapi-entity-servlet/src/test/java/io/datahubproject/openapi/util/OpenApiEntitiesUtilTest.java @@ -1,7 +1,7 @@ package io.datahubproject.openapi.util; import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor; -import com.linkedin.gms.factory.spring.YamlPropertySourceFactory; +import com.linkedin.metadata.spring.YamlPropertySourceFactory; import com.linkedin.metadata.models.registry.EntityRegistry; import io.datahubproject.openapi.config.OpenAPIEntityTestConfiguration; import io.datahubproject.openapi.dto.UpsertAspectRequest; diff --git a/metadata-service/restli-client/build.gradle b/metadata-service/restli-client/build.gradle index 45cf008d3c..b1b778b45c 100644 --- a/metadata-service/restli-client/build.gradle +++ b/metadata-service/restli-client/build.gradle @@ -7,6 +7,7 @@ dependencies { api project(path: ':metadata-service:restli-api', configuration: 'restClient') api project(':metadata-events:mxe-schemas') api project(':metadata-utils') + implementation project(':metadata-service:configuration') implementation externalDependency.slf4jApi compileOnly externalDependency.lombok diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/common/client/ClientCache.java b/metadata-service/restli-client/src/main/java/com/linkedin/common/client/ClientCache.java new file mode 100644 index 0000000000..8aa0984be5 --- /dev/null +++ b/metadata-service/restli-client/src/main/java/com/linkedin/common/client/ClientCache.java @@ -0,0 +1,134 @@ +package com.linkedin.common.client; + +import com.codahale.metrics.Gauge; +import com.github.benmanes.caffeine.cache.CacheLoader; +import com.github.benmanes.caffeine.cache.Caffeine; +import com.github.benmanes.caffeine.cache.Expiry; +import com.github.benmanes.caffeine.cache.LoadingCache; +import com.github.benmanes.caffeine.cache.Weigher; +import com.github.benmanes.caffeine.cache.stats.CacheStats; +import com.linkedin.metadata.config.cache.client.ClientCacheConfig; +import com.linkedin.metadata.utils.metrics.MetricUtils; +import lombok.Builder; +import lombok.NonNull; +import lombok.extern.slf4j.Slf4j; +import org.checkerframework.checker.nullness.qual.Nullable; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.function.BiFunction; +import java.util.function.Function; + +/** + * Generic cache with common configuration for limited weight, per item expiry, and batch loading + * @param key + * @param value + */ +@Slf4j +@Builder +public class ClientCache { + @NonNull + protected final C config; + @NonNull + protected final LoadingCache cache; + @NonNull + private final Function, Map> loadFunction; + @NonNull + private final Weigher weigher; + @NonNull + private final BiFunction ttlSecondsFunction; + + public @Nullable V get(@NonNull K key) { + return cache.get(key); + } + + public @NonNull Map<@NonNull K, @NonNull V> getAll(@NonNull Iterable keys) { + return cache.getAll(keys); + } + + public void refresh(@NonNull K key) { + cache.refresh(key); + } + + public static class ClientCacheBuilder { + + private ClientCacheBuilder cache(LoadingCache cache) { + return null; + } + private ClientCache build() { + return null; + } + + public ClientCache build(Class metricClazz) { + // loads data from entity client + CacheLoader loader = new CacheLoader<>() { + @Override + public V load(@NonNull K key) { + return loadAll(List.of(key)).get(key); + } + + @Override + @NonNull + public Map loadAll(@NonNull Iterable keys) { + return loadFunction.apply(keys); + } + }; + + // build cache + Caffeine caffeine = Caffeine.newBuilder() + .maximumWeight(config.getMaxBytes()) + // limit total size + .weigher(weigher) + .softValues() + // define per entity/aspect ttls + .expireAfter(new Expiry() { + public long expireAfterCreate(@NonNull K key, @NonNull V aspect, long currentTime) { + int ttlSeconds = ttlSecondsFunction.apply(config, key); + if (ttlSeconds < 0) { + ttlSeconds = Integer.MAX_VALUE; + } + return TimeUnit.SECONDS.toNanos(ttlSeconds); + } + public long expireAfterUpdate(@NonNull K key, @NonNull V aspect, + long currentTime, long currentDuration) { + return currentDuration; + } + public long expireAfterRead(@NonNull K key, @NonNull V aspect, + long currentTime, long currentDuration) { + return currentDuration; + } + }); + + if (config.isStatsEnabled()) { + caffeine.recordStats(); + } + + LoadingCache cache = caffeine.build(loader); + + if (config.isStatsEnabled()) { + ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1); + executor.scheduleAtFixedRate(() -> { + CacheStats cacheStats = cache.stats(); + + MetricUtils.gauge(metricClazz, "hitRate", () -> (Gauge) cacheStats::hitRate); + MetricUtils.gauge(metricClazz, "loadFailureRate", () -> + (Gauge) cacheStats::loadFailureRate); + MetricUtils.gauge(metricClazz, "evictionCount", () -> + (Gauge) cacheStats::evictionCount); + MetricUtils.gauge(metricClazz, "loadFailureCount", () -> + (Gauge) cacheStats::loadFailureCount); + MetricUtils.gauge(metricClazz, "averageLoadPenalty", () -> + (Gauge) cacheStats::averageLoadPenalty); + MetricUtils.gauge(metricClazz, "evictionWeight", () -> + (Gauge) cacheStats::evictionWeight); + + log.debug(metricClazz.getSimpleName() + ": " + cacheStats); + }, 0, config.getStatsIntervalSeconds(), TimeUnit.SECONDS); + } + + return new ClientCache<>(config, cache, loadFunction, weigher, ttlSecondsFunction); + } + } +} diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClientCache.java b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClientCache.java new file mode 100644 index 0000000000..3b35dc5289 --- /dev/null +++ b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClientCache.java @@ -0,0 +1,141 @@ +package com.linkedin.entity.client; + +import com.github.benmanes.caffeine.cache.LoadingCache; +import com.github.benmanes.caffeine.cache.Weigher; +import com.linkedin.common.client.ClientCache; +import com.linkedin.common.urn.Urn; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.EnvelopedAspect; +import com.linkedin.entity.EnvelopedAspectMap; +import com.linkedin.metadata.config.cache.client.EntityClientCacheConfig; +import com.linkedin.util.Pair; +import lombok.Builder; +import lombok.Data; +import lombok.NonNull; + +import javax.annotation.Nonnull; +import java.util.Collection; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.function.BiFunction; +import java.util.function.Function; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import java.util.stream.StreamSupport; + +import static com.linkedin.metadata.utils.PegasusUtils.urnToEntityName; + +@Builder +public class EntityClientCache { + @NonNull + private EntityClientCacheConfig config; + @NonNull + private final ClientCache cache; + @NonNull + private BiFunction, Set, Map> loadFunction; + + public EntityResponse getV2(@Nonnull final Urn urn, @Nonnull final Set aspectNames) { + return batchGetV2(Set.of(urn), aspectNames).get(urn); + } + + public Map batchGetV2(@Nonnull final Set urns, @Nonnull final Set aspectNames) { + final Map response; + + if (config.isEnabled()) { + Set keys = urns.stream() + .flatMap(urn -> aspectNames.stream() + .map(a -> Key.builder().urn(urn).aspectName(a).build())) + .collect(Collectors.toSet()); + Map envelopedAspects = cache.getAll(keys); + + Set responses = envelopedAspects.entrySet().stream() + .map(entry -> Pair.of(entry.getKey().getUrn(), entry.getValue())) + .collect(Collectors.groupingBy(Pair::getKey, Collectors.mapping(Pair::getValue, Collectors.toSet()))) + .entrySet().stream().map(e -> toEntityResponse(e.getKey(), e.getValue())) + .collect(Collectors.toSet()); + + response = responses.stream().collect(Collectors.toMap(EntityResponse::getUrn, Function.identity())); + } else { + response = loadFunction.apply(urns, aspectNames); + } + + return response; + } + + private static EntityResponse toEntityResponse(Urn urn, Collection envelopedAspects) { + final EntityResponse response = new EntityResponse(); + response.setUrn(urn); + response.setEntityName(urnToEntityName(urn)); + response.setAspects(new EnvelopedAspectMap( + envelopedAspects.stream() + .collect(Collectors.toMap(EnvelopedAspect::getName, aspect -> aspect)) + )); + return response; + } + + public static class EntityClientCacheBuilder { + + private EntityClientCacheBuilder cache(LoadingCache cache) { + return this; + } + + public EntityClientCache build(Class metricClazz) { + // estimate size + Weigher weighByEstimatedSize = (key, value) -> + value.getValue().data().values().parallelStream() + .mapToInt(o -> o.toString().getBytes().length) + .sum(); + + // batch loads data from entity client (restli or java) + Function, Map> loader = (Iterable keys) -> { + Map> keysByEntity = StreamSupport.stream(keys.spliterator(), true) + .collect(Collectors.groupingBy(Key::getEntityName, Collectors.toSet())); + + Stream> results = keysByEntity.entrySet().parallelStream() + .flatMap(entry -> { + Set urns = entry.getValue().stream() + .map(Key::getUrn) + .collect(Collectors.toSet()); + Set aspects = entry.getValue().stream() + .map(Key::getEntityName) + .collect(Collectors.toSet()); + return loadFunction.apply(urns, aspects).entrySet().stream(); + }) + .flatMap(resp -> resp.getValue().getAspects().values().stream() + .map(envAspect -> { + Key key = Key.builder().urn(resp.getKey()).aspectName(envAspect.getName()).build(); + return Map.entry(key, envAspect); + })); + + return results.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + }; + + // ideally the cache time comes from caching headers from service, but configuration driven for now + BiFunction ttlSeconds = (config, key) -> + Optional.ofNullable(config.getEntityAspectTTLSeconds()).orElse(Map.of()) + .getOrDefault(key.getEntityName(), Map.of()) + .getOrDefault(key.getAspectName(), config.getDefaultTTLSeconds()); + + cache = ClientCache.builder() + .weigher(weighByEstimatedSize) + .config(config) + .loadFunction(loader) + .ttlSecondsFunction(ttlSeconds) + .build(metricClazz); + + return new EntityClientCache(config, cache, loadFunction); + } + } + + @Data + @Builder + protected static class Key { + private final Urn urn; + private final String aspectName; + + public String getEntityName() { + return urn.getEntityType(); + } + } +} diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemEntityClient.java b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemEntityClient.java new file mode 100644 index 0000000000..94067abd0c --- /dev/null +++ b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemEntityClient.java @@ -0,0 +1,91 @@ +package com.linkedin.entity.client; + +import com.datahub.authentication.Authentication; +import com.linkedin.common.urn.Urn; +import com.linkedin.entity.EntityResponse; +import com.linkedin.metadata.config.cache.client.EntityClientCacheConfig; +import com.linkedin.mxe.MetadataChangeProposal; +import com.linkedin.mxe.PlatformEvent; +import com.linkedin.r2.RemoteInvocationException; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.net.URISyntaxException; +import java.util.Map; +import java.util.Set; + +/** + * Adds entity/aspect cache and assumes system authentication + */ +public interface SystemEntityClient extends EntityClient { + + EntityClientCache getEntityClientCache(); + Authentication getSystemAuthentication(); + + /** + * Builds the cache + * @param systemAuthentication system authentication + * @param cacheConfig cache configuration + * @return the cache + */ + default EntityClientCache buildEntityClientCache(Class metricClazz, Authentication systemAuthentication, EntityClientCacheConfig cacheConfig) { + return EntityClientCache.builder() + .config(cacheConfig) + .loadFunction((Set urns, Set aspectNames) -> { + try { + String entityName = urns.stream().findFirst().map(Urn::getEntityType).get(); + + if (urns.stream().anyMatch(urn -> !urn.getEntityType().equals(entityName))) { + throw new IllegalArgumentException("Urns must be of the same entity type. RestliEntityClient API limitation."); + } + + return batchGetV2(entityName, urns, aspectNames, systemAuthentication); + } catch (RemoteInvocationException | URISyntaxException e) { + throw new RuntimeException(e); + } + }).build(metricClazz); + } + + /** + * Get an entity by urn with the given aspects + * @param urn the id of the entity + * @param aspectNames aspects of the entity + * @return response object + * @throws RemoteInvocationException + * @throws URISyntaxException + */ + @Nullable + default EntityResponse getV2(@Nonnull Urn urn, @Nonnull Set aspectNames) + throws RemoteInvocationException, URISyntaxException { + return getEntityClientCache().getV2(urn, aspectNames); + } + + /** + * Batch get a set of aspects for a single entity type, multiple ids with the given aspects. + * + * @param urns the urns of the entities to batch get + * @param aspectNames the aspect names to batch get + * @throws RemoteInvocationException + */ + @Nonnull + default Map batchGetV2(@Nonnull Set urns, @Nonnull Set aspectNames) + throws RemoteInvocationException, URISyntaxException { + return getEntityClientCache().batchGetV2(urns, aspectNames); + } + + default void producePlatformEvent(@Nonnull String name, @Nullable String key, @Nonnull PlatformEvent event) throws Exception { + producePlatformEvent(name, key, event, getSystemAuthentication()); + } + + default boolean exists(@Nonnull Urn urn) throws RemoteInvocationException { + return exists(urn, getSystemAuthentication()); + } + + default String ingestProposal(@Nonnull final MetadataChangeProposal metadataChangeProposal, final boolean async) throws RemoteInvocationException { + return ingestProposal(metadataChangeProposal, getSystemAuthentication(), async); + } + + default void setWritable(boolean canWrite) throws RemoteInvocationException { + setWritable(canWrite, getSystemAuthentication()); + } +} diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemRestliEntityClient.java b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemRestliEntityClient.java new file mode 100644 index 0000000000..f3c3435342 --- /dev/null +++ b/metadata-service/restli-client/src/main/java/com/linkedin/entity/client/SystemRestliEntityClient.java @@ -0,0 +1,25 @@ +package com.linkedin.entity.client; + +import com.datahub.authentication.Authentication; +import com.linkedin.metadata.config.cache.client.EntityClientCacheConfig; +import com.linkedin.parseq.retry.backoff.BackoffPolicy; +import com.linkedin.restli.client.Client; +import lombok.Getter; + +import javax.annotation.Nonnull; + +/** + * Restli backed SystemEntityClient + */ +@Getter +public class SystemRestliEntityClient extends RestliEntityClient implements SystemEntityClient { + private final EntityClientCache entityClientCache; + private final Authentication systemAuthentication; + + public SystemRestliEntityClient(@Nonnull final Client restliClient, @Nonnull final BackoffPolicy backoffPolicy, int retryCount, + Authentication systemAuthentication, EntityClientCacheConfig cacheConfig) { + super(restliClient, backoffPolicy, retryCount); + this.systemAuthentication = systemAuthentication; + this.entityClientCache = buildEntityClientCache(SystemRestliEntityClient.class, systemAuthentication, cacheConfig); + } +} diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClient.java b/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClient.java index 47a15ccdd3..d2b8499615 100644 --- a/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClient.java +++ b/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClient.java @@ -5,6 +5,7 @@ import com.linkedin.common.EntityRelationships; import com.linkedin.common.WindowDuration; import com.linkedin.common.client.BaseClient; +import com.linkedin.metadata.config.cache.client.UsageClientCacheConfig; import com.linkedin.parseq.retry.backoff.BackoffPolicy; import com.linkedin.r2.RemoteInvocationException; import com.linkedin.restli.client.Client; @@ -17,19 +18,39 @@ public class UsageClient extends BaseClient { private static final UsageStatsRequestBuilders USAGE_STATS_REQUEST_BUILDERS = new UsageStatsRequestBuilders(); - public UsageClient(@Nonnull final Client restliClient, @Nonnull final BackoffPolicy backoffPolicy, int retryCount) { + private final UsageClientCache usageClientCache; + + public UsageClient(@Nonnull final Client restliClient, @Nonnull final BackoffPolicy backoffPolicy, int retryCount, + Authentication systemAuthentication, UsageClientCacheConfig cacheConfig) { super(restliClient, backoffPolicy, retryCount); + this.usageClientCache = UsageClientCache.builder() + .config(cacheConfig) + .loadFunction((String resource, UsageTimeRange range) -> { + try { + return getUsageStats(resource, range, systemAuthentication); + } catch (RemoteInvocationException | URISyntaxException e) { + throw new RuntimeException(e); + } + }).build(); + } + + /** + * Gets a specific version of downstream {@link EntityRelationships} for the given dataset. + * Using cache and system authentication. + * Validate permissions before use! + */ + @Nonnull + public UsageQueryResult getUsageStats(@Nonnull String resource, @Nonnull UsageTimeRange range) { + return usageClientCache.getUsageStats(resource, range); } /** * Gets a specific version of downstream {@link EntityRelationships} for the given dataset. */ @Nonnull - public UsageQueryResult getUsageStats( - @Nonnull String resource, - @Nonnull UsageTimeRange range, - @Nonnull Authentication authentication - ) throws RemoteInvocationException, URISyntaxException { + private UsageQueryResult getUsageStats(@Nonnull String resource, @Nonnull UsageTimeRange range, + @Nonnull Authentication authentication) + throws RemoteInvocationException, URISyntaxException { final UsageStatsDoQueryRangeRequestBuilder requestBuilder = USAGE_STATS_REQUEST_BUILDERS.actionQueryRange() .resourceParam(resource) .durationParam(WindowDuration.DAY) diff --git a/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClientCache.java b/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClientCache.java new file mode 100644 index 0000000000..a04c1e90fb --- /dev/null +++ b/metadata-service/restli-client/src/main/java/com/linkedin/usage/UsageClientCache.java @@ -0,0 +1,75 @@ +package com.linkedin.usage; + +import com.github.benmanes.caffeine.cache.LoadingCache; +import com.github.benmanes.caffeine.cache.Weigher; +import com.linkedin.common.client.ClientCache; +import com.linkedin.metadata.config.cache.client.UsageClientCacheConfig; +import lombok.Builder; +import lombok.Data; +import lombok.NonNull; + +import javax.annotation.Nonnull; +import java.util.Map; +import java.util.function.BiFunction; +import java.util.function.Function; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; + + +@Builder +public class UsageClientCache { + @NonNull + private UsageClientCacheConfig config; + @NonNull + private final ClientCache cache; + @NonNull + private BiFunction loadFunction; + + public UsageQueryResult getUsageStats(@Nonnull String resource, @Nonnull UsageTimeRange range) { + if (config.isEnabled()) { + return cache.get(Key.builder().resource(resource).range(range).build()); + } else { + return loadFunction.apply(resource, range); + } + } + + public static class UsageClientCacheBuilder { + + private UsageClientCacheBuilder cache(LoadingCache cache) { + return this; + } + + public UsageClientCache build() { + // estimate size + Weigher weighByEstimatedSize = (key, value) -> + value.data().values().parallelStream() + .mapToInt(o -> o.toString().getBytes().length) + .sum(); + + // batch loads data from usage client + Function, Map> loader = (Iterable keys) -> + StreamSupport.stream(keys.spliterator(), true) + .map(k -> Map.entry(k, loadFunction.apply(k.getResource(), k.getRange()))) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + + // default ttl only + BiFunction ttlSeconds = (config, key) -> config.getDefaultTTLSeconds(); + + cache = ClientCache.builder() + .weigher(weighByEstimatedSize) + .config(config) + .loadFunction(loader) + .ttlSecondsFunction(ttlSeconds) + .build(UsageClientCache.class); + + return new UsageClientCache(config, cache, loadFunction); + } + } + + @Data + @Builder + protected static class Key { + private final String resource; + private final UsageTimeRange range; + } +} diff --git a/metadata-utils/src/main/java/com/linkedin/metadata/utils/metrics/MetricUtils.java b/metadata-utils/src/main/java/com/linkedin/metadata/utils/metrics/MetricUtils.java index 3d90cba85b..9a8848e090 100644 --- a/metadata-utils/src/main/java/com/linkedin/metadata/utils/metrics/MetricUtils.java +++ b/metadata-utils/src/main/java/com/linkedin/metadata/utils/metrics/MetricUtils.java @@ -1,6 +1,7 @@ package com.linkedin.metadata.utils.metrics; import com.codahale.metrics.Counter; +import com.codahale.metrics.Gauge; import com.codahale.metrics.MetricRegistry; import com.codahale.metrics.SharedMetricRegistries; import com.codahale.metrics.Timer; @@ -48,4 +49,8 @@ public class MetricUtils { public static Timer timer(String metricName) { return REGISTRY.timer(MetricRegistry.name(metricName)); } + + public static > T gauge(Class clazz, String metricName, MetricRegistry.MetricSupplier supplier) { + return REGISTRY.gauge(MetricRegistry.name(clazz, metricName), supplier); + } }