feat: usage stats (part 1)

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

@@ -7,6 +7,7 @@
 import com.linkedin.lineage.client.Relationships;
 import com.linkedin.metadata.restli.DefaultRestliClientFactory;
 import com.linkedin.restli.client.Client;
+import com.linkedin.usage.UsageClient;
 import com.linkedin.util.Configuration;
 
 /**
@@ -36,6 +37,7 @@ public class GmsClientFactory {
     private static Relationships _relationships;
     private static EntityClient _entities;
     private static AspectClient _aspects;
+    private static UsageClient _usage;
 
 
     private GmsClientFactory() { }
@@ -94,4 +96,15 @@ public static AspectClient getAspectsClient() {
         }
         return _aspects;
     }
+
+    public static UsageClient getUsageClient() {
+        if (_usage == null) {
+            synchronized (GmsClientFactory.class) {
+                if (_usage == null) {
+                    _usage = new UsageClient(REST_CLIENT);
+                }
+            }
+        }
+        return _usage;
+    }
 }

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

@@ -13,9 +13,12 @@
 import com.linkedin.datahub.graphql.generated.RelatedDataset;
 import com.linkedin.datahub.graphql.generated.SearchResult;
 import com.linkedin.datahub.graphql.generated.InstitutionalMemoryMetadata;
+import com.linkedin.datahub.graphql.generated.UsageQueryResult;
+import com.linkedin.datahub.graphql.generated.UserUsageCounts;
 import com.linkedin.datahub.graphql.resolvers.load.AspectResolver;
 import com.linkedin.datahub.graphql.resolvers.load.EntityTypeResolver;
 import com.linkedin.datahub.graphql.resolvers.load.LoadableTypeBatchResolver;
+import com.linkedin.datahub.graphql.resolvers.load.UsageTypeResolver;
 import com.linkedin.datahub.graphql.resolvers.mutate.MutableTypeResolver;
 import com.linkedin.datahub.graphql.resolvers.type.AspectInterfaceTypeResolver;
 import com.linkedin.datahub.graphql.resolvers.type.HyperParameterValueTypeResolver;
@@ -54,6 +57,7 @@
 import com.linkedin.datahub.graphql.types.lineage.DataFlowDataJobsRelationshipsType;
 import com.linkedin.datahub.graphql.types.glossary.GlossaryTermType;
 
+import com.linkedin.datahub.graphql.types.usage.UsageType;
 import graphql.execution.DataFetcherResult;
 import graphql.schema.idl.RuntimeWiring;
 import org.apache.commons.io.IOUtils;
@@ -103,6 +107,7 @@ public class GmsGraphQLEngine {
     );
     public static final GlossaryTermType GLOSSARY_TERM_TYPE = new GlossaryTermType(GmsClientFactory.getEntitiesClient());
     public static final AspectType ASPECT_TYPE = new AspectType(GmsClientFactory.getAspectsClient());
+    public static final UsageType USAGE_TYPE = new UsageType(GmsClientFactory.getUsageClient());
 
     /**
      * Configures the graph objects that can be fetched primary key.
@@ -200,10 +205,11 @@ public static void configureRuntimeWiring(final RuntimeWiring.Builder builder) {
 
     public static GraphQLEngine.Builder builder() {
         return GraphQLEngine.builder()
-                .addSchema(schema())
-                .addDataLoaders(loaderSuppliers(LOADABLE_TYPES))
-                .addDataLoader("Aspect", (context) -> createAspectLoader(context))
-                .configureRuntimeWiring(GmsGraphQLEngine::configureRuntimeWiring);
+            .addSchema(schema())
+            .addDataLoaders(loaderSuppliers(LOADABLE_TYPES))
+            .addDataLoader("Aspect", (context) -> createAspectLoader(context))
+            .addDataLoader("UsageQueryResult", (context) -> createUsageLoader(context))
+            .configureRuntimeWiring(GmsGraphQLEngine::configureRuntimeWiring);
     }
 
     public static GraphQLEngine get() {
@@ -219,52 +225,52 @@ public static GraphQLEngine get() {
 
     private static void configureQueryResolvers(final RuntimeWiring.Builder builder) {
         builder.type("Query", typeWiring -> typeWiring
-                .dataFetcher("search", new AuthenticatedResolver<>(
-                        new SearchResolver(SEARCHABLE_TYPES)))
-                .dataFetcher("autoComplete", new AuthenticatedResolver<>(
-                        new AutoCompleteResolver(SEARCHABLE_TYPES)))
-                .dataFetcher("autoCompleteForAll", new AuthenticatedResolver<>(
-                        new AutoCompleteForAllResolver(SEARCHABLE_TYPES)))
-                .dataFetcher("browse", new AuthenticatedResolver<>(
-                        new BrowseResolver(BROWSABLE_TYPES)))
-                .dataFetcher("browsePaths", new AuthenticatedResolver<>(
-                        new BrowsePathsResolver(BROWSABLE_TYPES)))
-                .dataFetcher("dataset", new AuthenticatedResolver<>(
-                        new LoadableTypeResolver<>(
-                                DATASET_TYPE,
-                                (env) -> env.getArgument(URN_FIELD_NAME))))
-                .dataFetcher("corpUser", new AuthenticatedResolver<>(
-                        new LoadableTypeResolver<>(
-                                CORP_USER_TYPE,
-                                (env) -> env.getArgument(URN_FIELD_NAME))))
-                .dataFetcher("corpGroup", new AuthenticatedResolver<>(
-                        new LoadableTypeResolver<>(
-                                CORP_GROUP_TYPE,
-                                (env) -> env.getArgument(URN_FIELD_NAME))))
-                .dataFetcher("dashboard", new AuthenticatedResolver<>(
-                        new LoadableTypeResolver<>(
-                                DASHBOARD_TYPE,
-                                (env) -> env.getArgument(URN_FIELD_NAME))))
-                .dataFetcher("chart", new AuthenticatedResolver<>(
-                        new LoadableTypeResolver<>(
-                                CHART_TYPE,
-                                (env) -> env.getArgument(URN_FIELD_NAME))))
-                .dataFetcher("tag", new AuthenticatedResolver<>(
-                        new LoadableTypeResolver<>(
-                                TAG_TYPE,
-                                (env) -> env.getArgument(URN_FIELD_NAME))))
-                .dataFetcher("dataFlow", new AuthenticatedResolver<>(
-                        new LoadableTypeResolver<>(
-                                DATA_FLOW_TYPE,
-                                (env) -> env.getArgument(URN_FIELD_NAME))))
-                .dataFetcher("dataJob", new AuthenticatedResolver<>(
-                        new LoadableTypeResolver<>(
-                                DATA_JOB_TYPE,
-                                (env) -> env.getArgument(URN_FIELD_NAME))))
-                .dataFetcher("glossaryTerm", new AuthenticatedResolver<>(
-                        new LoadableTypeResolver<>(
-                                GLOSSARY_TERM_TYPE,
-                                (env) -> env.getArgument(URN_FIELD_NAME))))
+            .dataFetcher("search", new AuthenticatedResolver<>(
+                    new SearchResolver(SEARCHABLE_TYPES)))
+            .dataFetcher("autoComplete", new AuthenticatedResolver<>(
+                    new AutoCompleteResolver(SEARCHABLE_TYPES)))
+            .dataFetcher("autoCompleteForAll", new AuthenticatedResolver<>(
+                    new AutoCompleteForAllResolver(SEARCHABLE_TYPES)))
+            .dataFetcher("browse", new AuthenticatedResolver<>(
+                    new BrowseResolver(BROWSABLE_TYPES)))
+            .dataFetcher("browsePaths", new AuthenticatedResolver<>(
+                    new BrowsePathsResolver(BROWSABLE_TYPES)))
+            .dataFetcher("dataset", new AuthenticatedResolver<>(
+                    new LoadableTypeResolver<>(
+                            DATASET_TYPE,
+                            (env) -> env.getArgument(URN_FIELD_NAME))))
+            .dataFetcher("corpUser", new AuthenticatedResolver<>(
+                    new LoadableTypeResolver<>(
+                            CORP_USER_TYPE,
+                            (env) -> env.getArgument(URN_FIELD_NAME))))
+            .dataFetcher("corpGroup", new AuthenticatedResolver<>(
+                    new LoadableTypeResolver<>(
+                            CORP_GROUP_TYPE,
+                            (env) -> env.getArgument(URN_FIELD_NAME))))
+            .dataFetcher("dashboard", new AuthenticatedResolver<>(
+                    new LoadableTypeResolver<>(
+                            DASHBOARD_TYPE,
+                            (env) -> env.getArgument(URN_FIELD_NAME))))
+            .dataFetcher("chart", new AuthenticatedResolver<>(
+                    new LoadableTypeResolver<>(
+                            CHART_TYPE,
+                            (env) -> env.getArgument(URN_FIELD_NAME))))
+            .dataFetcher("tag", new AuthenticatedResolver<>(
+                    new LoadableTypeResolver<>(
+                            TAG_TYPE,
+                            (env) -> env.getArgument(URN_FIELD_NAME))))
+            .dataFetcher("dataFlow", new AuthenticatedResolver<>(
+                    new LoadableTypeResolver<>(
+                            DATA_FLOW_TYPE,
+                            (env) -> env.getArgument(URN_FIELD_NAME))))
+            .dataFetcher("dataJob", new AuthenticatedResolver<>(
+                    new LoadableTypeResolver<>(
+                            DATA_JOB_TYPE,
+                            (env) -> env.getArgument(URN_FIELD_NAME))))
+            .dataFetcher("glossaryTerm", new AuthenticatedResolver<>(
+                    new LoadableTypeResolver<>(
+                            GLOSSARY_TERM_TYPE,
+                            (env) -> env.getArgument(URN_FIELD_NAME))))
         );
     }
 
@@ -300,7 +306,8 @@ private static void configureDatasetResolvers(final RuntimeWiring.Builder builde
                                     UPSTREAM_LINEAGE_TYPE,
                                     (env) -> ((Entity) env.getSource()).getUrn()))
                     )
-                    .dataFetcher("schemaMetadata", new AuthenticatedResolver<>(
+                .dataFetcher("usageStats", new AuthenticatedResolver<>(new UsageTypeResolver()))
+                .dataFetcher("schemaMetadata", new AuthenticatedResolver<>(
                         new AspectResolver())
                     )
             )
@@ -311,6 +318,13 @@ private static void configureDatasetResolvers(final RuntimeWiring.Builder builde
                                     (env) -> ((Owner) env.getSource()).getOwner()))
                     )
             )
+            .type("UserUsageCounts", typeWiring -> typeWiring
+                .dataFetcher("user", new AuthenticatedResolver<>(
+                    new LoadableTypeResolver<>(
+                        CORP_USER_TYPE,
+                        (env) -> ((UserUsageCounts) env.getSource()).getUser().getUrn()))
+                )
+            )
             .type("RelatedDataset", typeWiring -> typeWiring
                     .dataFetcher("dataset", new AuthenticatedResolver<>(
                             new LoadableTypeResolver<>(
@@ -554,6 +568,18 @@ private static DataLoader<VersionedAspectKey, DataFetcherResult<Aspect>> createA
         }), loaderOptions);
     }
 
+    private static DataLoader<UsageStatsKey, DataFetcherResult<UsageQueryResult>> createUsageLoader(final QueryContext queryContext) {
+        BatchLoaderContextProvider contextProvider = () -> queryContext;
+        DataLoaderOptions loaderOptions = DataLoaderOptions.newOptions().setBatchLoaderContextProvider(contextProvider);
+        return DataLoader.newDataLoader((keys, context) -> CompletableFuture.supplyAsync(() -> {
+            try {
+                return USAGE_TYPE.batchLoad(keys, context.getContext());
+            } catch (Exception e) {
+                throw new RuntimeException(String.format("Failed to retrieve usage stats", e));
+            }
+        }), loaderOptions);
+    }
+
     private GmsGraphQLEngine() { }
 
 }

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

@@ -0,0 +1,16 @@
+package com.linkedin.datahub.graphql;
+
+import com.linkedin.usage.UsageTimeRange;
+import lombok.Data;
+
+
+@Data
+public class UsageStatsKey {
+  private String resource;
+  private UsageTimeRange range;
+
+  public UsageStatsKey(String resource, UsageTimeRange range) {
+    this.resource = resource;
+    this.range = range;
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/UsageTypeResolver.java

@@ -0,0 +1,37 @@
+package com.linkedin.datahub.graphql.resolvers.load;
+
+import com.linkedin.datahub.graphql.UsageStatsKey;
+
+import com.linkedin.datahub.graphql.types.LoadableType;
+import com.linkedin.pegasus2avro.usage.UsageQueryResult;
+import com.linkedin.usage.UsageTimeRange;
+import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
+import java.util.concurrent.CompletableFuture;
+import org.dataloader.DataLoader;
+
+
+/**
+ * Generic GraphQL resolver responsible for
+ *
+ *    1. Retrieving a single input urn.
+ *    2. Resolving a single {@link LoadableType}.
+ *
+ *  Note that this resolver expects that {@link DataLoader}s were registered
+ *  for the provided {@link LoadableType} under the name provided by {@link LoadableType#name()}
+ *
+ */
+public class UsageTypeResolver implements DataFetcher<CompletableFuture<UsageQueryResult>> {
+
+    @Override
+    public CompletableFuture<UsageQueryResult> get(DataFetchingEnvironment environment) {
+        final DataLoader<UsageStatsKey, UsageQueryResult> loader = environment.getDataLoaderRegistry().getDataLoader("UsageQueryResult");
+
+        String resource = environment.getArgument("resource");
+        UsageTimeRange duration = UsageTimeRange.valueOf(environment.getArgument("range"));
+
+        UsageStatsKey key = new UsageStatsKey(resource, duration);
+
+        return loader.load(key);
+    }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/usage/UsageAggregationMapper.java

@@ -0,0 +1,34 @@
+package com.linkedin.datahub.graphql.types.usage;
+
+import com.linkedin.datahub.graphql.generated.UsageAggregation;
+import com.linkedin.datahub.graphql.generated.WindowDuration;
+import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
+import javax.annotation.Nonnull;
+
+
+public class UsageAggregationMapper implements
+                                               ModelMapper<com.linkedin.usage.UsageAggregation, UsageAggregation> {
+
+  public static final UsageAggregationMapper INSTANCE = new UsageAggregationMapper();
+
+  public static UsageAggregation map(@Nonnull final com.linkedin.usage.UsageAggregation pdlUsageAggregation) {
+    return INSTANCE.apply(pdlUsageAggregation);
+  }
+
+  @Override
+  public UsageAggregation apply(@Nonnull final com.linkedin.usage.UsageAggregation pdlUsageAggregation) {
+    UsageAggregation result = new UsageAggregation();
+    result.setBucket(pdlUsageAggregation.getBucket());
+
+    if (pdlUsageAggregation.hasDuration()) {
+      result.setDuration(WindowDuration.valueOf(pdlUsageAggregation.getDuration().toString()));
+    }
+    if (pdlUsageAggregation.hasResource()) {
+      result.setResource(pdlUsageAggregation.getResource().toString());
+    }
+    if (pdlUsageAggregation.hasMetrics()) {
+      result.setMetrics(UsageAggregationMetricsMapper.map(pdlUsageAggregation.getMetrics()));
+    }
+    return result;
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/usage/UsageAggregationMetricsMapper.java

@@ -0,0 +1,33 @@
+package com.linkedin.datahub.graphql.types.usage;
+
+import com.linkedin.datahub.graphql.generated.UsageAggregationMetrics;
+import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
+import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
+
+
+public class UsageAggregationMetricsMapper implements
+                                               ModelMapper<com.linkedin.usage.UsageAggregationMetrics, UsageAggregationMetrics> {
+
+  public static final UsageAggregationMetricsMapper INSTANCE = new UsageAggregationMetricsMapper();
+
+  public static UsageAggregationMetrics map(@Nonnull final com.linkedin.usage.UsageAggregationMetrics usageAggregationMetrics) {
+    return INSTANCE.apply(usageAggregationMetrics);
+  }
+
+  @Override
+  public UsageAggregationMetrics apply(@Nonnull final com.linkedin.usage.UsageAggregationMetrics usageAggregationMetrics) {
+    UsageAggregationMetrics result = new UsageAggregationMetrics();
+    result.setTotalSqlQueries(usageAggregationMetrics.getTotalSqlQueries());
+    result.setUniqueUserCount(usageAggregationMetrics.getUniqueUserCount());
+    result.setTopSqlQueries(usageAggregationMetrics.getTopSqlQueries());
+    if (usageAggregationMetrics.hasUsers()) {
+      result.setUsers(usageAggregationMetrics.getUsers()
+          .stream()
+          .map(aggregation -> UserUsageCountsMapper.map(aggregation))
+          .collect(Collectors.toList()));
+    }
+
+    return result;
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/usage/UsageQueryResultAggregationMapper.java

@@ -0,0 +1,31 @@
+package com.linkedin.datahub.graphql.types.usage;
+
+import com.linkedin.datahub.graphql.generated.UsageQueryResultAggregations;
+import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
+import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
+
+
+public class UsageQueryResultAggregationMapper implements
+                                               ModelMapper<com.linkedin.usage.UsageQueryResultAggregations, UsageQueryResultAggregations> {
+
+  public static final UsageQueryResultAggregationMapper INSTANCE = new UsageQueryResultAggregationMapper();
+
+  public static UsageQueryResultAggregations map(@Nonnull final com.linkedin.usage.UsageQueryResultAggregations pdlUsageResultAggregations) {
+    return INSTANCE.apply(pdlUsageResultAggregations);
+  }
+
+  @Override
+  public UsageQueryResultAggregations apply(@Nonnull final com.linkedin.usage.UsageQueryResultAggregations pdlUsageResultAggregations) {
+    UsageQueryResultAggregations result = new UsageQueryResultAggregations();
+    result.setTotalSqlQueries(pdlUsageResultAggregations.getTotalSqlQueries());
+    result.setUniqueUserCount(pdlUsageResultAggregations.getUniqueUserCount());
+    if (pdlUsageResultAggregations.hasUsers()) {
+      result.setUsers(pdlUsageResultAggregations.getUsers()
+          .stream()
+          .map(aggregation -> UserUsageCountsMapper.map(aggregation))
+          .collect(Collectors.toList()));
+    }
+    return result;
+  }
+}