feat(entity-client): enable client side cache for entity-client and u…

…sage-client (datahub-project#8877)

datahub-frontend/app/auth/AuthModule.java

@@ -11,16 +11,19 @@
 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.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 @@ protected void configure() {
             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 @@ protected Authentication provideSystemAuthentication() {
 
     @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

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

@@ -13,7 +13,7 @@
 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 @@
 public class OidcCallbackLogic extends DefaultCallbackLogic<Result, PlayWebContext> {
 
   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;

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;
+}

datahub-frontend/app/controllers/SsoCallbackController.java

@@ -3,7 +3,7 @@
 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 SsoCallbackLogic implements CallbackLogic<Result, PlayWebContext> {
     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);
     }
 

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

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

@@ -302,6 +302,7 @@
 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 @@
 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 GmsGraphQLEngine(final GmsGraphQLEngineArgs args) {
         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;

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

@@ -11,6 +11,7 @@
 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 @@
 @Data
 public class GmsGraphQLEngineArgs {
     EntityClient entityClient;
+    SystemEntityClient systemEntityClient;
     GraphClient graphClient;
     UsageClient usageClient;
     AnalyticsService analyticsService;

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetStatsSummaryResolver.java

@@ -1,20 +1,24 @@
 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;
 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 CompletableFuture<DatasetStatsSummary> get(DataFetchingEnvironment enviro
 
       try {
 
+        if (!isAuthorized(resourceUrn, context)) {
+          log.debug("User {} is not authorized to view profile information for dataset {}",
+                  context.getActorUrn(),
+                  resourceUrn.toString());
+          return null;
+        }
+
         com.linkedin.usage.UsageQueryResult
-            usageQueryResult = usageClient.getUsageStats(resourceUrn.toString(), UsageTimeRange.MONTH, context.getAuthentication());
+            usageQueryResult = usageClient.getUsageStats(resourceUrn.toString(), UsageTimeRange.MONTH);
 
         final DatasetStatsSummary result = new DatasetStatsSummary();
         result.setQueryCountLast30Days(usageQueryResult.getAggregations().getTotalSqlQueries());
@@ -90,4 +101,10 @@ private CorpUser createPartialUser(final Urn userUrn) {
     result.setUrn(userUrn.toString());
     return result;
   }
+
+  private boolean isAuthorized(final Urn resourceUrn, final QueryContext context) {
+    return AuthorizationUtils.isAuthorized(context,
+            Optional.of(new ResourceSpec(resourceUrn.getEntityType(), resourceUrn.toString())),
+            PoliciesConfig.VIEW_DATASET_USAGE_PRIVILEGE);
+  }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetUsageStatsResolver.java

@@ -9,12 +9,10 @@
 import com.linkedin.datahub.graphql.generated.UsageQueryResult;
 import com.linkedin.datahub.graphql.types.usage.UsageQueryResultMapper;
 import com.linkedin.metadata.authorization.PoliciesConfig;
-import com.linkedin.r2.RemoteInvocationException;
 import com.linkedin.usage.UsageClient;
 import com.linkedin.usage.UsageTimeRange;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
-import java.net.URISyntaxException;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import lombok.extern.slf4j.Slf4j;
@@ -44,10 +42,10 @@ public CompletableFuture<UsageQueryResult> get(DataFetchingEnvironment environme
       }
       try {
         com.linkedin.usage.UsageQueryResult
-            usageQueryResult = usageClient.getUsageStats(resourceUrn.toString(), range, context.getAuthentication());
+            usageQueryResult = usageClient.getUsageStats(resourceUrn.toString(), range);
         return UsageQueryResultMapper.map(usageQueryResult);
-      } catch (RemoteInvocationException | URISyntaxException e) {
-        throw new RuntimeException(String.format("Failed to load Usage Stats for resource %s", resourceUrn.toString()), e);
+      } catch (Exception e) {
+        throw new RuntimeException(String.format("Failed to load Usage Stats for resource %s", resourceUrn), e);
       }
     });
   }

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/dashboard/DashboardStatsSummaryTest.java

@@ -117,8 +117,7 @@ public void testGetException() throws Exception {
     UsageClient mockClient = Mockito.mock(UsageClient.class);
     Mockito.when(mockClient.getUsageStats(
         Mockito.eq(TEST_DASHBOARD_URN),
-        Mockito.eq(UsageTimeRange.MONTH),
-        Mockito.any(Authentication.class)
+        Mockito.eq(UsageTimeRange.MONTH)
     )).thenThrow(RuntimeException.class);
 
     // Execute resolver

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetStatsSummaryResolverTest.java

@@ -1,6 +1,8 @@
 package com.linkedin.datahub.graphql.resolvers.dataset;
 
 import com.datahub.authentication.Authentication;
+import com.datahub.authorization.AuthorizationResult;
+import com.datahub.plugins.auth.authorization.Authorizer;
 import com.google.common.collect.ImmutableList;
 import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.datahub.graphql.QueryContext;
@@ -53,13 +55,18 @@ public void testGetSuccess() throws Exception {
     UsageClient mockClient = Mockito.mock(UsageClient.class);
     Mockito.when(mockClient.getUsageStats(
         Mockito.eq(TEST_DATASET_URN),
-        Mockito.eq(UsageTimeRange.MONTH),
-        Mockito.any(Authentication.class)
+        Mockito.eq(UsageTimeRange.MONTH)
     )).thenReturn(testResult);
 
     // Execute resolver
     DatasetStatsSummaryResolver resolver = new DatasetStatsSummaryResolver(mockClient);
     QueryContext mockContext = Mockito.mock(QueryContext.class);
+    Mockito.when(mockContext.getActorUrn()).thenReturn("urn:li:corpuser:test");
+    Authorizer mockAuthorizer = Mockito.mock(Authorizer.class);
+    AuthorizationResult mockAuthorizerResult = Mockito.mock(AuthorizationResult.class);
+    Mockito.when(mockAuthorizerResult.getType()).thenReturn(AuthorizationResult.Type.ALLOW);
+    Mockito.when(mockAuthorizer.authorize(Mockito.any())).thenReturn(mockAuthorizerResult);
+    Mockito.when(mockContext.getAuthorizer()).thenReturn(mockAuthorizer);
     Mockito.when(mockContext.getAuthentication()).thenReturn(Mockito.mock(Authentication.class));
     DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class);
     Mockito.when(mockEnv.getSource()).thenReturn(TEST_SOURCE);
@@ -79,8 +86,7 @@ public void testGetSuccess() throws Exception {
     newResult.setAggregations(new UsageQueryResultAggregations());
     Mockito.when(mockClient.getUsageStats(
         Mockito.eq(TEST_DATASET_URN),
-        Mockito.eq(UsageTimeRange.MONTH),
-        Mockito.any(Authentication.class)
+        Mockito.eq(UsageTimeRange.MONTH)
     )).thenReturn(newResult);
 
     // Then verify that the new result is _not_ returned (cache hit)
@@ -116,8 +122,7 @@ public void testGetException() throws Exception {
     UsageClient mockClient = Mockito.mock(UsageClient.class);
     Mockito.when(mockClient.getUsageStats(
         Mockito.eq(TEST_DATASET_URN),
-        Mockito.eq(UsageTimeRange.MONTH),
-        Mockito.any(Authentication.class)
+        Mockito.eq(UsageTimeRange.MONTH)
     )).thenThrow(RuntimeException.class);
 
     // Execute resolver

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/common/steps/GMSDisableWriteModeStep.java

@@ -1,20 +1,18 @@
 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 GMSDisableWriteModeStep implements UpgradeStep {
 
-  private final Authentication _systemAuthentication;
-  private final RestliEntityClient _entityClient;
+  private final SystemRestliEntityClient _entityClient;
 
   @Override
   public String id() {
@@ -30,7 +28,7 @@ public int retryCount() {
   public Function<UpgradeContext, UpgradeStepResult> 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");